wafer-core 0.1.26__py3-none-any.whl → 0.1.28__py3-none-any.whl

wafer_core/targets/providers/baremetal.py

@@ -0,0 +1,72 @@
+"""Baremetal provider — degenerate case with no cloud API.
+
+Baremetal targets have no provisioning lifecycle. The "resource" is just the
+SSH endpoint from the spec. list_targets returns nothing (no API to query),
+and provision/terminate are errors.
+"""
+
+from __future__ import annotations
+
+from wafer_core.targets.types import Target, TargetSpec
+from wafer_core.utils.kernel_utils.targets.config import BaremetalTarget, VMTarget
+
+
+def target_from_ssh_spec(spec: BaremetalTarget | VMTarget) -> Target:
+    """Build a Target from a baremetal/VM spec's SSH info.
+
+    Since there's no cloud API, the resource_id is synthetic:
+    "baremetal:{host}:{port}" to make it unique and stable.
+    """
+    # Parse user@host:port
+    ssh_target = spec.ssh_target
+    assert ":" in ssh_target, f"ssh_target must include port, got: {ssh_target}"
+
+    user_host, port_str = ssh_target.rsplit(":", 1)
+    if "@" in user_host:
+        user, host = user_host.split("@", 1)
+    else:
+        user = "root"
+        host = user_host
+
+    port = int(port_str)
+
+    return Target(
+        resource_id=f"baremetal:{host}:{port}",
+        provider="baremetal",
+        status="running",  # Assumed running; TCP check happens elsewhere
+        public_ip=host,
+        ssh_port=port,
+        ssh_username=user,
+        gpu_type=spec.gpu_type,
+        name=spec.name,
+        spec_name=spec.name,
+    )
+
+
+class BaremetalProvider:
+    """Baremetal implementation of TargetProvider.
+
+    Baremetal has no cloud API. list_targets returns empty (no remote state
+    to query). Use target_from_ssh_spec() to build a Target from a spec
+    when you already know which spec you want.
+    """
+
+    async def list_targets(self) -> list[Target]:
+        """Baremetal has no API to list. Returns empty."""
+        return []
+
+    async def get_target(self, resource_id: str) -> Target | None:
+        """Baremetal has no API to query. Returns None."""
+        return None
+
+    async def provision(self, spec: TargetSpec) -> Target:
+        """Baremetal targets cannot be provisioned — they already exist."""
+        assert isinstance(spec, (BaremetalTarget, VMTarget)), (
+            f"BaremetalProvider.provision requires BaremetalTarget or VMTarget, "
+            f"got {type(spec).__name__}"
+        )
+        return target_from_ssh_spec(spec)
+
+    async def terminate(self, resource_id: str) -> bool:
+        """Baremetal targets cannot be terminated via API."""
+        return False

wafer_core/targets/providers/digitalocean.py

@@ -0,0 +1,164 @@
+"""DigitalOcean provider — adapts existing DO REST API to TargetProvider protocol."""
+
+from __future__ import annotations
+
+import logging
+import time
+from datetime import datetime, timezone
+
+from wafer_core.targets.digitalocean import (
+    DigitalOceanError,
+    _api_request_async,
+    _get_ssh_key_ids,
+    _wait_for_ssh,
+)
+from wafer_core.targets.digitalocean import (
+    terminate_droplet as _terminate_droplet,
+)
+from wafer_core.targets.types import Target, TargetSpec
+from wafer_core.utils.kernel_utils.targets.config import DigitalOceanTarget
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_droplet_to_target(droplet: dict) -> Target:
+    """Parse a DigitalOcean API droplet response into a Target."""
+    droplet_id = str(droplet.get("id", ""))
+    droplet_name = droplet.get("name", "")
+    status_raw = droplet.get("status", "").lower()
+
+    # Map DO statuses to our values
+    # DO: new, active, off, archive
+    status_map = {
+        "new": "pending",
+        "active": "running",
+        "off": "stopped",
+        "archive": "terminated",
+    }
+    status = status_map.get(status_raw, status_raw)
+
+    # Extract public IP
+    public_ip = None
+    networks = droplet.get("networks", {})
+    for net in networks.get("v4", []):
+        if net.get("type") == "public":
+            public_ip = net.get("ip_address")
+            break
+
+    # Infer spec_name from naming convention: wafer-{spec_name}-{timestamp}
+    spec_name = None
+    if droplet_name.startswith("wafer-"):
+        parts = droplet_name.split("-")
+        if len(parts) >= 3:
+            spec_name = "-".join(parts[1:-1])
+
+    created_at = droplet.get("created_at")
+
+    # Extract GPU type from size slug
+    size = droplet.get("size", {})
+    size_slug = (
+        size.get("slug", "") if isinstance(size, dict) else str(droplet.get("size_slug", ""))
+    )
+    gpu_type = "MI300X" if "mi300x" in size_slug.lower() else "unknown"
+
+    return Target(
+        resource_id=droplet_id,
+        provider="digitalocean",
+        status=status,
+        public_ip=public_ip,
+        ssh_port=22,
+        ssh_username="root",
+        gpu_type=gpu_type,
+        name=droplet_name or None,
+        created_at=created_at,
+        spec_name=spec_name,
+    )
+
+
+class DigitalOceanProvider:
+    """DigitalOcean implementation of TargetProvider.
+
+    Wraps existing REST API calls for droplet management.
+    """
+
+    async def list_targets(self) -> list[Target]:
+        """List all droplets on the DigitalOcean account."""
+        try:
+            response = await _api_request_async("GET", "/droplets", params={"per_page": "200"})
+        except DigitalOceanError:
+            raise
+        except Exception as e:
+            logger.warning(f"Failed to list DigitalOcean droplets: {e}")
+            return []
+
+        droplets = response.get("droplets", [])
+        return [_parse_droplet_to_target(d) for d in droplets]
+
+    async def get_target(self, resource_id: str) -> Target | None:
+        """Get a specific droplet by ID."""
+        try:
+            response = await _api_request_async("GET", f"/droplets/{resource_id}")
+        except Exception as e:
+            logger.warning(f"Failed to get DigitalOcean droplet {resource_id}: {e}")
+            return None
+
+        droplet = response.get("droplet")
+        if not droplet:
+            return None
+
+        return _parse_droplet_to_target(droplet)
+
+    async def provision(self, spec: TargetSpec) -> Target:
+        """Provision a new DigitalOcean droplet from a spec.
+
+        Blocks until SSH is ready.
+        """
+        assert isinstance(spec, DigitalOceanTarget), (
+            f"DigitalOceanProvider.provision requires DigitalOceanTarget, got {type(spec).__name__}"
+        )
+
+        droplet_name = f"wafer-{spec.name}-{int(time.time())}"
+
+        ssh_key_ids = await _get_ssh_key_ids()
+        if not ssh_key_ids:
+            logger.warning("No SSH keys found - droplet may not be accessible")
+
+        create_data = {
+            "name": droplet_name,
+            "region": spec.region,
+            "size": spec.size_slug,
+            "image": spec.image,
+            "ssh_keys": ssh_key_ids,
+            "backups": False,
+            "ipv6": True,
+            "monitoring": True,
+        }
+
+        logger.info(f"Provisioning DigitalOcean droplet: {droplet_name}")
+        response = await _api_request_async("POST", "/droplets", data=create_data)
+
+        if not response or "droplet" not in response:
+            raise DigitalOceanError(f"Failed to create droplet: {response}")
+
+        droplet = response["droplet"]
+        droplet_id = str(droplet["id"])
+        logger.info(f"Droplet created: {droplet_id}")
+
+        public_ip = await _wait_for_ssh(droplet_id, spec.provision_timeout)
+
+        return Target(
+            resource_id=droplet_id,
+            provider="digitalocean",
+            status="running",
+            public_ip=public_ip,
+            ssh_port=22,
+            ssh_username="root",
+            gpu_type=spec.gpu_type,
+            name=droplet_name,
+            created_at=datetime.now(timezone.utc).isoformat(),
+            spec_name=spec.name,
+        )
+
+    async def terminate(self, resource_id: str) -> bool:
+        """Terminate a DigitalOcean droplet."""
+        return await _terminate_droplet(resource_id)

wafer_core/targets/providers/runpod.py

@@ -0,0 +1,250 @@
+"""RunPod provider — adapts existing RunPod GraphQL API to TargetProvider protocol."""
+
+from __future__ import annotations
+
+import logging
+import time
+from datetime import datetime, timezone
+
+from wafer_core.targets.runpod import (
+    RunPodError,
+    _graphql_request_async,
+    _wait_for_ssh,
+)
+from wafer_core.targets.runpod import (
+    terminate_pod as _terminate_pod,
+)
+from wafer_core.targets.types import Target, TargetSpec
+from wafer_core.utils.kernel_utils.targets.config import RunPodTarget
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_pod_to_target(pod: dict) -> Target | None:
+    """Parse a RunPod API pod response into a Target.
+
+    Returns None if the pod has no usable SSH info.
+    """
+    pod_id = pod.get("id", "")
+    pod_name = pod.get("name", "")
+    status_raw = pod.get("desiredStatus", "").lower()
+
+    # Map RunPod statuses to our status values
+    status = status_raw if status_raw else "unknown"
+
+    # Extract SSH info from runtime ports
+    public_ip = None
+    ssh_port = None
+    runtime = pod.get("runtime")
+    if runtime:
+        for port in runtime.get("ports") or []:
+            if port.get("privatePort") == 22 and port.get("isIpPublic"):
+                ip = port.get("ip")
+                # Skip proxy SSH (ssh.runpod.io), want direct IP
+                if ip and ip != "ssh.runpod.io":
+                    public_ip = ip
+                    ssh_port = port.get("publicPort")
+                    break
+
+    # Infer spec_name from pod naming convention: wafer-{spec_name}-{timestamp}
+    spec_name = None
+    if pod_name.startswith("wafer-"):
+        parts = pod_name.split("-")
+        if len(parts) >= 3:
+            spec_name = "-".join(parts[1:-1])
+
+    # Extract GPU type
+    gpu_type = ""
+    machine = pod.get("machine")
+    if machine:
+        gpu_type_info = machine.get("gpuType")
+        if gpu_type_info:
+            gpu_type = gpu_type_info.get("displayName", "")
+
+    cost = pod.get("costPerHr")
+
+    return Target(
+        resource_id=pod_id,
+        provider="runpod",
+        status=status,
+        public_ip=public_ip,
+        ssh_port=ssh_port,
+        ssh_username="root",
+        gpu_type=gpu_type,
+        name=pod_name or None,
+        created_at=None,  # RunPod API doesn't expose creation time in list query
+        spec_name=spec_name,
+        price_per_hour=float(cost) if cost else None,
+    )
+
+
+class RunPodProvider:
+    """RunPod implementation of TargetProvider.
+
+    Wraps existing GraphQL API calls:
+    - list_targets: myself { pods { ... } }
+    - get_target: pod(input: { podId }) { ... }
+    - provision: podFindAndDeployOnDemand
+    - terminate: podTerminate
+    """
+
+    async def list_targets(self) -> list[Target]:
+        """List all running pods on the RunPod account."""
+        query = """
+        query {
+            myself {
+                pods {
+                    id
+                    name
+                    desiredStatus
+                    costPerHr
+                    machine {
+                        podHostId
+                        gpuType {
+                            displayName
+                        }
+                    }
+                    runtime {
+                        ports {
+                            ip
+                            isIpPublic
+                            privatePort
+                            publicPort
+                            type
+                        }
+                    }
+                }
+            }
+        }
+        """
+
+        try:
+            data = await _graphql_request_async(query)
+        except RunPodError:
+            raise
+        except Exception as e:
+            logger.warning(f"Failed to list RunPod pods: {e}")
+            return []
+
+        pods = data.get("myself", {}).get("pods", [])
+        targets = []
+
+        for pod in pods:
+            target = _parse_pod_to_target(pod)
+            if target is not None:
+                targets.append(target)
+
+        return targets
+
+    async def get_target(self, resource_id: str) -> Target | None:
+        """Get a specific pod by ID."""
+        query = """
+        query pod($input: PodFilter!) {
+            pod(input: $input) {
+                id
+                name
+                desiredStatus
+                costPerHr
+                machine {
+                    podHostId
+                    gpuType {
+                        displayName
+                    }
+                }
+                runtime {
+                    ports {
+                        ip
+                        isIpPublic
+                        privatePort
+                        publicPort
+                        type
+                    }
+                }
+            }
+        }
+        """
+        variables = {"input": {"podId": resource_id}}
+
+        try:
+            data = await _graphql_request_async(query, variables)
+        except Exception as e:
+            logger.warning(f"Failed to get RunPod pod {resource_id}: {e}")
+            return None
+
+        pod = data.get("pod")
+        if not pod:
+            return None
+
+        return _parse_pod_to_target(pod)
+
+    async def provision(self, spec: TargetSpec) -> Target:
+        """Provision a new RunPod pod from a spec.
+
+        Blocks until SSH is ready.
+        """
+        assert isinstance(spec, RunPodTarget), (
+            f"RunPodProvider.provision requires RunPodTarget, got {type(spec).__name__}"
+        )
+
+        pod_name = f"wafer-{spec.name}-{int(time.time())}"
+
+        mutation = """
+        mutation podFindAndDeployOnDemand($input: PodFindAndDeployOnDemandInput!) {
+            podFindAndDeployOnDemand(input: $input) {
+                id
+                machineId
+                machine {
+                    podHostId
+                }
+            }
+        }
+        """
+
+        pod_input: dict = {
+            "gpuTypeId": spec.gpu_type_id,
+            "gpuCount": spec.gpu_count,
+            "cloudType": "SECURE",
+            "name": pod_name,
+            "supportPublicIp": True,
+            "containerDiskInGb": spec.container_disk_gb,
+            "minVcpuCount": 1,
+            "minMemoryInGb": 4,
+            "ports": "22/tcp",
+            "startSsh": True,
+            "startJupyter": False,
+            "env": [],
+        }
+
+        if spec.template_id:
+            pod_input["templateId"] = spec.template_id
+        else:
+            pod_input["imageName"] = spec.image
+
+        logger.info(f"Provisioning RunPod pod: {pod_name}")
+        data = await _graphql_request_async(mutation, {"input": pod_input})
+
+        pod_data = data.get("podFindAndDeployOnDemand")
+        if not pod_data:
+            raise RunPodError("No pod returned from deployment")
+
+        pod_id = pod_data["id"]
+        logger.info(f"Pod created: {pod_id}")
+
+        public_ip, ssh_port, ssh_username = await _wait_for_ssh(pod_id, spec.provision_timeout)
+
+        return Target(
+            resource_id=pod_id,
+            provider="runpod",
+            status="running",
+            public_ip=public_ip,
+            ssh_port=ssh_port,
+            ssh_username=ssh_username,
+            gpu_type=spec.gpu_type,
+            name=pod_name,
+            created_at=datetime.now(timezone.utc).isoformat(),
+            spec_name=spec.name,
+        )
+
+    async def terminate(self, resource_id: str) -> bool:
+        """Terminate a RunPod pod."""
+        return await _terminate_pod(resource_id)

wafer_core/targets/reconcile.py

@@ -0,0 +1,90 @@
+"""Reconciliation: compare TargetSpecs to live Targets.
+
+Pure function — no API calls, no side effects. Takes specs and targets as
+inputs, returns a ReconcileResult describing what's bound, what's orphaned,
+and what's unprovisioned.
+"""
+
+from __future__ import annotations
+
+from wafer_core.targets.types import ReconcileResult, Target, TargetSpec
+from wafer_core.utils.kernel_utils.targets.config import (
+    BaremetalTarget,
+    DigitalOceanTarget,
+    RunPodTarget,
+    VMTarget,
+)
+
+
+def _is_cloud_spec(spec: TargetSpec) -> bool:
+    """Check if a spec represents a cloud-provisioned resource.
+
+    Baremetal and VM specs don't have cloud-managed lifecycles,
+    so they're excluded from "unprovisioned" checks.
+    """
+    return isinstance(spec, (RunPodTarget, DigitalOceanTarget))
+
+
+def _spec_provider(spec: TargetSpec) -> str | None:
+    """Get the provider name for a spec, or None if not cloud-managed."""
+    if isinstance(spec, RunPodTarget):
+        return "runpod"
+    if isinstance(spec, DigitalOceanTarget):
+        return "digitalocean"
+    if isinstance(spec, (BaremetalTarget, VMTarget)):
+        return "baremetal"
+    return None
+
+
+def reconcile(
+    specs: list[TargetSpec],
+    targets: list[Target],
+    binding_hints: dict[str, str] | None = None,
+) -> ReconcileResult:
+    """Compare specs to live targets and classify each.
+
+    Matching rules (in priority order):
+    1. Target.spec_name matches Spec.name exactly (set by naming convention
+       or explicit binding).
+    2. binding_hints maps resource_id → spec_name (from local cache).
+    3. No match → target is unbound (orphan).
+
+    A cloud spec with no matching target is "unprovisioned".
+    Baremetal/VM specs are never "unprovisioned" (they don't have a cloud
+    lifecycle — the machine is always there or it isn't).
+
+    Args:
+        specs: All known TargetSpecs (loaded from TOML files).
+        targets: All live Targets (fetched from provider APIs).
+        binding_hints: Optional resource_id → spec_name cache for targets
+            whose spec_name can't be inferred from naming convention.
+
+    Returns:
+        ReconcileResult with bound, unbound, and unprovisioned lists.
+    """
+    hints = binding_hints or {}
+    spec_by_name = {s.name: s for s in specs}
+    claimed_spec_names: set[str] = set()
+
+    bound: list[tuple[TargetSpec, Target]] = []
+    unbound: list[Target] = []
+
+    for target in targets:
+        # Try to find the spec this target belongs to
+        resolved_spec_name = target.spec_name or hints.get(target.resource_id)
+
+        if resolved_spec_name and resolved_spec_name in spec_by_name:
+            spec = spec_by_name[resolved_spec_name]
+            bound.append((spec, target))
+            claimed_spec_names.add(resolved_spec_name)
+        else:
+            unbound.append(target)
+
+    # Cloud specs with no bound target are unprovisioned
+    unprovisioned = [s for s in specs if s.name not in claimed_spec_names and _is_cloud_spec(s)]
+
+    return ReconcileResult(
+        bound=bound,
+        unbound=unbound,
+        unprovisioned=unprovisioned,
+    )