chore(container): cache large pypi packages in layers (#4690)

* chore(container): preheat large pypi package * fix
bentoml · May 6, 2024 · 61b731d · 61b731d
1 parent dd94c93
commit 61b731d
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 0 deletions.
diff --git a/src/bentoml/_internal/container/frontend/dockerfile/templates/base.j2 b/src/bentoml/_internal/container/frontend/dockerfile/templates/base.j2
@@ -60,6 +60,13 @@ WORKDIR $BENTO_PATH
 
 # Block SETUP_BENTO_COMPONENTS
 {% block SETUP_BENTO_COMPONENTS %}
+
+{% if __pip_preheat_packages__ %}
+{% for value in __pip_preheat_packages__ -%}
+{% call common.RUN(__enable_buildkit__) -%} {{ __pip_cache__ }} {% endcall -%} bash -c "pip3 install {{value}} || true"
+{% endfor -%}
+{% endif -%}
+
 COPY --chown={{ bento__user }}:{{ bento__user }} ./env/python ./env/python/
 {% set __install_python_scripts__ = expands_bento_path("env", "python", "install.sh", bento_path=bento__path) %}
 {% set __pip_cache__ = common.mount_cache("/root/.cache/pip") %}

diff --git a/src/bentoml/_internal/container/generate.py b/src/bentoml/_internal/container/generate.py
@@ -57,6 +57,7 @@ def get_templates_variables(
     conda: CondaOptions,
     bento_fs: FS,
     *,
+    python_packages: dict[str, str] | None = None,
     _is_cuda: bool = False,
     **bento_env: str | bool,
 ) -> dict[str, t.Any]:
@@ -97,13 +98,20 @@ def get_templates_variables(
     if bento_env:
         default_env.update(bento_env)
 
+    PREHEAT_PIP_PACKAGES = ["torch", "vllm"]
+
     return {
         **{to_options_field(k): v for k, v in docker.to_dict().items()},
         **{to_bento_field(k): v for k, v in default_env.items()},
         "__prometheus_port__": BentoMLContainer.grpc.metrics.port.get(),
         "__base_image__": base_image,
         "__conda_python_version__": conda_python_version,
         "__is_cuda__": _is_cuda,
+        "__pip_preheat_packages__": [
+            python_packages[k]
+            for k in PREHEAT_PIP_PACKAGES
+            if python_packages and python_packages.get(k)
+        ],
     }
 
 
@@ -190,12 +198,46 @@ def generate_containerfile(
             globals={"bento_base_template": template, **J2_FUNCTION},
         )
 
+    requirement_file = bento_fs.getsyspath("env/python/requirements.lock.txt")
+    if not os.path.exists(requirement_file):
+        requirement_file = bento_fs.getsyspath("env/python/requirements.txt")
+    if os.path.exists(requirement_file):
+        python_packages = _resolve_package_versions(requirement_file)
+    else:
+        python_packages = {}
+
     return template.render(
         **get_templates_variables(
             docker,
             conda,
             bento_fs,
+            python_packages=python_packages,
             _is_cuda=release_type == "cuda",
             **override_bento_env,
         )
     )
+
+
+def _resolve_package_versions(requirement: str) -> dict[str, str]:
+    from pip_requirements_parser import RequirementsFile
+
+    requirements_txt = RequirementsFile.from_file(
+        requirement,
+        include_nested=True,
+    )
+    deps: dict[str, str] = {}
+    for req in requirements_txt.requirements:
+        if (
+            req.is_editable
+            or req.is_local_path
+            or req.is_url
+            or req.is_wheel
+            or not req.name
+        ):
+            continue
+        for sp in req.specifier:
+            if sp.operator == "==":
+                assert req.line is not None
+                deps[req.name] = req.line
+                break
+    return deps