poly-hammer-worker 0.1.1.dev9__tar.gz → 0.1.1.dev11__tar.gz

{poly_hammer_worker-0.1.1.dev9 → poly_hammer_worker-0.1.1.dev11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: poly-hammer-worker
-Version: 0.1.1.dev9
+Version: 0.1.1.dev11
 Summary: Poly Hammer self-hosted GPU worker agent
 License-Expression: MIT
 Requires-Python: >=3.13

{poly_hammer_worker-0.1.1.dev9 → poly_hammer_worker-0.1.1.dev11}/poly_hammer_worker/docker_runner.py

@@ -37,15 +37,6 @@ class DockerRunError(Exception):
         super().__init__(f"Container exited with code {exit_code}: {stderr}")
 
 
-def resolve_image_tag(docker_image: str, cuda_version: str) -> str:
-    """Build the fully-qualified image reference.
-
-    The backend stores bare image names (e.g. ``ph-hy-motion``).
-    This appends the CUDA version tag matching the local GPU.
-    """
-    return f"{docker_image}:cuda{cuda_version}"
-
-
 class InsufficientDiskError(Exception):
     """Raised when there isn't enough disk space to pull an image."""
 
@@ -300,6 +291,7 @@ def run_inference_container(
     job_payload: dict,
     on_progress: Callable[[float], None] | None = None,
     on_status: Callable[[str], None] | None = None,
+    on_log: Callable[[str], None] | None = None,
 ) -> dict:
     """Run an inference container and return the result dict.
 
@@ -315,6 +307,8 @@ def run_inference_container(
         Called with progress float (0.0–1.0) as the container reports progress.
     on_status:
         Called with human-readable status strings (pulling, running, etc.).
+    on_log:
+        Called with each line of container output (both stdout and stderr).
 
     Returns
     -------
@@ -324,14 +318,12 @@ def run_inference_container(
     ------
     DockerRunError  — If the container exits non-zero or emits an error message.
     """
-    stdin_bytes = json.dumps(job_payload).encode()
-
     if on_status:
         on_status(f"Creating container from {image_ref}")
 
     container = client.containers.create(
         image_ref,
-        stdin_open=True,
+        environment={"JOB_PAYLOAD": json.dumps(job_payload)},
         detach=True,
         device_requests=[
             docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]]),
@@ -344,26 +336,26 @@ def run_inference_container(
     try:
         container.start()
 
-        # Attach and pipe job JSON into stdin
-        sock = container.attach_socket(params={"stdin": 1, "stream": 1})
-        sock._sock.sendall(stdin_bytes)
-        sock._sock.shutdown(1)  # EOF on stdin
-        sock.close()
-
         if on_status:
             on_status("Inference running …")
 
-        # Stream stdout for JSON-line progress messages
+        # Stream both stdout and stderr for JSON-line progress + container logs
+        stderr_lines: list[str] = []
         for raw_line in container.logs(
-            stream=True, follow=True, stdout=True, stderr=False
+            stream=True, follow=True, stdout=True, stderr=True
         ):
             line = raw_line.decode("utf-8", errors="replace").strip()
             if not line:
                 continue
+
+            if on_log:
+                on_log(line)
+
             try:
                 msg = json.loads(line)
             except json.JSONDecodeError:
-                logger.debug("Non-JSON output: %s", line[:200])
+                # Non-JSON output is stderr or unstructured stdout
+                stderr_lines.append(line)
                 continue
 
             msg_type = msg.get("type")
@@ -384,10 +376,8 @@ def run_inference_container(
             raise DockerRunError(exit_code, error_msg)
 
         if exit_code != 0:
-            stderr_logs = container.logs(stdout=False, stderr=True).decode(
-                "utf-8", errors="replace"
-            )[-2000:]
-            raise DockerRunError(exit_code, stderr_logs)
+            stderr_tail = "\n".join(stderr_lines)[-2000:]
+            raise DockerRunError(exit_code, stderr_tail)
 
         if result_msg is None:
             raise DockerRunError(

{poly_hammer_worker-0.1.1.dev9 → poly_hammer_worker-0.1.1.dev11}/poly_hammer_worker/image_discovery.py

@@ -10,7 +10,12 @@ Required labels (image is skipped if any are missing):
   - ai.poly-hammer.portal-compatible = "true"
   - ai.poly-hammer.service-id
   - ai.poly-hammer.model-ids       (comma-separated model IDs)
-  - ai.poly-hammer.cuda-version
+
+Optional labels:
+  - ai.poly-hammer.min-cuda-version   Minimum host CUDA version required.
+        When set **and** *host_cuda_version* is passed to
+        :func:`discover_local_images`, images whose minimum exceeds the
+        host capability are skipped with a warning.
 
 Per-model labels (optional, looked up for each model ID):
   - ai.poly-hammer.model.<id>.name
@@ -32,17 +37,22 @@ _REQUIRED_LABELS = frozenset(
         "portal-compatible",
         "service-id",
         "model-ids",
-        "cuda-version",
     }
 )
 
 
-def discover_local_images(client: docker.DockerClient) -> list[dict[str, Any]]:
+def discover_local_images(
+    client: docker.DockerClient,
+    host_cuda_version: str | None = None,
+) -> list[dict[str, Any]]:
     """Scan local Docker images for portal-compatible OCI labels.
 
     Returns a list of dicts matching the ``WorkerImageInfo`` schema — one
     entry per Docker image. Each entry contains a ``models`` list with
     per-model metadata (name, hardware requirements).
+
+    If *host_cuda_version* is provided (e.g. ``"13.1"``), images whose
+    ``ai.poly-hammer.min-cuda-version`` label exceeds it are skipped.
     """
     results: list[dict[str, Any]] = []
 
@@ -77,8 +87,23 @@ def discover_local_images(client: docker.DockerClient) -> list[dict[str, Any]]:
             continue
 
         image_ref = _best_tag(image)
+        if not image_ref:
+            logger.debug("Skipping untagged image %s", image.short_id)
+            continue
+
         service_id = ph_labels["service-id"]
-        cuda_version = ph_labels["cuda-version"]
+        min_cuda_version: str | None = ph_labels.get("min-cuda-version")
+
+        # Host compatibility check: skip images that need a newer CUDA
+        if min_cuda_version and host_cuda_version:
+            if _parse_version(min_cuda_version) > _parse_version(host_cuda_version):
+                logger.warning(
+                    "Image %s requires CUDA >= %s but host has %s — skipping",
+                    image_ref,
+                    min_cuda_version,
+                    host_cuda_version,
+                )
+                continue
 
         # Image size in GB (Docker reports bytes)
         image_size_gb = round(image.attrs.get("Size", 0) / (1024**3), 2)
@@ -106,7 +131,7 @@ def discover_local_images(client: docker.DockerClient) -> list[dict[str, Any]]:
             {
                 "image_ref": image_ref,
                 "service_id": service_id,
-                "cuda_version": cuda_version,
+                "min_cuda_version": min_cuda_version,
                 "image_size_gb": image_size_gb,
                 "models": models,
             }
@@ -121,13 +146,12 @@ def discover_local_images(client: docker.DockerClient) -> list[dict[str, Any]]:
     return results
 
 
-def _best_tag(image: Any) -> str:
-    """Return the most useful tag for an image, or fall back to short ID."""
+def _best_tag(image: Any) -> str | None:
+    """Return the most useful tag for an image, or None if untagged."""
     tags = image.tags
     if tags:
         return tags[0]
-    short_id = image.short_id
-    return short_id.replace("sha256:", "") if short_id else "<untagged>"
+    return None
 
 
 def _parse_int(value: str) -> int:
@@ -136,3 +160,29 @@ def _parse_int(value: str) -> int:
         return int(value)
     except (ValueError, TypeError):
         return 0
+
+
+def _parse_version(version: str) -> tuple[int, ...]:
+    """Parse a dotted version string into a tuple of ints for comparison."""
+    try:
+        return tuple(int(p) for p in version.split("."))
+    except (ValueError, TypeError):
+        return (0,)
+
+
+def resolve_image_for_model(
+    available_images: list[dict[str, Any]], model_id: str
+) -> str | None:
+    """Find the local Docker image reference that supports a given model ID.
+
+    Searches the list returned by :func:`discover_local_images` for the
+    first image whose ``models`` list contains the requested *model_id*.
+
+    Returns the ``image_ref`` (e.g. ``"hy-motion:1.0"``) or ``None`` if
+    no matching image is found.
+    """
+    for img in available_images:
+        for model in img.get("models", []):
+            if model.get("model_id") == model_id:
+                return img.get("image_ref")
+    return None

{poly_hammer_worker-0.1.1.dev9 → poly_hammer_worker-0.1.1.dev11}/poly_hammer_worker/worker_loop.py

@@ -22,10 +22,10 @@ from poly_hammer_worker.docker_runner import (
     InsufficientDiskError,
     ensure_image,
     get_free_disk_gb,
-    resolve_image_tag,
     run_inference_container,
 )
 from poly_hammer_worker.hardware import detect_hardware
+from poly_hammer_worker.image_discovery import resolve_image_for_model
 
 logger = logging.getLogger(__name__)
 
@@ -87,7 +87,9 @@ def run_worker_loop(
     docker_client = docker.from_env()
 
     # ---- Image discovery (initial) -------------------------------------------
-    available_images = discover_local_images(docker_client)
+    available_images = discover_local_images(
+        docker_client, host_cuda_version=cuda_version
+    )
     if available_images:
         total_models = sum(len(img["models"]) for img in available_images)
         console.print(
@@ -132,7 +134,9 @@ def run_worker_loop(
         nonlocal _last_disk_check, available_images
         while not shutdown.is_set():
             # Re-discover images on each heartbeat (cheap — Docker API list)
-            available_images = discover_local_images(docker_client)
+            available_images = discover_local_images(
+                docker_client, host_cuda_version=cuda_version
+            )
 
             # Throttled disk space check
             now = time.monotonic()
@@ -217,7 +221,7 @@ def run_worker_loop(
             portal=portal,
             docker_client=docker_client,
             job=job,
-            cuda_version=cuda_version,
+            available_images=available_images,
             auto_cleanup=auto_cleanup,
             registry_url=registry_url,
             console=console,
@@ -234,21 +238,26 @@ def _execute_job(
     portal: PortalClient,
     docker_client: docker.DockerClient,
     job: dict,
-    cuda_version: str,
+    available_images: list[dict],
     console: Console,
     auto_cleanup: bool = False,
     registry_url: str | None = None,
 ) -> None:
     """Ensure image (local-first, optional registry pull), run container, report outcome."""
     job_id = str(job["job_id"])
-    docker_image = job.get("docker_image", "")
+    model_id = job.get("model_id", "")
 
-    if not docker_image:
-        portal.fail_job(job_id, "Job missing docker_image field", "NO_IMAGE")
-        console.print("[red]  Job missing docker_image — skipped[/red]")
+    image_ref = resolve_image_for_model(available_images, model_id)
+    if not image_ref:
+        portal.fail_job(
+            job_id,
+            f"No local Docker image found for model '{model_id}'. "
+            "Build the appropriate image and restart the worker.",
+            "NO_IMAGE",
+        )
+        console.print(f"[red]  No local image for model '{model_id}' — skipped[/red]")
         return
 
-    image_ref = resolve_image_tag(docker_image, cuda_version)
     start_time = time.monotonic()
 
     try:
@@ -262,22 +271,42 @@ def _execute_job(
         )
 
         # ---- Run container ----------------------------------------------------
+        container_logger = logging.getLogger(f"{__name__}.container")
+
+        def _on_log(line: str) -> None:
+            container_logger.info(line)
+            console.print(f"  [dim]│ {line}[/dim]", highlight=False)
+
         result = run_inference_container(
             docker_client,
             image_ref,
             job_payload=job,
             on_progress=lambda p: _on_progress(portal, job_id, p, console),
             on_status=lambda s: console.print(f"  [dim]{s}[/dim]"),
+            on_log=_on_log,
         )
 
         elapsed = time.monotonic() - start_time
 
         # ---- Report completion ------------------------------------------------
+        # The container uploads to S3 via the presigned URL but doesn't
+        # echo the key back in its result JSON.  Fall back to the s3_key
+        # from the job payload so the backend can store it.
+        result_url = result.get("result_url") or job.get("s3_key", "")
+
+        # Build result_metadata: use explicit field if present, otherwise
+        # collect useful top-level fields the container emitted.
+        result_metadata = result.get("result_metadata") or {
+            k: v
+            for k, v in result.items()
+            if k not in ("type", "status", "result_url", "result_metadata")
+        }
+
         portal.complete_job(
             job_id,
             elapsed_seconds=elapsed,
-            result_url=result.get("result_url"),
-            result_metadata=result.get("result_metadata"),
+            result_url=result_url,
+            result_metadata=result_metadata,
         )
         console.print(f"  [green]Completed in {elapsed:.1f}s[/green]")
 

{poly_hammer_worker-0.1.1.dev9 → poly_hammer_worker-0.1.1.dev11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "poly-hammer-worker"
-version = "0.1.1.dev9"
+version = "0.1.1.dev11"
 description = "Poly Hammer self-hosted GPU worker agent"
 readme = "README.md"
 requires-python = ">=3.13"