OpenLLM/tools/dependencies.py

#!/usr/bin/env python3
# Copyright 2023 BentoML Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from __future__ import annotations

import dataclasses
import os
import typing as t

import inflection
import tomlkit

import openllm

if t.TYPE_CHECKING:
    from tomlkit.items import Array
    from tomlkit.items import Table


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


@dataclasses.dataclass(frozen=True)
class Classifier:
    identifier: t.Dict[str, str] = dataclasses.field(
        default_factory=lambda: {
            "status": "Development Status",
            "environment": "Environment",
            "license": "License",
            "topic": "Topic",
            "os": "Operating System",
            "audience": "Intended Audience",
            "typing": "Typing",
            "language": "Programming Language",
        }
    )

    joiner: str = " :: "

    @staticmethod
    def status() -> dict[int, str]:
        return {
            v: status
            for v, status in zip(
                range(1, 8),
                [
                    "1 - Planning",
                    "2 - Pre-Alpha",
                    "3 - Alpha",
                    "4 - Beta",
                    "5 - Production/Stable",
                    "6 - Mature",
                    "7 - Inactive",
                ],
            )
        }

    @staticmethod
    def apache() -> str:
        return Classifier.create_classifier("license", "OSI Approved", "Apache Software License")

    @staticmethod
    def create_classifier(identifier: str, *decls: t.Any) -> str:
        cls_ = Classifier()
        if identifier not in cls_.identifier:
            raise ValueError(f"{identifier} is not yet supported (supported alias: {Classifier.identifier})")
        return cls_.joiner.join([cls_.identifier[identifier], *decls])

    @staticmethod
    def create_python_classifier(
        implementation: list[str] | None = None, supported_version: list[str] | None = None
    ) -> list[str]:
        if supported_version is None:
            supported_version = ["3.8", "3.9", "3.10", "3.11", "3.12"]
        if implementation is None:
            implementation = ["CPython", "PyPy"]
        base = [
            Classifier.create_classifier("language", "Python"),
            Classifier.create_classifier("language", "Python", "3"),
        ]
        base.append(Classifier.create_classifier("language", "Python", "3", "Only"))
        base.extend([Classifier.create_classifier("language", "Python", version) for version in supported_version])
        base.extend(
            [Classifier.create_classifier("language", "Python", "Implementation", impl) for impl in implementation]
        )
        return base

    @staticmethod
    def create_status_classifier(level: int) -> str:
        return Classifier.create_classifier("status", Classifier.status()[level])


@dataclasses.dataclass(frozen=True)
class Dependencies:
    name: str
    git_repo_url: t.Optional[str] = None
    branch: t.Optional[str] = None
    extensions: t.Optional[t.List[str]] = None
    subdirectory: t.Optional[str] = None
    requires_gpu: bool = False
    lower_constraint: t.Optional[str] = None
    upper_constraint: t.Optional[str] = None
    platform: t.Optional[t.Tuple[t.Literal["Linux", "Windows", "Darwin"], t.Literal["eq", "ne"]]] = None

    def with_options(self, **kwargs: t.Any) -> Dependencies:
        return dataclasses.replace(self, **kwargs)

    @property
    def has_constraint(self) -> bool:
        return self.lower_constraint is not None or self.upper_constraint is not None

    @property
    def pypi_extensions(self) -> str:
        return "" if self.extensions is None else f"[{','.join(self.extensions)}]"

    @staticmethod
    def platform_restriction(platform: t.LiteralString, op: t.Literal["eq", "ne"] = "eq") -> str:
        return f'platform_system{"==" if op == "eq" else "!="}"{platform}"'

    def to_str(self) -> str:
        deps: list[str] = []
        if self.lower_constraint is not None and self.upper_constraint is not None:
            dep = f"{self.name}{self.pypi_extensions}>={self.lower_constraint},<{self.upper_constraint}"
        elif self.lower_constraint is not None:
            dep = f"{self.name}{self.pypi_extensions}>={self.lower_constraint}"
        elif self.upper_constraint is not None:
            dep = f"{self.name}{self.pypi_extensions}<{self.upper_constraint}"
        elif self.subdirectory is not None:
            dep = f"{self.name}{self.pypi_extensions} @ git+https://github.com/{self.git_repo_url}.git#subdirectory={self.subdirectory}"
        elif self.branch is not None:
            dep = f"{self.name}{self.pypi_extensions} @ git+https://github.com/{self.git_repo_url}.git@{self.branch}"
        else:
            dep = f"{self.name}{self.pypi_extensions}"

        deps.append(dep)

        if self.platform:
            deps.append(self.platform_restriction(*self.platform))

        return ";".join(deps)

    @classmethod
    def from_tuple(cls, *decls: t.Any) -> Dependencies:
        return cls(*decls)


_BENTOML_EXT = ["grpc", "io"]
_TRANSFORMERS_EXT = ["torch", "tokenizers", "accelerate"]

_BASE_DEPENDENCIES = [
    Dependencies(name="bentoml", extensions=_BENTOML_EXT, lower_constraint="1.0.25"),
    Dependencies(name="transformers", extensions=_TRANSFORMERS_EXT, lower_constraint="4.29.0"),
    Dependencies(name="safetensors"),
    Dependencies(name="optimum"),
    Dependencies(name="attrs", lower_constraint="23.1.0"),
    Dependencies(name="cattrs", lower_constraint="23.1.0"),
    Dependencies(name="orjson"),
    Dependencies(name="inflection"),
    Dependencies(name="tabulate", extensions=["widechars"], lower_constraint="0.9.0"),
    Dependencies(name="httpx"),
    Dependencies(name="typing_extensions"),
    Dependencies(name="cuda-python", platform=("Darwin", "ne")),
    Dependencies(name="bitsandbytes", upper_constraint="0.40"),  # Currently only <0.40 works with CUDA 11.8
]

_NIGHTLY_MAPPING: dict[str, Dependencies] = {
    "bentoml": Dependencies.from_tuple("bentoml", "bentoml/bentoml", "main", _BENTOML_EXT),
    "peft": Dependencies.from_tuple("peft", "huggingface/peft", "main", None),
    "transformers": Dependencies.from_tuple("transformers", "huggingface/transformers", "main", _TRANSFORMERS_EXT),
    "optimum": Dependencies.from_tuple("optimum", "huggingface/optimum", "main", None),
    "accelerate": Dependencies.from_tuple("accelerate", "huggingface/accelerate", "main", None),
    "bitsandbytes": Dependencies.from_tuple("bitsandbytes", "TimDettmers/bitsandbytes", "main", None),
    "trl": Dependencies.from_tuple("trl", "lvwerra/trl", "main", None),
    "vllm": Dependencies.from_tuple("vllm", "vllm-project/vllm", "main", None, None, True, None),
}

_ALL_RUNTIME_DEPS = ["flax", "jax", "jaxlib", "tensorflow", "keras"]
FINE_TUNE_DEPS = ["peft", "bitsandbytes", "datasets", "accelerate", "deepspeed", "trl"]
FLAN_T5_DEPS = _ALL_RUNTIME_DEPS
OPT_DEPS = _ALL_RUNTIME_DEPS
OPENAI_DEPS = ["openai", "tiktoken"]
AGENTS_DEPS = ["transformers[agents]>=4.30", "diffusers", "soundfile"]
PLAYGROUND_DEPS = ["jupyter", "notebook", "ipython", "jupytext", "nbformat"]
GGML_DEPS = ["ctransformers"]
GPTQ_DEPS = ["auto-gptq[triton]"]
VLLM_DEPS = ["vllm", "ray"]

_base_requirements = {
    inflection.dasherize(name): config_cls.__openllm_requirements__
    for name, config_cls in openllm.CONFIG_MAPPING.items()
    if config_cls.__openllm_requirements__
}

# shallow copy from locals()
_locals = locals().copy()

# NOTE: update this table when adding new external dependencies
# sync with openllm.utils.OPTIONAL_DEPENDENCIES
_base_requirements.update(
    {v: _locals.get(f"{inflection.underscore(v).upper()}_DEPS") for v in openllm.utils.OPTIONAL_DEPENDENCIES}
)

_base_requirements = {k: v for k, v in sorted(_base_requirements.items())}

fname = f"{os.path.basename(os.path.dirname(__file__))}/{os.path.basename(__file__)}"


def create_classifiers() -> Array:
    arr = tomlkit.array()
    arr.extend(
        [
            Classifier.create_status_classifier(5),
            Classifier.create_classifier("environment", "GPU", "NVIDIA CUDA"),
            Classifier.create_classifier("environment", "GPU", "NVIDIA CUDA", "12"),
            Classifier.create_classifier("environment", "GPU", "NVIDIA CUDA", "11.8"),
            Classifier.create_classifier("environment", "GPU", "NVIDIA CUDA", "11.7"),
            Classifier.apache(),
            Classifier.create_classifier("topic", "Scientific/Engineering", "Artificial Intelligence"),
            Classifier.create_classifier("topic", "Software Development", "Libraries"),
            Classifier.create_classifier("os", "OS Independent"),
            Classifier.create_classifier("audience", "Developers"),
            Classifier.create_classifier("audience", "Science/Research"),
            Classifier.create_classifier("audience", "System Administrators"),
            Classifier.create_classifier("typing", "Typed"),
            *Classifier.create_python_classifier(),
        ]
    )
    return arr.multiline(True)


def create_optional_table() -> Table:
    all_array = tomlkit.array()
    all_array.extend([f"openllm[{k}]" for k in _base_requirements])

    table = tomlkit.table(is_super_table=True)
    _base_requirements.update({"all": all_array.multiline(True)})
    table.update({k: v for k, v in sorted(_base_requirements.items())})
    table.add(tomlkit.nl())

    return table


def create_url_table() -> Table:
    table = tomlkit.table()
    _urls = {
        "Blog": "https://modelserving.com",
        "Discord": "https://l.bentoml.com/join-openllm-discord",
        "Documentation": "https://github.com/bentoml/openllm#readme",
        "GitHub": "https://github.com/bentoml/openllm",
        "History": "https://github.com/bentoml/openllm/blob/main/CHANGELOG.md",
        "Homepage": "https://bentoml.com",
        "Tracker": "https://github.com/bentoml/openllm/issues",
        "Twitter": "https://twitter.com/bentomlai",
    }
    table.update({k: v for k, v in sorted(_urls.items())})
    return table


def main() -> int:
    with open(os.path.join(ROOT, "pyproject.toml"), "r") as f:
        pyproject = tomlkit.parse(f.read())

    dependencies_array = tomlkit.array()
    dependencies_array.extend([v.to_str() for v in _BASE_DEPENDENCIES])

    pyproject["project"]["urls"] = create_url_table()
    pyproject["project"]["classifiers"] = create_classifiers()
    pyproject["project"]["optional-dependencies"] = create_optional_table()
    pyproject["project"]["scripts"] = {"openllm": "openllm.cli:cli"}
    pyproject["project"]["dependencies"] = dependencies_array.multiline(True)

    with open(os.path.join(ROOT, "pyproject.toml"), "w") as f:
        f.write(tomlkit.dumps(pyproject))

    with open(os.path.join(ROOT, "nightly-requirements.txt"), "w") as f:
        f.write(f"# This file is generated by `{fname}`. DO NOT EDIT\n-e .[playground,flan-t5]\n")
        f.writelines([f"{v.to_str()}\n" for v in _NIGHTLY_MAPPING.values() if not v.requires_gpu])
    with open(os.path.join(ROOT, "nightly-requirements-gpu.txt"), "w") as f:
        f.write(f"# This file is generated by `{fname}`. # DO NOT EDIT\n")
        f.write(
            "# For Jax, Flax, Tensorflow, PyTorch CUDA support, please refers to their official installation for your specific setup.\n"
        )
        f.write("-r nightly-requirements.txt\n-e .[all]\n")
        f.writelines([f"{v.to_str()}\n" for v in _NIGHTLY_MAPPING.values() if v.requires_gpu])

    return 0


if __name__ == "__main__":
    raise SystemExit(main())