diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 3011544dec..1e82f050a8 100644 --- a/flytekit/bin/entrypoint.py +++ b/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) diff --git a/flytekit/core/constants.py b/flytekit/core/constants.py index 2ed491ce60..4ea432bc96 100644 --- a/flytekit/core/constants.py +++ b/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" diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index aa0327299b..b9e4fe2992 100644 --- a/flytekit/core/python_auto_container.py +++ b/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, diff --git a/flytekit/image_spec/__init__.py b/flytekit/image_spec/__init__.py index c06c22c62b..c2f596bc6e 100644 --- a/flytekit/image_spec/__init__.py +++ b/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) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 41383c7e5f..ec93b028d9 100644 --- a/flytekit/image_spec/image_spec.py +++ b/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.""" diff --git a/flytekit/image_spec/noop_builder.py b/flytekit/image_spec/noop_builder.py new file mode 100644 index 0000000000..8bbf7d89a8 --- /dev/null +++ b/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 diff --git a/pydoclint-errors-baseline.txt b/pydoclint-errors-baseline.txt index c6c02f2413..b2db60b96c 100644 --- a/pydoclint-errors-baseline.txt +++ b/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 diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py index 10438d7e01..ddbcfe7b74 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/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"] diff --git a/tests/flytekit/unit/core/image_spec/test_noop_builder.py b/tests/flytekit/unit/core/image_spec/test_noop_builder.py new file mode 100644 index 0000000000..c9812ba0dc --- /dev/null +++ b/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)