diff --git a/FEATURES.md b/FEATURES.md index 717ddf09e..eda45a15c 100644 --- a/FEATURES.md +++ b/FEATURES.md @@ -98,7 +98,7 @@ The Flyte CLI follows a **verb noun** structure. Full reference: [CLI Docs](http flyte run hello.py main --numbers '[1,2,3]' # Run a task flyte serve serving.py env # Serve an app flyte deploy my_workflow.py # Deploy environments -flyte build my_workflow.py --push # Build and push images +flyte build my_workflow.py # Build and push images flyte get logs # Get logs for a run flyte abort run # Abort a run ``` diff --git a/examples/genai/handoff/README.md b/examples/genai/handoff/README.md index 8957aaf8b..ee684aed2 100644 --- a/examples/genai/handoff/README.md +++ b/examples/genai/handoff/README.md @@ -211,7 +211,7 @@ flyte deploy agent_handoff.py ```bash # Build and push image -flyte build agent_handoff.py --push +flyte build agent_handoff.py # Deploy to Flyte cluster flyte deploy agent_handoff.py --domain production diff --git a/examples/image/ci_build_image.py b/examples/image/ci_build_image.py index 08659195e..9fdf75c3d 100644 --- a/examples/image/ci_build_image.py +++ b/examples/image/ci_build_image.py @@ -1,76 +1,16 @@ -"""Build and push an image to a user-specified target from CI. +"""Build and push a custom image from CI. -Takes a source image and re-tags/pushes it to a target registry/name:tag. -Works with both the local Docker builder and the remote builder. +Define your image below, then run: -Usage:: - - # Re-tag an existing image and push to ECR - python examples/image/ci_build_image.py \ - --from ghcr.io/flyteorg/flyte:py3.12-v2.0.0b56 \ - --to 123456789.dkr.ecr.us-west-2.amazonaws.com/myorg/myimage:v1.0.0 - - # Use the remote builder - python examples/image/ci_build_image.py \ - --from ghcr.io/flyteorg/flyte:py3.12-v2.0.0b56 \ - --to 123456789.dkr.ecr.us-west-2.amazonaws.com/myorg/myimage:v1.0.0 \ - --builder remote - - # Force rebuild even if the target already exists - python examples/image/ci_build_image.py \ - --from ghcr.io/flyteorg/flyte:py3.12-v2.0.0b56 \ - --to 123456789.dkr.ecr.us-west-2.amazonaws.com/myorg/myimage:v1.0.0 \ - --force - - # Force rebuild using the remote builder - python examples/image/ci_build_image.py \ - --from ghcr.io/flyteorg/flyte:py3.12-v2.0.0b56 \ - --to 123456789.dkr.ecr.us-west-2.amazonaws.com/myorg/myimage:v1.0.0 \ - --builder remote --force + flyte build examples/image/ci_build_image.py env """ -import argparse -import asyncio - import flyte -from flyte import Image -from flyte.extend import ImageBuildEngine - - -def parse_target(target: str) -> tuple[str, str, str]: - """Parse a target image string into (registry, name, tag). - - Example: - >>> parse_target("123456789.dkr.ecr.us-west-2.amazonaws.com/myorg/myimage:v1.0.0") - ('123456789.dkr.ecr.us-west-2.amazonaws.com/myorg', 'myimage', 'v1.0.0') - """ - if ":" not in target: - raise ValueError(f"Target '{target}' must contain a tag (e.g., myregistry/myimage:v1.0)") - image_path, tag = target.rsplit(":", 1) - if "/" not in image_path: - raise ValueError(f"Target '{target}' must contain a registry (e.g., myregistry/myimage:v1.0)") - registry, name = image_path.rsplit("/", 1) - return registry, name, tag - - -async def build_and_push(from_image: str, to_target: str, builder: str = "local", force: bool = False) -> str: - """Build an image from a base and push it to a target registry/name:tag.""" - registry, name, tag = parse_target(to_target) - image = Image.from_base(from_image).clone(registry=registry, name=name) - object.__setattr__(image, "_tag", tag) - result = await ImageBuildEngine.build(image, builder=builder, force=force) - return result.uri - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Build and push an image to a target registry/name:tag.") - parser.add_argument("--from", dest="from_image", required=True, help="Source image URI") - parser.add_argument("--to", dest="to_target", required=True, help="Target image as registry/name:tag") - parser.add_argument("--builder", choices=["local", "remote"], default="local", help="Image builder to use") - parser.add_argument("--force", action="store_true", help="Skip existence check, always rebuild") - - args = parser.parse_args() - flyte.init_from_config() - uri = asyncio.run(build_and_push(args.from_image, args.to_target, args.builder, args.force)) - print(uri) +env = flyte.Environment( + name="env", + image=flyte.Image.from_debian_base( + registry="ghcr.io/myorg", + name="myimage", + ), +) diff --git a/maint_tools/build_default_image.py b/maint_tools/build_default_image.py index bfb418261..d236c4b74 100644 --- a/maint_tools/build_default_image.py +++ b/maint_tools/build_default_image.py @@ -70,7 +70,7 @@ async def build_flyte_connector_image( suffix = __version__.replace("+", "-") python_version = _detect_python_version() tag = f"py{python_version[0]}.{python_version[1]}-{suffix}" - object.__setattr__(default_image, "_tag", tag) + default_image = default_image.clone(tag=tag) await ImageBuildEngine.build(default_image, builder=builder) diff --git a/src/flyte/_deploy.py b/src/flyte/_deploy.py index ca6d77f9b..241055f47 100644 --- a/src/flyte/_deploy.py +++ b/src/flyte/_deploy.py @@ -349,7 +349,7 @@ def _update_interface_inputs_and_outputs_docstring( return updated_interface -async def _build_image_bg(env_name: str, image: Image) -> Tuple[str, str, Optional[Any]]: +async def _build_image_bg(env_name: str, image: Image, force: bool = False) -> Tuple[str, str, Optional[Any]]: """ Build the image in the background and return the environment name, the built image URI, and the RunIdentifierData (if built by the remote image builder). @@ -358,7 +358,7 @@ async def _build_image_bg(env_name: str, image: Image) -> Tuple[str, str, Option from ._internal.imagebuild.image_builder import RunIdentifierData status.step(f"Building image {image.name} for environment {env_name}") - result = await build.aio(image) + result = await build.aio(image, force=force) assert result.uri is not None, "Image build result URI is None, make sure to wait for the build to complete" run_id_data = None if result.remote_run: @@ -367,7 +367,9 @@ async def _build_image_bg(env_name: str, image: Image) -> Tuple[str, str, Option return env_name, result.uri, run_id_data -async def _build_images(deployment: DeploymentPlan, image_refs: Dict[str, str] | None = None) -> ImageCache: +async def _build_images( + deployment: DeploymentPlan, image_refs: Dict[str, str] | None = None, force: bool = False +) -> ImageCache: """ Build the images for the given deployment plan and update the environment with the built image. """ @@ -397,7 +399,7 @@ async def _build_images(deployment: DeploymentPlan, image_refs: Dict[str, str] | image_identifier_map[env_name] = image_uri continue logger.debug(f"Building Image for environment {env_name}, image: {env.image}") - images.append(_build_image_bg(env_name, env.image)) + images.append(_build_image_bg(env_name, env.image, force=force)) elif env.image == "auto" and "auto" not in image_identifier_map: if _DEFAULT_IMAGE_REF_NAME in image_refs: @@ -406,7 +408,7 @@ async def _build_images(deployment: DeploymentPlan, image_refs: Dict[str, str] | image_identifier_map[env_name] = image_uri continue auto_image = Image.from_debian_base() - images.append(_build_image_bg(env_name, auto_image)) + images.append(_build_image_bg(env_name, auto_image, force=force)) if images: with status.group(f"Building {len(images)} image{'s' if len(images) > 1 else ''}..."): @@ -603,7 +605,7 @@ async def deploy( @syncify -async def build_images(envs: Environment) -> ImageCache: +async def build_images(envs: Environment, force: bool = False) -> ImageCache: """ Build the images for the given environments. :param envs: Environment to build images for. @@ -612,4 +614,4 @@ async def build_images(envs: Environment) -> ImageCache: cfg = get_init_config() images = cfg.images if cfg else {} deployment = plan_deploy(envs) - return await _build_images(deployment[0], images) + return await _build_images(deployment[0], images, force=force) diff --git a/src/flyte/_image.py b/src/flyte/_image.py index 46a352da5..ca92206b5 100644 --- a/src/flyte/_image.py +++ b/src/flyte/_image.py @@ -483,8 +483,8 @@ class Image: # Layers to be added to the image. In init, because frozen, but users shouldn't access, so underscore. _layers: Tuple[Layer, ...] = field(default_factory=tuple) - # Only settable internally. - _tag: Optional[str] = field(default=None, init=False) + # Explicitly set tag — overrides the content-hash default in _final_tag when provided. + tag: Optional[str] = field(default=None) _DEFAULT_IMAGE_PREFIXES: ClassVar = { PYTHON_3_10: "py3.10-", @@ -545,7 +545,7 @@ def _get_default_image_for( preset_tag = f"py{python_version[0]}.{python_version[1]}-{suffix}" image = Image._new( base_image=f"python:{python_version[0]}.{python_version[1]}-slim-bookworm", - registry=_BASE_REGISTRY, + registry=None, name=_DEFAULT_IMAGE_NAME, python_version=python_version, platform=("linux/amd64", "linux/arm64") if platform is None else platform, @@ -581,8 +581,13 @@ def _get_default_image_for( image = image.with_pip_packages(f"flyte=={flyte_version}", pre=True) else: image = image.with_pip_packages(f"flyte=={flyte_version}") + # Set the registry last so internal clone() calls during construction + # (above) don't inherit _BASE_REGISTRY — clone() strips it whenever + # self.registry == _BASE_REGISTRY, so we defer stamping it until the + # image is fully assembled. + object.__setattr__(image, "registry", _BASE_REGISTRY) if not dev_mode: - object.__setattr__(image, "_tag", preset_tag) + object.__setattr__(image, "tag", preset_tag) return image @@ -623,7 +628,9 @@ def from_debian_base( ) if registry or name: - return base_image.clone(registry=registry, name=name, registry_secret=registry_secret, extendable=True) + return base_image.clone( + registry=registry, name=name, registry_secret=registry_secret, extendable=True + ) return base_image @@ -686,7 +693,6 @@ def from_uv_script( version :param script: path to the uv script :param platform: architecture to use for the image, default is linux/amd64, use tuple for multiple values - :param python_version: Python version for the image, if not specified, will use the current Python version :param index_url: index url to use for pip install, default is None :param extra_index_urls: extra index urls to use for pip install, default is True :param pre: whether to allow pre-release versions, default is False @@ -694,9 +700,6 @@ def from_uv_script( :param secret_mounts: Secret mounts to use for the image, default is None. :return: Image - - Args: - secret_mounts: """ ll = UVScript( script=Path(script), @@ -727,6 +730,7 @@ def clone( python_version: Optional[Tuple[int, int]] = None, addl_layer: Optional[Layer] = None, extendable: Optional[bool] = None, + tag: Optional[str] = None, ) -> Image: """ Use this method to clone the current image and change the registry and name @@ -741,6 +745,7 @@ def clone( image for other images, and additional layers can be added on top of it. If False, the image cannot be used as a base image for other images, and additional layers cannot be added on top of it. If None (default), defaults to False for safety. + :param tag: Explicit tag for the cloned image. If omitted, a content-hash tag is used. :return: """ from flyte import Secret @@ -757,7 +762,17 @@ def clone( "Flyte current cannot add additional layers to a Dockerfile-based Image." " Please amend the dockerfile directly." ) - registry = registry or self.registry + # Registry inheritance: only carry forward a registry that the caller + # actually owns. _BASE_REGISTRY ("ghcr.io/flyteorg") is an internal + # constant used as the home of the SDK's own prebuilt base images — + # virtually no user has write access to it. Inheriting it into a + # user-defined clone would silently produce a URI like + # "ghcr.io/flyteorg/my-image:tag" that can never be pushed. + # When no explicit registry is provided and the parent carries + # _BASE_REGISTRY, treat the clone as registry-less (None) so the + # build system assigns the correct target registry at build time. + parent_registry = None if self.registry == _BASE_REGISTRY else self.registry + registry = registry or parent_registry name = name or self.name registry_secret = registry_secret or self._image_registry_secret base_image = base_image or self.base_image @@ -771,6 +786,7 @@ def clone( dockerfile=self.dockerfile, registry=registry, name=name, + tag=tag or None, platform=self.platform, python_version=python_version or self.python_version, extendable=extendable if extendable is not None else self.extendable, @@ -783,7 +799,12 @@ def clone( @classmethod def from_dockerfile( - cls, file: Path, registry: str, name: str, platform: Union[Architecture, Tuple[Architecture, ...], None] = None + cls, + file: Path, + registry: str, + name: str, + platform: Union[Architecture, Tuple[Architecture, ...], None] = None, + tag: Optional[str] = None, ) -> Image: """ Use this method to create a new image with the specified dockerfile. Note you cannot use additional layers @@ -794,10 +815,11 @@ def from_dockerfile( context for the builder will be the directory where the dockerfile is located. :param file: path to the dockerfile - :param name: name of the image :param registry: registry to use for the image + :param name: name of the image :param platform: architecture to use for the image, default is linux/amd64, use tuple for multiple values Example: ("linux/amd64", "linux/arm64") + :param tag: Explicit tag for the built image. If omitted, a content-hash tag is used. :return: """ @@ -806,13 +828,12 @@ def from_dockerfile( "dockerfile": file, "registry": registry, "name": name, + "tag": tag or None, "extendable": False, # Dockerfile-based images cannot have additional layers } if platform: kwargs["platform"] = platform - img = cls._new(**kwargs) - - return img + return cls._new(**kwargs) def _get_hash_digest(self) -> str: """ @@ -835,7 +856,7 @@ def _get_hash_digest(self) -> str: @property def _final_tag(self) -> str: - t = self._tag or self._get_hash_digest() + t = self.tag or self._get_hash_digest() return t or "latest" @cached_property diff --git a/src/flyte/cli/_build.py b/src/flyte/cli/_build.py index 982b448c6..76e8008c6 100644 --- a/src/flyte/cli/_build.py +++ b/src/flyte/cli/_build.py @@ -23,6 +23,18 @@ class BuildArguments: ) }, ) + force: bool = field( + default=False, + metadata={ + "click.option": click.Option( + ["--force"], + is_flag=True, + type=bool, + default=False, + help="Force rebuild and push even if the image already exists in the registry.", + ) + }, + ) @classmethod def from_dict(cls, d: Dict[str, Any]) -> "BuildArguments": @@ -50,7 +62,7 @@ def invoke(self, ctx: click.Context): status.step(f"Building environment: {self.obj_name}") obj.init() with common.cli_status(obj.output_format, "Building...", spinner="dots"): - image_cache = flyte.build_images(self.obj) + image_cache = flyte.build_images(self.obj, force=self.build_args.force) status.success(f"Environment {self.obj_name} built") common.print_output(common.format("Images", image_cache.repr(), obj.output_format), obj.output_format) diff --git a/tests/cli/test_build.py b/tests/cli/test_build.py new file mode 100644 index 000000000..fa4f4ab67 --- /dev/null +++ b/tests/cli/test_build.py @@ -0,0 +1,73 @@ +from unittest.mock import Mock, patch + +import pytest +from click.testing import CliRunner + +from flyte._internal.imagebuild.image_builder import ImageCache +from flyte.cli._build import build + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def mock_cli_config(): + cfg = Mock() + cfg.output_format = "table-simple" + cfg.log_level = None + cfg.init.return_value = None + return cfg + + +def test_build_force_flag_passed_to_build_images(runner, tmp_path, mock_cli_config): + """--force is forwarded to flyte.build_images.""" + images_file = tmp_path / "images.py" + # Variable name "my_env" is used as the subcommand, not the Environment's name attribute. + images_file.write_text( + "import flyte\n" + "my_env = flyte.Environment(name='my_env',\n" + " image=flyte.Image.from_base('python:3.12').clone(\n" + " registry='reg', name='img'))\n" + ) + + mock_cache = ImageCache(image_lookup={"my_env": "reg/img:abc123"}) + + with patch("flyte.build_images", return_value=mock_cache) as mock_build: + result = runner.invoke( + build, + ["--force", str(images_file), "my_env"], + obj=mock_cli_config, + ) + + assert result.exit_code == 0, result.output + mock_build.assert_called_once() + call_kwargs = mock_build.call_args[1] + assert call_kwargs.get("force") is True + + +def test_build_force_defaults_to_false(runner, tmp_path, mock_cli_config): + """force defaults to False when --force is not passed.""" + images_file = tmp_path / "images.py" + # Variable name "my_env" is used as the subcommand, not the Environment's name attribute. + images_file.write_text( + "import flyte\n" + "my_env = flyte.Environment(name='my_env',\n" + " image=flyte.Image.from_base('python:3.12').clone(\n" + " registry='reg', name='img'))\n" + ) + + mock_cache = ImageCache(image_lookup={"my_env": "reg/img:abc123"}) + + with patch("flyte.build_images", return_value=mock_cache) as mock_build: + result = runner.invoke( + build, + [str(images_file), "my_env"], + obj=mock_cli_config, + ) + + assert result.exit_code == 0, result.output + mock_build.assert_called_once() + call_kwargs = mock_build.call_args[1] + assert not call_kwargs.get("force", False) diff --git a/tests/flyte/test_deploy.py b/tests/flyte/test_deploy.py index 48b8aded1..3ad80132a 100644 --- a/tests/flyte/test_deploy.py +++ b/tests/flyte/test_deploy.py @@ -354,3 +354,55 @@ async def test_build_images_no_build_run_urls_for_local_build(): assert cache.image_lookup["my-env"] == "registry/my-image:sha256abc" assert cache.build_run_ids == {} + + +@pytest.mark.asyncio +async def test_build_image_bg_passes_force_to_build(): + """force=True is forwarded to build.aio.""" + image = flyte.Image.from_base("python:3.10") + mock_result = ImageBuild(uri="registry/my-image:abc", remote_run=None) + + with patch("flyte._build.build") as mock_build: + mock_build.aio = AsyncMock(return_value=mock_result) + await _build_image_bg("my-env", image, force=True) + mock_build.aio.assert_called_once_with(image, force=True) + + +@pytest.mark.asyncio +async def test_build_image_bg_force_defaults_to_false(): + """force defaults to False when not specified.""" + image = flyte.Image.from_base("python:3.10") + mock_result = ImageBuild(uri="registry/my-image:abc", remote_run=None) + + with patch("flyte._build.build") as mock_build: + mock_build.aio = AsyncMock(return_value=mock_result) + await _build_image_bg("my-env", image) + mock_build.aio.assert_called_once_with(image, force=False) + + +@pytest.mark.asyncio +async def test_build_images_passes_force_to_bg(): + """force=True is forwarded through _build_images to _build_image_bg.""" + image = flyte.Image.from_base("python:3.10") + env = flyte.TaskEnvironment(name="my-env", image=image) + plan = DeploymentPlan(envs={"my-env": env}) + mock_result = ImageBuild(uri="registry/my-image:abc", remote_run=None) + + with patch("flyte._build.build") as mock_build: + mock_build.aio = AsyncMock(return_value=mock_result) + await _build_images(plan, force=True) + mock_build.aio.assert_called_once_with(image, force=True) + + +@pytest.mark.asyncio +async def test_build_images_force_defaults_to_false(): + """force defaults to False when not specified.""" + image = flyte.Image.from_base("python:3.10") + env = flyte.TaskEnvironment(name="my-env", image=image) + plan = DeploymentPlan(envs={"my-env": env}) + mock_result = ImageBuild(uri="registry/my-image:abc", remote_run=None) + + with patch("flyte._build.build") as mock_build: + mock_build.aio = AsyncMock(return_value=mock_result) + await _build_images(plan) + mock_build.aio.assert_called_once_with(image, force=False) diff --git a/tests/flyte/test_image.py b/tests/flyte/test_image.py index a82b2b183..21455b987 100644 --- a/tests/flyte/test_image.py +++ b/tests/flyte/test_image.py @@ -199,6 +199,35 @@ def test_base_image_cloned(): assert cloned_default_image.uri.startswith("ghcr.io/flyteorg/flyte-clone") +def test_clone_strips_base_registry(): + # Any clone of a from_debian_base() image without an explicit registry should + # have registry=None — the clone is a user-owned derivative, not the SDK's + # ghcr.io/flyteorg image. + img = Image.from_debian_base(python_version=(3, 13)).clone(name="my-image") + assert img.registry is None + assert img.uri.startswith("my-image:") + + # Same via from_debian_base(name=...) + img2 = Image.from_debian_base(python_version=(3, 13), name="my-image") + assert img2.registry is None + assert img2.uri.startswith("my-image:") + + # with_* chains also strip _BASE_REGISTRY — they produce user-owned derivatives. + img3 = Image.from_debian_base(python_version=(3, 13)).with_pip_packages("numpy") + assert img3.registry is None + assert img3.uri.startswith("flyte:") + + # Explicit registry always wins regardless. + img4 = Image.from_debian_base(python_version=(3, 13)).clone( + registry="myregistry.io", name="my-image" + ) + assert img4.registry == "myregistry.io" + + # The unmodified base image itself still carries _BASE_REGISTRY. + base = Image.from_debian_base(python_version=(3, 13)) + assert base.registry == "ghcr.io/flyteorg" + + def test_base_image_clone_same(): default_image = Image.from_debian_base(python_version=(3, 13)) cloned_default_image = Image.from_debian_base(python_version=(3, 13)).clone( @@ -440,6 +469,21 @@ def test_from_dockerfile_is_not_extendable(): dockerfile_path.unlink() +def test_from_dockerfile_with_explicit_tag(tmp_path): + dockerfile = tmp_path / "Dockerfile" + dockerfile.write_text("FROM python:3.12-slim\n") + img = Image.from_dockerfile(file=dockerfile, registry="reg", name="my-img", tag="v2.0.0") + assert img.tag == "v2.0.0" + assert img.uri == "reg/my-img:v2.0.0" + + +def test_from_dockerfile_empty_string_tag_falls_back_to_content_hash(tmp_path): + dockerfile = tmp_path / "Dockerfile" + dockerfile.write_text("FROM python:3.12-slim\n") + img = Image.from_dockerfile(file=dockerfile, registry="reg", name="my-img", tag="") + assert img.tag is None + + def test_with_code_bundle_defaults(): """with_code_bundle() creates a CodeBundleLayer with default values.""" img = Image.from_debian_base(registry="localhost", name="test-image").with_code_bundle() @@ -525,3 +569,28 @@ def test_resolve_code_bundle_loaded_modules_copy_style_none(tmp_path): assert bundle_layers[0].root_dir == tmp_path assert bundle_layers[0].copy_style == "loaded_modules" assert bundle_layers[0].dst == "." + + +def test_clone_with_explicit_tag(): + img = Image.from_debian_base(registry="reg", name="img", python_version=(3, 12)) + cloned = img.clone(registry="other-reg", name="other-img", tag="v1.2.3") + assert cloned.tag == "v1.2.3" + assert cloned.uri == "other-reg/other-img:v1.2.3" + + +def test_clone_tag_not_inherited_from_source(): + # clone() never inherits tag from source — always fresh content hash unless tag= given + img = Image.from_debian_base(registry="reg", name="img", python_version=(3, 12)) + # default image has tag set (version-based); clone without tag= should content-hash + cloned = img.clone(registry="other-reg", name="other-img") + assert cloned.tag is None + assert cloned.uri.startswith("other-reg/other-img:") + assert cloned.uri != img.uri # different because hash differs (no version tag) + + +def test_clone_empty_string_tag_falls_back_to_content_hash(): + img = Image.from_debian_base(registry="reg", name="img", python_version=(3, 12)) + cloned = img.clone(registry="reg", name="img", tag="") + assert cloned.tag is None # empty string normalized to None + +