more-compute 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

morecompute/services/pod_manager.py

@@ -1,35 +1,44 @@
 import asyncio
 import subprocess
 import os
-import sys
 import tempfile
 import tarfile
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from .prime_intellect import PrimeIntellectService, PodResponse
+from .providers.base_provider import BaseGPUProvider
+from ..models.api_models import PodResponse
 
 if TYPE_CHECKING:
     from ..execution.executor import NextZmqExecutor
 
+# Type alias for supported provider services
+ProviderService = BaseGPUProvider
+
+
 class PodKernelManager:
     """
-    Manages remote GPU pod connections (currently PI as provider, hope to provide other providers in the future)
-    and SSH tunnels for ZMQ execution
+    Manages remote GPU pod connections and SSH tunnels for ZMQ execution.
+
+    Supports multiple GPU providers:
+    - RunPod
+    - Lambda Labs
+    - Vast.ai
     """
-    pi_service: PrimeIntellectService
+    provider_service: ProviderService
+    provider_type: str
     pod: PodResponse | None
     ssh_tunnel_proc: subprocess.Popen[bytes] | None
     local_cmd_port: int
     local_pub_port: int
-    remote_cmd_port : int
+    remote_cmd_port: int
     remote_pub_port: int
     executor: "NextZmqExecutor | None"
     _ssh_key_cache: str | None
 
     def __init__(
         self,
-        pi_service: PrimeIntellectService,
+        provider_service: ProviderService,
         local_cmd_port: int = 15555,
         local_pub_port: int = 15556,
         remote_cmd_port: int = 5555,
@@ -39,13 +48,15 @@ class PodKernelManager:
         Initialize pod manager
 
         args:
-            pi_service : Prime Intellect API service
+            provider_service: GPU provider service implementing BaseGPUProvider
             local_cmd_port: Local port for REQ/REP tunnel
             local_pub_port: Local port for PUB/SUB tunnel
             remote_cmd_port: Remote port for REQ/REP socket
             remote_pub_port: Remote port for PUB/SUB socket
         """
-        self.pi_service = pi_service
+        self.provider_service = provider_service
+        self.provider_type = getattr(provider_service, 'PROVIDER_NAME', 'unknown')
+
         self.pod = None
         self.ssh_tunnel_proc = None
         self.local_cmd_port = local_cmd_port
@@ -71,9 +82,8 @@ class PodKernelManager:
                 self._ssh_key_cache = expanded
                 return expanded
 
-        # Try common SSH key paths (including Prime Intellect's recommended name)
+        # Try common SSH key paths
         common_keys = [
-            "~/.ssh/primeintellect_ed25519",  # Prime Intellect's recommended name
             "~/.ssh/id_ed25519",
             "~/.ssh/id_rsa",
         ]
@@ -85,6 +95,84 @@ class PodKernelManager:
 
         return None
 
+    def _is_key_encrypted(self, key_path: str) -> bool:
+        """Check if an SSH private key is encrypted with a passphrase."""
+        try:
+            with open(key_path, 'r') as f:
+                content = f.read(500)  # Read first 500 bytes
+                # OpenSSH encrypted keys contain these markers
+                return 'aes256-ctr' in content or 'aes128-ctr' in content or 'bcrypt' in content
+        except Exception:
+            return False
+
+    def _is_key_in_agent(self, key_path: str) -> bool:
+        """Check if the SSH key is loaded in the ssh-agent."""
+        try:
+            result = subprocess.run(
+                ["ssh-add", "-l"],
+                capture_output=True,
+                text=True,
+                timeout=5
+            )
+            if result.returncode != 0:
+                return False
+            # Check if the key fingerprint is in the agent
+            # Get the fingerprint of our key
+            fp_result = subprocess.run(
+                ["ssh-keygen", "-lf", key_path],
+                capture_output=True,
+                text=True,
+                timeout=5
+            )
+            if fp_result.returncode == 0:
+                # Extract fingerprint (e.g., SHA256:xxx)
+                parts = fp_result.stdout.split()
+                if len(parts) >= 2:
+                    fingerprint = parts[1]
+                    return fingerprint in result.stdout
+            return False
+        except Exception:
+            return False
+
+    def _get_ssh_setup_instructions(self) -> str:
+        """Get provider-specific SSH setup instructions."""
+        provider_instructions = {
+            "runpod": (
+                "SSH authentication failed. Please add your SSH public key to RunPod:\n"
+                "1. Visit https://www.runpod.io/console/user/settings\n"
+                "2. Go to 'SSH Public Keys' section\n"
+                "3. Add your public key (~/.ssh/id_ed25519.pub or ~/.ssh/id_rsa.pub)\n"
+                "4. Try connecting again"
+            ),
+            "lambda_labs": (
+                "SSH authentication failed. Lambda Labs SSH key mismatch.\n\n"
+                "IMPORTANT: Lambda Labs assigns SSH keys at instance creation time.\n"
+                "Your current instance may have been created with a different key.\n\n"
+                "To fix this:\n"
+                "1. Terminate the current instance\n"
+                "2. Visit https://cloud.lambdalabs.com/ssh-keys\n"
+                "3. Add your public key (~/.ssh/id_ed25519.pub) if not already there\n"
+                "4. Create a new instance - it will use your registered key\n\n"
+                "To view your public key, run: cat ~/.ssh/id_ed25519.pub"
+            ),
+            "vastai": (
+                "SSH authentication failed. Please add your SSH public key to Vast.ai:\n"
+                "1. Visit https://cloud.vast.ai/account/\n"
+                "2. Go to 'SSH Keys' section\n"
+                "3. Add your public key (~/.ssh/id_ed25519.pub or ~/.ssh/id_rsa.pub)\n"
+                "4. Try connecting again"
+            ),
+        }
+
+        return provider_instructions.get(
+            self.provider_type,
+            (
+                "SSH authentication failed. Please ensure your SSH public key is added to your GPU provider.\n"
+                "Upload your public key (~/.ssh/id_ed25519.pub or ~/.ssh/id_rsa.pub) to your provider's dashboard.\n"
+                "Then try connecting again."
+            )
+        )
+
     async def connect_to_pod(self, pod_id:str) -> dict[str, object]:
         """
         Connects to existing pod and set up ssh tunnel
@@ -94,8 +182,6 @@ class PodKernelManager:
         Response:
             dict with connection status
         """
-        import sys
-
         # Check if already connected to this pod
         if self.pod and self.pod.id == pod_id:
             # Check if tunnel is still alive
@@ -105,18 +191,13 @@ class PodKernelManager:
                     "message": f"Already connected to pod {pod_id}"
                 }
             # Tunnel died, clean up and reconnect
-            print(f"[POD MANAGER] Existing tunnel dead, reconnecting...", file=sys.stderr, flush=True)
             await self.disconnect()
 
         # If connected to different pod, disconnect first
         if self.pod and self.pod.id != pod_id:
-            print(f"[POD MANAGER] Disconnecting from {self.pod.id} to connect to {pod_id}", file=sys.stderr, flush=True)
             await self.disconnect()
 
-        self.pod = await self.pi_service.get_pod(pod_id)
-
-        print(f"[POD MANAGER] Pod status: {self.pod.status}", file=sys.stderr, flush=True)
-        print(f"[POD MANAGER] SSH connection: {self.pod.sshConnection}", file=sys.stderr, flush=True)
+        self.pod = await self.provider_service.get_pod(pod_id)
 
         if not self.pod.sshConnection:
             return{
@@ -132,7 +213,7 @@ class PodKernelManager:
             }
 
         # Parse SSH connection string
-        # Format can be: "ssh root@ip -p port" OR "root@ip -p port"
+        # Format can be: "ssh root@ip -p port" OR "root@ip -p port" OR "ssh ubuntu@ip"
         ssh_parts = self.pod.sshConnection.split()
 
         # Find the part containing @ (user@host)
@@ -148,8 +229,8 @@ class PodKernelManager:
                 "message": f"Invalid SSH connection format (no user@host found): {self.pod.sshConnection}"
             }
 
-        # Extract host from user@host
-        ssh_host = host_part.split("@")[1]
+        # Extract user and host from user@host
+        ssh_user, ssh_host = host_part.split("@")
         ssh_port = "22"
 
         # Extract port if specified with -p flag
@@ -158,33 +239,22 @@ class PodKernelManager:
             if port_idx + 1 < len(ssh_parts):
                 ssh_port = ssh_parts[port_idx + 1]
 
-        print(f"[POD MANAGER] Parsed SSH host: {ssh_host}, port: {ssh_port}", file=sys.stderr, flush=True)
-
         #deploy worker code to pod
-        print(f"[POD MANAGER] Deploying worker code to pod...", file=sys.stderr, flush=True)
-        deploy_result = await self._deploy_worker(ssh_host, ssh_port)
-        print(f"[POD MANAGER] Deploy result: {deploy_result}", file=sys.stderr, flush=True)
+        deploy_result = await self._deploy_worker(ssh_user, ssh_host, ssh_port)
         if deploy_result.get("status") ==  "error":
             return deploy_result
 
         #create ssh tunnel for ZMQ ports
-        print(f"[POD MANAGER] Creating SSH tunnel...", file=sys.stderr, flush=True)
-        tunnel_result = await self._create_ssh_tunnel(ssh_host, ssh_port)
-        print(f"[POD MANAGER] Tunnel result: {tunnel_result}", file=sys.stderr, flush=True)
+        tunnel_result = await self._create_ssh_tunnel(ssh_user, ssh_host, ssh_port)
         if tunnel_result.get("status") ==  "error":
             return tunnel_result
 
         #start remote worker
-        worker_result = await self._start_remote_worker(ssh_host, ssh_port)
+        worker_result = await self._start_remote_worker(ssh_user, ssh_host, ssh_port)
         if worker_result.get("status") ==  "error":
             await self.disconnect()
             return worker_result
 
-        # Note: Worker may take a few seconds to start and install matplotlib
-        # The connection should work even if verification fails
-        print(f"[POD MANAGER] Remote worker is starting (matplotlib install may take a few seconds)", file=sys.stderr, flush=True)
-        print(f"[POD MANAGER] Connection established - try running code in ~5 seconds", file=sys.stderr, flush=True)
-
         return {
             "status": "ok",
             "message": f"Connected to pod {pod_id}",
@@ -196,11 +266,12 @@ class PodKernelManager:
             }
         }
 
-    async def _deploy_worker(self, ssh_host: str, ssh_port: str) -> dict[str,object]:
+    async def _deploy_worker(self, ssh_user: str, ssh_host: str, ssh_port: str) -> dict[str,object]:
         """
         Deploy worker code to remote pod via Secure Copy Protocol.
 
         args:
+            ssh_user: SSH username (e.g., 'root' or 'ubuntu')
             ssh_host: SSH host address
             ssh_port: SSH port
 
@@ -208,6 +279,21 @@ class PodKernelManager:
             dict with deployment status
         """
         try:
+            # Check if SSH key is encrypted and not in agent
+            ssh_key = self._get_ssh_key()
+            if ssh_key and self._is_key_encrypted(ssh_key):
+                if not self._is_key_in_agent(ssh_key):
+                    key_name = os.path.basename(ssh_key)
+                    return {
+                        "status": "error",
+                        "message": (
+                            f"Your SSH key ({key_name}) is protected with a passphrase but not loaded in ssh-agent.\n\n"
+                            f"To fix this, run:\n"
+                            f"  ssh-add {ssh_key}\n\n"
+                            f"Enter your passphrase when prompted, then try connecting again."
+                        )
+                    }
+
             # Create temporary tarball of morecompute package
             project_root = Path(__file__).parent.parent.parent
             morecompute_dir = project_root / "morecompute"
@@ -231,7 +317,7 @@ class PodKernelManager:
                 "-o", "BatchMode=yes",  # Prevent password prompts, fail fast if key auth doesn't work
                 "-o", "ConnectTimeout=10",
                 tmp_path,
-                f"root@{ssh_host}:/tmp/morecompute.tar.gz"
+                f"{ssh_user}@{ssh_host}:/tmp/morecompute.tar.gz"
             ])
 
             result = subprocess.run(
@@ -244,14 +330,11 @@ class PodKernelManager:
             if result.returncode != 0:
                 error_msg = result.stderr.lower()
                 if "permission denied" in error_msg or "publickey" in error_msg:
+                    # Get provider-specific SSH setup instructions
+                    ssh_help = self._get_ssh_setup_instructions()
                     return {
                         "status": "error",
-                        "message": (
-                            "SSH authentication failed. Please add your SSH public key to Prime Intellect:\n"
-                            "1. Visit https://app.primeintellect.ai/dashboard/tokens\n"
-                            "2. Upload your public key (~/.ssh/id_ed25519.pub or ~/.ssh/id_rsa.pub)\n"
-                            "3. Try connecting again"
-                        )
+                        "message": ssh_help
                     }
                 elif "host key verification failed" in error_msg:
                     return {
@@ -265,6 +348,9 @@ class PodKernelManager:
                     }
 
             # Extract on remote and install dependencies
+            # Use sudo for non-root users to run pip install
+            pip_cmd = "pip install --quiet pyzmq matplotlib" if ssh_user == "root" else "sudo pip install --quiet pyzmq matplotlib"
+
             ssh_cmd = ["ssh", "-p", ssh_port]
 
             if ssh_key:
@@ -275,11 +361,11 @@ class PodKernelManager:
                 "-o", "UserKnownHostsFile=/dev/null",
                 "-o", "BatchMode=yes",
                 "-o", "ConnectTimeout=10",
-                f"root@{ssh_host}",
+                f"{ssh_user}@{ssh_host}",
                 (
                     "cd /tmp && "
                     "tar -xzf morecompute.tar.gz && "
-                    "pip install --quiet pyzmq matplotlib && "
+                    f"{pip_cmd} && "
                     "echo 'Deployment complete'"
                 )
             ])
@@ -308,11 +394,12 @@ class PodKernelManager:
                 "message": f"Deployment error: {str(e)}"
             }
 
-    async def _create_ssh_tunnel(self, ssh_host: str, ssh_port: str) -> dict[str, object]:
+    async def _create_ssh_tunnel(self, ssh_user: str, ssh_host: str, ssh_port: str) -> dict[str, object]:
         """
         Create SSH tunnel for ZMQ ports.
 
         args:
+            ssh_user: SSH username (e.g., 'root' or 'ubuntu')
             ssh_host: SSH host address
             ssh_port: SSH port
 
@@ -336,7 +423,7 @@ class PodKernelManager:
                 "-N",  # No command execution
                 "-L", f"{self.local_cmd_port}:localhost:{self.remote_cmd_port}",
                 "-L", f"{self.local_pub_port}:localhost:{self.remote_pub_port}",
-                f"root@{ssh_host}"
+                f"{ssh_user}@{ssh_host}"
             ])
 
             self.ssh_tunnel_proc = subprocess.Popen(
@@ -369,17 +456,17 @@ class PodKernelManager:
             }
 
         except Exception as e:
-            print(f"[POD MANAGER] Exception creating tunnel: {e}", file=sys.stderr, flush=True)
             return {
                 "status": "error",
                 "message": f"Tunnel creation error: {str(e)}"
             }
 
-    async def _start_remote_worker(self, ssh_host: str, ssh_port: str) -> dict[str, object]:
+    async def _start_remote_worker(self, ssh_user: str, ssh_host: str, ssh_port: str) -> dict[str, object]:
         """
         Start ZMQ worker on remote pod.
 
         args:
+            ssh_user: SSH username (e.g., 'root' or 'ubuntu')
             ssh_host: SSH host address
             ssh_port: SSH port
 
@@ -387,31 +474,46 @@ class PodKernelManager:
             dict with worker start status
         """
         try:
-            print(f"[POD MANAGER] Starting remote worker on {ssh_host}:{ssh_port}", file=sys.stderr, flush=True)
-
             # Start worker in background on remote pod
             # Use 'python3' instead of sys.executable since remote pod may have different Python path
+            # Use sudo for non-root users
+            python_cmd = "python3" if ssh_user == "root" else "sudo python3"
+
             ssh_key = self._get_ssh_key()
             worker_cmd = ["ssh", "-p", ssh_port]
 
             if ssh_key:
                 worker_cmd.extend(["-i", ssh_key])
 
+            # Build the command - for non-root users, we need to pass env vars through sudo
+            if ssh_user == "root":
+                remote_cmd = (
+                    f"cd /tmp && "
+                    f"MC_ZMQ_CMD_ADDR=tcp://0.0.0.0:{self.remote_cmd_port} "
+                    f"MC_ZMQ_PUB_ADDR=tcp://0.0.0.0:{self.remote_pub_port} "
+                    f"setsid python3 -u /tmp/morecompute/execution/worker.py "
+                    f"</dev/null >/tmp/worker.log 2>&1 & "
+                    f"echo $!"
+                )
+            else:
+                # For non-root, use sudo with env vars passed via sudo's env mechanism
+                remote_cmd = (
+                    f"cd /tmp && "
+                    f"sudo MC_ZMQ_CMD_ADDR=tcp://0.0.0.0:{self.remote_cmd_port} "
+                    f"MC_ZMQ_PUB_ADDR=tcp://0.0.0.0:{self.remote_pub_port} "
+                    f"setsid python3 -u /tmp/morecompute/execution/worker.py "
+                    f"</dev/null >/tmp/worker.log 2>&1 & "
+                    f"echo $!"
+                )
+
             worker_cmd.extend([
                 "-o", "StrictHostKeyChecking=no",
                 "-o", "UserKnownHostsFile=/dev/null",
                 "-o", "BatchMode=yes",
                 "-o", "ConnectTimeout=10",
-                f"root@{ssh_host}",
+                f"{ssh_user}@{ssh_host}",
                 "sh", "-c",
-                (
-                    f"'cd /tmp && "
-                    f"MC_ZMQ_CMD_ADDR=tcp://0.0.0.0:{self.remote_cmd_port} "
-                    f"MC_ZMQ_PUB_ADDR=tcp://0.0.0.0:{self.remote_pub_port} "
-                    f"setsid python3 -u /tmp/morecompute/execution/worker.py "
-                    f"</dev/null >/tmp/worker.log 2>&1 & "
-                    f"echo $!'"
-                )
+                f"'{remote_cmd}'"
             ])
 
             result = subprocess.run(
@@ -428,13 +530,10 @@ class PodKernelManager:
                 }
 
             remote_pid = result.stdout.strip()
-            print(f"[POD MANAGER] Remote worker PID: {remote_pid}", file=sys.stderr, flush=True)
 
             # Wait for worker to be ready
             await asyncio.sleep(2)
 
-            print(f"[POD MANAGER] Remote worker should be ready now", file=sys.stderr, flush=True)
-
             return {
                 "status": "ok",
                 "message": "Remote worker started",
@@ -580,7 +679,7 @@ class PodKernelManager:
 
         # Get updated pod info
         try:
-            updated_pod = await self.pi_service.get_pod(pod.id)
+            updated_pod = await self.provider_service.get_pod(pod.id)
             pod_status = updated_pod.status
         except Exception:
             pod_status = "unknown"
@@ -594,12 +693,14 @@ class PodKernelManager:
                 "gpu_type": pod.gpuName,
                 "gpu_count": pod.gpuCount,
                 "price_hr": pod.priceHr,
-                "ssh_connection": pod.sshConnection
+                "ssh_connection": pod.sshConnection,
+                "provider": self.provider_type
             },
             "tunnel": {
                 "alive": tunnel_alive,
                 "local_cmd_port": self.local_cmd_port,
                 "local_pub_port": self.local_pub_port
             },
-            "executor_attached": self.executor is not None
+            "executor_attached": self.executor is not None,
+            "provider": self.provider_type
         }

morecompute/services/pod_monitor.py

@@ -1,36 +1,55 @@
 """Service for monitoring GPU pod status updates."""
 
 import asyncio
-import sys
-from typing import Callable, Awaitable
+from typing import Callable, Awaitable, Union
 from cachetools import TTLCache
 
 from .prime_intellect import PrimeIntellectService
+from .providers.base_provider import BaseGPUProvider
 
 
 PodUpdateCallback = Callable[[dict], Awaitable[None]]
 
+# Type alias for supported provider services
+ProviderService = Union[PrimeIntellectService, BaseGPUProvider]
+
 
 class PodMonitor:
-    """Monitors GPU pod status and broadcasts updates."""
+    """Monitors GPU pod status and broadcasts updates.
+
+    Supports monitoring pods from any GPU provider that implements
+    the BaseGPUProvider interface.
+    """
 
     POLL_INTERVAL_SECONDS = 5
 
     def __init__(
         self,
-        prime_intellect: PrimeIntellectService,
         pod_cache: TTLCache,
-        update_callback: PodUpdateCallback
+        update_callback: PodUpdateCallback,
+        prime_intellect: PrimeIntellectService | None = None,
+        provider_service: BaseGPUProvider | None = None,
     ):
         """
         Initialize pod monitor.
 
         Args:
-            prime_intellect: Prime Intellect API service
             pod_cache: Cache to clear on updates
             update_callback: Async callback for broadcasting updates
+            prime_intellect: Legacy Prime Intellect API service (deprecated, use provider_service)
+            provider_service: GPU provider service implementing BaseGPUProvider
         """
-        self.pi_service = prime_intellect
+        # Support both old and new interface
+        if provider_service is not None:
+            self.provider = provider_service
+            self.provider_name = provider_service.PROVIDER_NAME
+        elif prime_intellect is not None:
+            # Backwards compatibility
+            self.provider = prime_intellect
+            self.provider_name = "prime_intellect"
+        else:
+            raise ValueError("Either prime_intellect or provider_service must be provided")
+
         self.pod_cache = pod_cache
         self.update_callback = update_callback
         self.monitoring_tasks: dict[str, asyncio.Task] = {}
@@ -44,12 +63,10 @@ class PodMonitor:
         """
         # 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:
         """
@@ -65,7 +82,26 @@ class PodMonitor:
                 await task
             except asyncio.CancelledError:
                 pass
-        print(f"[POD MONITOR] Stopped monitoring pod {pod_id}", file=sys.stderr, flush=True)
+
+    def _normalize_status(self, status: str) -> str:
+        """Normalize status across different providers."""
+        # Common status normalization
+        status_map = {
+            # Common statuses
+            "running": "ACTIVE",
+            "active": "ACTIVE",
+            "ready": "ACTIVE",
+            "starting": "STARTING",
+            "pending": "PENDING",
+            "stopped": "STOPPED",
+            "terminated": "TERMINATED",
+            "error": "ERROR",
+            # Provider-specific
+            "exited": "TERMINATED",
+            "loading": "STARTING",
+            "booting": "STARTING",
+        }
+        return status_map.get(status.lower(), status.upper())
 
     async def _monitor_loop(self, pod_id: str) -> None:
         """
@@ -78,61 +114,53 @@ class PodMonitor:
             while True:
                 try:
                     # Fetch current pod status
-                    pod = await self.pi_service.get_pod(pod_id)
+                    pod = await self.provider.get_pod(pod_id)
 
-                    print(
-                        f"[POD MONITOR] Pod {pod_id} status: {pod.status}",
-                        file=sys.stderr,
-                        flush=True
-                    )
+                    # Normalize the status
+                    normalized_status = self._normalize_status(pod.status)
 
                     # Clear cache to force fresh data
                     self.pod_cache.clear()
 
-                    # Broadcast update
+                    # Broadcast update with provider info
                     await self.update_callback({
                         "type": "pod_status_update",
                         "data": {
                             "pod_id": pod_id,
                             "name": pod.name,
-                            "status": pod.status,
+                            "status": normalized_status,
                             "ssh_connection": pod.sshConnection,
                             "ip": pod.ip,
                             "gpu_name": pod.gpuName,
-                            "price_hr": pod.priceHr
+                            "gpu_count": pod.gpuCount,
+                            "price_hr": pod.priceHr,
+                            "provider": self.provider_name
                         }
                     })
 
                     # 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
-                        )
+                    if normalized_status in {"ERROR", "TERMINATED"}:
                         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
+                    # Note: Modal doesn't support SSH, so we just check for ACTIVE
+                    if normalized_status == "ACTIVE":
+                        if pod.sshConnection or self.provider_name == "modal":
+                            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
-                    )
+                except Exception:
                     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)
+
+    def stop_all(self) -> None:
+        """Stop monitoring all pods."""
+        for pod_id in list(self.monitoring_tasks.keys()):
+            task = self.monitoring_tasks.pop(pod_id, None)
+            if task and not task.done():
+                task.cancel()

morecompute/services/prime_intellect.py

@@ -97,9 +97,7 @@ class PrimeIntellectService:
         """
         Create a new pod
         """
-        import sys
         payload = pod_request.model_dump(exclude_none=True)
-        print(f"[PI SERVICE] Creating pod with payload: {payload}", file=sys.stderr, flush=True)
 
         response = await self._make_request(
             "POST",
@@ -294,9 +292,7 @@ class PrimeIntellectService:
             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)