fix(container): preheat packages support markers

bentoml · May 16, 2024 · 6d005da · 6d005da
1 parent 078b137
commit 6d005da
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/src/bentoml/_internal/bento/build_config.py b/src/bentoml/_internal/bento/build_config.py
@@ -518,12 +518,14 @@ def is_empty(self) -> bool:
 
     @property
     def _jinja_environment(self) -> jinja2.Environment:
-        return jinja2.Environment(
+        env = jinja2.Environment(
             extensions=["jinja2.ext.debug"],
             variable_start_string="<<",
             variable_end_string=">>",
             loader=jinja2.FileSystemLoader(os.path.dirname(__file__), followlinks=True),
         )
+        env.filters["bash_quote"] = shlex.quote
+        return env
 
     def write_to_bento(self, bento_fs: FS, build_ctx: str) -> None:
         from .bentoml_builder import build_bentoml_sdist

diff --git a/src/bentoml/_internal/container/frontend/dockerfile/templates/base.j2 b/src/bentoml/_internal/container/frontend/dockerfile/templates/base.j2
@@ -63,7 +63,7 @@ WORKDIR $BENTO_PATH
 
 {% if __pip_preheat_packages__ %}
 {% for value in __pip_preheat_packages__ -%}
-{% call common.RUN(__enable_buildkit__) -%} {{ __pip_cache__ }} {% endcall -%} bash -c "pip3 install {{value}} || true"
+{% call common.RUN(__enable_buildkit__) -%} {{ __pip_cache__ }} {% endcall -%} pip3 install {{value | bash_quote}} ; exit 0
 {% endfor -%}
 {% endif -%}
 

diff --git a/src/bentoml/_internal/container/generate.py b/src/bentoml/_internal/container/generate.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import shlex
 import os
 import typing as t
 from typing import TYPE_CHECKING
@@ -166,6 +167,7 @@ def generate_containerfile(
         lstrip_blocks=True,
         loader=FileSystemLoader(TEMPLATES_PATH, followlinks=True),
     )
+    ENVIRONMENT.filters["bash_quote"] = shlex.quote
 
     if docker.cuda_version is not None:
         release_type = "cuda"