agent-hitch 0.2.6 → 0.2.8

package/integrations/harbor/hitch_harbor_environment.py

@@ -0,0 +1,469 @@
+"""Harbor Docker environment that stamps lease ownership on Compose resources."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import re
+import tempfile
+from pathlib import Path
+from typing import Any, Mapping
+
+import yaml
+from harbor.constants import MAIN_SERVICE_NAME
+from harbor.environments.docker.docker import DockerEnvironment
+
+_LABEL_ROOT = "io.hitch.root-id"
+_LABEL_PROVIDER = "io.hitch.provider"
+_LABEL_EVAL = "io.hitch.eval-id"
+_LABEL_WORK = "io.hitch.work-id"
+_LABEL_LEASE = "io.hitch.lease-id"
+_LABEL_EPOCH = "io.hitch.lease-epoch"
+_LABEL_TASK = "io.hitch.task-id"
+_REQUIRED_LABELS = {
+    _LABEL_ROOT,
+    _LABEL_PROVIDER,
+    _LABEL_EVAL,
+    _LABEL_WORK,
+    _LABEL_LEASE,
+    _LABEL_EPOCH,
+}
+_ALLOWED_LABELS = _REQUIRED_LABELS | {_LABEL_TASK}
+_COMPOSE_RESET_MARKER = "__HITCH_COMPOSE_RESET_NULL__"
+_GPU_OVERRIDE_PREFIX = "__HITCH_GPU_OVERRIDE_"
+
+
+class HitchHarborDockerEnvironment(DockerEnvironment):
+    """Use Harbor's Docker semantics with a final ownership-label overlay."""
+
+    _hitch_static_no_network = False
+
+    def __init__(
+        self,
+        *args: Any,
+        hitch_ownership_labels: Mapping[str, str] | None = None,
+        hitch_service_resource_limits: Mapping[str, Mapping[str, int]] | None = None,
+        hitch_resolved_images: Mapping[str, str] | None = None,
+        hitch_prebuilt_task_image: str | None = None,
+        hitch_model_proxy_host_gateway: bool = False,
+        hitch_main_gpu_count: int = 0,
+        **kwargs: Any,
+    ) -> None:
+        self._hitch_ownership_labels = _validate_labels(hitch_ownership_labels)
+        self._hitch_service_resource_limits = _validate_resource_limits(
+            hitch_service_resource_limits
+        )
+        self._hitch_resolved_images = _validate_resolved_images(hitch_resolved_images)
+        self._hitch_prebuilt_task_image = _validate_prebuilt_task_image(
+            hitch_prebuilt_task_image
+        )
+        if not isinstance(hitch_model_proxy_host_gateway, bool):
+            raise ValueError("Hitch model proxy host gateway flag is invalid")
+        self._hitch_model_proxy_host_gateway = hitch_model_proxy_host_gateway
+        if isinstance(hitch_main_gpu_count, bool) or not isinstance(hitch_main_gpu_count, int) or hitch_main_gpu_count < 0:
+            raise ValueError("Hitch main GPU count is invalid")
+        self._hitch_main_gpu_count = hitch_main_gpu_count
+        self._hitch_ownership_temp_dir: tempfile.TemporaryDirectory[str] | None = None
+        self._hitch_ownership_compose_path: Path | None = None
+        self._hitch_phase_compose_path: Path | None = None
+        self._hitch_candidate_recycler = None
+        positional = list(args)
+        task_env_config = kwargs.get("task_env_config")
+        task_env_position: int | None = None
+        if task_env_config is None and len(positional) >= 5:
+            task_env_position = 4
+            task_env_config = positional[task_env_position]
+        self._hitch_static_no_network = _static_no_network(positional, kwargs, task_env_config)
+        self._hitch_static_network_receipt = None
+        if task_env_config is not None:
+            requested = getattr(task_env_config, "docker_image", None)
+            resolved = self._hitch_prebuilt_task_image or self._hitch_resolved_images.get(requested)
+            if resolved is not None:
+                updated = task_env_config.model_copy(update={"docker_image": resolved})
+                if task_env_position is None:
+                    kwargs["task_env_config"] = updated
+                else:
+                    positional[task_env_position] = updated
+        super().__init__(*positional, **kwargs)
+        self._hitch_benchmark = None
+        self._hitch_benchmark_platform = None
+        directory = getattr(self, "environment_dir", None)
+        if directory and (directory.parent / ".hitch-benchmark.json").is_file():
+            self._hitch_benchmark_platform = "linux/amd64"
+        if directory and directory.name == "environment" and (directory.parent / ".hitch-benchmark.json").is_file():
+            from hitch_benchmark import BenchmarkSession, descriptor
+            self._hitch_benchmark = BenchmarkSession(self, descriptor(directory))
+        if (
+            self._hitch_ownership_labels
+            or self._hitch_service_resource_limits
+            or self._hitch_resolved_images
+            or self._hitch_prebuilt_task_image
+            or self._hitch_model_proxy_host_gateway
+            or self._hitch_main_gpu_count > 0
+            or self._hitch_benchmark_platform
+            or self._hitch_static_no_network
+        ):
+            self._hitch_ownership_compose_path = self._write_ownership_overlay()
+
+    @property
+    def _docker_compose_paths(self) -> list[Path]:
+        if self._hitch_static_no_network and (self._environment_docker_compose_path.exists() or self.extra_docker_compose_paths):
+            raise ValueError("static no-network requires a single Dockerfile/image environment without Compose additions")
+        paths = list(super()._docker_compose_paths)
+        if self._hitch_ownership_compose_path is not None:
+            paths.append(self._hitch_ownership_compose_path)
+        if self._hitch_phase_compose_path is not None:
+            paths.append(self._hitch_phase_compose_path)
+        return paths
+
+    def _requires_egress_control(self, *, startup_network_policy, phase_network_policies):
+        if self._hitch_static_no_network:
+            return False
+        return super()._requires_egress_control(startup_network_policy=startup_network_policy, phase_network_policies=phase_network_policies)
+
+    @property
+    def capabilities(self):
+        result = super().capabilities
+        if self._hitch_static_no_network:
+            return result.model_copy(update={"disable_internet": True, "dynamic_network_policy": False})
+        return result
+
+    async def _apply_network_policy(self, network_policy):
+        if self._hitch_static_no_network:
+            if network_policy != self.network_policy:
+                raise ValueError("static no-network cannot change runtime network policy")
+            return
+        await super()._apply_network_policy(network_policy)
+
+    async def _verify_static_network(self):
+        result = await self._run_docker_compose_command(["ps", "--all", "--quiet", MAIN_SERVICE_NAME])
+        identity = result.stdout.strip()
+        if not re.fullmatch(r"[a-f0-9]{64}", identity):
+            raise RuntimeError("static no-network requires exactly one owned main container")
+        template = '{"mode":{{json .HostConfig.NetworkMode}},"networks":{{json .NetworkSettings.Networks}},"ports":{{json .HostConfig.PortBindings}}}'
+        process = await asyncio.create_subprocess_exec("docker", "inspect", "--format", template, identity,
+                                                       stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL)
+        try:
+            stdout, _ = await asyncio.wait_for(process.communicate(), 15)
+        finally:
+            if process.returncode is None:
+                process.kill(); await process.wait()
+        if process.returncode != 0 or len(stdout) > 65536:
+            raise RuntimeError("could not verify static no-network container")
+        inspected = json.loads(stdout)
+        if inspected.get("mode") != "none" or set(inspected.get("networks") or {}) != {"none"} or inspected.get("ports"):
+            raise RuntimeError("Docker did not enforce the static no-network declaration")
+        self._hitch_static_network_receipt = {"container_id": identity, "network_mode": "none", "host_ports": False}
+
+    async def recycle_candidate_phase(self, phase_index: int) -> dict[str, Any]:
+        """Replace only main after the supervisor has exported the retired run.
+
+        Does not start a model or prepare a tool binding. The native task VM and
+        service containers remain live. Failures require whole-trial cleanup.
+        """
+        if self._hitch_candidate_recycler is None:
+            from hitch_candidate_recycle import CandidateRecycler
+            self._hitch_candidate_recycler = CandidateRecycler(self)
+        return await self._hitch_candidate_recycler.recycle(phase_index)
+
+    async def start(self, force_build: bool):
+        try:
+            await super().start(force_build)
+            if self._hitch_static_no_network:
+                await self._verify_static_network()
+            if self._hitch_benchmark:
+                await self._hitch_benchmark.prepare()
+        except BaseException:
+            # Harbor may abort setup before entering its normal teardown path.
+            try:
+                await self.stop(delete=True)
+            except BaseException:
+                pass
+            raise
+
+    async def stop_service(self, service: str) -> None:
+        try:
+            await super().stop_service(service)
+            if service == MAIN_SERVICE_NAME and self._hitch_benchmark:
+                await self._hitch_benchmark.snapshot()
+        except BaseException as error:
+            if self._hitch_benchmark:
+                self._hitch_benchmark.failure = {"phase": "quiesce", "error": type(error).__name__}
+                self._hitch_benchmark.write_journal()
+            raise
+
+    async def stop(self, delete: bool):
+        try:
+            if self._hitch_benchmark and self._hitch_benchmark.started:
+                try:
+                    await super().stop_service(MAIN_SERVICE_NAME)
+                finally:
+                    await self._hitch_benchmark.cleanup()
+        finally:
+            await super().stop(delete)
+
+    def _write_ownership_overlay(self) -> Path:
+        services = {MAIN_SERVICE_NAME}
+        if self._enable_egress_control:
+            services.add(self._EGRESS_CONTROL_SERVICE_NAME)
+        networks = {"default"}
+        volumes: set[str] = set()
+        external_networks: set[str] = set()
+        external_volumes: set[str] = set()
+        build_services: set[str] = set()
+        declared_platforms: set[str] = set()
+        image_overlays: dict[str, str] = {}
+        sources = [self._environment_docker_compose_path, *self.extra_docker_compose_paths]
+        for source in sources:
+            if not source.exists():
+                continue
+            document = yaml.safe_load(source.read_text(encoding="utf-8"))
+            if document is None:
+                continue
+            if not isinstance(document, dict):
+                raise ValueError(f"Docker Compose ownership source must be a mapping: {source}")
+            services.update(_mapping_names(document.get("services"), "services", source))
+            for name, config in (document.get("services") or {}).items():
+                if isinstance(config, dict):
+                    if 'platform' in config:
+                        if config['platform'] is not None:
+                            declared_platforms.add(name)
+                        else:
+                            declared_platforms.discard(name)
+                    if "build" in config and config.get("build") is not None:
+                        build_services.add(name)
+                    reference = config.get("image")
+                    resolved = self._hitch_resolved_images.get(reference)
+                    if resolved is not None:
+                        image_overlays[name] = resolved
+                    elif reference is not None:
+                        image_overlays.pop(name, None)
+            _collect_resources(document.get("networks"), networks, external_networks, "networks", source)
+            _collect_resources(document.get("volumes"), volumes, external_volumes, "volumes", source)
+
+        labels = dict(self._hitch_ownership_labels)
+        service_overlays: dict[str, dict[str, Any]] = {}
+        for name in sorted(services):
+            config: dict[str, Any] = {}
+            # The benchmark default must not replace explicit per-service
+            # architecture, e.g. an ARM QEMU sidecar beside an AMD64 candidate.
+            if getattr(self, "_hitch_benchmark_platform", None) and name not in declared_platforms:
+                config["platform"] = self._hitch_benchmark_platform
+            if labels:
+                config["labels"] = labels
+            if name == MAIN_SERVICE_NAME and self._hitch_static_no_network:
+                config["network_mode"] = "none"
+            resolved_image = image_overlays.get(name)
+            if name == MAIN_SERVICE_NAME and self._hitch_prebuilt_task_image:
+                resolved_image = self._hitch_prebuilt_task_image
+            if resolved_image is not None:
+                config["image"] = resolved_image
+            if (
+                name == MAIN_SERVICE_NAME
+                and self._hitch_prebuilt_task_image
+                and name in build_services
+            ):
+                # Compose's normal merge semantics retain an earlier build stanza
+                # when an overlay only supplies image (or even build: null).  The
+                # !reset tag is the supported way to make a successful Hitch
+                # prebuild authoritative and prevent a second mutable-context build.
+                config["build"] = _COMPOSE_RESET_MARKER
+            if name == MAIN_SERVICE_NAME and self._hitch_model_proxy_host_gateway:
+                config["extra_hosts"] = ["host.docker.internal:host-gateway"]
+            if name == MAIN_SERVICE_NAME and self._hitch_main_gpu_count > 0:
+                config["deploy"] = {
+                    "resources": {
+                        "reservations": {
+                            "devices": f"{_GPU_OVERRIDE_PREFIX}{self._hitch_main_gpu_count}__"
+                        }
+                    }
+                }
+            if name != MAIN_SERVICE_NAME and self._hitch_service_resource_limits:
+                limits = self._hitch_service_resource_limits.get(name)
+                if limits is None:
+                    raise ValueError(
+                        f"Docker Compose sidecar has no Hitch resource limit: {name}"
+                    )
+                config.update({
+                    "cpus": limits["cpu_millis"] / 1000,
+                    "mem_limit": limits["memory_bytes"],
+                })
+                if limits.get("gpu_count", 0) > 0:
+                    config["deploy"] = {
+                        "resources": {
+                            "reservations": {
+                                "devices": f"{_GPU_OVERRIDE_PREFIX}{limits['gpu_count']}__"
+                            }
+                        }
+                    }
+            service_overlays[name] = config
+        overlay = {
+            "services": service_overlays,
+            "networks": {
+                name: {"labels": labels}
+                for name in sorted(networks - external_networks)
+            },
+            "volumes": {
+                name: {"labels": labels}
+                for name in sorted(volumes - external_volumes)
+            },
+        }
+        self._hitch_ownership_temp_dir = tempfile.TemporaryDirectory(
+            prefix="hitch-harbor-ownership-"
+        )
+        target = Path(self._hitch_ownership_temp_dir.name) / "docker-compose-hitch-ownership.yaml"
+        contents = json.dumps(overlay, indent=2, sort_keys=True)
+        contents = contents.replace(
+            f'"build": "{_COMPOSE_RESET_MARKER}"',
+            '"build": !reset null',
+        )
+        contents = re.sub(
+            rf'"devices": "{_GPU_OVERRIDE_PREFIX}([1-9][0-9]*)__"',
+            lambda match: '"devices": !override [{"capabilities": ["gpu"], "count": ' + match.group(1) + '}]',
+            contents,
+        )
+        target.write_text(contents, encoding="utf-8")
+        return target
+
+
+def _static_no_network(positional, kwargs, task_env_config) -> bool:
+    """Use Docker's isolated namespace only when no phase needs connectivity.
+
+    Compose tasks and dynamic policies retain Harbor's egress-control path.
+    A single image/Dockerfile has no sibling-service network to preserve.
+    """
+    policy = kwargs.get("network_policy", positional[6] if len(positional) > 6 else None)
+    phases = kwargs.get("phase_network_policies", positional[7] if len(positional) > 7 else ()) or ()
+    directory = kwargs.get("environment_dir", positional[0] if positional else None)
+    return bool(
+        getattr(getattr(task_env_config, "os", None), "value", None) == "linux"
+        and getattr(getattr(policy, "network_mode", None), "value", None) == "no-network"
+        and all(phase == policy for phase in phases)
+        and directory is not None
+        and not (Path(directory) / "docker-compose.yaml").exists()
+        and not kwargs.get("extra_docker_compose")
+    )
+
+
+def _mapping_names(value: Any, label: str, source: Path) -> set[str]:
+    if value is None:
+        return set()
+    if not isinstance(value, dict) or any(not isinstance(name, str) or not name for name in value):
+        raise ValueError(f"Docker Compose {label} must be a named mapping: {source}")
+    return set(value)
+
+
+def _collect_resources(
+    value: Any,
+    names: set[str],
+    external: set[str],
+    label: str,
+    source: Path,
+) -> None:
+    for name in _mapping_names(value, label, source):
+        names.add(name)
+        config = value[name]
+        if isinstance(config, dict) and bool(config.get("external")):
+            external.add(name)
+
+
+def _validate_labels(value: Mapping[str, str] | None) -> dict[str, str]:
+    if value is None:
+        return {}
+    if not isinstance(value, Mapping) or set(value) - _ALLOWED_LABELS or not _REQUIRED_LABELS <= set(value):
+        raise ValueError("Hitch Docker ownership label fields are invalid")
+    labels = dict(value)
+    if (
+        not re.fullmatch(r"[a-f0-9]{24}", labels.get(_LABEL_ROOT, ""))
+        or labels.get(_LABEL_PROVIDER) != "local-docker"
+        or not re.fullmatch(r"eval_[a-f0-9]{32}", labels.get(_LABEL_EVAL, ""))
+        or not re.fullmatch(r"work_[a-f0-9]{32}", labels.get(_LABEL_WORK, ""))
+        or not re.fullmatch(r"lease_[a-f0-9]{32}", labels.get(_LABEL_LEASE, ""))
+        or not re.fullmatch(r"[1-9][0-9]*", labels.get(_LABEL_EPOCH, ""))
+        or any(not isinstance(entry, str) or not entry or len(entry) > 4096 or "\x00" in entry or "\n" in entry or "\r" in entry for entry in labels.values())
+    ):
+        raise ValueError("Hitch Docker ownership label values are invalid")
+    return dict(sorted(labels.items()))
+
+
+def _validate_resource_limits(
+    value: Mapping[str, Mapping[str, int]] | None,
+) -> dict[str, dict[str, int]]:
+    if value is None:
+        return {}
+    if not isinstance(value, Mapping):
+        raise ValueError("Hitch Docker service resource limits are invalid")
+    result: dict[str, dict[str, int]] = {}
+    for name, resources in value.items():
+        if (
+            not isinstance(name, str)
+            or not name
+            or name == MAIN_SERVICE_NAME
+            or len(name) > 255
+            or any(character in name for character in ("\x00", "\n", "\r"))
+            or not isinstance(resources, Mapping)
+            or set(resources) - {"cpu_millis", "memory_bytes", "gpu_count"}
+            or not {"cpu_millis", "memory_bytes"} <= set(resources)
+            or isinstance(resources.get("cpu_millis"), bool)
+            or not isinstance(resources.get("cpu_millis"), int)
+            or resources["cpu_millis"] < 1
+            or isinstance(resources.get("memory_bytes"), bool)
+            or not isinstance(resources.get("memory_bytes"), int)
+            or resources["memory_bytes"] < 1
+            or ("gpu_count" in resources and (
+                isinstance(resources.get("gpu_count"), bool)
+                or not isinstance(resources.get("gpu_count"), int)
+                or resources["gpu_count"] < 1
+            ))
+        ):
+            raise ValueError("Hitch Docker service resource limits are invalid")
+        result[name] = {
+            "cpu_millis": resources["cpu_millis"],
+            "memory_bytes": resources["memory_bytes"],
+            **({"gpu_count": resources["gpu_count"]} if "gpu_count" in resources else {}),
+        }
+    return dict(sorted(result.items()))
+
+
+def _validate_resolved_images(value: Mapping[str, str] | None) -> dict[str, str]:
+    if value is None:
+        return {}
+    if not isinstance(value, Mapping):
+        raise ValueError("Hitch resolved image mapping is invalid")
+    result: dict[str, str] = {}
+    for requested, resolved in value.items():
+        if (
+            not _valid_image_reference(requested)
+            or not _valid_image_reference(resolved)
+            or re.search(r"@sha256:[a-f0-9]{64}$", resolved) is None
+            or _image_repository(requested) != _image_repository(resolved)
+        ):
+            raise ValueError("Hitch resolved image mapping is invalid")
+        result[requested] = resolved
+    return dict(sorted(result.items()))
+
+
+def _validate_prebuilt_task_image(value: str | None) -> str | None:
+    if value is None:
+        return None
+    if not isinstance(value, str) or re.fullmatch(r"sha256:[a-f0-9]{64}", value) is None:
+        raise ValueError("Hitch prebuilt task image is invalid")
+    return value
+
+
+def _valid_image_reference(value: Any) -> bool:
+    return (
+        isinstance(value, str)
+        and bool(value)
+        and len(value) <= 1024
+        and not any(character.isspace() or character == "\x00" for character in value)
+        and "://" not in value
+        and "$" not in value
+    )
+
+
+def _image_repository(reference: str) -> str:
+    without_digest = reference.split("@", 1)[0]
+    slash = without_digest.rfind("/")
+    colon = without_digest.rfind(":")
+    return without_digest[:colon] if colon > slash else without_digest