more-compute 0.1.3__py3-none-any.whl → 0.2.0__py3-none-any.whl

morecompute/services/pod_monitor.py

@@ -0,0 +1,138 @@
+"""Service for monitoring GPU pod status updates."""
+
+import asyncio
+import sys
+from typing import Callable, Awaitable
+from cachetools import TTLCache
+
+from .prime_intellect import PrimeIntellectService
+
+
+PodUpdateCallback = Callable[[dict], Awaitable[None]]
+
+
+class PodMonitor:
+    """Monitors GPU pod status and broadcasts updates."""
+
+    POLL_INTERVAL_SECONDS = 5
+
+    def __init__(
+        self,
+        prime_intellect: PrimeIntellectService,
+        pod_cache: TTLCache,
+        update_callback: PodUpdateCallback
+    ):
+        """
+        Initialize pod monitor.
+
+        Args:
+            prime_intellect: Prime Intellect API service
+            pod_cache: Cache to clear on updates
+            update_callback: Async callback for broadcasting updates
+        """
+        self.pi_service = prime_intellect
+        self.pod_cache = pod_cache
+        self.update_callback = update_callback
+        self.monitoring_tasks: dict[str, asyncio.Task] = {}
+
+    async def start_monitoring(self, pod_id: str) -> None:
+        """
+        Start monitoring a pod's status.
+
+        Args:
+            pod_id: ID of pod to monitor
+        """
+        # Don't start duplicate monitors
+        if pod_id in self.monitoring_tasks:
+            print(f"[POD MONITOR] Already monitoring pod {pod_id}", file=sys.stderr, flush=True)
+            return
+
+        task = asyncio.create_task(self._monitor_loop(pod_id))
+        self.monitoring_tasks[pod_id] = task
+        print(f"[POD MONITOR] Started monitoring pod {pod_id}", file=sys.stderr, flush=True)
+
+    async def stop_monitoring(self, pod_id: str) -> None:
+        """
+        Stop monitoring a pod.
+
+        Args:
+            pod_id: ID of pod to stop monitoring
+        """
+        task = self.monitoring_tasks.pop(pod_id, None)
+        if task and not task.done():
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+        print(f"[POD MONITOR] Stopped monitoring pod {pod_id}", file=sys.stderr, flush=True)
+
+    async def _monitor_loop(self, pod_id: str) -> None:
+        """
+        Main monitoring loop for a pod.
+
+        Args:
+            pod_id: ID of pod to monitor
+        """
+        try:
+            while True:
+                try:
+                    # Fetch current pod status
+                    pod = await self.pi_service.get_pod(pod_id)
+
+                    print(
+                        f"[POD MONITOR] Pod {pod_id} status: {pod.status}",
+                        file=sys.stderr,
+                        flush=True
+                    )
+
+                    # Clear cache to force fresh data
+                    self.pod_cache.clear()
+
+                    # Broadcast update
+                    await self.update_callback({
+                        "type": "pod_status_update",
+                        "data": {
+                            "pod_id": pod_id,
+                            "name": pod.name,
+                            "status": pod.status,
+                            "ssh_connection": pod.sshConnection,
+                            "ip": pod.ip,
+                            "gpu_name": pod.gpuName,
+                            "price_hr": pod.priceHr
+                        }
+                    })
+
+                    # Stop monitoring if ERROR or TERMINATED
+                    if pod.status in {"ERROR", "TERMINATED"}:
+                        print(
+                            f"[POD MONITOR] Pod {pod_id} reached terminal state: {pod.status}",
+                            file=sys.stderr,
+                            flush=True
+                        )
+                        break
+
+                    # If ACTIVE and has SSH connection, pod is fully ready - stop monitoring
+                    if pod.status == "ACTIVE" and pod.sshConnection:
+                        print(
+                            f"[POD MONITOR] Pod {pod_id} is ACTIVE with SSH connection: {pod.sshConnection}",
+                            file=sys.stderr,
+                            flush=True
+                        )
+                        break
+
+                    # Wait before next check
+                    await asyncio.sleep(self.POLL_INTERVAL_SECONDS)
+
+                except Exception as e:
+                    print(
+                        f"[POD MONITOR] Error checking pod {pod_id}: {e}",
+                        file=sys.stderr,
+                        flush=True
+                    )
+                    await asyncio.sleep(self.POLL_INTERVAL_SECONDS)
+
+        finally:
+            # Clean up
+            self.monitoring_tasks.pop(pod_id, None)
+            print(f"[POD MONITOR] Stopped monitoring pod {pod_id}", file=sys.stderr, flush=True)

morecompute/services/prime_intellect.py

@@ -1,69 +1,14 @@
-from pydantic import BaseModel
-from datetime import datetime
 import httpx
 from fastapi import HTTPException
 
-class EnvVar(BaseModel):
-    key: str
-    value: str
-
-class PodConfig(BaseModel):
-    # Required fields
-    name: str
-    cloudId: str
-    gpuType: str
-    socket: str
-    gpuCount: int = 1
-
-    # Optional
-    diskSize: int | None = None
-    vcpus: int | None = None
-    memory: int | None = None
-    maxPrice: float | None = None
-    image: str | None = None
-    customTemplateId: str | None = None
-    dataCenterId: str | None = None
-    country: str | None = None
-    security: str | None = None
-    envVars: list[EnvVar] | None = None
-    jupyterPassword: str | None = None
-    autoRestart: bool | None = None
-
-
-class ProviderConfig(BaseModel):
-    type: str = "runpod"
-
-
-class TeamConfig(BaseModel):
-    teamId: str | None = None
-
-
-class CreatePodRequest(BaseModel):
-    pod: PodConfig
-    provider: ProviderConfig
-    team: TeamConfig | None = None
-
-
-class PodResponse(BaseModel):
-    id: str
-    userId: str
-    teamId: str | None
-    name: str
-    status: str
-    gpuName: str
-    gpuCount: int
-    priceHr: float
-    sshConnection: str | None
-    ip: str | None
-    createdAt: datetime
-    updatedAt: datetime
-
-
-class AvailabilityQuery(BaseModel):
-    regions: list[str] | None = None
-    gpu_count: int | None = None
-    gpu_type: str | None = None
-    security: str | None = None
+from ..models.api_models import (
+    CreatePodRequest,
+    CreateDiskRequest,
+    PodResponse,
+    DiskResponse,
+)
+
+# Helper methods for https://docs.primeintellect.ai/api-reference/
 
 
 class PrimeIntellectService:
@@ -314,3 +259,82 @@ class PrimeIntellectService:
             Dict with confirmation
         """
         return await self._make_request("PATCH", f"/ssh-keys/{key_id}/primary")
+
+
+    async def get_disks(
+        self,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> dict[str, object]:
+        """
+        Get list of all disks
+
+        args:
+            limit : max number of results (default : 100)
+            offset : for Pagination
+
+        returns:
+            dict with:
+                - data: list of disk objects
+                - total_count : total number of disks
+                - offset : current offset
+                - limit: current limit
+        """
+        params : dict[str, str | int | float | list[str]] = {"limit":limit, "offset":offset}
+        return await self._make_request("GET", "/disks/", params=params)
+
+    async def create_disks(self, disk_request :CreateDiskRequest) -> DiskResponse:
+        """
+        Create new persistent disk
+
+        Args:
+            disk_reqeust: disk configuration request
+
+        Returns:
+            Disk response with disk details
+
+        """
+        import sys
+        payload = disk_request.model_dump(exclude_none=True)
+        print(f"[PI SERVICE] Creating disk with payload: + {payload}", file=sys.stderr, flush=True)
+        response = await self._make_request("POST", "/disks/", json_data=payload)
+        return DiskResponse.model_validate(response)
+
+    async def get_disk(self, disk_id:str) -> DiskResponse:
+        """
+        Get details of specific disk
+
+        args:
+            disk_id: disk identifier
+
+        returns:
+            DiskResponse
+        """
+        response = await self._make_request("GET", f"/disks/{disk_id}")
+        return DiskResponse.model_validate(response)
+
+    async def update_disk(self, disk_id:str, name:str) -> dict[str,object]:
+        """
+        update disk name
+
+        args:
+            disk_id: disk identifier
+            name: new name of disk
+
+        returns:
+            dict with update confirmation
+        """
+        data: dict[str, object] = {"name":name}
+        return await self._make_request("PATCH", f"/disks/{disk_id}", json_data=data)
+
+    async def delete_disk(self, disk_id: str) -> dict[str, object]:
+        """
+        Delete a disk.
+
+        Args:
+            disk_id: The disk identifier
+
+        Returns:
+            Dict with deletion confirmation
+        """
+        return await self._make_request("DELETE", f"/disks/{disk_id}")

morecompute/utils/config_util.py

@@ -0,0 +1,59 @@
+"""Configuration utilities for managing API keys and environment variables."""
+
+from pathlib import Path
+import os
+
+
+def load_api_key_from_env(env_var: str, env_file_path: Path | None = None) -> str | None:
+    """
+    Load API key from environment variable or .env file.
+
+    Args:
+        env_var: Environment variable name to check
+        env_file_path: Path to .env file (optional)
+
+    Returns:
+        API key string or None if not found
+    """
+    api_key = os.getenv(env_var)
+    if api_key:
+        return api_key
+
+    if env_file_path and env_file_path.exists():
+        try:
+            with env_file_path.open("r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if line.startswith(f"{env_var}="):
+                        return line.split("=", 1)[1].strip().strip('"').strip("'")
+        except Exception:
+            pass
+
+    return None
+
+
+def save_api_key_to_env(env_var: str, api_key: str, env_file_path: Path) -> None:
+    """
+    Save API key to .env file, replacing existing value if present.
+
+    Args:
+        env_var: Environment variable name
+        api_key: API key value to save
+        env_file_path: Path to .env file
+
+    Raises:
+        ValueError: If API key is empty
+        IOError: If file cannot be written
+    """
+    if not api_key.strip():
+        raise ValueError("API key cannot be empty")
+
+    existing_lines = []
+    if env_file_path.exists():
+        with env_file_path.open("r", encoding="utf-8") as f:
+            existing_lines = f.readlines()
+
+    new_lines = [line for line in existing_lines if not line.strip().startswith(f"{env_var}=")]
+    new_lines.append(f"{env_var}={api_key}\n")
+    with env_file_path.open("w", encoding="utf-8") as f:
+        f.writelines(new_lines)

morecompute/utils/special_commands.py

@@ -122,11 +122,17 @@ class AsyncSpecialCommandHandler:
             # Check if command failed
             if return_code != 0:
                 result["status"] = "error"
-                result["error"] = {
-                    "ename": "ShellCommandError",
-                    "evalue": f"Command failed with return code {return_code}",
-                    "traceback": [f"Shell command failed: {command}"]
-                }
+                # Only add generic error if we don't already have detailed stderr output
+                # The detailed stderr is already in result["outputs"] from streaming
+                has_stderr = any(o.get("name") == "stderr" and o.get("text", "").strip()
+                               for o in result.get("outputs", []))
+                if not has_stderr:
+                    # No detailed error output, add generic error
+                    result["error"] = {
+                        "ename": "ShellCommandError",
+                        "evalue": f"Command failed with return code {return_code}",
+                        "traceback": [f"Shell command failed: {command}"]
+                    }
 
         except Exception as e:
             result["status"] = "error"

morecompute/utils/zmq_util.py

@@ -0,0 +1,51 @@
+"""Utilities for ZMQ socket management and reconnection."""
+
+import zmq
+import os
+
+
+def reconnect_zmq_sockets(
+    executor: any,
+    cmd_addr: str | None = None,
+    pub_addr: str | None = None,
+    is_remote: bool = False
+) -> None:
+    """
+    Reconnect ZMQ executor sockets to new addresses.
+
+    Args:
+        executor: Executor instance with ZMQ sockets
+        cmd_addr: Command socket address (defaults to local)
+        pub_addr: Publish socket address (defaults to local)
+        is_remote: True if connecting to remote worker, False for local
+    """
+    # Use provided addresses or fall back to defaults
+    final_cmd_addr = cmd_addr or os.getenv('MC_ZMQ_CMD_ADDR', 'tcp://127.0.0.1:5555')
+    final_pub_addr = pub_addr or os.getenv('MC_ZMQ_PUB_ADDR', 'tcp://127.0.0.1:5556')
+
+    # Update executor addresses
+    executor.cmd_addr = final_cmd_addr
+    executor.pub_addr = final_pub_addr
+    executor.is_remote = is_remote
+
+    # Reconnect command socket (REQ)
+    executor.req.close(0)  # type: ignore[reportAttributeAccessIssue]
+    executor.req = executor.ctx.socket(zmq.REQ)  # type: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
+    executor.req.connect(executor.cmd_addr)  # type: ignore[reportAttributeAccessIssue]
+
+    # Reconnect publish socket (SUB)
+    executor.sub.close(0)  # type: ignore[reportAttributeAccessIssue]
+    executor.sub = executor.ctx.socket(zmq.SUB)  # type: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
+    executor.sub.connect(executor.pub_addr)  # type: ignore[reportAttributeAccessIssue]
+    executor.sub.setsockopt_string(zmq.SUBSCRIBE, '')  # type: ignore[reportAttributeAccessIssue]
+
+
+def reset_to_local_zmq(executor: any) -> None:
+    """
+    Reset executor to local ZMQ addresses.
+
+    Args:
+        executor: Executor instance to reset
+    """
+    executor.is_remote = False
+    reconnect_zmq_sockets(executor)

frontend/components/MarkdownRenderer.tsx

@@ -1,84 +0,0 @@
-'use client';
-
-import React from 'react';
-
-interface MarkdownRendererProps {
-  source: string;
-  onClick?: () => void;
-}
-
-const MarkdownRenderer: React.FC<MarkdownRendererProps> = ({ source, onClick }) => {
-  const escapeHtml = (text: string) => {
-    const div = document.createElement('div');
-    div.textContent = text;
-    return div.innerHTML;
-  };
-
-  const renderMarkdown = (text: string) => {
-    let html = text;
-
-    // Code blocks (must be processed first)
-    html = html.replace(/```([\s\S]*?)```/g, (match, code) => {
-      return `<pre><code>${escapeHtml(code.trim())}</code></pre>`;
-    });
-
-    // Language-specific code blocks
-    html = html.replace(/```(\w+)\n([\s\S]*?)```/g, (match, lang, code) => {
-      return `<pre><code class="language-${lang}">${escapeHtml(code.trim())}</code></pre>`;
-    });
-
-    // Headers
-    html = html.replace(/^### (.*$)/gim, '<h3>$1</h3>');
-    html = html.replace(/^## (.*$)/gim, '<h2>$1</h2>');
-    html = html.replace(/^# (.*$)/gim, '<h1>$1</h1>');
-
-    // Bold and Italic
-    html = html.replace(/\*\*\*([^*]+)\*\*\*/g, '<strong><em>$1</em></strong>');
-    html = html.replace(/\*\*([^*]+)\*\*/g, '<strong>$1</strong>');
-    html = html.replace(/\*([^*]+)\*/g, '<em>$1</em>');
-
-    // Strikethrough
-    html = html.replace(/~~([^~]+)~~/g, '<del>$1</del>');
-    
-    // Links
-    html = html.replace(/\[([^\]]+)\]\(([^)]+)\)/g, '<a href="$2" target="_blank" rel="noopener noreferrer">$1</a>');
-
-    // Unordered lists
-    html = html.replace(/^\s*[-*+] (.+)$/gim, (match, item) => `<li>${item}</li>`);
-    html = html.replace(/(<li>[\s\S]*?<\/li>)/g, '<ul>$1</ul>').replace(/<\/ul>\s*<ul>/g, '');
-
-    // Ordered lists
-    html = html.replace(/^\s*\d+\. (.+)$/gim, (match, item) => `<ol><li>${item}</li></ol>`).replace(/<\/ol>\s*<ol>/g, '');
-
-    // Inline code
-    html = html.replace(/`([^`]+)`/g, '<code>$1</code>');
-
-    // Blockquotes
-    html = html.replace(/^> (.+)$/gim, '<blockquote>$1</blockquote>');
-    
-    // Handle HTML br tags first - convert them to proper line breaks
-    html = html.replace(/<br\s*\/?>/gi, '\n');
-
-    // Line breaks and paragraphs - handle paragraph separation properly
-    // Split by double newlines to create paragraphs, but preserve single line breaks within paragraphs
-    html = html.replace(/\n\s*\n/g, '</p><p>');  // Paragraph breaks (double newlines with optional whitespace)
-    html = html.replace(/\n(?!<\/p>)/g, '<br>');  // Line breaks within paragraphs (but not paragraph breaks)
-
-    // Ensure content is wrapped in paragraph tags if not already in a block element
-    if (!html.match(/^<(h[1-6]|ul|ol|pre|blockquote|hr|p)/) && html.trim()) {
-        html = '<p>' + html + '</p>';
-    }
-    
-    return html;
-  };
-
-  return (
-    <div 
-      className="markdown-rendered" 
-      onClick={onClick}
-      dangerouslySetInnerHTML={{ __html: renderMarkdown(source) }}
-    />
-  );
-};
-
-export default MarkdownRenderer;