Add task image IDs and compile-time image overrides

packages/tangle-cli/src/tangle_cli/__init__.py

@@ -14,6 +14,6 @@
 try:
     __version__ = metadata_version("tangle-cli")
 except PackageNotFoundError:
-    __version__ = "0.1.5"
+    __version__ = "0.1.6"
 
 __all__ = ["TangleDynamicDiscoveryClient", "__version__"]

packages/tangle-cli/src/tangle_cli/pipeline_compiler.py

@@ -72,6 +72,40 @@
 from .utils import dump_yaml
 
 
+IMAGE_IDS: dict[str, str] = {}
+"""Logical image-id registry used by ``@task(image_id=...)``.
+
+OSS ships this empty. Downstream distributions may register environment-specific
+image defaults without hard-coding their image literals into the open-source
+package; CLI ``--image ID=REF`` overrides still win over registered defaults.
+"""
+
+
+def register_image_id(image_id: str, image: str) -> None:
+    """Register a default image reference for ``@task(image_id=...)``."""
+    if not image_id:
+        raise ValueError("image_id must be a non-empty string")
+    if not image:
+        raise ValueError("image must be a non-empty string")
+    IMAGE_IDS[image_id] = image
+
+
+def get_image_id(image_id: str) -> str | None:
+    """Return the registered default image for ``image_id``, if any."""
+    return IMAGE_IDS.get(image_id)
+
+
+def resolve_image_id(
+    image_id: str,
+    image_overrides: Mapping[str, str] | None = None,
+) -> str | None:
+    """Resolve ``image_id`` using compile-time overrides, then registered defaults."""
+    overrides = image_overrides or {}
+    if image_id in overrides:
+        return overrides[image_id]
+    return get_image_id(image_id)
+
+
 @dataclass
 class CompileResult:
     """Outcome of a :func:`compile_pipeline` call.
@@ -102,6 +136,7 @@ def compile_pipeline(
     *,
     pipeline_name: str | None = None,
     emit_components_sidecar: bool = True,
+    image_overrides: Mapping[str, str] | None = None,
 ) -> CompileResult:
     """Compile ``script`` to a single pipeline YAML at ``output``.
 
@@ -132,6 +167,9 @@ def compile_pipeline(
             compilation fails with a :class:`CompileError` (a valid
             dehydrated pipeline cannot be produced without the sidecar).
             Pipelines that use only ``ref(url=...)`` are unaffected.
+        image_overrides: Compile-time image-id overrides from ``--image
+            ID=REF``. They apply only to ``@task(image_id=...)`` refs that do
+            not also set an explicit ``image=``.
 
     Returns:
         A :class:`CompileResult`. ``components_path`` is the sidecar path
@@ -143,6 +181,7 @@ def compile_pipeline(
             unreachable ``@task`` source/dependency files).
     """
     overrides = dict(overrides or {})
+    image_overrides = dict(image_overrides or {})
 
     # 1. Validate the script path.
     script_path = Path(script).resolve()
@@ -197,6 +236,7 @@ def compile_pipeline(
             root_overrides=overrides,
             emit_components_sidecar=emit_components_sidecar,
             source_dirs=purge_dirs,
+            image_overrides=image_overrides,
         )
 
         # 5. Compile the root (and, recursively, all children) into in-memory
@@ -436,7 +476,11 @@ def _compile_pipeline_fn(
         )
     if task_refs:
         components_path = output_path.with_name(output_path.stem + ".components.yaml")
-        components_entries = _build_local_from_python_components(task_refs, components_yaml_dir=components_path.parent)
+        components_entries = _build_local_from_python_components(
+            task_refs,
+            components_yaml_dir=components_path.parent,
+            image_overrides=ctx.image_overrides,
+        )
         _rewrite_task_componentref_urls(
             body_dict=body_dict,
             task_refs=task_refs,
@@ -1427,7 +1471,10 @@ def _relocate_relative_local_url(url: str, source_dir: Path, sidecar_dir: Path)
 
 
 def _build_local_from_python_components(
-    task_refs: list[tuple[str, CallableRef]], *, components_yaml_dir: Path
+    task_refs: list[tuple[str, CallableRef]],
+    *,
+    components_yaml_dir: Path,
+    image_overrides: Mapping[str, str] | None = None,
 ) -> dict[str, Any]:
     """Build the ``<stem>.components.yaml`` content for @task refs.
 
@@ -1492,6 +1539,16 @@ def _build_local_from_python_components(
         local_from_python: dict[str, Any] = {}
         if ref._task_image is not None:
             local_from_python["image"] = ref._task_image
+        elif ref._task_image_id is not None:
+            resolved_image = resolve_image_id(ref._task_image_id, image_overrides)
+            if resolved_image is None:
+                raise CompileError(
+                    f"@task image_id={ref._task_image_id!r} on function "
+                    f"{ref._task_function_name!r} did not resolve to an image. "
+                    f"Pass --image {ref._task_image_id}=IMAGE to `tangle sdk pipelines compile`, "
+                    f"or register a default with register_image_id({ref._task_image_id!r}, IMAGE)."
+                )
+            local_from_python["image"] = resolved_image
         # Always pin the function name. Without it the hydrator defaults
         # to the file stem and extracts the wrong symbol.
         assert ref._task_function_name is not None
@@ -2233,6 +2290,7 @@ def compile_file(
         overrides: Mapping[str, str] | None = None,
         pipeline_name: str | None = None,
         emit_components_sidecar: bool = True,
+        image_overrides: Mapping[str, str] | None = None,
     ) -> CompileResult:
         """Compile ``script`` to a single dehydrated pipeline YAML at ``output``.
 
@@ -2251,6 +2309,7 @@ def compile_file(
             overrides,
             pipeline_name=pipeline_name,
             emit_components_sidecar=emit_components_sidecar,
+            image_overrides=image_overrides,
         )
         self.log.info(f"wrote {result.pipeline_path}")
         if result.components_path is not None:

packages/tangle-cli/src/tangle_cli/pipelines.py

@@ -258,6 +258,7 @@ def compile_pipeline_file(
     overrides: Mapping[str, str] | None = None,
     pipeline_name: str | None = None,
     emit_components_sidecar: bool = True,
+    image_overrides: Mapping[str, str] | None = None,
     logger: Any | None = None,
 ) -> CompileResult:
     """Compile a Python-authored pipeline to a dehydrated YAML bundle.
@@ -284,6 +285,7 @@ def compile_pipeline_file(
             overrides=dict(overrides) if overrides else None,
             pipeline_name=pipeline_name,
             emit_components_sidecar=emit_components_sidecar,
+            image_overrides=dict(image_overrides) if image_overrides else None,
         )
     except (CompileError, SchemaValidationError) as exc:
         raise PipelineValidationError(str(exc)) from exc

packages/tangle-cli/src/tangle_cli/pipelines_cli.py

@@ -138,6 +138,18 @@ def _parse_overrides(values: list[str] | None) -> dict[str, str]:
     return parsed
 
 
+def _parse_image_overrides(values: list[str] | None) -> dict[str, str]:
+    parsed: dict[str, str] = {}
+    for value in values or []:
+        if "=" not in value:
+            raise SystemExit("--image entries must use ID=REF syntax")
+        image_id, image_ref = value.split("=", 1)
+        if not image_id or not image_ref:
+            raise SystemExit("--image entries must use ID=REF syntax")
+        parsed[image_id] = image_ref
+    return parsed
+
+
 @app.command(name="hydrate")
 def pipelines_hydrate(
     pipeline_path: pathlib.Path,
@@ -266,6 +278,17 @@ def pipelines_compile(
             negative_iterable=(),
         ),
     ] = None,
+    image: Annotated[
+        list[str] | None,
+        Parameter(
+            name="--image",
+            help=(
+                "Compile-time image-id override as ID=REF for @task(image_id=ID). "
+                "Repeat for multiple IDs."
+            ),
+            negative_iterable=(),
+        ),
+    ] = None,
     log_type: LogTypeOption = "console",
 ) -> None:
     """Compile a Python-authored pipeline to a dehydrated YAML bundle."""
@@ -276,6 +299,7 @@ def pipelines_compile(
             pipeline_path,
             output,
             overrides=_parse_overrides(override),
+            image_overrides=_parse_image_overrides(image),
             pipeline_name=pipeline,
             logger=logger,
         )

packages/tangle-cli/src/tangle_cli/python_pipeline/compiler_context.py

@@ -252,6 +252,7 @@ class CompileContext:
     subgraph_dir: Path
     root_overrides: dict[str, str] = field(default_factory=dict)
     emit_components_sidecar: bool = True
+    image_overrides: dict[str, str] = field(default_factory=dict)
     max_depth: int = 32
     # Compiled CHILD artifacts keyed by compile key (Decision M dedup).
     # The root is NOT stored here; it is returned directly.

packages/tangle-cli/src/tangle_cli/python_pipeline/ref.py

@@ -49,6 +49,7 @@ class CallableRef:
     _task_source_path: Path | None = None
     _task_function_name: str | None = None
     _task_image: str | None = None
+    _task_image_id: str | None = None
     _task_dependencies_from: Path | None = None
     _task_mode: str | None = None
     _task_resolve_root: Path | None = None

packages/tangle-cli/src/tangle_cli/python_pipeline/task.py

@@ -65,6 +65,7 @@ def task(
     *,
     env: TaskEnv | None = None,
     image: str | None = None,
+    image_id: str | None = None,
     dependencies_from: str | Path | None = None,
     mode: str | None = None,
     resolve_root: str | Path | None = None,
@@ -95,6 +96,10 @@ def task(
             written verbatim into the emitted
             ``components.yaml#local_from_python.image`` field. Overrides
             ``env.image`` when both are given.
+        image_id: Logical image identifier resolved at compile time via
+            registered defaults or ``tangle sdk pipelines compile --image
+            ID=REF`` overrides. Ignored when explicit ``image=`` supplies
+            an image; otherwise it also overrides ``env.image``.
         dependencies_from: Path to a ``pyproject.toml`` (or any file
             the hydrator understands) that declares pip
             dependencies. Resolved relative to the caller's source
@@ -138,7 +143,11 @@ def task(
     # overrides the corresponding ``env`` field; otherwise the env value
     # (if any) is used. This mirrors YAML anchor semantics: start from the
     # declared-once defaults, override locally where needed.
-    effective_image = image if image is not None else (env.image if env else None)
+    effective_image = (
+        image
+        if image is not None
+        else (None if image_id is not None else (env.image if env else None))
+    )
     effective_deps_raw = (
         dependencies_from
         if dependencies_from is not None
@@ -202,6 +211,7 @@ def decorator(fn: Callable[..., Any]) -> CallableRef:
             _task_source_path=source_path,
             _task_function_name=function_name,
             _task_image=effective_image,
+            _task_image_id=image_id,
             _task_dependencies_from=deps_path,
             _task_mode=mode,
             _task_resolve_root=resolve_root_path,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tangle-cli"
-version = "0.1.5"
+version = "0.1.6"
 description = "CLI for Tangle, the open-source ML pipeline orchestration platform"
 readme = "README.md"
 authors = [

tests/test_component_from_func.py

@@ -50,6 +50,22 @@
 # ``sys.modules``) is provided suite-wide by the autouse
 # ``downstream_authoring_surface`` fixture in ``tests/conftest.py``.
 
+# ============================================================================
+# @task image-id authoring metadata
+# ============================================================================
+
+
+def test_task_decorator_records_image_id_without_explicit_image():
+    from tangle_cli.python_pipeline import task
+
+    @task(image_id="eval-slim")
+    def uses_image_id() -> str:
+        return "ok"
+
+    assert uses_image_id._task_image is None
+    assert uses_image_id._task_image_id == "eval-slim"
+
+
 # ============================================================================
 # Type resolution tests
 # ============================================================================

tests/test_packaging.py

@@ -178,7 +178,7 @@ def test_tangle_cli_wheel_supports_expert_no_deps_import_path_without_tangle_api
     requires_dist = [line for line in metadata.splitlines() if line.startswith("Requires-Dist: ")]
     assert not any(name.startswith("tangle_api/") for name in names)
     assert "tangle_cli/openapi/openapi.json" not in names
-    assert "Version: 0.1.5" in metadata
+    assert "Version: 0.1.6" in metadata
     assert "Requires-Dist: tangle-api==0.1.1" in requires_dist
     assert not any("extra == 'native'" in line for line in requires_dist)
     assert "Provides-Extra: native" in metadata