From 32b2fa72d19b48ce865a35700fb1831b820f8d1c Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Fri, 18 Apr 2025 01:23:23 -0400 Subject: [PATCH 1/2] Adds ImageSpec.with_runtime_packages (#3231) * Adds ImageSpec.with_dev_dependencies Signed-off-by: Thomas J. Fan * Fix Signed-off-by: Thomas J. Fan * Add tests for noop builder Signed-off-by: Thomas J. Fan * Use runtime_packages Signed-off-by: Thomas J. Fan * Add docs abount how to use runtime packages Signed-off-by: Thomas J. Fan * Less diffs Signed-off-by: Thomas J. Fan * Fix formatting Signed-off-by: Thomas J. Fan * Fix docstring Signed-off-by: Thomas J. Fan * Dix docstring Signed-off-by: Thomas J. Fan * Let pip default to user by itself to be more compatible Signed-off-by: Thomas J. Fan --------- Signed-off-by: Thomas J. Fan --- flytekit/bin/entrypoint.py | 31 +++++++++++++------ flytekit/core/constants.py | 3 ++ flytekit/core/python_auto_container.py | 7 +++++ flytekit/image_spec/__init__.py | 3 ++ flytekit/image_spec/image_spec.py | 16 +++++++++- flytekit/image_spec/noop_builder.py | 17 ++++++++++ pydoclint-errors-baseline.txt | 2 +- .../unit/core/image_spec/test_image_spec.py | 6 ++++ .../unit/core/image_spec/test_noop_builder.py | 24 ++++++++++++++ 9 files changed, 98 insertions(+), 11 deletions(-) create mode 100644 flytekit/image_spec/noop_builder.py create mode 100644 tests/flytekit/unit/core/image_spec/test_noop_builder.py 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 808a93731e..6f02063681 100644 --- a/pydoclint-errors-baseline.txt +++ b/pydoclint-errors-baseline.txt @@ -174,7 +174,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) From 0505eabf183c46df740dc06fa16df411c9014c69 Mon Sep 17 00:00:00 2001 From: Michael Hotan Date: Wed, 23 Apr 2025 03:15:49 +1000 Subject: [PATCH 2/2] Image spec builder options (#3233) * Image spec builder options Provide the ability to specify image `builder` specific options per Image Spec. Signed-off-by: Mike Hotan * Add builder_options validation Signed-off-by: Mike Hotan * updates Signed-off-by: Mike Hotan --------- Signed-off-by: Mike Hotan --- flytekit/image_spec/image_spec.py | 97 ++++++++++++------- pydoclint-errors-baseline.txt | 2 - .../unit/core/image_spec/test_image_spec.py | 37 ++++++- 3 files changed, 94 insertions(+), 42 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index ec93b028d9..ec3ef4a2ab 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -11,7 +11,7 @@ from dataclasses import asdict, dataclass from functools import cached_property, lru_cache from importlib import metadata -from typing import Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union import click import requests @@ -30,56 +30,57 @@ class ImageSpec: """ This class is used to specify the docker image that will be used to run the task. - Args: - name: name of the image. - python_version: python version of the image. Use default python in the base image if None. - builder: Type of plugin to build the image. Use envd by default. - source_root: source root of the image. - env: environment variables of the image. - registry: registry of the image. - packages: list of python packages to install. - conda_packages: list of conda packages to install. - conda_channels: list of conda channels. - requirements: path to the requirements.txt file. - apt_packages: list of apt packages to install. - cuda: version of cuda to install. - cudnn: version of cudnn to install. - base_image: base image of the image. - platform: Specify the target platforms for the build output (for example, windows/amd64 or linux/amd64,darwin/arm64 - pip_index: Specify the custom pip index url - pip_extra_index_url: Specify one or more pip index urls as a list - pip_secret_mounts: Specify a list of tuples to mount secret for pip install. Each tuple should contain the path to + Attributes: + name (str): Name of the image. + python_version (str): Python version of the image. Use default python in the base image if None. + builder (Optional[str]): Type of plugin to build the image. Use envd by default. + source_root (Optional[str]): Source root of the image. + env (Optional[Dict[str, str]]): Environment variables of the image. + registry (Optional[str]): Registry of the image. + packages (Optional[List[str]]): List of python packages to install. + conda_packages (Optional[List[str]]): List of conda packages to install. + conda_channels (Optional[List[str]]): List of conda channels. + requirements (Optional[str]): Path to the requirements.txt file. + apt_packages (Optional[List[str]]): List of apt packages to install. + cuda (Optional[str]): Version of cuda to install. + cudnn (Optional[str]): Version of cudnn to install. + base_image (Optional[Union[str, 'ImageSpec']]): Base image of the image. + platform (str): Specify the target platforms for the build output (for example, windows/amd64 or linux/amd64,darwin/arm64). + pip_index (Optional[str]): Specify the custom pip index url. + pip_extra_index_url (Optional[List[str]]): Specify one or more pip index urls as a list. + pip_secret_mounts (Optional[List[Tuple[str, str]]]): Specify a list of tuples to mount secret for pip install. Each tuple should contain the path to the secret file and the mount path. For example, [(".gitconfig", "/etc/gitconfig")]. This is experimental and the interface may change in the future. Configuring this should not change the built image. - pip_extra_args: Specify one or more extra pip install arguments as a space-delimited string - registry_config: Specify the path to a JSON registry config file - entrypoint: List of strings to overwrite the entrypoint of the base image with, set to [] to remove the entrypoint. - commands: Command to run during the building process - tag_format: Custom string format for image tag. The ImageSpec hash passed in as `spec_hash`. For example, - to add a "dev" suffix to the image tag, set `tag_format="{spec_hash}-dev"` - source_copy_mode: This option allows the user to specify which source files to copy from the local host, into the image. + pip_extra_args (Optional[str]): Specify one or more extra pip install arguments as a space-delimited string. + registry_config (Optional[str]): Specify the path to a JSON registry config file. + entrypoint (Optional[List[str]]): List of strings to overwrite the entrypoint of the base image with, set to [] to remove the entrypoint. + commands (Optional[List[str]]): Command to run during the building process. + tag_format (Optional[str]): Custom string format for image tag. The ImageSpec hash passed in as `spec_hash`. For example, + to add a "dev" suffix to the image tag, set `tag_format="{spec_hash}-dev"`. + source_copy_mode (Optional[CopyFileDetection]): This option allows the user to specify which source files to copy from the local host, into the image. Not setting this option means to use the default flytekit behavior. The default behavior is: - if fast register is used, source files are not copied into the image (because they're already copied into the fast register tar layer). - if fast register is not used, then the LOADED_MODULES (aka 'auto') option is used to copy loaded Python files into the image. - 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 + copy (Optional[List[str]]): List of files/directories to copy to /root. e.g. ["src/file1.txt", "src/file2.txt"]. + python_exec (Optional[str]): Python executable to use for install packages. + runtime_packages (Optional[List[str]]): 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"])` + use `builder="noop"`: `ImageSpec(base_image="ghcr.io/name/my-custom-image", builder="noop").with_runtime_packages(["numpy"])`. + builder_options (Optional[Dict[str, Any]]): Additional options for the builder. This is a dictionary that will be passed to the builder. + The options are builder-specific and may not be supported by all builders. """ name: str = "flytekit" python_version: str = None # Use default python in the base image if None. builder: Optional[str] = None source_root: Optional[str] = None # a.txt:auto - env: Optional[typing.Dict[str, str]] = None + env: Optional[Dict[str, str]] = None registry: Optional[str] = None packages: Optional[List[str]] = None conda_packages: Optional[List[str]] = None @@ -102,6 +103,7 @@ class ImageSpec: copy: Optional[List[str]] = None python_exec: Optional[str] = None runtime_packages: Optional[List[str]] = None + builder_options: Optional[Dict[str, Any]] = None def __post_init__(self): self.name = self.name.lower() @@ -153,6 +155,9 @@ def __post_init__(self): error_msg = "pip_secret_mounts must be a list of tuples of two strings or None" raise ValueError(error_msg) + if self.builder_options is not None and not isinstance(self.builder_options, dict): + raise ValueError("builder_options must be a dictionary or None") + @cached_property def id(self) -> str: """ @@ -318,16 +323,30 @@ def exist(self) -> Optional[bool]: click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red") return None - def _update_attribute(self, attr_name: str, values: Union[str, List[str]]) -> "ImageSpec": + def _update_attribute(self, attr_name: str, values: Union[str, List[str], Dict[str, Any]]) -> "ImageSpec": """ - Generic method to update a specified list attribute, either appending or extending. + Generic method to update a specified attribute, handling strings, lists, and dictionaries. """ - current_value = copy.deepcopy(getattr(self, attr_name)) or [] + current_value = copy.deepcopy(getattr(self, attr_name)) + + if current_value is None: + if isinstance(values, dict): + current_value = {} + else: + current_value = [] if isinstance(values, str): + if not isinstance(current_value, list): + raise TypeError(f"Cannot append string to non-list attribute {attr_name}") current_value.append(values) elif isinstance(values, list): + if not isinstance(current_value, list): + raise TypeError(f"Cannot extend non-list attribute {attr_name}") current_value.extend(values) + elif isinstance(values, dict): + if not isinstance(current_value, dict): + raise TypeError(f"Cannot update non-dict attribute {attr_name}") + current_value.update(values) return dataclasses.replace(self, **{attr_name: current_value}) @@ -372,6 +391,12 @@ def with_runtime_packages(self, runtime_packages: List[str]) -> "ImageSpec": """ return self._update_attribute("runtime_packages", runtime_packages) + def with_builder_options(self, builder_options: Dict[str, Any]) -> "ImageSpec": + """ + Builder that returns a new image spec with additional builder options. + """ + return self._update_attribute("builder_options", builder_options) + @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/pydoclint-errors-baseline.txt b/pydoclint-errors-baseline.txt index 6f02063681..95fdb70a68 100644 --- a/pydoclint-errors-baseline.txt +++ b/pydoclint-errors-baseline.txt @@ -173,8 +173,6 @@ flytekit/extras/tensorflow/record.py DOC603: Class `TFRecordDatasetConfig`: 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: [buffer_size: Optional[int], compression_type: Optional[str], name: Optional[str], num_parallel_reads: Optional[int]]. (Please read https://jsh9.github.io/pydoclint/checking_class_attributes.html on how to correctly document class attributes.) -------------------- 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], 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 ddbcfe7b74..f9dfa63a9b 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -5,13 +5,14 @@ import mock import pytest +from flytekit.configuration import (FastSerializationSettings, ImageConfig, + SerializationSettings) +from flytekit.constants import CopyFileDetection from flytekit.core import context_manager from flytekit.core.context_manager import ExecutionState +from flytekit.core.python_auto_container import update_image_spec_copy_handling from flytekit.image_spec import ImageSpec from flytekit.image_spec.image_spec import _F_IMG_ID, ImageBuildEngine -from flytekit.core.python_auto_container import update_image_spec_copy_handling -from flytekit.configuration import SerializationSettings, FastSerializationSettings, ImageConfig -from flytekit.constants import CopyFileDetection REQUIREMENT_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "requirements.txt") REGISTRY_CONFIG_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "registry_config.json") @@ -33,7 +34,8 @@ def test_image_spec(mock_image_spec_builder, monkeypatch): requirements=REQUIREMENT_FILE, registry_config=REGISTRY_CONFIG_FILE, entrypoint=["/bin/bash"], - copy=["/src/file1.txt"] + copy=["/src/file1.txt"], + builder_options={"builder_option": "builder_option_value"}, ) assert image_spec._is_force_push is False @@ -62,6 +64,7 @@ def test_image_spec(mock_image_spec_builder, monkeypatch): assert image_spec._is_force_push is True assert image_spec.entrypoint == ["/bin/bash"] assert image_spec.copy == ["/src/file1.txt", "/src", "/src/file2.txt"] + assert image_spec.builder_options == {"builder_option": "builder_option_value"} assert image_spec.image_name() == f"localhost:30001/flytekit:{image_spec.tag}" ctx = context_manager.FlyteContext.current_context() @@ -307,3 +310,29 @@ 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"] + +def test_invalid_builder_options(): + msg = "builder_options must be a dictionary or None" + with pytest.raises(ValueError, match=msg): + ImageSpec(name="localhost:30000/flytekit:0.1.5", builder_options="invalid_builder_option") + with pytest.raises(ValueError, match=msg): + ImageSpec(name="localhost:30000/flytekit:0.1.5", + builder_options=["invalid_builder_option"]) + +def test_with_builder_options(): + image_spec = ImageSpec( + name="localhost:30000/flytekit:0.1.5", + builder_options={ + "existing_builder_option_1": "existing_builder_option_value_1", + } + ) + new_image_spec = image_spec.with_builder_options( + {"new_builder_option_1": "new_builder_option_value_1"}) + + assert image_spec.builder_options == { + "existing_builder_option_1": "existing_builder_option_value_1", + } + assert new_image_spec.builder_options == { + "existing_builder_option_1": "existing_builder_option_value_1", + "new_builder_option_1": "new_builder_option_value_1" + }