hypercli-sdk 0.4.2__py3-none-any.whl

c3/jobs.py

@@ -0,0 +1,285 @@
+"""Jobs API"""
+import base64
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Iterator
+
+if TYPE_CHECKING:
+    from .http import HTTPClient
+
+
+@dataclass
+class Job:
+    job_id: str
+    job_key: str
+    state: str
+    gpu_type: str
+    gpu_count: int
+    region: str
+    interruptible: bool
+    price_per_hour: float
+    price_per_second: float
+    docker_image: str
+    runtime: int
+    hostname: str | None = None
+    created_at: float | None = None
+    started_at: float | None = None
+    completed_at: float | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Job":
+        return cls(
+            job_id=data.get("job_id", ""),
+            job_key=data.get("job_key", ""),
+            state=data.get("state", ""),
+            gpu_type=data.get("gpu_type", ""),
+            gpu_count=data.get("gpu_count", 1),
+            region=data.get("region", ""),
+            interruptible=data.get("interruptible", True),
+            price_per_hour=data.get("price_per_hour", 0),
+            price_per_second=data.get("price_per_second", 0),
+            docker_image=data.get("docker_image", ""),
+            runtime=data.get("runtime", 0),
+            hostname=data.get("hostname"),
+            created_at=data.get("created_at"),
+            started_at=data.get("started_at"),
+            completed_at=data.get("completed_at"),
+        )
+
+
+@dataclass
+class GPUMetrics:
+    index: int
+    name: str
+    utilization: float
+    memory_used: float
+    memory_total: float
+    temperature: int
+    power_draw: float
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "GPUMetrics":
+        return cls(
+            index=data.get("index", 0),
+            name=data.get("name", ""),
+            utilization=data.get("utilization_gpu_percent", 0),
+            memory_used=data.get("memory_used_mb", 0),
+            memory_total=data.get("memory_total_mb", 0),
+            temperature=data.get("temperature_c", 0),
+            power_draw=data.get("power_draw_w", 0),
+        )
+
+
+@dataclass
+class SystemMetrics:
+    cpu_percent: float
+    cpu_cores: float
+    cpu_unix_percent: float
+    memory_used: float
+    memory_limit: float
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "SystemMetrics":
+        return cls(
+            cpu_percent=data.get("cpu_percent", 0),
+            cpu_cores=data.get("cpu_cores", 1),
+            cpu_unix_percent=data.get("cpu_unix_percent", data.get("cpu_percent", 0)),
+            memory_used=data.get("memory_used_mb", 0),
+            memory_limit=data.get("memory_limit_mb", 0),
+        )
+
+
+@dataclass
+class JobMetrics:
+    gpus: list[GPUMetrics] = field(default_factory=list)
+    system: SystemMetrics | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "JobMetrics":
+        system_data = data.get("system")
+        return cls(
+            gpus=[GPUMetrics.from_dict(g) for g in data.get("gpus", [])],
+            system=SystemMetrics.from_dict(system_data) if system_data else None,
+        )
+
+
+class Jobs:
+    """Jobs API wrapper"""
+
+    def __init__(self, http: "HTTPClient"):
+        self._http = http
+
+    def list(self, state: str = None) -> list[Job]:
+        """List all jobs"""
+        params = {"state": state} if state else None
+        data = self._http.get("/api/jobs", params=params)
+        # API returns {"jobs": [...], "total_count": ...}
+        jobs = data.get("jobs", []) if isinstance(data, dict) else data
+        return [Job.from_dict(j) for j in jobs]
+
+    def get(self, job_id: str) -> Job:
+        """Get job details"""
+        data = self._http.get(f"/api/jobs/{job_id}")
+        return Job.from_dict(data)
+
+    def create(
+        self,
+        image: str,
+        command: str = None,
+        gpu_type: str = "l40s",
+        gpu_count: int = 1,
+        region: str = None,
+        runtime: int = None,
+        interruptible: bool = True,
+        env: dict[str, str] = None,
+        ports: dict[str, int] = None,
+        auth: bool = False,
+    ) -> Job:
+        """Create a new job.
+
+        Args:
+            image: Docker image to run
+            command: Command to execute (base64 encoded internally)
+            gpu_type: GPU type (e.g., "l40s", "a100")
+            gpu_count: Number of GPUs
+            region: Region to run in
+            runtime: Max runtime in seconds
+            interruptible: Allow spot/preemptible instances
+            env: Environment variables
+            ports: Ports to expose. Use {"lb": port} for HTTPS load balancer
+            auth: Enable Bearer token auth on load balancer (use with ports={"lb": port})
+        """
+        payload = {
+            "docker_image": image,
+            "gpu_type": gpu_type,
+            "gpu_count": gpu_count,
+            "interruptible": interruptible,
+            "command": base64.b64encode((command or "").encode()).decode(),
+        }
+        if region:
+            payload["region"] = region
+        if runtime:
+            payload["runtime"] = runtime
+        if env:
+            payload["env_vars"] = env
+        if ports:
+            payload["ports"] = ports
+        if auth:
+            payload["auth"] = auth
+
+        data = self._http.post("/api/jobs", json=payload)
+        return Job.from_dict(data)
+
+    def cancel(self, job_id: str) -> dict:
+        """Cancel a job"""
+        return self._http.delete(f"/api/jobs/{job_id}")
+
+    def extend(self, job_id: str, runtime: int) -> Job:
+        """Extend job runtime"""
+        data = self._http.patch(f"/api/jobs/{job_id}", json={"runtime": runtime})
+        return Job.from_dict(data)
+
+    def logs(self, job_id: str) -> str:
+        """Get job logs"""
+        data = self._http.get(f"/api/jobs/{job_id}/logs")
+        return data.get("logs", "")
+
+    def metrics(self, job_id: str) -> JobMetrics:
+        """Get job GPU metrics"""
+        data = self._http.get(f"/api/jobs/{job_id}/metrics")
+        return JobMetrics.from_dict(data)
+
+    def token(self, job_id: str) -> str:
+        """Get job auth token"""
+        data = self._http.get(f"/api/jobs/{job_id}/token")
+        return data.get("token", "")
+
+
+# Utility functions for finding jobs
+
+
+def is_uuid(s: str) -> bool:
+    """Check if string looks like a UUID (job ID)"""
+    return "-" in s and len(s) > 30
+
+
+def find_by_id(jobs: Jobs, job_id: str) -> Job | None:
+    """Find job by UUID via direct API call.
+
+    Args:
+        jobs: Jobs API instance
+        job_id: Full job UUID
+
+    Returns:
+        Job if found, None if not found or error
+    """
+    try:
+        return jobs.get(job_id)
+    except Exception:
+        return None
+
+
+def find_by_hostname(job_list: list[Job], hostname: str) -> Job | None:
+    """Find job by hostname (exact or prefix match).
+
+    Args:
+        job_list: List of Job objects to search
+        hostname: Hostname to match (can be partial prefix)
+
+    Returns:
+        First matching Job or None
+    """
+    for job in job_list:
+        if job.hostname and (job.hostname == hostname or job.hostname.startswith(hostname)):
+            return job
+    return None
+
+
+def find_by_ip(job_list: list[Job], ip: str) -> Job | None:
+    """Find job by IP address (extracted from hostname).
+
+    Args:
+        job_list: List of Job objects to search
+        ip: IP address to match
+
+    Returns:
+        First matching Job or None
+    """
+    import socket
+
+    for job in job_list:
+        if not job.hostname:
+            continue
+        try:
+            job_ip = socket.gethostbyname(job.hostname)
+            if job_ip == ip:
+                return job
+        except socket.gaierror:
+            continue
+    return None
+
+
+def find_job(jobs: Jobs, identifier: str, state: str = None) -> Job | None:
+    """Find a job by UUID, hostname, or IP address.
+
+    Args:
+        jobs: Jobs API instance
+        identifier: Job UUID, hostname (partial match), or IP address
+        state: Optional state filter for listing jobs
+
+    Returns:
+        Matching Job or None
+    """
+    # Try UUID first (direct API call)
+    if is_uuid(identifier):
+        return find_by_id(jobs, identifier)
+
+    # Get job list for hostname/IP search
+    job_list = jobs.list(state=state)
+
+    # Try hostname match
+    job = find_by_hostname(job_list, identifier)
+    if job:
+        return job
+
+    # Try IP match (slower, requires DNS lookup)
+    return find_by_ip(job_list, identifier)

c3/logs.py

@@ -0,0 +1,273 @@
+"""Async log streaming for jobs"""
+import asyncio
+import json
+from collections import deque
+from typing import TYPE_CHECKING, AsyncIterator, Callable
+
+import websockets
+
+from .config import get_ws_url, WS_LOGS_PATH
+
+if TYPE_CHECKING:
+    from .client import C3
+
+
+# Default limits to prevent memory blowup
+DEFAULT_MAX_INITIAL_LINES = 1000  # Max lines to fetch on initial REST call
+DEFAULT_MAX_BUFFER = 5000  # Max lines to keep in memory buffer
+
+
+def fetch_logs(c3: "C3", job_id: str, tail: int = None) -> list[str]:
+    """Fetch logs via REST API (one-time call).
+
+    Args:
+        c3: C3 client
+        job_id: Job ID
+        tail: Only return last N lines (default: all)
+
+    Returns:
+        List of log lines
+    """
+    try:
+        logs = c3.jobs.logs(job_id)
+        if not logs:
+            return []
+        lines = logs.strip().split("\n")
+        if tail and len(lines) > tail:
+            return lines[-tail:]
+        return lines
+    except Exception:
+        return []
+
+
+class LogStream:
+    """Async log streamer - websocket streaming with optional initial fetch.
+
+    Usage:
+        stream = LogStream(c3, job)
+        await stream.connect()
+        async for line in stream:
+            print(line)
+        await stream.close()
+
+    This class guarantees:
+    - Initial logs fetched ONCE on connect (limited to max_initial_lines)
+    - All subsequent logs via websocket (NO polling)
+    - Bounded buffer to prevent memory blowup
+    - Proper cleanup on close
+    """
+
+    def __init__(
+        self,
+        c3: "C3",
+        job_id: str,
+        job_key: str = None,
+        fetch_initial: bool = True,
+        max_initial_lines: int = DEFAULT_MAX_INITIAL_LINES,
+        max_buffer: int = DEFAULT_MAX_BUFFER,
+    ):
+        """
+        Args:
+            c3: C3 client
+            job_id: Job ID for REST log fetch
+            job_key: Job key for websocket (if None, fetched from job)
+            fetch_initial: Whether to fetch existing logs on connect
+            max_initial_lines: Max lines to fetch initially (prevents huge fetch)
+            max_buffer: Max lines to keep in buffer (oldest dropped)
+        """
+        self.c3 = c3
+        self.job_id = job_id
+        self.job_key = job_key
+        self.fetch_initial = fetch_initial
+        self.max_initial_lines = max_initial_lines
+        self.max_buffer = max_buffer
+
+        self._ws = None
+        self._buffer: deque[str] = deque(maxlen=max_buffer)
+        self._initial_fetched = False
+        self._connected = False
+        self._closed = False
+
+    @property
+    def status(self) -> str:
+        """Connection status: disconnected, connecting, connected, closed"""
+        if self._closed:
+            return "closed"
+        if self._connected:
+            return "connected"
+        if self._ws:
+            return "connecting"
+        return "disconnected"
+
+    async def connect(self) -> list[str]:
+        """Connect to log stream.
+
+        Returns initial logs (if fetch_initial=True).
+        After this, iterate with `async for line in stream`.
+        """
+        if self._closed:
+            raise RuntimeError("LogStream is closed")
+
+        initial_lines = []
+
+        # Fetch initial logs ONCE (bounded)
+        if self.fetch_initial and not self._initial_fetched:
+            initial_lines = fetch_logs(self.c3, self.job_id, tail=self.max_initial_lines)
+            for line in initial_lines:
+                self._buffer.append(line)
+            self._initial_fetched = True
+
+        # Get job_key if not provided
+        if not self.job_key:
+            job = self.c3.jobs.get(self.job_id)
+            self.job_key = job.job_key
+
+        # Connect websocket
+        if self.job_key and not self._ws:
+            ws_url = get_ws_url()
+            full_url = f"{ws_url}{WS_LOGS_PATH}/{self.job_key}"
+            self._ws = await websockets.connect(full_url)
+            self._connected = True
+
+        return initial_lines
+
+    async def close(self):
+        """Close the websocket connection"""
+        self._closed = True
+        self._connected = False
+        if self._ws:
+            await self._ws.close()
+            self._ws = None
+
+    def get_buffer(self) -> list[str]:
+        """Get current buffer contents (bounded, oldest may be dropped)"""
+        return list(self._buffer)
+
+    def clear_buffer(self):
+        """Clear the buffer"""
+        self._buffer.clear()
+
+    async def __aiter__(self) -> AsyncIterator[str]:
+        """Async iterate over NEW log lines from websocket.
+
+        Note: This does NOT yield initial logs. Call connect() first
+        and handle the returned initial lines separately.
+        """
+        if not self._ws:
+            raise RuntimeError("Not connected. Call connect() first.")
+
+        try:
+            async for message in self._ws:
+                if self._closed:
+                    break
+                try:
+                    data = json.loads(message)
+                    if data.get("event") == "log" and data.get("log"):
+                        for line in data["log"].splitlines():
+                            if line:
+                                self._buffer.append(line)
+                                yield line
+                except json.JSONDecodeError:
+                    continue
+        except websockets.ConnectionClosed:
+            self._connected = False
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+
+async def stream_logs(
+    c3: "C3",
+    job_id: str,
+    on_line: Callable[[str], None],
+    until_state: set[str] = None,
+    poll_state_interval: float = 2.0,
+    fetch_initial: bool = True,
+    fetch_final: bool = True,
+    max_initial_lines: int = DEFAULT_MAX_INITIAL_LINES,
+) -> None:
+    """Stream logs until job reaches a terminal state.
+
+    Args:
+        c3: C3 client
+        job_id: Job ID to stream logs from
+        on_line: Callback for each log line (called immediately, no buffering)
+        until_state: States to stop on (default: terminal states)
+        poll_state_interval: How often to check job STATE (NOT log polling!)
+        fetch_initial: Fetch existing logs on start
+        fetch_final: Fetch logs one more time after job terminates
+        max_initial_lines: Max lines to fetch initially
+
+    This function:
+    - Fetches initial logs ONCE (bounded)
+    - Streams via websocket (NO log polling)
+    - Polls job STATE only (to detect termination)
+    - Optionally fetches final logs ONCE when job terminates
+    """
+    if until_state is None:
+        until_state = {"succeeded", "failed", "canceled", "terminated"}
+
+    job = c3.jobs.get(job_id)
+    initial_fetched = False
+    ws = None
+
+    try:
+        # Wait for job to be assigned/running
+        while job.state in ("pending", "queued"):
+            await asyncio.sleep(poll_state_interval)
+            job = c3.jobs.get(job_id)
+
+        # Check for immediate terminal state
+        if job.state in until_state:
+            if fetch_final:
+                for line in fetch_logs(c3, job_id, tail=max_initial_lines):
+                    on_line(line)
+            return
+
+        # Fetch initial logs ONCE when running (bounded)
+        if fetch_initial and job.state == "running" and not initial_fetched:
+            for line in fetch_logs(c3, job_id, tail=max_initial_lines):
+                on_line(line)
+            initial_fetched = True
+
+        # Connect websocket
+        if job.job_key:
+            ws_url = get_ws_url()
+            full_url = f"{ws_url}{WS_LOGS_PATH}/{job.job_key}"
+            ws = await websockets.connect(full_url)
+
+            # Stream logs while checking job state periodically
+            while True:
+                try:
+                    # Wait for message with timeout to allow state checks
+                    message = await asyncio.wait_for(ws.recv(), timeout=poll_state_interval)
+                    try:
+                        data = json.loads(message)
+                        if data.get("event") == "log" and data.get("log"):
+                            for line in data["log"].splitlines():
+                                if line:
+                                    on_line(line)
+                    except json.JSONDecodeError:
+                        continue
+                except asyncio.TimeoutError:
+                    # Check job state (NOT polling logs!)
+                    job = c3.jobs.get(job_id)
+                    if job.state in until_state:
+                        break
+                except websockets.ConnectionClosed:
+                    break
+
+        # Fetch final logs ONCE (may have missed some during shutdown)
+        if fetch_final:
+            # Small delay to let final logs flush
+            await asyncio.sleep(0.5)
+            for line in fetch_logs(c3, job_id, tail=max_initial_lines):
+                on_line(line)
+
+    finally:
+        if ws:
+            await ws.close()