flowmesh-sdk-stack 0.1.0__py3-none-any.whl

flowmesh_stack/images.py

@@ -0,0 +1,112 @@
+"""FlowMesh Docker image reference management.
+
+Provides the canonical mapping from build targets to image references,
+used by both the CLI dev commands and programmatic build/deploy scripts.
+"""
+
+BUILD_TARGETS: dict[str, str] = {
+    "flowmesh_server": "{registry}/flowmesh_server:{version}",
+    "flowmesh_worker_cpu": "{registry}/flowmesh_worker:{version}-cpu",
+    "flowmesh_worker_gpu_builder": "{registry}/flowmesh_worker_builder:{version}-gpu",
+    "flowmesh_worker_gpu": "{registry}/flowmesh_worker:{version}-gpu",
+    "flowmesh_ssh_cpu": "{registry}/flowmesh_ssh:{version}-cpu",
+    "flowmesh_ssh_gpu": "{registry}/flowmesh_ssh:{version}-gpu",
+}
+"""Mapping from build target name to image reference format string."""
+
+CACHE_TARGETS: dict[str, str] = {
+    "flowmesh_server": "{registry}/flowmesh_server:{scope}",
+    "flowmesh_worker_cpu": "{registry}/flowmesh_worker:{scope}-cpu",
+    "flowmesh_worker_gpu_builder": "{registry}/flowmesh_worker_builder:{scope}-gpu",
+    "flowmesh_worker_gpu": "{registry}/flowmesh_worker:{scope}-gpu",
+    "flowmesh_ssh_cpu": "{registry}/flowmesh_ssh:{scope}-cpu",
+    "flowmesh_ssh_gpu": "{registry}/flowmesh_ssh:{scope}-gpu",
+}
+"""Mapping from build target name to registry cache reference format string."""
+
+BUILD_GROUPS: dict[str, list[str]] = {
+    "server": ["flowmesh_server"],
+    "workers": [
+        "flowmesh_worker_cpu",
+        "flowmesh_worker_gpu",
+        "flowmesh_ssh_cpu",
+        "flowmesh_ssh_gpu",
+    ],
+    "builders": ["flowmesh_worker_gpu_builder"],
+}
+"""Mapping from group name to list of build targets."""
+BUILD_GROUPS["default"] = [
+    target for group in ("server", "workers") for target in BUILD_GROUPS[group]
+]
+
+BUILD_DEPENDENCIES: dict[str, list[str]] = {
+    "flowmesh_worker_gpu": ["flowmesh_worker_gpu_builder"],
+}
+"""Auxiliary build targets that must be configured alongside a selected target."""
+
+PUSH_PLATFORMS: dict[str, str] = {
+    target: "linux/amd64,linux/arm64" for target in BUILD_TARGETS
+}
+"""Default platform matrix to publish for each build target."""
+
+
+def get_image_ref(registry: str, version: str, target: str) -> str:
+    """Resolve a Docker image reference for a build target.
+
+    Args:
+        registry: Container registry (e.g. ``ghcr.io/mlsys-io``).
+        version: Image version tag (e.g. ``dev``, ``0.1.0``).
+        target: Build target name (must be a key in :data:`BUILD_TARGETS`).
+
+    Raises:
+        ValueError: If the target is unknown.
+    """
+    if target not in BUILD_TARGETS:
+        raise ValueError(f"Unknown build target: {target}")
+    return BUILD_TARGETS[target].format(registry=registry, version=version)
+
+
+def get_cache_ref(registry: str, scope: str, target: str) -> str:
+    """Resolve a Docker registry cache reference for a build target."""
+
+    if target not in CACHE_TARGETS:
+        raise ValueError(f"Unknown build target: {target}")
+    normalized_scope = (scope or "").strip()
+    if not normalized_scope or normalized_scope == "cache":
+        cache_scope = "cache"
+    elif normalized_scope.startswith("cache-"):
+        cache_scope = normalized_scope
+    else:
+        cache_scope = f"cache-{normalized_scope}"
+    return CACHE_TARGETS[target].format(registry=registry, scope=cache_scope)
+
+
+def expand_build_targets(targets: list[str]) -> list[str]:
+    """Expand explicit targets with any dependent helper targets.
+
+    The returned list preserves first-seen order and inserts helper targets ahead
+    of any target that requires them.
+    """
+
+    expanded: list[str] = []
+    seen: set[str] = set()
+
+    def _visit(target: str) -> None:
+        if target in seen:
+            return
+        for dep in BUILD_DEPENDENCIES.get(target, []):
+            _visit(dep)
+        seen.add(target)
+        expanded.append(target)
+
+    for target in targets:
+        _visit(target)
+    return expanded
+
+
+def get_push_platforms(target: str) -> str:
+    """Resolve the default push platform matrix for a build target."""
+
+    if target not in PUSH_PLATFORMS:
+        raise ValueError(f"Unknown build target: {target}")
+    return PUSH_PLATFORMS[target]

flowmesh_stack/node_client.py

@@ -0,0 +1,173 @@
+"""Node direct API client."""
+
+import os
+from typing import Any
+
+import httpx
+from flowmesh.exceptions import FlowMeshConnectionError, FlowMeshError
+
+_DEFAULT_TIMEOUT = 30.0
+_WORKER_CREATE_TIMEOUT = 600.0
+
+
+def _raise_for_status(response: httpx.Response, method: str) -> None:
+    if response.status_code < 400:
+        return
+    try:
+        body = response.json()
+    except Exception:
+        body = response.text
+    message = ""
+    if isinstance(body, dict):
+        message = body.get("detail", "") or body.get("message", "")
+    if not message:
+        message = str(body)
+    raise FlowMeshError(
+        f"{method} {response.url} returned {response.status_code}: {message}"
+    )
+
+
+class NodeClient:
+    """Client for a node's worker management HTTP API.
+
+    Used by the CLI stack sub-package for worker lifecycle management.
+
+    Args:
+        base_url: Node HTTP endpoint (e.g. ``http://localhost:8000``).
+        token: Optional bearer token for authentication.
+        timeout: Request timeout in seconds.
+    """
+
+    def __init__(
+        self,
+        base_url: str | None = None,
+        token: str | None = None,
+        timeout: float = _DEFAULT_TIMEOUT,
+    ) -> None:
+        resolved_url = base_url or _default_node_url()
+        resolved_token = token or os.getenv("FLOWMESH_API_KEY") or None
+        headers: dict[str, str] = {"Accept": "application/json"}
+        if resolved_token:
+            headers["Authorization"] = f"Bearer {resolved_token}"
+        self._base_url = resolved_url.rstrip("/")
+        self._http = httpx.Client(
+            base_url=self._base_url,
+            headers=headers,
+            timeout=httpx.Timeout(timeout),
+        )
+
+    # -- Workers --------------------------------------------------------- #
+
+    def list_workers(self) -> list[dict[str, Any]]:
+        """List all workers managed by this node."""
+        return self._request("GET", "/api/v1/stack/workers")
+
+    def create_worker(self, config: str | dict[str, Any]) -> dict[str, Any]:
+        """Create a worker from a JSON/YAML config.
+
+        Args:
+            config: Worker init config as a JSON string or dict.
+        """
+        if isinstance(config, dict):
+            return self._request(
+                "POST",
+                "/api/v1/stack/workers",
+                json_body=config,
+                timeout=_WORKER_CREATE_TIMEOUT,
+            )
+        return self._request(
+            "POST",
+            "/api/v1/stack/workers",
+            data=config,
+            headers={"Content-Type": "application/json"},
+            timeout=_WORKER_CREATE_TIMEOUT,
+        )
+
+    def start_worker(self, name: str) -> None:
+        """Start a stopped worker."""
+        self._request("POST", f"/api/v1/stack/workers/{name}/start")
+
+    def stop_worker(self, name: str) -> None:
+        """Stop a running worker."""
+        self._request("POST", f"/api/v1/stack/workers/{name}/stop")
+
+    def destroy_worker(self, name: str) -> None:
+        """Destroy a single worker, removing its container."""
+        self._request("DELETE", f"/api/v1/stack/workers/{name}")
+
+    def destroy_all_workers(self, *, ignore_unreachable: bool = False) -> bool:
+        """Destroy all workers managed by this node.
+
+        Returns ``True`` on success, ``False`` when ``ignore_unreachable=True``
+        and the FlowMesh server was unreachable. Other errors propagate.
+        """
+        try:
+            self._request(
+                "DELETE",
+                "/api/v1/stack/workers",
+                headers={"Content-Type": "application/json"},
+            )
+        except FlowMeshConnectionError:
+            if not ignore_unreachable:
+                raise
+            return False
+        return True
+
+    def worker_names(self) -> list[str]:
+        """Return a list of all worker names."""
+        data = self.list_workers()
+        names: list[str] = []
+        for item in data:
+            if isinstance(item, str):
+                names.append(item)
+            elif isinstance(item, dict):
+                name = item.get("name")
+                if isinstance(name, str) and name:
+                    names.append(name)
+        return names
+
+    # -- Transport ------------------------------------------------------- #
+
+    def _request(
+        self,
+        method: str,
+        path: str,
+        json_body: Any = None,
+        data: str | bytes | None = None,
+        headers: dict[str, str] | None = None,
+        timeout: float | None = None,
+    ) -> Any:
+        url = path
+        kwargs: dict[str, Any] = {}
+        if json_body is not None:
+            kwargs["json"] = json_body
+        if data is not None:
+            kwargs["content"] = data
+        if headers:
+            kwargs["headers"] = headers
+        if timeout is not None:
+            kwargs["timeout"] = timeout
+        try:
+            response = self._http.request(method, url, **kwargs)
+        except httpx.ConnectError as exc:
+            raise FlowMeshConnectionError(
+                f"Failed to connect to {self._base_url}{path}: {exc}"
+            )
+        _raise_for_status(response, method)
+        if not response.content:
+            return None
+        return response.json()
+
+    def close(self) -> None:
+        self._http.close()
+
+    def __enter__(self) -> "NodeClient":
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.close()
+
+
+def _default_node_url() -> str:
+    port = os.getenv("SERVER_HTTP_PORT", os.getenv("SERVER_APP_PORT", "8000"))
+    return f"http://localhost:{port}"

flowmesh_stack/paths.py

@@ -0,0 +1,21 @@
+"""Filesystem path helpers shared by FlowMesh tooling."""
+
+from pathlib import Path
+
+
+def resolve_path(value: str, default: str, base_dir: Path) -> Path:
+    """Resolve a possibly relative path against a base directory."""
+    raw = value.strip() or default
+    path = Path(raw).expanduser()
+    return path if path.is_absolute() else (base_dir / path).resolve()
+
+
+def ensure_dir(path: Path) -> None:
+    """Create a directory if it does not exist."""
+    path.mkdir(parents=True, exist_ok=True)
+
+
+def ensure_file(path: Path) -> None:
+    """Create a file and its parent directory if needed."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.touch(exist_ok=True)

flowmesh_stack/workers.py

@@ -0,0 +1,278 @@
+"""Worker lifecycle helpers for direct node operations."""
+
+import json
+import subprocess
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Any
+
+import yaml
+from flowmesh.exceptions import FlowMeshError
+
+from .docker import ensure_docker_available
+from .node_client import NodeClient
+
+_MAX_PARALLEL_REQUESTS = 16
+
+
+def operate_workers(
+    client: NodeClient,
+    names: list[str],
+    operation: str,
+) -> list[str]:
+    """Apply a start/stop/destroy operation to one or more workers."""
+    if not names:
+        return []
+    if "all" in names:
+        if len(names) != 1:
+            raise FlowMeshError("Use either 'all' or worker names, not both.")
+        names = client.worker_names()
+        if not names:
+            return []
+
+    def _apply(name: str) -> str:
+        match operation:
+            case "start":
+                client.start_worker(name)
+            case "stop":
+                client.stop_worker(name)
+            case "destroy":
+                client.destroy_worker(name)
+            case _:
+                raise FlowMeshError(f"Unsupported worker operation: {operation}")
+        return name
+
+    successes: list[str] = []
+    errors: list[str] = []
+    max_workers = min(len(names), _MAX_PARALLEL_REQUESTS)
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = {executor.submit(_apply, name): name for name in names}
+        for future in as_completed(futures):
+            name = futures[future]
+            try:
+                successes.append(future.result())
+            except Exception as exc:
+                errors.append(f"{name}: {exc}")
+    if errors:
+        raise FlowMeshError("; ".join(errors))
+    return successes
+
+
+def create_workers(
+    client: NodeClient,
+    kind: str = "cpu",
+    count: int = 1,
+    targets: str = "all",
+    config_paths: list[Path] | None = None,
+    config_raw: list[str] | None = None,
+) -> list[tuple[str, dict[str, Any]]]:
+    """Create node workers from configs or built-in cpu/gpu presets.
+
+    When ``kind`` is "gpu", if ``count`` is equal to 1, a single worker with the
+    specified GPU targets will be created. If ``count`` is greater than 1, one worker
+    will be created per GPU target, and the number of targets must match ``count``.
+    """
+    payloads = _payloads_for_worker_create(
+        kind=kind,
+        count=count,
+        targets=targets,
+        config_paths=config_paths,
+        config_raw=config_raw,
+    )
+    if not payloads:
+        return []
+
+    created: list[tuple[str, dict[str, Any]]] = []
+    errors: list[str] = []
+    max_workers = min(len(payloads), _MAX_PARALLEL_REQUESTS)
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = {
+            executor.submit(client.create_worker, payload): label
+            for payload, label in payloads
+        }
+        for future in as_completed(futures):
+            label = futures[future]
+            try:
+                created.append((label, future.result()))
+            except Exception as exc:
+                errors.append(f"{label}: {exc}")
+    if errors:
+        raise FlowMeshError("; ".join(errors))
+    return created
+
+
+def select_worker_images(
+    kinds: list[str],
+    images: dict[str, str],
+    builder_images: dict[str, str] | None = None,
+    builder: bool = False,
+) -> list[str]:
+    """Resolve requested worker image kinds into concrete image refs."""
+    normalized = [kind.strip().lower() for kind in kinds if kind.strip()]
+    if not normalized:
+        raise FlowMeshError("worker pull expects cpu|gpu|ssh-cpu|ssh-gpu|all")
+
+    selected_images = builder_images if builder else images
+    valid_source = builder_images if builder and builder_images is not None else images
+    valid = set([*valid_source, "all"])
+    invalid = [kind for kind in normalized if kind not in valid]
+    if invalid:
+        unique = ", ".join(sorted(set(invalid)))
+        raise FlowMeshError(f"Invalid kind(s): {unique}")
+
+    if "all" in normalized:
+        return list((selected_images or {}).values())
+
+    requested = set(normalized)
+    if builder and builder_images is not None:
+        invalid_builder = [kind for kind in requested if kind not in builder_images]
+        if invalid_builder:
+            unique = ", ".join(sorted(invalid_builder))
+            raise FlowMeshError(f"No builder image for: {unique}")
+        return [builder_images[kind] for kind in builder_images if kind in requested]
+
+    return [images[kind] for kind in images if kind in requested]
+
+
+def pull_images(images: list[str]) -> None:
+    """Pull one or more worker images via docker."""
+    ensure_docker_available()
+    for image in images:
+        result = subprocess.run(["docker", "pull", image], text=True, check=False)
+        if result.returncode != 0:
+            raise FlowMeshError(f"Failed to pull image: {image}")
+
+
+def detect_gpu_targets(targets: str) -> list[str]:
+    """Resolve requested GPU ids into a concrete target list."""
+    if targets != "all":
+        return [item.strip() for item in targets.split(",") if item.strip()]
+
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+    except FileNotFoundError:
+        return []
+    if result.returncode != 0:
+        return []
+    return [item.strip() for item in result.stdout.splitlines() if item.strip()]
+
+
+def _payloads_for_worker_create(
+    kind: str,
+    count: int,
+    targets: str,
+    config_paths: list[Path] | None,
+    config_raw: list[str] | None,
+) -> list[tuple[str, str]]:
+    if count < 1:
+        raise FlowMeshError("Worker count must be at least 1.")
+    if config_paths is not None or config_raw is not None:
+        payloads: list[tuple[str, str]] = []
+        for config_path in config_paths or []:
+            if not config_path.exists():
+                raise FlowMeshError(f"Config not found: {config_path}")
+            _extend_payloads(
+                payloads, f"worker from {config_path.name}", config_path.read_text()
+            )
+        for idx, raw in enumerate(config_raw or []):
+            _extend_payloads(payloads, f"worker from raw#{idx}", raw)
+        return payloads
+
+    if kind == "cpu":
+        return [
+            (
+                json.dumps(
+                    {
+                        "provider": "docker",
+                        "init_on_start": True,
+                        "worker_config": {
+                            "worker_type": "cpu",
+                            "worker_alias": f"worker_cpu_{idx}",
+                        },
+                    }
+                ),
+                "CPU worker",
+            )
+            for idx in range(count)
+        ]
+
+    if kind == "gpu":
+        raw_gpu_ids = detect_gpu_targets(targets)
+        gpu_ids: list[int] = []
+        for raw_gpu_id in raw_gpu_ids:
+            if not raw_gpu_id.isdigit():
+                raise FlowMeshError(f"Invalid GPU id '{raw_gpu_id}'")
+            gpu_ids.append(int(raw_gpu_id))
+
+        if count > 1:
+            if count != len(gpu_ids):
+                raise FlowMeshError(
+                    f"GPU worker count {count} does not match "
+                    f"detected GPU targets: {gpu_ids}. "
+                    f"Consider setting count={len(gpu_ids)} or specifying exactly "
+                    f"{count} GPU targets."
+                )
+            gpu_payloads = [
+                (
+                    json.dumps(
+                        {
+                            "provider": "docker",
+                            "init_on_start": True,
+                            "worker_config": {
+                                "worker_type": "gpu",
+                                "cuda_devices": [gpu_id],
+                                "worker_alias": f"worker_gpu_{gpu_id}",
+                            },
+                        }
+                    ),
+                    f"GPU worker for GPU {gpu_id}",
+                )
+                for gpu_id in gpu_ids
+            ]
+        else:
+            worker_suffix = "all" if targets == "all" else "_".join(raw_gpu_ids)
+            gpu_payloads = [
+                (
+                    json.dumps(
+                        {
+                            "provider": "docker",
+                            "init_on_start": True,
+                            "worker_config": {
+                                "worker_type": "gpu",
+                                "cuda_devices": gpu_ids,
+                                "worker_alias": f"worker_gpu_{worker_suffix}",
+                            },
+                        }
+                    ),
+                    f"GPU worker for GPUs {', '.join(raw_gpu_ids)}",
+                )
+            ]
+
+        if not gpu_payloads:
+            raise FlowMeshError("No GPUs detected or specified.")
+        return gpu_payloads
+
+    raise FlowMeshError("worker up expects kind cpu|gpu or use --config")
+
+
+def _extend_payloads(
+    payloads: list[tuple[str, str]], label_prefix: str, payload_text: str
+) -> None:
+    try:
+        payload_obj = yaml.safe_load(payload_text)
+    except Exception:
+        payload_obj = None
+
+    if isinstance(payload_obj, list):
+        for idx, item in enumerate(payload_obj):
+            payloads.append((json.dumps(item), f"{label_prefix}#{idx}"))
+        return
+    if isinstance(payload_obj, dict):
+        payloads.append((json.dumps(payload_obj), label_prefix))
+        return
+    payloads.append((payload_text, label_prefix))