perf: faster LLM loading

using attrs for faster class creation opposed to metaclass

Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>

Author: aarnphm

src/openllm/_configuration.py

@@ -66,7 +66,7 @@ class GenerationConfig:
 import openllm
 
 from .exceptions import ForbiddenAttributeError, GpuNotAvailableError, OpenLLMException
-from .utils import DEBUG, LazyType, bentoml_cattr, dantic, first_not_none, lenient_issubclass
+from .utils import DEBUG, LazyType, bentoml_cattr, codegen, dantic, first_not_none, lenient_issubclass
 
 if hasattr(t, "Required"):
     from typing import Required
@@ -85,7 +85,7 @@ class GenerationConfig:
     import tensorflow as tf
     import torch
     import transformers
-    from attr import _CountingAttr, _make_init, _make_method, _make_repr, _transform_attrs
+    from attr import _CountingAttr, _make_init, _make_repr, _transform_attrs  # type: ignore
     from transformers.generation.beam_constraints import Constraint
 
     from ._types import ClickFunctionWrapper, F, O_co, P
@@ -103,7 +103,7 @@ class GenerationConfig:
     ItemgetterAny = itemgetter
     # NOTE: Using internal API from attr here, since we are actually
     # allowing subclass of openllm.LLMConfig to become 'attrs'-ish
-    from attr._make import _CountingAttr, _make_init, _make_method, _make_repr, _transform_attrs
+    from attr._make import _CountingAttr, _make_init, _make_repr, _transform_attrs
 
     transformers = openllm.utils.LazyLoader("transformers", globals(), "transformers")
     torch = openllm.utils.LazyLoader("torch", globals(), "torch")
@@ -389,89 +389,10 @@ def _populate_value_from_env_var(
     return os.environ.get(key, fallback)
 
 
-# sentinel object for unequivocal object() getattr
-_sentinel = object()
-
-
 def _field_env_key(model_name: str, key: str, suffix: str | None = None) -> str:
     return "_".join(filter(None, map(str.upper, ["OPENLLM", model_name, suffix.strip("_") if suffix else "", key])))
 
 
-def _has_own_attribute(cls: type[t.Any], attrib_name: t.Any):
-    """
-    Check whether *cls* defines *attrib_name* (and doesn't just inherit it).
-    """
-    attr = getattr(cls, attrib_name, _sentinel)
-    if attr is _sentinel:
-        return False
-
-    for base_cls in cls.__mro__[1:]:
-        a = getattr(base_cls, attrib_name, None)
-        if attr is a:
-            return False
-
-    return True
-
-
-def _get_annotations(cls: type[t.Any]) -> DictStrAny:
-    """
-    Get annotations for *cls*.
-    """
-    if _has_own_attribute(cls, "__annotations__"):
-        return cls.__annotations__
-
-    return DictStrAny()
-
-
-_classvar_prefixes = (
-    "typing.ClassVar",
-    "t.ClassVar",
-    "ClassVar",
-    "typing_extensions.ClassVar",
-)
-
-
-def _is_class_var(annot: str | t.Any) -> bool:
-    """
-    Check whether *annot* is a typing.ClassVar.
-
-    The string comparison hack is used to avoid evaluating all string
-    annotations which would put attrs-based classes at a performance
-    disadvantage compared to plain old classes.
-    """
-    annot = str(annot)
-
-    # Annotation can be quoted.
-    if annot.startswith(("'", '"')) and annot.endswith(("'", '"')):
-        annot = annot[1:-1]
-
-    return annot.startswith(_classvar_prefixes)
-
-
-def _add_method_dunders(cls: type[t.Any], method_or_cls: _T, _overwrite_doc: str | None = None) -> _T:
-    """
-    Add __module__ and __qualname__ to a *method* if possible.
-    """
-    try:
-        method_or_cls.__module__ = cls.__module__
-    except AttributeError:
-        pass
-
-    try:
-        method_or_cls.__qualname__ = ".".join((cls.__qualname__, method_or_cls.__name__))
-    except AttributeError:
-        pass
-
-    try:
-        method_or_cls.__doc__ = (
-            _overwrite_doc or "Method or class generated by LLMConfig for class " f"{cls.__qualname__}."
-        )
-    except AttributeError:
-        pass
-
-    return method_or_cls
-
-
 # cached it here to save one lookup per assignment
 _object_getattribute = object.__getattribute__
 
@@ -506,8 +427,8 @@ class ModelSettings(t.TypedDict, total=False):
     generation_class: t.Type[GenerationConfig]
 
 
-_ModelSettings: type[attr.AttrsInstance] = _add_method_dunders(
-    type("__internal__", (ModelSettings,), {"__module__": "openllm._configuration"}),
+_ModelSettings: type[attr.AttrsInstance] = codegen.add_method_dunders(
+    type("__openllm_internal__", (ModelSettings,), {"__module__": "openllm._configuration"}),
     attr.make_class(
         "ModelSettings",
         {
@@ -563,7 +484,7 @@ def structure_settings(cl_: type[LLMConfig], cls: type[t.Any]):
     partialed = functools.partial(_field_env_key, model_name=model_name, suffix="generation")
 
     def auto_env_transformers(_: t.Any, fields: list[attr.Attribute[t.Any]]) -> list[attr.Attribute[t.Any]]:
-        _has_own_gen = _has_own_attribute(cl_, "GenerationConfig")
+        _has_own_gen = codegen.has_own_attribute(cl_, "GenerationConfig")
         return [
             f.evolve(
                 default=_populate_value_from_env_var(
@@ -576,20 +497,20 @@ def auto_env_transformers(_: t.Any, fields: list[attr.Attribute[t.Any]]) -> list
             for f in fields
         ]
 
-    _target: DictStrAny = {
-        "default_id": settings["default_id"],
-        "model_ids": settings["model_ids"],
-        "url": settings.get("url", ""),
-        "requires_gpu": settings.get("requires_gpu", False),
-        "trust_remote_code": settings.get("trust_remote_code", False),
-        "requirements": settings.get("requirements", None),
-        "name_type": name_type,
-        "model_name": model_name,
-        "start_name": start_name,
-        "env": openllm.utils.ModelEnv(model_name),
-        "timeout": settings.get("timeout", 3600),
-        "workers_per_resource": settings.get("workers_per_resource", 1),
-        "generation_class": attr.make_class(
+    return cls(
+        default_id=settings["default_id"],
+        model_ids=settings["model_ids"],
+        url=settings.get("url", ""),
+        requires_gpu=settings.get("requires_gpu", False),
+        trust_remote_code=settings.get("trust_remote_code", False),
+        requirements=settings.get("requirements", None),
+        name_type=name_type,
+        model_name=model_name,
+        start_name=start_name,
+        env=openllm.utils.ModelEnv(model_name),
+        timeout=settings.get("timeout", 3600),
+        workers_per_resource=settings.get("workers_per_resource", 1),
+        generation_class=attr.make_class(
             f"{_cl_name}GenerationConfig",
             [],
             bases=(GenerationConfig,),
@@ -599,18 +520,12 @@ def auto_env_transformers(_: t.Any, fields: list[attr.Attribute[t.Any]]) -> list
             repr=True,
             field_transformer=auto_env_transformers,
         ),
-    }
-
-    return cls(**_target)
+    )
 
 
 bentoml_cattr.register_structure_hook(_ModelSettings, structure_settings)
 
 
-def _generate_unique_filename(cls: type[t.Any], func_name: str):
-    return f"<LLMConfig generated {func_name} {cls.__module__}." f"{getattr(cls, '__qualname__', cls.__name__)}>"
-
-
 def _setattr_class(attr_name: str, value_var: t.Any, add_dunder: bool = False):
     """
     Use the builtin setattr to set *attr_name* to *value_var*.
@@ -632,7 +547,7 @@ def _make_assignment_script(cls: type[LLMConfig], attributes: attr.AttrsInstance
         "cls": cls,
         "_cached_attribute": attributes,
         "_cached_getattribute_get": _object_getattribute.__get__,
-        "__add_dunder": _add_method_dunders,
+        "__add_dunder": codegen.add_method_dunders,
     }
     annotations: DictStrAny = {"return": None}
 
@@ -643,19 +558,9 @@ def _make_assignment_script(cls: type[LLMConfig], attributes: attr.AttrsInstance
         lines.append(_setattr_class(arg_name, attr_name, add_dunder=attr_name in _dunder_add))
         annotations[attr_name] = field.type
 
-    script = "def __assign_attr(cls, %s):\n    %s\n" % (", ".join(args), "\n    ".join(lines) if lines else "pass")
-    assign_method = _make_method(
-        "__assign_attr",
-        script,
-        _generate_unique_filename(cls, "__assign_attr"),
-        globs,
+    return codegen.generate_function(
+        cls, "__assign_attr", lines, args=("cls", *args), globs=globs, annotations=annotations
     )
-    assign_method.__annotations__ = annotations
-
-    if DEBUG:
-        logger.info("Generated script:\n%s", script)
-
-    return assign_method
 
 
 _reserved_namespace = {"__config__", "GenerationConfig"}
@@ -841,7 +746,7 @@ def __init_subclass__(cls):
         _make_assignment_script(cls, bentoml_cattr.structure(cls, _ModelSettings))(cls)
         # process a fields under cls.__dict__ and auto convert them with dantic.Field
         cd = cls.__dict__
-        anns = _get_annotations(cls)
+        anns = codegen.get_annotations(cls)
         partialed = functools.partial(_field_env_key, model_name=cls.__openllm_model_name__)
 
         def auto_config_env(_: type[LLMConfig], attrs: list[attr.Attribute[t.Any]]) -> list[attr.Attribute[t.Any]]:
@@ -861,7 +766,7 @@ def auto_config_env(_: type[LLMConfig], attrs: list[attr.Attribute[t.Any]]) -> l
         these: dict[str, _CountingAttr[t.Any]] = {}
         annotated_names: set[str] = set()
         for attr_name, typ in anns.items():
-            if _is_class_var(typ):
+            if codegen.is_class_var(typ):
                 continue
             annotated_names.add(attr_name)
             val = cd.get(attr_name, attr.NOTHING)
@@ -907,7 +812,7 @@ def auto_config_env(_: type[LLMConfig], attrs: list[attr.Attribute[t.Any]]) -> l
         cls.__attrs_attrs__ = attrs
         # generate a __attrs_init__ for the subclass, since we will
         # implement a custom __init__
-        cls.__attrs_init__ = _add_method_dunders(
+        cls.__attrs_init__ = codegen.add_method_dunders(
             cls,
             _make_init(
                 cls,  # cls (the attrs-decorated class)
@@ -924,7 +829,7 @@ def auto_config_env(_: type[LLMConfig], attrs: list[attr.Attribute[t.Any]]) -> l
             ),
         )
         # __repr__ function with the updated fields.
-        cls.__repr__ = _add_method_dunders(cls, _make_repr(cls.__attrs_attrs__, None, cls))
+        cls.__repr__ = codegen.add_method_dunders(cls, _make_repr(cls.__attrs_attrs__, None, cls))
         # Traverse the MRO to collect existing slots
         # and check for an existing __weakref__.
         existing_slots: DictStrAny = dict()