antioch-py 2.2.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

common/core/container.py

@@ -0,0 +1,261 @@
+import contextlib
+import time
+from enum import Enum
+
+import docker
+from docker.errors import APIError
+from docker.models.containers import Container
+
+from common.ark import Ark as ArkDefinition, Environment
+from common.ark.module import ModuleImage, ModuleReady, ModuleStart
+from common.constants import ANTIOCH_API_URL
+from common.core.auth import AuthHandler
+from common.core.rome import RomeClient
+from common.utils.comms import CommsSession
+from common.utils.time import now_us
+
+# Container naming prefix for all Antioch module containers
+CONTAINER_PREFIX = "antioch-module-"
+
+# Synchronization paths for module coordination
+ARK_MODULE_READY_PATH = "_ark/module_ready"
+ARK_MODULE_START_PATH = "_ark/module_start"
+
+
+class ContainerSource(str, Enum):
+    """
+    Source location for container images.
+    """
+
+    LOCAL = "Local"
+    REMOTE = "Remote"
+
+
+class ContainerManagerError(Exception):
+    """
+    Raised when container management operations fail.
+    """
+
+    pass
+
+
+class ContainerManager:
+    """
+    Manages Docker containers for Ark modules.
+
+    Handles launching, coordination, and cleanup of module containers.
+    Uses host networking for Zenoh communication.
+    """
+
+    def __init__(self) -> None:
+        """
+        Create a new container manager.
+
+        Initializes Docker client and Zenoh communication session.
+        """
+
+        self._comms = CommsSession()
+        self._client = docker.from_env()
+        self._containers: dict[str, Container] = {}
+
+    def launch_ark(
+        self,
+        ark: ArkDefinition,
+        source: ContainerSource = ContainerSource.LOCAL,
+        environment: Environment = Environment.SIM,
+        debug: bool = False,
+        timeout: float = 30.0,
+    ) -> int:
+        """
+        Launch all module containers for an Ark.
+
+        Stops any existing module containers first to ensure idempotent behavior.
+
+        :param ark: Ark definition to launch.
+        :param source: Container image source (local or remote).
+        :param environment: Environment to run in (sim or real).
+        :param debug: Enable debug mode.
+        :param timeout: Timeout in seconds for modules to become ready.
+        :return: Global start time in microseconds.
+        :raises ContainerManagerError: If environment is incompatible or launch fails.
+        """
+
+        # Validate environment compatibility
+        if ark.capability == Environment.SIM and environment == Environment.REAL:
+            raise ContainerManagerError(f"Ark '{ark.name}' has sim capability but requested for real")
+        if ark.capability == Environment.REAL and environment == Environment.SIM:
+            raise ContainerManagerError(f"Ark '{ark.name}' has real capability but requested for sim")
+
+        # Stop all existing module containers (idempotent)
+        self._stop_all()
+
+        # Get GAR credentials if pulling from remote
+        gar_auth = self._get_gar_auth() if source == ContainerSource.REMOTE else None
+
+        # Build container configs
+        configs: list[tuple[str, str, str]] = []
+        for module in ark.modules:
+            image = self._get_image(module.image, environment)
+            if image is None:
+                raise ContainerManagerError(f"No image for module '{module.name}' in {environment}")
+            if gar_auth is not None:
+                image = f"{gar_auth['registry_host']}/{gar_auth['repository']}/{image}"
+            container_name = f"{CONTAINER_PREFIX}{ark.name.replace('_', '-')}-{module.name.replace('_', '-')}"
+            configs.append((module.name, container_name, image))
+
+        # Pull images if remote
+        if gar_auth is not None:
+            self._pull_images([c[2] for c in configs], gar_auth)
+
+        # Set up ready subscriber before launching
+        ready_sub = self._comms.declare_async_subscriber(ARK_MODULE_READY_PATH)
+
+        # Launch containers
+        ark_json = ark.model_dump_json()
+        for module_name, container_name, image in configs:
+            self._launch(module_name, container_name, image, ark_json, environment, debug)
+
+        # Wait for all modules to be ready
+        pending = {m.name for m in ark.modules}
+        start = time.time()
+        while pending:
+            if time.time() - start > timeout:
+                raise ContainerManagerError(f"Timeout waiting for modules: {', '.join(sorted(pending))}")
+            msg = ready_sub.recv_timeout(ModuleReady, timeout=0.1)
+            if msg is not None:
+                print(f"Module ready: {msg.module_name}")
+                pending.discard(msg.module_name)
+
+        # Broadcast global start time (2s in future for sync)
+        global_start_us = ((now_us() // 1_000_000) + 2) * 1_000_000
+        self._comms.declare_publisher(ARK_MODULE_START_PATH).publish(ModuleStart(global_start_time_us=global_start_us))
+        return global_start_us
+
+    def stop(self, timeout: float = 10.0) -> None:
+        """
+        Stop all module containers.
+
+        :param timeout: Timeout in seconds for container stop operation.
+        """
+
+        self._stop_all(timeout)
+
+    def close(self, timeout: float = 10.0) -> None:
+        """
+        Close the container manager and clean up resources.
+
+        Stops all module containers and closes the Zenoh session.
+
+        :param timeout: Timeout in seconds for container stop operation.
+        """
+
+        self._stop_all(timeout)
+        self._comms.close()
+
+    def _launch(
+        self,
+        module_name: str,
+        container_name: str,
+        image: str,
+        ark_json: str,
+        environment: Environment,
+        debug: bool,
+    ) -> None:
+        """
+        Launch a single module container.
+
+        :param module_name: Name of the module.
+        :param container_name: Docker container name.
+        :param image: Docker image to use.
+        :param ark_json: Serialized Ark definition.
+        :param environment: Environment (sim or real).
+        :param debug: Enable debug mode.
+        :raises ContainerManagerError: If container launch fails.
+        """
+
+        try:
+            container = self._client.containers.run(
+                image=image,
+                name=container_name,
+                environment={
+                    "_MODULE_NAME": module_name,
+                    "_ARK": ark_json,
+                    "_ENVIRONMENT": str(environment.value),
+                    "_DEBUG": str(debug).lower(),
+                },
+                network_mode="host",
+                ipc_mode="host",
+                detach=True,
+                remove=False,
+            )
+            self._containers[container_name] = container
+            print(f"Launched container: {container_name}")
+        except APIError as e:
+            raise ContainerManagerError(f"Failed to launch '{container_name}': {e}") from e
+
+    def _stop_all(self, timeout: float = 10.0) -> None:
+        """
+        Stop all Antioch module containers.
+
+        Finds all containers with the antioch-module- prefix and stops them.
+
+        :param timeout: Timeout in seconds for stop operation.
+        """
+
+        with contextlib.suppress(APIError):
+            for container in self._client.containers.list(all=True):
+                if container.name and container.name.startswith(CONTAINER_PREFIX):
+                    print(f"Stopping container: {container.name}")
+                    self._stop_container(container, timeout)
+
+        self._containers.clear()
+
+    def _stop_container(self, container: Container, timeout: float) -> None:
+        """
+        Stop and remove a single container.
+
+        :param container: Docker container to stop.
+        :param timeout: Timeout in seconds for stop operation.
+        """
+
+        with contextlib.suppress(APIError):
+            container.stop(timeout=int(timeout))
+        with contextlib.suppress(APIError):
+            container.remove(force=True)
+
+    def _get_image(self, image: str | ModuleImage, environment: Environment) -> str | None:
+        """
+        Get the image name for the given environment.
+
+        :param image: Image specification (string or ModuleImage).
+        :param environment: Environment to get image for.
+        :return: Image name or None if not available.
+        """
+
+        if isinstance(image, str):
+            return image
+        return image.sim if environment == Environment.SIM else image.real
+
+    def _get_gar_auth(self) -> dict:
+        """
+        Get Google Artifact Registry authentication credentials.
+
+        :return: Dictionary containing registry host, repository, and access token.
+        """
+
+        auth = AuthHandler()
+        rome = RomeClient(ANTIOCH_API_URL, auth.get_token())
+        return rome.get_gar_token()
+
+    def _pull_images(self, images: list[str], gar_auth: dict) -> None:
+        """
+        Pull container images from the registry.
+
+        :param images: List of image names to pull.
+        :param gar_auth: GAR authentication credentials.
+        """
+
+        auth_config = {"username": "oauth2accesstoken", "password": gar_auth["access_token"]}
+        for image in set(images):
+            print(f"Pulling image: {image}")
+            self._client.images.pull(image, auth_config=auth_config)

common/core/registry.py

@@ -1,12 +1,15 @@
 import json
+import os
+import tempfile
 from collections import defaultdict
 from datetime import datetime
 from pathlib import Path
 
-from common.ark import Ark as ArkDefinition, ArkReference, ArkVersionReference, AssetReference, AssetVersionReference
+from common.ark import Ark as ArkDefinition
 from common.constants import ANTIOCH_API_URL, get_ark_dir, get_asset_dir
-from common.core.auth import AuthError, AuthHandler
-from common.rome import RomeClient
+from common.core.auth import AuthHandler
+from common.core.rome import RomeClient
+from common.core.types import ArkReference, ArkVersionReference, AssetReference, AssetVersionReference
 
 
 def list_local_arks() -> list[ArkReference]:
@@ -17,17 +20,43 @@ def list_local_arks() -> list[ArkReference]:
     """
 
     arks_dir = get_ark_dir()
-    files_by_name = defaultdict(list)
-    for file_path in arks_dir.iterdir():
-        if file_path.is_file() and (file_path.name.endswith(":ark.json") or file_path.name.endswith(":asset.usdz")):
-            name = file_path.name.split(":")[0]
-            files_by_name[name].append(file_path)
 
+    # Group files by ark name and version
+    # File format: {name}:{version}:ark.json or {name}:{version}:asset.usdz
+    files_by_name_version: dict[str, dict[str, dict[str, Path]]] = defaultdict(lambda: defaultdict(dict))
+    for file_path in arks_dir.iterdir():
+        if not file_path.is_file():
+            continue
+        if file_path.name.endswith(":ark.json"):
+            name, version, _ = file_path.name.rsplit(":", 2)
+            files_by_name_version[name][version]["ark"] = file_path
+        elif file_path.name.endswith(":asset.usdz"):
+            name, version, _ = file_path.name.rsplit(":", 2)
+            files_by_name_version[name][version]["asset"] = file_path
+
+    # Build references for each ark
     results = []
-    for name, files in files_by_name.items():
-        ref = _build_ark_reference(name, files)
-        if ref is not None:
-            results.append(ref)
+    for name, versions in files_by_name_version.items():
+        version_refs = []
+        for version, files in versions.items():
+            ark_file = files.get("ark")
+            if ark_file is None:
+                continue
+            asset_file = files.get("asset")
+            ark_stat = ark_file.stat()
+            version_refs.append(
+                ArkVersionReference(
+                    version=version,
+                    full_path=str(ark_file),
+                    asset_path=str(asset_file) if asset_file else None,
+                    size_bytes=ark_stat.st_size,
+                    created_at=datetime.fromtimestamp(ark_stat.st_ctime).isoformat(),
+                    updated_at=datetime.fromtimestamp(ark_stat.st_mtime).isoformat(),
+                    asset_size_bytes=asset_file.stat().st_size if asset_file else None,
+                )
+            )
+        if version_refs:
+            results.append(build_ark_reference_from_versions(name, version_refs))
 
     return results
 
@@ -60,8 +89,8 @@ def get_ark_version_reference(name: str, version: str) -> ArkVersionReference:
             for version_ref in ark_ref.versions:
                 if version_ref.version == version:
                     return version_ref
-            raise FileNotFoundError(f"Version {version} of Ark {name} not found in local storage")
-    raise FileNotFoundError(f"No versions of Ark {name} found in local storage")
+            raise FileNotFoundError(f"Version {version} of Ark {name} not found locally. Please pull the Ark first.")
+    raise FileNotFoundError(f"No versions of Ark {name} found locally. Please pull the Ark first.")
 
 
 def get_asset_path(name: str, version: str, extension: str = "usdz", assert_exists: bool = True) -> Path:
@@ -94,19 +123,34 @@ def list_local_assets() -> list[AssetReference]:
     if not assets_dir.exists():
         return []
 
-    files_by_name: dict[str, list[Path]] = defaultdict(list)
+    # Group files by asset name and version
+    # File format: {name}:{version}:file.{extension}
+    files_by_name_version: dict[str, dict[str, Path]] = defaultdict(dict)
     for file_path in assets_dir.iterdir():
-        if file_path.is_file():
-            # Parse filename format: {name}:{version}:file.{extension}
-            parts = file_path.stem.split(":")
-            if len(parts) == 3 and parts[-1] == "file":
-                files_by_name[parts[0]].append(file_path)
+        if not file_path.is_file():
+            continue
+        parts = file_path.stem.split(":")
+        if len(parts) == 3 and parts[-1] == "file":
+            name, version = parts[0], parts[1]
+            files_by_name_version[name][version] = file_path
 
+    # Build references for each asset
     results = []
-    for name, files in files_by_name.items():
-        ref = _build_asset_reference(name, files)
-        if ref is not None:
-            results.append(ref)
+    for name, versions in files_by_name_version.items():
+        version_refs = []
+        for version, asset_file in versions.items():
+            asset_stat = asset_file.stat()
+            version_refs.append(
+                AssetVersionReference(
+                    version=version,
+                    full_path=str(asset_file),
+                    size_bytes=asset_stat.st_size,
+                    created_at=datetime.fromtimestamp(asset_stat.st_ctime).isoformat(),
+                    updated_at=datetime.fromtimestamp(asset_stat.st_mtime).isoformat(),
+                )
+            )
+        if version_refs:
+            results.append(build_asset_reference_from_versions(name, version_refs))
 
     return results
 
@@ -121,13 +165,8 @@ def list_remote_arks() -> list[ArkReference]:
     :raises AuthError: If not authenticated.
     """
 
-    # Get auth token
     auth = AuthHandler()
     token = auth.get_token()
-    if token is None:
-        raise AuthError("User not authenticated. Please login first")
-
-    # Create Rome client and list arks
     rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
     return rome_client.list_arks()
 
@@ -136,6 +175,8 @@ def pull_remote_ark(name: str, version: str, overwrite: bool = False) -> ArkDefi
     """
     Pull an Ark from remote registry to local storage.
 
+    Downloads the Ark config (ark.json) and asset (asset.usdz) if present.
+
     Requires authentication.
 
     :param name: Name of the Ark.
@@ -148,30 +189,26 @@ def pull_remote_ark(name: str, version: str, overwrite: bool = False) -> ArkDefi
     # Check if Ark already exists locally
     arks_dir = get_ark_dir()
     ark_json_path = arks_dir / f"{name}:{version}:ark.json"
+    ark_asset_path = arks_dir / f"{name}:{version}:asset.usdz"
     if ark_json_path.exists() and not overwrite:
         return load_local_ark(name, version)
 
-    # Get auth token
     auth = AuthHandler()
     token = auth.get_token()
-    if not token:
-        raise AuthError("User not authenticated. Please login first")
-
-    # Create Rome client and fetch Ark definition and save to local storage
     rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
-    ark = rome_client.get_ark(name=name, version=version)
-
-    # Save Ark JSON
-    with open(ark_json_path, "wb") as f:
-        f.write(json.dumps(ark).encode("utf-8"))
 
-    # Download asset usdz only if ark has asset_hash (hardware modules exist)
-    if ark.get("metadata", {}).get("asset_hash") is not None:
-        asset_content = rome_client.download_ark_assets(name=name, version=version)
-        with open(arks_dir / f"{name}:{version}:asset.usdz", "wb") as f:
-            f.write(asset_content)
-
-    return ArkDefinition(**ark)
+    print(f"Pulling {name} v{version}")
+    downloaded_asset = rome_client.pull_ark(
+        name=name,
+        version=version,
+        config_output_path=str(ark_json_path),
+        asset_output_path=str(ark_asset_path),
+    )
+    print("  ✓ Config downloaded")
+    if downloaded_asset:
+        print("  ✓ Asset downloaded")
+    print(f"✓ Ark {name} v{version} pulled successfully")
+    return load_local_ark(name, version)
 
 
 def list_remote_assets() -> list[AssetReference]:
@@ -184,17 +221,12 @@ def list_remote_assets() -> list[AssetReference]:
     :raises AuthError: If not authenticated.
     """
 
-    # Get auth token
     token = AuthHandler().get_token()
-    if token is None:
-        raise AuthError("User not authenticated. Please login first")
-
-    # Create Rome client and list assets
     rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
     return rome_client.list_assets()
 
 
-def pull_remote_asset(name: str, version: str, overwrite: bool = False, show_progress: bool = True) -> Path:
+def pull_remote_asset(name: str, version: str, overwrite: bool = False) -> Path:
     """
     Pull an asset from remote registry to local storage.
 
@@ -203,129 +235,76 @@ def pull_remote_asset(name: str, version: str, overwrite: bool = False, show_pro
     :param name: Name of the asset.
     :param version: Version of the asset.
     :param overwrite: Overwrite local asset if it already exists.
-    :param show_progress: Show download progress bar.
     :return: Path to the downloaded asset file.
     :raises AuthError: If not authenticated.
     """
 
-    # Get auth token
-    token = AuthHandler().get_token()
-    if token is None:
-        raise AuthError("User not authenticated. Please login first")
-
-    # Create Rome client and get asset metadata to determine extension
-    rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
-    metadata = rome_client.get_asset_metadata(name=name, version=version)
-    extension = metadata.get("extension", "usdz")
-
     # Check if asset already exists locally
-    asset_file_path = get_asset_path(name=name, version=version, extension=extension, assert_exists=False)
+    # NOTE: Only checks USDZ assets for now
+    asset_file_path = get_asset_path(name=name, version=version, extension="usdz", assert_exists=False)
     if asset_file_path.exists() and not overwrite:
-        print(f"Asset {name}:{version} already exists locally, skipping download")
         return asset_file_path
 
-    # Download the asset file
-    rome_client.download_asset(name=name, version=version, output_path=str(asset_file_path), show_progress=show_progress)
-    return asset_file_path
+    token = AuthHandler().get_token()
+    rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
+    temp_path: str | None = None
+
+    try:
+        # Download to a temp file in the destination directory so publishing can be atomic
+        # This avoids EXDEV when the asset directory is a separate mount (common in Kubernetes)
+        asset_file_path.parent.mkdir(parents=True, exist_ok=True)
+        safe_prefix = f".{name.replace(':', '_')}.{version.replace(':', '_')}."
+        fd, temp_path = tempfile.mkstemp(prefix=safe_prefix, dir=str(asset_file_path.parent))
+        os.close(fd)
+
+        # Pull asset - metadata comes back from response body
+        metadata = rome_client.pull_asset(name=name, version=version, output_path=temp_path)
+        extension = metadata.get("extension", "usdz")
+
+        # Get final path with correct extension
+        asset_file_path = get_asset_path(name=name, version=version, extension=extension, assert_exists=False)
+        asset_file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Publish atomically on the same filesystem
+        Path(temp_path).replace(asset_file_path)
+        return asset_file_path
+    finally:
+        # Clean up if download fails or publish raises
+        if temp_path is not None:
+            Path(temp_path).unlink(missing_ok=True)
 
 
-def _build_ark_reference(name: str, files: list[Path]) -> ArkReference | None:
+def build_ark_reference_from_versions(name: str, version_refs: list[ArkVersionReference]) -> ArkReference | None:
     """
-    Create an ArkReference from a list of files for a given ark.
+    Build an ArkReference from version references.
 
-    :param name: The name of the ark.
-    :param files: List of file paths (ark JSON and asset USDZ files).
-    :return: ArkReference object or None if no valid versions found.
+    :param name: The Ark name.
+    :param version_refs: List of ArkVersionReference instances.
+    :return: ArkReference or None if no versions exist.
     """
 
-    file_stats = [f.stat() for f in files]
-    created_at = min(datetime.fromtimestamp(stat.st_ctime) for stat in file_stats).isoformat()
-    updated_at = max(datetime.fromtimestamp(stat.st_mtime) for stat in file_stats).isoformat()
-
-    # Group files by version - parse from {name}:{version}:ark.json or {name}:{version}:asset.usdz
-    files_by_version: dict[str, list[Path]] = defaultdict(list)
-    for file_path in files:
-        files_by_version[file_path.name.split(":")[1]].append(file_path)
-
-    # Create an ArkVersionReference for each version
-    version_refs = []
-    for version, version_files in files_by_version.items():
-        ark_file = None
-        asset_file = None
-        for file_path in version_files:
-            if file_path.name.endswith(":ark.json"):
-                ark_file = file_path
-            elif file_path.name.endswith(":asset.usdz"):
-                asset_file = file_path
-        if ark_file is None:
-            continue
-
-        ark_stat = ark_file.stat()
-        version_refs.append(
-            ArkVersionReference(
-                version=version,
-                full_path=str(ark_file),
-                asset_path=str(asset_file) if asset_file else None,
-                size_bytes=ark_stat.st_size,
-                created_at=datetime.fromtimestamp(ark_stat.st_ctime).isoformat(),
-                updated_at=datetime.fromtimestamp(ark_stat.st_mtime).isoformat(),
-                asset_size_bytes=asset_file.stat().st_size if asset_file else None,
-            )
-        )
-
     if not version_refs:
         return None
 
-    return ArkReference(
-        name=name,
-        versions=version_refs,
-        created_at=created_at,
-        updated_at=updated_at,
-    )
+    # Aggregate timestamps from versions, fallback to empty string if all are missing
+    created_at = min((v.created_at for v in version_refs if v.created_at), default="")
+    updated_at = max((v.updated_at for v in version_refs if v.updated_at), default="")
+    return ArkReference(name=name, versions=version_refs, created_at=created_at, updated_at=updated_at)
 
 
-def _build_asset_reference(name: str, files: list[Path]) -> AssetReference | None:
+def build_asset_reference_from_versions(name: str, version_refs: list[AssetVersionReference]) -> AssetReference | None:
     """
-    Create an AssetReference from a list of files for a given asset.
+    Build an AssetReference from version references.
 
-    :param name: The name of the asset.
-    :param files: List of file paths (.usdz files).
-    :return: AssetReference object or None if no valid versions found.
+    :param name: The Asset name.
+    :param version_refs: List of AssetVersionReference instances.
+    :return: AssetReference or None if no versions exist.
     """
 
-    file_stats = [f.stat() for f in files]
-    created_at = min(datetime.fromtimestamp(stat.st_ctime) for stat in file_stats).isoformat()
-    updated_at = max(datetime.fromtimestamp(stat.st_mtime) for stat in file_stats).isoformat()
-
-    # Group files by version
-    files_by_version: dict[str, list[Path]] = defaultdict(list)
-    for file_path in files:
-        # Parse filename format: {name}:{version}:file.usdz
-        version = file_path.name.split(":")[1]
-        files_by_version[version].append(file_path)
-
-    # Create an AssetVersionReference for each version
-    version_refs = []
-    for version, version_files in files_by_version.items():
-        # Should only be one file per version, but take first if multiple
-        asset_file = version_files[0]
-        asset_stat = asset_file.stat()
-        version_refs.append(
-            AssetVersionReference(
-                version=version,
-                full_path=str(asset_file),
-                size_bytes=asset_stat.st_size,
-                created_at=datetime.fromtimestamp(asset_stat.st_ctime).isoformat(),
-                updated_at=datetime.fromtimestamp(asset_stat.st_mtime).isoformat(),
-            )
-        )
-
     if not version_refs:
         return None
 
-    return AssetReference(
-        name=name,
-        versions=version_refs,
-        created_at=created_at,
-        updated_at=updated_at,
-    )
+    # Aggregate timestamps from versions, fallback to empty string if all are missing
+    created_at = min((v.created_at for v in version_refs if v.created_at), default="")
+    updated_at = max((v.updated_at for v in version_refs if v.updated_at), default="")
+    return AssetReference(name=name, versions=version_refs, created_at=created_at, updated_at=updated_at)