Adds ImageSpec.with_runtime_packages

flytekit/bin/entrypoint.py

@@ -30,7 +30,7 @@
 from flytekit.core import utils
 from flytekit.core.base_task import IgnoreOutputs, PythonTask
 from flytekit.core.checkpointer import SyncCheckpoint
-from flytekit.core.constants import FLYTE_FAIL_ON_ERROR
+from flytekit.core.constants import FLYTE_FAIL_ON_ERROR, RUNTIME_PACKAGES_ENV_NAME
 from flytekit.core.context_manager import (
     ExecutionParameters,
     ExecutionState,
@@ -63,6 +63,18 @@ def get_version_message():
     return f"Welcome to Flyte! Version: {flytekit.__version__}"
 
 
+def _run_subprocess(cmd: List[str], env: Optional[dict] = None) -> int:
+    """Run cmd with proper SIGTERM handling."""
+    p = subprocess.Popen(cmd, env=env)
+
+    def handle_sigterm(signum, frame):
+        logger.info(f"passing signum {signum} [frame={frame}] to subprocess")
+        p.send_signal(signum)
+
+    signal.signal(signal.SIGTERM, handle_sigterm)
+    return p.wait()
+
+
 def _compute_array_job_index():
     """
     Computes the absolute index of the current array job. This is determined by summing the compute-environment-specific
@@ -432,6 +444,14 @@ def setup_execution(
 
     compressed_serialization_settings = os.environ.get(SERIALIZED_CONTEXT_ENV_VAR, "")
 
+    if runtime_packages := os.getenv(RUNTIME_PACKAGES_ENV_NAME):
+        import importlib
+        import site
+
+        dev_packages_list = runtime_packages.split(" ")
+        _run_subprocess([sys.executable, "-m", "pip", "install", *dev_packages_list])
+        importlib.reload(site)
+
     ctx = FlyteContextManager.current_context()
     # Create directories
     user_workspace_dir = ctx.file_access.get_random_local_directory()
@@ -751,14 +771,7 @@ def fast_execute_task_cmd(additional_distribution: str, dest_dir: str, task_exec
             env["PYTHONPATH"] += os.pathsep + dest_dir_resolved
         else:
             env["PYTHONPATH"] = dest_dir_resolved
-    p = subprocess.Popen(cmd, env=env)
-
-    def handle_sigterm(signum, frame):
-        logger.info(f"passing signum {signum} [frame={frame}] to subprocess")
-        p.send_signal(signum)
-
-    signal.signal(signal.SIGTERM, handle_sigterm)
-    returncode = p.wait()
+    returncode = _run_subprocess(cmd, env)
     exit(returncode)
 
 

flytekit/core/constants.py

@@ -44,3 +44,6 @@
 # Shared memory mount name and path
 SHARED_MEMORY_MOUNT_NAME = "flyte-shared-memory"
 SHARED_MEMORY_MOUNT_PATH = "/dev/shm"
+
+# Packages to be installed at the beginning of runtime
+RUNTIME_PACKAGES_ENV_NAME = "_F_RUNTIME_PACKAGES"

flytekit/core/python_auto_container.py

@@ -12,6 +12,7 @@
 from flytekit.configuration import ImageConfig, SerializationSettings
 from flytekit.constants import CopyFileDetection
 from flytekit.core.base_task import PythonTask, TaskMetadata, TaskResolverMixin
+from flytekit.core.constants import RUNTIME_PACKAGES_ENV_NAME
 from flytekit.core.context_manager import FlyteContextManager
 from flytekit.core.pod_template import PodTemplate
 from flytekit.core.resources import Resources, ResourceSpec, construct_extended_resources
@@ -231,6 +232,12 @@ def _get_container(self, settings: SerializationSettings) -> _task_model.Contain
         for elem in (settings.env, self.environment):
             if elem:
                 env.update(elem)
+
+        # Add runtime dependencies into environment
+        if isinstance(self.container_image, ImageSpec) and self.container_image.runtime_packages:
+            runtime_packages = " ".join(self.container_image.runtime_packages)
+            env[RUNTIME_PACKAGES_ENV_NAME] = runtime_packages
+
         return _get_container_definition(
             image=self.get_image(settings),
             resource_spec=self.resources,

flytekit/image_spec/__init__.py

@@ -17,6 +17,9 @@
 
 from .default_builder import DefaultImageBuilder
 from .image_spec import ImageBuildEngine, ImageSpec
+from .noop_builder import NoOpBuilder
 
 # Set this to a lower priority compared to `envd` to maintain backward compatibility
 ImageBuildEngine.register(DefaultImageBuilder.builder_type, DefaultImageBuilder(), priority=1)
+# Lower priority compared to Default.
+ImageBuildEngine.register(NoOpBuilder.builder_type, NoOpBuilder(), priority=0)

flytekit/image_spec/image_spec.py

@@ -67,6 +67,12 @@ class ImageSpec:
             If the option is set by the user, then that option is of course used.
         copy: List of files/directories to copy to /root. e.g. ["src/file1.txt", "src/file2.txt"]
         python_exec: Python executable to use for install packages
+        runtime_packages: List of packages to be installed during runtime. `runtime_packages` requires `pip` to be installed
+            in your base image.
+                - If you are using an ImageSpec as your base image, please include `pip` into your packages:
+                  `ImageSpec(..., packages=["pip"])`.
+                - If you want to install runtime packages into a fixed base_image and not use an image builder, you can
+                  use `builder="noop"`: `ImageSpec(base_image="ghcr.io/name/my-custom-image", builder="noop").with_runtime_packages(["numpy"])`
     """
 
     name: str = "flytekit"
@@ -95,6 +101,7 @@ class ImageSpec:
     source_copy_mode: Optional[CopyFileDetection] = None
     copy: Optional[List[str]] = None
     python_exec: Optional[str] = None
+    runtime_packages: Optional[List[str]] = None
 
     def __post_init__(self):
         self.name = self.name.lower()
@@ -127,6 +134,7 @@ def __post_init__(self):
             "pip_extra_index_url",
             "entrypoint",
             "commands",
+            "runtime_packages",
         ]
         for parameter in parameters_str_list:
             attr = getattr(self, parameter)
@@ -160,7 +168,7 @@ def id(self) -> str:
 
         :return: a unique identifier of the ImageSpec
         """
-        parameters_to_exclude = ["pip_secret_mounts", "builder"]
+        parameters_to_exclude = ["pip_secret_mounts", "builder", "runtime_packages"]
         # Only get the non-None values in the ImageSpec to ensure the hash is consistent across different Flytekit versions.
         image_spec_dict = asdict(
             self, dict_factory=lambda x: {k: v for (k, v) in x if v is not None and k not in parameters_to_exclude}
@@ -358,6 +366,12 @@ def force_push(self) -> "ImageSpec":
 
         return copied_image_spec
 
+    def with_runtime_packages(self, runtime_packages: List[str]) -> "ImageSpec":
+        """
+        Builder that returns a new image spec with runtime packages. Dev packages will be installed during runtime.
+        """
+        return self._update_attribute("runtime_packages", runtime_packages)
+
     @classmethod
     def from_env(cls, *, pinned_packages: Optional[List[str]] = None, **kwargs) -> "ImageSpec":
         """Create ImageSpec with the environment's Python version and packages pinned to the ones in the environment."""

flytekit/image_spec/noop_builder.py

@@ -0,0 +1,17 @@
+from flytekit.image_spec.image_spec import ImageSpec, ImageSpecBuilder
+
+
+class NoOpBuilder(ImageSpecBuilder):
+    """Noop image builder."""
+
+    builder_type = "noop"
+
+    def build_image(self, image_spec: ImageSpec) -> str:
+        if not isinstance(image_spec.base_image, str):
+            msg = "base_image must be a string to use the noop image builder"
+            raise ValueError(msg)
+
+        import click
+
+        click.secho(f"Using image: {image_spec.base_image}", fg="blue")
+        return image_spec.base_image

pydoclint-errors-baseline.txt

@@ -170,7 +170,7 @@ flytekit/extras/tensorflow/record.py
 --------------------
 flytekit/image_spec/image_spec.py
     DOC601: Class `ImageSpec`: Class docstring contains fewer class attributes than actual class attributes.  (Please read https://jsh9.github.io/pydoclint/checking_class_attributes.html on how to correctly document class attributes.)
-    DOC603: Class `ImageSpec`: Class docstring attributes are different from actual class attributes. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Attributes in the class definition but not in the docstring: [apt_packages: Optional[List[str]], base_image: Optional[Union[str, 'ImageSpec']], builder: Optional[str], commands: Optional[List[str]], conda_channels: Optional[List[str]], conda_packages: Optional[List[str]], copy: Optional[List[str]], cuda: Optional[str], cudnn: Optional[str], entrypoint: Optional[List[str]], env: Optional[typing.Dict[str, str]], name: str, packages: Optional[List[str]], pip_extra_args: Optional[str], pip_extra_index_url: Optional[List[str]], pip_index: Optional[str], pip_secret_mounts: Optional[List[Tuple[str, str]]], platform: str, python_exec: Optional[str], python_version: str, registry: Optional[str], registry_config: Optional[str], requirements: Optional[str], source_copy_mode: Optional[CopyFileDetection], source_root: Optional[str], tag_format: Optional[str]]. (Please read https://jsh9.github.io/pydoclint/checking_class_attributes.html on how to correctly document class attributes.)
+    DOC603: Class `ImageSpec`: Class docstring attributes are different from actual class attributes. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Attributes in the class definition but not in the docstring: [apt_packages: Optional[List[str]], base_image: Optional[Union[str, 'ImageSpec']], builder: Optional[str], commands: Optional[List[str]], conda_channels: Optional[List[str]], conda_packages: Optional[List[str]], copy: Optional[List[str]], cuda: Optional[str], cudnn: Optional[str], entrypoint: Optional[List[str]], env: Optional[typing.Dict[str, str]], name: str, packages: Optional[List[str]], pip_extra_args: Optional[str], pip_extra_index_url: Optional[List[str]], pip_index: Optional[str], pip_secret_mounts: Optional[List[Tuple[str, str]]], platform: str, python_exec: Optional[str], python_version: str, registry: Optional[str], registry_config: Optional[str], requirements: Optional[str], runtime_packages: Optional[List[str]], source_copy_mode: Optional[CopyFileDetection], source_root: Optional[str], tag_format: Optional[str]]. (Please read https://jsh9.github.io/pydoclint/checking_class_attributes.html on how to correctly document class attributes.)
     DOC109: Method `ImageSpecBuilder.build_image`: The option `--arg-type-hints-in-docstring` is `True` but there are no type hints in the docstring arg list
     DOC110: Method `ImageSpecBuilder.build_image`: The option `--arg-type-hints-in-docstring` is `True` but not all args in the docstring arg list have type hints
     DOC105: Method `ImageSpecBuilder.build_image`: Argument names match, but type hints in these args do not match: image_spec

tests/flytekit/unit/core/image_spec/test_image_spec.py

@@ -301,3 +301,9 @@ def test_image_spec_same_id_and_tag_with_builder():
     image_spec_with_builder = ImageSpec(name="my_image", builder="envd")
     assert image_spec.id == image_spec_with_builder.id
     assert image_spec.tag == image_spec_with_builder.tag
+
+
+def test_dev_packages():
+    image_spec = ImageSpec(name="localhost:30000/flytekit:0.1.5")
+    new_image_spec = image_spec.with_runtime_packages(["my-new-package"])
+    assert new_image_spec.runtime_packages == ["my-new-package"]

tests/flytekit/unit/core/image_spec/test_noop_builder.py

@@ -0,0 +1,24 @@
+import pytest
+from flytekit.image_spec.noop_builder import NoOpBuilder
+from flytekit import ImageSpec
+
+
+def test_noop_builder():
+    builder = NoOpBuilder()
+
+    image_spec = ImageSpec(base_image="localhost:30000/flytekit")
+    image = builder.build_image(image_spec)
+    assert image == "localhost:30000/flytekit"
+
+
+@pytest.mark.parametrize("base_image", [
+    None,
+    ImageSpec(base_image="another_none")
+])
+def test_noop_builder_error(base_image):
+    builder = NoOpBuilder()
+
+    msg = "base_image must be a string to use the noop image builder"
+    image_spec = ImageSpec(base_image=base_image)
+    with pytest.raises(ValueError, match=msg):
+        builder.build_image(image_spec)