From 865cf7b779b7505d094bcbc0923ba5aab72f0dbe Mon Sep 17 00:00:00 2001 From: "adriano@exa.ai" Date: Thu, 30 Apr 2026 03:19:46 +0000 Subject: [PATCH 1/7] feat(nix-builder): push images from remote builders Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- flytekit/image_spec/default_builder.py | 244 +++++++++++++++++- .../core/image_spec/test_default_builder.py | 122 ++++++++- 2 files changed, 358 insertions(+), 8 deletions(-) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index d21370f52c..a491baef72 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -10,7 +10,8 @@ from pathlib import Path from string import Template from subprocess import run -from typing import ClassVar, List, NamedTuple +from typing import ClassVar, List, NamedTuple, Optional +from urllib.parse import parse_qs, urlparse import click import toml @@ -569,6 +570,235 @@ class _PythonInstallTemplate(NamedTuple): extra_path: str +class _NixRemoteBuilder(NamedTuple): + store_uri: str + system: str + ssh_host: str + ssh_key: Optional[str] + + +def _split_nix_config_assignments(config_text: str, key: str) -> List[str]: + values: List[str] = [] + pattern = re.compile(rf"^\s*{re.escape(key)}\s*=\s*(.+?)\s*$") + for line in config_text.splitlines(): + match = pattern.match(line) + if match: + values.append(match.group(1)) + return values + + +def _machine_file_paths_from_config(config_text: str) -> List[Path]: + paths = [] + for builders_value in _split_nix_config_assignments(config_text, "builders"): + for builder_ref in builders_value.split(): + if builder_ref.startswith("@"): + paths.append(Path(builder_ref[1:]).expanduser()) + return paths + + +def _nix_config_paths() -> List[Path]: + paths = [] + xdg_config_home = os.environ.get("XDG_CONFIG_HOME") + if xdg_config_home: + paths.append(Path(xdg_config_home) / "nix" / "nix.conf") + else: + home = os.environ.get("HOME") + if home: + paths.append(Path(home) / ".config" / "nix" / "nix.conf") + paths.append(Path("/etc/nix/nix.conf")) + return paths + + +def _machine_file_paths() -> List[Path]: + paths = [] + explicit_path = os.environ.get("FLYTEKIT_NIX_REMOTE_BUILDERS_FILE") + if explicit_path: + paths.append(Path(explicit_path).expanduser()) + + nix_config = os.environ.get("NIX_CONFIG") + if nix_config: + paths.extend(_machine_file_paths_from_config(nix_config)) + + for config_path in _nix_config_paths(): + if config_path.exists(): + paths.extend(_machine_file_paths_from_config(config_path.read_text())) + + unique_paths = [] + seen = set() + for path in paths: + path_key = os.fspath(path) + if path_key not in seen: + seen.add(path_key) + unique_paths.append(path) + return unique_paths + + +def _ssh_host_from_store_uri(store_uri: str) -> Optional[str]: + parsed_uri = urlparse(store_uri) + if parsed_uri.scheme not in {"ssh", "ssh-ng"} or not parsed_uri.netloc: + return None + return parsed_uri.netloc + + +def _ssh_key_from_store_uri(store_uri: str) -> Optional[str]: + parsed_uri = urlparse(store_uri) + query = parse_qs(parsed_uri.query) + ssh_key_values = query.get("ssh-key") + if not ssh_key_values: + return None + return ssh_key_values[0] + + +def _local_ssh_key_path(ssh_key: Optional[str]) -> Optional[str]: + if not ssh_key or ssh_key == "-": + return None + + home = os.environ.get("HOME") + if home: + home_key = Path(home) / ".ssh" / Path(ssh_key).name + if home_key.exists(): + return os.fspath(home_key) + + return ssh_key + + +def _store_uri_with_ssh_key(builder: _NixRemoteBuilder) -> str: + if not builder.ssh_key or "ssh-key=" in builder.store_uri: + return builder.store_uri + + separator = "&" if "?" in builder.store_uri else "?" + return f"{builder.store_uri}{separator}ssh-key={builder.ssh_key}" + + +def _parse_nix_machine_line(line: str) -> Optional[_NixRemoteBuilder]: + stripped = line.strip() + if not stripped or stripped.startswith("#"): + return None + + parts = stripped.split() + if len(parts) < 2: + return None + + store_uri = parts[0] + ssh_host = _ssh_host_from_store_uri(store_uri) + if ssh_host is None: + return None + + ssh_key = _ssh_key_from_store_uri(store_uri) + if ssh_key is None and len(parts) >= 3: + ssh_key = parts[2] + + return _NixRemoteBuilder( + store_uri=store_uri, + system=parts[1], + ssh_host=ssh_host, + ssh_key=_local_ssh_key_path(ssh_key), + ) + + +def _configured_nix_remote_builders() -> List[_NixRemoteBuilder]: + builders = [] + inline_builders = os.environ.get("FLYTEKIT_NIX_REMOTE_BUILDERS") + if inline_builders: + for line in inline_builders.splitlines(): + builder = _parse_nix_machine_line(line) + if builder: + builders.append(builder) + + for machine_file_path in _machine_file_paths(): + if not machine_file_path.exists(): + continue + for line in machine_file_path.read_text().splitlines(): + builder = _parse_nix_machine_line(line) + if builder: + builders.append(builder) + + return builders + + +def _select_nix_remote_builder(nix_system: str) -> Optional[_NixRemoteBuilder]: + if os.environ.get("FLYTEKIT_NIX_REMOTE_PUSH", "").lower() in {"0", "false", "no"}: + return None + + for builder in _configured_nix_remote_builders(): + if builder.system == nix_system: + return builder + return None + + +def _ssh_command(builder: _NixRemoteBuilder, remote_command: List[str]) -> List[str]: + command = [ + "ssh", + "-o", "StrictHostKeyChecking=no", + "-o", "UserKnownHostsFile=/dev/null", + "-o", "BatchMode=yes", + ] + if builder.ssh_key: + command.extend(["-i", builder.ssh_key]) + command.append(builder.ssh_host) + command.extend(remote_command) + return command + + +def _remote_nix_copy_to_ecr( + *, + tmp_dir: str, + nix_system: str, + image_name: str, + ecr_token: str, + builder: _NixRemoteBuilder, +) -> None: + push_attr = f"packages.{nix_system}.push-to-ecr" + build_command = [ + "nix", "build", + "--no-link", + "--print-out-paths", + "--eval-store", "auto", + "--store", _store_uri_with_ssh_key(builder), + "--builders", "", + "--builders-use-substitutes", + "--system", nix_system, + f"path:{tmp_dir}#{push_attr}", + ] + click.secho(f"Remote nix2container build on {builder.ssh_host}: {' '.join(build_command)}", fg="blue") + build_env = { + **os.environ, + "NIX_SSHOPTS": "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o BatchMode=yes", + } + build_result = run(build_command, capture_output=True, text=True, env=build_env) + if build_result.returncode != 0: + raise RuntimeError( + f"Remote nix2container build failed with exit code {build_result.returncode}: " + f"{' '.join(build_command)}\n{build_result.stderr}" + ) + + output_paths = [line.strip() for line in build_result.stdout.splitlines() if line.strip()] + if not output_paths: + raise RuntimeError(f"Remote nix2container build produced no output path: {' '.join(build_command)}") + + push_to_ecr_bin = f"{output_paths[-1]}/bin/push-to-ecr" + push_command = _ssh_command( + builder, + [ + "env", + f"IMAGE_NAME={image_name}", + f"ECR_TOKEN={ecr_token}", + push_to_ecr_bin, + ], + ) + log_push_command = list(push_command) + for i, arg in enumerate(log_push_command): + if arg.startswith("ECR_TOKEN="): + log_push_command[i] = "ECR_TOKEN=[REDACTED]" + click.secho(f"Remote nix2container push on {builder.ssh_host}: {' '.join(log_push_command)}", fg="blue") + push_result = run(push_command) + if push_result.returncode != 0: + raise RuntimeError( + f"Remote nix2container push failed with exit code {push_result.returncode}: " + f"{' '.join(log_push_command)}" + ) + + def prepare_python_executable(image_spec: ImageSpec) -> _PythonInstallTemplate: if image_spec.python_exec: if image_spec.conda_channels: @@ -837,7 +1067,17 @@ def _build_image(self, image_spec: ImageSpec, *, push: bool = True) -> str: ["aws", "ecr", "get-login-password", "--region", "us-west-2"], capture_output=True, text=True, check=True, ).stdout.strip() - if is_cross_build: + remote_builder = _select_nix_remote_builder(nix_system) + if remote_builder: + _remote_nix_copy_to_ecr( + tmp_dir=tmp_dir, + nix_system=nix_system, + image_name=image_spec.image_name(), + ecr_token=ecr_token, + builder=remote_builder, + ) + return image_spec.image_name() + elif is_cross_build: docker_attr = f"packages.{local_system}.docker-{nix_system}.copyTo" click.secho(f"Cross-build: {nix_system} image via {local_system} n2c", fg="yellow") else: diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index a50b567230..6af55857d1 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -1,15 +1,25 @@ -import re import os -from unittest.mock import patch, Mock +import re +import tempfile +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock import pytest import flytekit -from flytekit.image_spec import ImageSpec -from flytekit.image_spec.default_builder import DefaultImageBuilder, create_docker_context from flytekit.constants import CopyFileDetection -from pathlib import Path -import tempfile +from flytekit.image_spec import ImageSpec +from flytekit.image_spec.default_builder import ( + DefaultImageBuilder, + _configured_nix_remote_builders, + _NixRemoteBuilder, + _parse_nix_machine_line, + _remote_nix_copy_to_ecr, + _store_uri_with_ssh_key, + create_docker_context, +) + def test_create_docker_context(tmp_path): docker_context_path = tmp_path / "builder_root" @@ -358,3 +368,103 @@ def test_python_exec_errors(tmp_path, key, value): msg = f"{key} is not supported with python_exec" with pytest.raises(ValueError, match=msg): create_docker_context(image_spec, docker_context_path) + + +def test_parse_nix_machine_line_uses_local_key_from_home(monkeypatch, tmp_path): + key_path = tmp_path / ".ssh" / "nix-runner-key" + key_path.parent.mkdir() + key_path.write_text("key") + monkeypatch.setenv("HOME", os.fspath(tmp_path)) + + builder = _parse_nix_machine_line( + "ssh-ng://root@10.0.0.1 x86_64-linux /root/.ssh/nix-runner-key 64 64 big-parallel,kvm - -" + ) + + assert builder + assert builder.store_uri == "ssh-ng://root@10.0.0.1" + assert builder.system == "x86_64-linux" + assert builder.ssh_host == "root@10.0.0.1" + assert builder.ssh_key == os.fspath(key_path) + + +def test_configured_nix_remote_builders_from_nix_config(monkeypatch, tmp_path): + machine_path = tmp_path / "machines" + machine_path.write_text("ssh-ng://root@10.0.0.2 aarch64-linux - 64 64 big-parallel - -\n") + monkeypatch.setenv("NIX_CONFIG", f"builders = @{machine_path}") + monkeypatch.delenv("FLYTEKIT_NIX_REMOTE_BUILDERS", raising=False) + monkeypatch.delenv("FLYTEKIT_NIX_REMOTE_BUILDERS_FILE", raising=False) + monkeypatch.setenv("XDG_CONFIG_HOME", os.fspath(tmp_path / "empty-config")) + + builders = _configured_nix_remote_builders() + + assert len(builders) == 1 + assert builders[0].store_uri == "ssh-ng://root@10.0.0.2" + assert builders[0].system == "aarch64-linux" + assert builders[0].ssh_host == "root@10.0.0.2" + assert builders[0].ssh_key is None + + +def test_store_uri_with_ssh_key_uses_local_key(): + builder = _NixRemoteBuilder( + store_uri="ssh-ng://root@10.0.0.3", + system="x86_64-linux", + ssh_host="root@10.0.0.3", + ssh_key="/home/runner/.ssh/nix-runner-key", + ) + + assert _store_uri_with_ssh_key(builder) == "ssh-ng://root@10.0.0.3?ssh-key=/home/runner/.ssh/nix-runner-key" + + +def test_remote_nix_copy_to_ecr_builds_remote_store_and_pushes_over_ssh(monkeypatch): + builder = _NixRemoteBuilder( + store_uri="ssh-ng://root@10.0.0.4", + system="x86_64-linux", + ssh_host="root@10.0.0.4", + ssh_key="/home/runner/.ssh/nix-runner-key", + ) + calls = [] + + def fake_run(command, **kwargs): + calls.append((command, kwargs)) + if command[0] == "nix": + return SimpleNamespace(returncode=0, stdout="/nix/store/copy-to\n", stderr="") + return SimpleNamespace(returncode=0) + + monkeypatch.setattr("flytekit.image_spec.default_builder.run", fake_run) + + _remote_nix_copy_to_ecr( + tmp_dir="/build-context", + nix_system="x86_64-linux", + image_name="472386928882.dkr.ecr.us-west-2.amazonaws.com/example:tag", + ecr_token="secret-token", + builder=builder, + ) + + assert calls[0][0] == [ + "nix", "build", + "--no-link", + "--print-out-paths", + "--eval-store", "auto", + "--store", "ssh-ng://root@10.0.0.4?ssh-key=/home/runner/.ssh/nix-runner-key", + "--builders", "", + "--builders-use-substitutes", + "--system", "x86_64-linux", + "path:/build-context#packages.x86_64-linux.push-to-ecr", + ] + assert calls[0][1]["capture_output"] is True + assert calls[0][1]["text"] is True + assert calls[0][1]["env"]["NIX_SSHOPTS"] == ( + "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o BatchMode=yes" + ) + assert calls[1][0] == [ + "ssh", + "-o", "StrictHostKeyChecking=no", + "-o", "UserKnownHostsFile=/dev/null", + "-o", "BatchMode=yes", + "-i", "/home/runner/.ssh/nix-runner-key", + "root@10.0.0.4", + "env", + "IMAGE_NAME=472386928882.dkr.ecr.us-west-2.amazonaws.com/example:tag", + "ECR_TOKEN=secret-token", + "/nix/store/copy-to/bin/push-to-ecr", + ] From 87ad1a7e741c814eb411aa8d2fd0ad53433773d1 Mon Sep 17 00:00:00 2001 From: "adriano@exa.ai" Date: Thu, 30 Apr 2026 03:30:05 +0000 Subject: [PATCH 2/7] fix(nix-builder): match multi-system remote builders Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- flytekit/image_spec/default_builder.py | 2 +- .../unit/core/image_spec/test_default_builder.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index a491baef72..da8a53d0fc 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -721,7 +721,7 @@ def _select_nix_remote_builder(nix_system: str) -> Optional[_NixRemoteBuilder]: return None for builder in _configured_nix_remote_builders(): - if builder.system == nix_system: + if nix_system in builder.system.split(","): return builder return None diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index 6af55857d1..c6f3ebfa99 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -16,6 +16,7 @@ _NixRemoteBuilder, _parse_nix_machine_line, _remote_nix_copy_to_ecr, + _select_nix_remote_builder, _store_uri_with_ssh_key, create_docker_context, ) @@ -404,6 +405,21 @@ def test_configured_nix_remote_builders_from_nix_config(monkeypatch, tmp_path): assert builders[0].ssh_key is None +def test_select_nix_remote_builder_matches_comma_separated_systems(monkeypatch): + monkeypatch.setenv( + "FLYTEKIT_NIX_REMOTE_BUILDERS", + "ssh-ng://root@10.0.0.5 x86_64-linux,aarch64-linux - 64 64 big-parallel - -", + ) + monkeypatch.delenv("FLYTEKIT_NIX_REMOTE_BUILDERS_FILE", raising=False) + monkeypatch.setenv("XDG_CONFIG_HOME", os.fspath(Path.cwd() / "missing-config")) + + builder = _select_nix_remote_builder("aarch64-linux") + + assert builder + assert builder.store_uri == "ssh-ng://root@10.0.0.5" + assert builder.system == "x86_64-linux,aarch64-linux" + + def test_store_uri_with_ssh_key_uses_local_key(): builder = _NixRemoteBuilder( store_uri="ssh-ng://root@10.0.0.3", From 7a637584a7adc141aaf9a075b6904270bc8f5720 Mon Sep 17 00:00:00 2001 From: "adriano@exa.ai" Date: Thu, 30 Apr 2026 03:54:21 +0000 Subject: [PATCH 3/7] fix(image-spec): use absolute ssh for remote nix push Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- flytekit/image_spec/default_builder.py | 31 +++++++++++++++++-- .../core/image_spec/test_default_builder.py | 10 +++--- 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index da8a53d0fc..d58bf2af48 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -670,6 +670,28 @@ def _store_uri_with_ssh_key(builder: _NixRemoteBuilder) -> str: return f"{builder.store_uri}{separator}ssh-key={builder.ssh_key}" +def _ssh_binary() -> Optional[str]: + candidates = [ + shutil.which("ssh"), + "/run/current-system/sw/bin/ssh", + "/usr/bin/ssh", + "/bin/ssh", + ] + for candidate in candidates: + if candidate and os.access(candidate, os.X_OK): + return candidate + return None + + +def _nix_sshopts() -> str: + ssh_options = [ + "-o", "StrictHostKeyChecking=no", + "-o", "UserKnownHostsFile=/dev/null", + "-o", "BatchMode=yes", + ] + return " ".join(ssh_options) + + def _parse_nix_machine_line(line: str) -> Optional[_NixRemoteBuilder]: stripped = line.strip() if not stripped or stripped.startswith("#"): @@ -728,7 +750,7 @@ def _select_nix_remote_builder(nix_system: str) -> Optional[_NixRemoteBuilder]: def _ssh_command(builder: _NixRemoteBuilder, remote_command: List[str]) -> List[str]: command = [ - "ssh", + _ssh_binary() or "ssh", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null", "-o", "BatchMode=yes", @@ -761,9 +783,14 @@ def _remote_nix_copy_to_ecr( f"path:{tmp_dir}#{push_attr}", ] click.secho(f"Remote nix2container build on {builder.ssh_host}: {' '.join(build_command)}", fg="blue") + ssh_binary = _ssh_binary() + path = os.environ.get("PATH", "") + if ssh_binary: + path = f"{Path(ssh_binary).parent}{os.pathsep}{path}" build_env = { **os.environ, - "NIX_SSHOPTS": "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o BatchMode=yes", + "NIX_SSHOPTS": _nix_sshopts(), + "PATH": path, } build_result = run(build_command, capture_output=True, text=True, env=build_env) if build_result.returncode != 0: diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index c6f3ebfa99..891d5d0705 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -469,11 +469,11 @@ def fake_run(command, **kwargs): ] assert calls[0][1]["capture_output"] is True assert calls[0][1]["text"] is True - assert calls[0][1]["env"]["NIX_SSHOPTS"] == ( - "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o BatchMode=yes" - ) - assert calls[1][0] == [ - "ssh", + assert "-o StrictHostKeyChecking=no" in calls[0][1]["env"]["NIX_SSHOPTS"] + assert "-o UserKnownHostsFile=/dev/null" in calls[0][1]["env"]["NIX_SSHOPTS"] + assert "-o BatchMode=yes" in calls[0][1]["env"]["NIX_SSHOPTS"] + assert calls[1][0][0].endswith("/ssh") or calls[1][0][0] == "ssh" + assert calls[1][0][1:] == [ "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null", "-o", "BatchMode=yes", From a89dd1e05a2694ea40d543bdc914dea2e82723ad Mon Sep 17 00:00:00 2001 From: "adriano@exa.ai" Date: Thu, 30 Apr 2026 04:03:50 +0000 Subject: [PATCH 4/7] fix(image-spec): replace embedded nix ssh key Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- flytekit/image_spec/default_builder.py | 19 ++++++++++++++++--- .../core/image_spec/test_default_builder.py | 14 ++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index d58bf2af48..cb1308660b 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -663,11 +663,24 @@ def _local_ssh_key_path(ssh_key: Optional[str]) -> Optional[str]: def _store_uri_with_ssh_key(builder: _NixRemoteBuilder) -> str: - if not builder.ssh_key or "ssh-key=" in builder.store_uri: + if not builder.ssh_key: return builder.store_uri - separator = "&" if "?" in builder.store_uri else "?" - return f"{builder.store_uri}{separator}ssh-key={builder.ssh_key}" + parsed_uri = urlparse(builder.store_uri) + query_parts = parsed_uri.query.split("&") if parsed_uri.query else [] + next_query_parts = [] + replaced_ssh_key = False + for query_part in query_parts: + if query_part.split("=", 1)[0] == "ssh-key": + next_query_parts.append(f"ssh-key={builder.ssh_key}") + replaced_ssh_key = True + else: + next_query_parts.append(query_part) + + if not replaced_ssh_key: + next_query_parts.append(f"ssh-key={builder.ssh_key}") + + return parsed_uri._replace(query="&".join(next_query_parts)).geturl() def _ssh_binary() -> Optional[str]: diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index 891d5d0705..4a53ac0d79 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -431,6 +431,20 @@ def test_store_uri_with_ssh_key_uses_local_key(): assert _store_uri_with_ssh_key(builder) == "ssh-ng://root@10.0.0.3?ssh-key=/home/runner/.ssh/nix-runner-key" +def test_store_uri_with_ssh_key_replaces_embedded_key(): + builder = _NixRemoteBuilder( + store_uri="ssh-ng://root@10.0.0.3?ssh-key=/root/.ssh/nix-runner-key&compress=true", + system="x86_64-linux", + ssh_host="root@10.0.0.3", + ssh_key="/home/runner/.ssh/nix-runner-key", + ) + + assert ( + _store_uri_with_ssh_key(builder) + == "ssh-ng://root@10.0.0.3?ssh-key=/home/runner/.ssh/nix-runner-key&compress=true" + ) + + def test_remote_nix_copy_to_ecr_builds_remote_store_and_pushes_over_ssh(monkeypatch): builder = _NixRemoteBuilder( store_uri="ssh-ng://root@10.0.0.4", From 4356aed1a70a5814c48ce3d8157f2c6e9994a45a Mon Sep 17 00:00:00 2001 From: "adriano@exa.ai" Date: Thu, 30 Apr 2026 18:27:37 +0000 Subject: [PATCH 5/7] test(image-spec): cover nix remote push selection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../core/image_spec/test_default_builder.py | 175 ++++++++++++++++++ 1 file changed, 175 insertions(+) diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index 4a53ac0d79..71a8c1fad0 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -1,5 +1,6 @@ import os import re +import subprocess import tempfile from pathlib import Path from types import SimpleNamespace @@ -498,3 +499,177 @@ def fake_run(command, **kwargs): "ECR_TOKEN=secret-token", "/nix/store/copy-to/bin/push-to-ecr", ] + + +def _simple_nix_image_spec(platform_value="linux/amd64"): + return ImageSpec( + name="remote-nix-test", + registry="472386928882.dkr.ecr.us-west-2.amazonaws.com/flytekit-tests", + nix=True, + platform=platform_value, + ) + + +def _patch_nix_builder_basics(monkeypatch, *, machine="x86_64", system="Linux", aws_token="secret-token"): + monkeypatch.setattr("flytekit.image_spec.default_builder.create_docker_context", lambda image_spec, tmp_path: None) + monkeypatch.setattr( + "flytekit.image_spec.default_builder.shutil.which", + lambda binary: "/usr/bin/nix" if binary == "nix" else None, + ) + monkeypatch.setattr("flytekit.image_spec.default_builder.platform.system", lambda: system) + monkeypatch.setattr("flytekit.image_spec.default_builder.platform.machine", lambda: machine) + + aws_run = Mock(return_value=SimpleNamespace(stdout=f"{aws_token}\n")) + monkeypatch.setattr("flytekit.image_spec.default_builder.subprocess.run", aws_run) + return aws_run + + +def test_build_image_with_remote_nix_builder_pushes_from_remote_store(monkeypatch): + image_spec = _simple_nix_image_spec() + builder = _NixRemoteBuilder( + store_uri="ssh-ng://root@10.0.0.10", + system="x86_64-linux", + ssh_host="root@10.0.0.10", + ssh_key="/home/runner/.ssh/nix-runner-key", + ) + aws_run = _patch_nix_builder_basics(monkeypatch) + local_run = Mock(return_value=SimpleNamespace(returncode=0)) + remote_calls = [] + + def fake_remote_nix_copy_to_ecr(**kwargs): + remote_calls.append(kwargs) + + monkeypatch.setattr("flytekit.image_spec.default_builder._select_nix_remote_builder", lambda nix_system: builder) + monkeypatch.setattr("flytekit.image_spec.default_builder._remote_nix_copy_to_ecr", fake_remote_nix_copy_to_ecr) + monkeypatch.setattr("flytekit.image_spec.default_builder.run", local_run) + + result = DefaultImageBuilder()._build_image(image_spec, push=True) + + assert result == image_spec.image_name() + aws_run.assert_called_once_with( + ["aws", "ecr", "get-login-password", "--region", "us-west-2"], + capture_output=True, + text=True, + check=True, + ) + assert len(remote_calls) == 1 + assert remote_calls[0]["nix_system"] == "x86_64-linux" + assert remote_calls[0]["image_name"] == image_spec.image_name() + assert remote_calls[0]["ecr_token"] == "secret-token" + assert remote_calls[0]["builder"] == builder + local_run.assert_not_called() + + +def test_build_image_without_remote_nix_builder_falls_back_to_local_copyto(monkeypatch): + image_spec = _simple_nix_image_spec() + _patch_nix_builder_basics(monkeypatch) + commands = [] + + def fake_run(command, **kwargs): + commands.append((command, kwargs)) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr("flytekit.image_spec.default_builder._select_nix_remote_builder", lambda nix_system: None) + remote_push = Mock() + monkeypatch.setattr("flytekit.image_spec.default_builder._remote_nix_copy_to_ecr", remote_push) + monkeypatch.setattr("flytekit.image_spec.default_builder.run", fake_run) + + result = DefaultImageBuilder()._build_image(image_spec, push=True) + + assert result == image_spec.image_name() + assert len(commands) == 1 + command = commands[0][0] + assert command[:2] == ["nix", "run"] + assert command[2].startswith("path:") + assert command[2].endswith("#packages.x86_64-linux.docker.copyTo") + assert command[3:] == [ + "--", + f"docker://{image_spec.image_name()}", + "--dest-creds", + "AWS:secret-token", + "--image-parallel-copies", + "32", + ] + remote_push.assert_not_called() + + +def test_build_image_without_remote_nix_builder_cross_builds_locally(monkeypatch): + image_spec = _simple_nix_image_spec(platform_value="linux/arm64") + _patch_nix_builder_basics(monkeypatch, machine="x86_64") + commands = [] + + def fake_run(command, **kwargs): + commands.append((command, kwargs)) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr("flytekit.image_spec.default_builder._select_nix_remote_builder", lambda nix_system: None) + monkeypatch.setattr("flytekit.image_spec.default_builder.run", fake_run) + + result = DefaultImageBuilder()._build_image(image_spec, push=True) + + assert result == image_spec.image_name() + assert len(commands) == 1 + command = commands[0][0] + assert command[:2] == ["nix", "run"] + assert command[2].startswith("path:") + assert command[2].endswith("#packages.x86_64-linux.docker-aarch64-linux.copyTo") + assert command[3:] == [ + "--", + f"docker://{image_spec.image_name()}", + "--dest-creds", + "AWS:secret-token", + "--image-parallel-copies", + "32", + ] + + +def test_build_image_nix_without_push_builds_local_docker_package(monkeypatch): + image_spec = _simple_nix_image_spec() + aws_run = _patch_nix_builder_basics(monkeypatch) + commands = [] + + def fake_run(command, **kwargs): + commands.append((command, kwargs)) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr("flytekit.image_spec.default_builder._select_nix_remote_builder", Mock()) + monkeypatch.setattr("flytekit.image_spec.default_builder.run", fake_run) + + result = DefaultImageBuilder()._build_image(image_spec, push=False) + + assert result == image_spec.image_name() + aws_run.assert_not_called() + assert len(commands) == 1 + command = commands[0][0] + assert command[:2] == ["nix", "build"] + assert command[2].startswith("path:") + assert command[2].endswith("#packages.x86_64-linux.docker") + + +def test_build_image_nix_errors_when_nix_is_missing(monkeypatch): + image_spec = _simple_nix_image_spec() + monkeypatch.setattr("flytekit.image_spec.default_builder.create_docker_context", lambda image_spec, tmp_path: None) + monkeypatch.setattr("flytekit.image_spec.default_builder.shutil.which", lambda binary: None) + + with pytest.raises(RuntimeError, match="Nix is not installed"): + DefaultImageBuilder()._build_image(image_spec, push=True) + + +def test_build_image_nix_errors_on_unsupported_platform(monkeypatch): + image_spec = _simple_nix_image_spec(platform_value="linux/riscv64") + aws_run = _patch_nix_builder_basics(monkeypatch) + + with pytest.raises(RuntimeError, match="Unsupported platform for nix builds: linux/riscv64"): + DefaultImageBuilder()._build_image(image_spec, push=True) + + aws_run.assert_not_called() + + +def test_build_image_nix_propagates_ecr_token_failure(monkeypatch): + image_spec = _simple_nix_image_spec() + _patch_nix_builder_basics(monkeypatch) + aws_error = subprocess.CalledProcessError(returncode=1, cmd=["aws", "ecr", "get-login-password"]) + monkeypatch.setattr("flytekit.image_spec.default_builder.subprocess.run", Mock(side_effect=aws_error)) + + with pytest.raises(subprocess.CalledProcessError): + DefaultImageBuilder()._build_image(image_spec, push=True) From d840a085eabd7fbc0727d7b7660c778fd985b75e Mon Sep 17 00:00:00 2001 From: "adriano@exa.ai" Date: Thu, 30 Apr 2026 21:35:43 +0000 Subject: [PATCH 6/7] test(image-spec): cover remote nix push selection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- flytekit/image_spec/default_builder.py | 20 +++++- .../core/image_spec/test_default_builder.py | 71 ++++++++++++++----- 2 files changed, 72 insertions(+), 19 deletions(-) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index cb1308660b..9bdd424a0e 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -236,6 +236,10 @@ def _copy_local_packages_and_update_lock(image_spec: ImageSpec, tmp_dir: Path): # Copy each local package from the lock file and update its path for package in lock_data["package"]: + if "source" not in package: + non_vendored_packages.append(package) + continue + source = package["source"] if "directory" in source: @@ -487,7 +491,7 @@ def _copy_local_packages_and_update_lock(image_spec: ImageSpec, tmp_dir: Path): # Export requirements from uv.lock to requirements.txt format # This excludes editable installs (-e) and local relative path dependencies requirements_export_cmd = rf"uv export --format requirements-txt {image_spec.uv_export_args} | grep -v '^\(-e\|\.\./\)' > {requirements_path}" - subprocess.run(requirements_export_cmd, shell=True, check=True) + subprocess.run(requirements_export_cmd, shell=True, check=True, cwd=tmp_dir) # Write local packages file local_packages_path = tmp_dir / "local_packages.txt" @@ -499,8 +503,18 @@ def _copy_lock_files_into_context(image_spec: ImageSpec, lock_file: str, tmp_dir msg = f"Support for {lock_file} files and packages is mutually exclusive" raise ValueError(msg) - # Copy and update local packages first - _copy_local_packages_and_update_lock(image_spec, tmp_dir) + if lock_file == "uv.lock": + _copy_local_packages_and_update_lock(image_spec, tmp_dir) + return + + lock_path = Path(image_spec.requirements) + lock_dir = lock_path.parent + pyproject_path = lock_dir / "pyproject.toml" + if not pyproject_path.exists(): + raise ValueError(f"pyproject.toml must exist in the same directory as {lock_file}") + + shutil.copy2(lock_path, tmp_dir / lock_file) + shutil.copy2(pyproject_path, tmp_dir / "pyproject.toml") def prepare_uv_lock_command(image_spec: ImageSpec, pip_install_args: List[str], tmp_dir: Path) -> str: diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index 71a8c1fad0..f6bc89a47b 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -23,6 +23,33 @@ ) +def _write_minimal_uv_project(tmp_path: Path) -> Path: + uv_lock_file = tmp_path / "uv.lock" + uv_lock_file.write_text( + """ +version = 1 +requires-python = ">=3.12" + +[[package]] +name = "flytekit" +version = "1.0.0" +""" + ) + + pyproject_file = tmp_path / "pyproject.toml" + pyproject_file.write_text( + """ +[project] +name = "test-project" +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [] +""" + ) + + return uv_lock_file + + def test_create_docker_context(tmp_path): docker_context_path = tmp_path / "builder_root" docker_context_path.mkdir() @@ -216,13 +243,13 @@ def test_should_push_env(monkeypatch, push_image_spec): image_spec = ImageSpec(name="my_flytekit", python_version="3.12", registry="localhost:30000") monkeypatch.setenv("FLYTE_PUSH_IMAGE_SPEC", push_image_spec) - run_mock = Mock() + run_mock = Mock(return_value=SimpleNamespace(returncode=0, stderr="")) monkeypatch.setattr("flytekit.image_spec.default_builder.run", run_mock) builder = DefaultImageBuilder() builder.build_image(image_spec) - run_mock.assert_called_once() + assert run_mock.call_count == 2 call_args = run_mock.call_args.args if push_image_spec == "0": @@ -231,15 +258,12 @@ def test_should_push_env(monkeypatch, push_image_spec): assert "--push" in call_args[0] -def test_create_docker_context_uv_lock(tmp_path): +def test_create_docker_context_uv_lock(monkeypatch, tmp_path): docker_context_path = tmp_path / "builder_root" docker_context_path.mkdir() - uv_lock_file = tmp_path / "uv.lock" - uv_lock_file.write_text("this is a lock file") - - pyproject_file = tmp_path / "pyproject.toml" - pyproject_file.write_text("this is a pyproject.toml file") + uv_lock_file = _write_minimal_uv_project(tmp_path) + subprocess_run = Mock(return_value=SimpleNamespace(returncode=0)) image_spec = ImageSpec( name="FLYTEKIT", @@ -251,6 +275,7 @@ def test_create_docker_context_uv_lock(tmp_path): ) warning_msg = "uv.lock support is experimental" + monkeypatch.setattr("flytekit.image_spec.default_builder.subprocess.run", subprocess_run) with pytest.warns(UserWarning, match=warning_msg): create_docker_context(image_spec, docker_context_path) @@ -258,11 +283,13 @@ def test_create_docker_context_uv_lock(tmp_path): assert dockerfile_path.exists() dockerfile_content = dockerfile_path.read_text() - assert ( - "uv sync --index-url https://url.com --extra-index-url " - "https://extra-url.com --no-install-package library-to-skip " - "--locked --no-dev --no-install-project" - ) in dockerfile_content + assert "uv venv && uv pip sync requirements.txt" in dockerfile_content + assert "uv pip install" not in dockerfile_content + subprocess_run.assert_called_once() + export_command = subprocess_run.call_args.args[0] + assert "uv export --format requirements-txt" in export_command + assert f"> {docker_context_path / 'requirements.txt'}" in export_command + assert subprocess_run.call_args.kwargs == {"shell": True, "check": True, "cwd": docker_context_path} @pytest.mark.parametrize("lock_file", ["uv.lock", "poetry.lock"]) @@ -272,7 +299,11 @@ def test_lock_errors_no_pyproject_toml(monkeypatch, tmp_path, lock_file): monkeypatch.setattr("flytekit.image_spec.default_builder.run", run_mock) lock_file_path = tmp_path / lock_file - lock_file_path.write_text("this is a lock file") + if lock_file == "uv.lock": + lock_file_path = _write_minimal_uv_project(tmp_path) + (tmp_path / "pyproject.toml").unlink() + else: + lock_file_path.write_text("this is a lock file") image_spec = ImageSpec( name="FLYTEKIT", @@ -282,7 +313,7 @@ def test_lock_errors_no_pyproject_toml(monkeypatch, tmp_path, lock_file): builder = DefaultImageBuilder() - with pytest.raises(ValueError, match="a pyproject.toml file must be in the same"): + with pytest.raises(ValueError, match="pyproject.toml must exist in the same directory"): builder.build_image(image_spec) @@ -317,7 +348,15 @@ def test_create_poetry_lock(tmp_path): poetry_lock.write_text("this is a lock file") pyproject_file = tmp_path / "pyproject.toml" - pyproject_file.write_text("this is a pyproject.toml file") + pyproject_file.write_text( + """ +[tool.poetry] +name = "test-project" +version = "0.1.0" +description = "" +authors = [] +""" + ) image_spec = ImageSpec( name="FLYTEKIT", From bd1193acf821ef60546215392ac82abe80ef7b97 Mon Sep 17 00:00:00 2001 From: "adriano@exa.ai" Date: Thu, 30 Apr 2026 22:24:33 +0000 Subject: [PATCH 7/7] fix(image-spec): freeze uv export for lock contexts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- flytekit/image_spec/default_builder.py | 2 +- .../core/image_spec/test_default_builder.py | 89 ++++++++++++++++++- 2 files changed, 89 insertions(+), 2 deletions(-) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index 9bdd424a0e..9878d1ebd9 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -490,7 +490,7 @@ def _copy_local_packages_and_update_lock(image_spec: ImageSpec, tmp_dir: Path): # Export requirements from uv.lock to requirements.txt format # This excludes editable installs (-e) and local relative path dependencies - requirements_export_cmd = rf"uv export --format requirements-txt {image_spec.uv_export_args} | grep -v '^\(-e\|\.\./\)' > {requirements_path}" + requirements_export_cmd = rf"uv export --frozen --format requirements-txt {image_spec.uv_export_args} | grep -v '^\(-e\|\.\./\)' > {requirements_path}" subprocess.run(requirements_export_cmd, shell=True, check=True, cwd=tmp_dir) # Write local packages file diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index f6bc89a47b..c8e3f9d4ad 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -287,11 +287,98 @@ def test_create_docker_context_uv_lock(monkeypatch, tmp_path): assert "uv pip install" not in dockerfile_content subprocess_run.assert_called_once() export_command = subprocess_run.call_args.args[0] - assert "uv export --format requirements-txt" in export_command + assert "uv export --frozen --format requirements-txt" in export_command assert f"> {docker_context_path / 'requirements.txt'}" in export_command assert subprocess_run.call_args.kwargs == {"shell": True, "check": True, "cwd": docker_context_path} +def test_create_docker_context_uv_lock_with_local_editables_uses_frozen_export(monkeypatch, tmp_path): + source_root = tmp_path / "repo" + source_root.mkdir() + (source_root / ".git").mkdir() + project_dir = source_root / "app" + project_dir.mkdir() + local_package_dir = source_root / "shared" / "local_package" + local_package_dir.mkdir(parents=True) + (local_package_dir / "local_package").mkdir() + (local_package_dir / "local_package" / "__init__.py").write_text("") + (local_package_dir / "pyproject.toml").write_text( + """ +[project] +name = "local-package" +version = "0.1.0" +""" + ) + + (project_dir / "pyproject.toml").write_text( + """ +[project] +name = "test-project" +version = "0.1.0" +requires-python = ">=3.12" +dependencies = ["local-package"] + +[tool.uv.sources] +local-package = { path = "../shared/local_package", editable = true } +""" + ) + uv_lock_file = project_dir / "uv.lock" + uv_lock_file.write_text( + """ +version = 1 +requires-python = ">=3.12" + +[[package]] +name = "local-package" +version = "0.1.0" +source = { editable = "../shared/local_package" } + +[[package]] +name = "test-project" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "local-package" }, +] + +[package.metadata] +requires-dist = [ + { name = "local-package", editable = "../shared/local_package" }, +] +""" + ) + docker_context_path = tmp_path / "builder_root" + docker_context_path.mkdir() + calls = [] + + def subprocess_run(command, **kwargs): + calls.append((command, kwargs)) + if command[0] == "git": + return SimpleNamespace(returncode=0, stdout=b"") + return SimpleNamespace(returncode=0) + + monkeypatch.setattr("flytekit.image_spec.default_builder.subprocess.run", subprocess_run) + + image_spec = ImageSpec( + name="FLYTEKIT", + python_version="3.12", + requirements=os.fspath(uv_lock_file), + vendor_local=False, + ) + + with pytest.warns(UserWarning, match="uv.lock support is experimental"): + create_docker_context(image_spec, docker_context_path) + + export_command = next( + command for command, _ in calls if isinstance(command, str) and command.startswith("uv export") + ) + assert "uv export --frozen --format requirements-txt" in export_command + assert f"> {docker_context_path / 'requirements.txt'}" in export_command + assert 'editable = "local_packages/shared/local_package"' in (docker_context_path / "uv.lock").read_text() + assert 'path = "local_packages/shared/local_package"' in (docker_context_path / "pyproject.toml").read_text() + assert (docker_context_path / "local_packages" / "shared" / "local_package" / "pyproject.toml").exists() + + @pytest.mark.parametrize("lock_file", ["uv.lock", "poetry.lock"]) @pytest.mark.filterwarnings("ignore::UserWarning") def test_lock_errors_no_pyproject_toml(monkeypatch, tmp_path, lock_file):