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

morecompute/services/lsp_service.py

@@ -0,0 +1,335 @@
+"""
+Language Server Protocol (LSP) service for Python autocomplete.
+Manages Pyright language server for providing IntelliSense features.
+"""
+
+import asyncio
+import json
+import subprocess
+import sys
+from typing import Optional, Dict, Any, List
+from pathlib import Path
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class LSPService:
+    """Manages Pyright language server for Python code intelligence."""
+
+    def __init__(self, workspace_root: Path):
+        self.workspace_root = workspace_root
+        self.process: Optional[subprocess.Popen] = None
+        self.msg_id = 0
+        self.pending_requests: Dict[int, asyncio.Future] = {}
+        self.initialized = False
+        self.documents: Dict[str, str] = {}  # Track open documents
+        self._reader_task: Optional[asyncio.Task] = None
+
+    async def start(self):
+        """Start the Pyright language server process."""
+        try:
+            # Start Pyright in LSP mode
+            self.process = subprocess.Popen(
+                ["pyright-langserver", "--stdio"],
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                bufsize=0,
+            )
+
+            logger.info("Pyright language server started")
+
+            # Start reading responses in background
+            self._reader_task = asyncio.create_task(self._read_responses())
+
+            # Initialize the language server
+            init_result = await self._send_request("initialize", {
+                "processId": None,
+                "rootUri": f"file://{self.workspace_root}",
+                "capabilities": {
+                    "textDocument": {
+                        "completion": {
+                            "completionItem": {
+                                "snippetSupport": True,
+                                "documentationFormat": ["markdown", "plaintext"],
+                                "resolveSupport": {
+                                    "properties": ["documentation", "detail"]
+                                }
+                            },
+                            "contextSupport": True
+                        },
+                        "hover": {
+                            "contentFormat": ["markdown", "plaintext"]
+                        },
+                        "signatureHelp": {
+                            "signatureInformation": {
+                                "documentationFormat": ["markdown", "plaintext"]
+                            }
+                        }
+                    }
+                },
+                "initializationOptions": {
+                    "python": {
+                        "analysis": {
+                            "autoSearchPaths": True,
+                            "useLibraryCodeForTypes": True,
+                            "diagnosticMode": "openFilesOnly"
+                        }
+                    }
+                }
+            })
+
+            # Send initialized notification
+            await self._send_notification("initialized", {})
+            self.initialized = True
+            logger.info("Pyright language server initialized")
+
+        except Exception as e:
+            logger.error(f"Failed to start LSP service: {e}")
+            raise
+
+    async def get_completions(self, cell_id: str, source: str, line: int, character: int) -> List[Dict[str, Any]]:
+        """Get code completions at a specific position."""
+        if not self.initialized:
+            return []
+
+        try:
+            # Create a virtual file URI for the cell
+            file_uri = f"file://{self.workspace_root}/cell_{cell_id}.py"
+
+            # Update the document
+            if file_uri in self.documents:
+                # Document already open, send change notification
+                await self._send_notification("textDocument/didChange", {
+                    "textDocument": {
+                        "uri": file_uri,
+                        "version": self.documents[file_uri]["version"] + 1
+                    },
+                    "contentChanges": [{"text": source}]
+                })
+                self.documents[file_uri]["version"] += 1
+                self.documents[file_uri]["text"] = source
+            else:
+                # Open new document
+                await self._send_notification("textDocument/didOpen", {
+                    "textDocument": {
+                        "uri": file_uri,
+                        "languageId": "python",
+                        "version": 1,
+                        "text": source
+                    }
+                })
+                self.documents[file_uri] = {"version": 1, "text": source}
+
+            # Request completions
+            result = await self._send_request("textDocument/completion", {
+                "textDocument": {"uri": file_uri},
+                "position": {"line": line, "character": character},
+                "context": {"triggerKind": 1}  # Invoked
+            })
+
+            if not result:
+                return []
+
+            # Handle both list and CompletionList formats
+            items = result.get("items", []) if isinstance(result, dict) else result
+
+            return items if isinstance(items, list) else []
+
+        except Exception as e:
+            logger.error(f"Error getting completions: {e}")
+            return []
+
+    async def get_hover(self, cell_id: str, source: str, line: int, character: int) -> Optional[Dict[str, Any]]:
+        """Get hover information at a specific position."""
+        if not self.initialized:
+            return None
+
+        try:
+            file_uri = f"file://{self.workspace_root}/cell_{cell_id}.py"
+
+            # Ensure document is open
+            if file_uri not in self.documents:
+                await self._send_notification("textDocument/didOpen", {
+                    "textDocument": {
+                        "uri": file_uri,
+                        "languageId": "python",
+                        "version": 1,
+                        "text": source
+                    }
+                })
+                self.documents[file_uri] = {"version": 1, "text": source}
+
+            result = await self._send_request("textDocument/hover", {
+                "textDocument": {"uri": file_uri},
+                "position": {"line": line, "character": character}
+            })
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Error getting hover info: {e}")
+            return None
+
+    async def _send_request(self, method: str, params: Dict[str, Any]) -> Any:
+        """Send a JSON-RPC request and wait for response."""
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("LSP process not running")
+
+        self.msg_id += 1
+        msg_id = self.msg_id
+
+        message = {
+            "jsonrpc": "2.0",
+            "id": msg_id,
+            "method": method,
+            "params": params
+        }
+
+        # Create future for response
+        future: asyncio.Future = asyncio.Future()
+        self.pending_requests[msg_id] = future
+
+        # Send message
+        content = json.dumps(message)
+        headers = f"Content-Length: {len(content)}\r\n\r\n"
+        try:
+            self.process.stdin.write((headers + content).encode())
+            self.process.stdin.flush()
+        except Exception as e:
+            del self.pending_requests[msg_id]
+            raise RuntimeError(f"Failed to send LSP request: {e}")
+
+        # Wait for response with timeout
+        try:
+            result = await asyncio.wait_for(future, timeout=5.0)
+            return result
+        except asyncio.TimeoutError:
+            del self.pending_requests[msg_id]
+            logger.warning(f"LSP request timeout for method: {method}")
+            return None
+
+    async def _send_notification(self, method: str, params: Dict[str, Any]):
+        """Send a notification (no response expected)."""
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("LSP process not running")
+
+        message = {
+            "jsonrpc": "2.0",
+            "method": method,
+            "params": params
+        }
+
+        content = json.dumps(message)
+        headers = f"Content-Length: {len(content)}\r\n\r\n"
+        self.process.stdin.write((headers + content).encode())
+        self.process.stdin.flush()
+
+    async def _read_responses(self):
+        """Background task to read LSP responses."""
+        if not self.process or not self.process.stdout:
+            return
+
+        buffer = b""
+
+        try:
+            while self.process.poll() is None:
+                # Read data
+                chunk = await asyncio.get_event_loop().run_in_executor(
+                    None, self.process.stdout.read, 1024
+                )
+
+                if not chunk:
+                    break
+
+                buffer += chunk
+
+                # Process complete messages
+                while b"\r\n\r\n" in buffer:
+                    header_end = buffer.index(b"\r\n\r\n")
+                    headers = buffer[:header_end].decode('utf-8')
+                    buffer = buffer[header_end + 4:]
+
+                    # Parse Content-Length
+                    content_length = 0
+                    for line in headers.split("\r\n"):
+                        if line.startswith("Content-Length:"):
+                            content_length = int(line.split(":")[1].strip())
+                            break
+
+                    # Wait for complete message
+                    while len(buffer) < content_length:
+                        chunk = await asyncio.get_event_loop().run_in_executor(
+                            None, self.process.stdout.read, content_length - len(buffer)
+                        )
+                        if not chunk:
+                            break
+                        buffer += chunk
+
+                    # Parse message
+                    message_data = buffer[:content_length]
+                    buffer = buffer[content_length:]
+
+                    try:
+                        message = json.loads(message_data.decode('utf-8'))
+                        await self._handle_message(message)
+                    except json.JSONDecodeError as e:
+                        logger.error(f"Failed to parse LSP message: {e}")
+
+        except Exception as e:
+            logger.error(f"Error reading LSP responses: {e}")
+
+    async def _handle_message(self, message: Dict[str, Any]):
+        """Handle incoming LSP message."""
+        # Response to our request
+        if "id" in message and message["id"] in self.pending_requests:
+            future = self.pending_requests.pop(message["id"])
+            if "result" in message:
+                future.set_result(message["result"])
+            elif "error" in message:
+                future.set_exception(RuntimeError(message["error"]))
+            else:
+                future.set_result(None)
+
+        # Notification from server (e.g., diagnostics)
+        elif "method" in message:
+            # We can handle server notifications here if needed
+            pass
+
+    async def shutdown(self):
+        """Shutdown the language server."""
+        if not self.initialized:
+            return
+
+        try:
+            # Close all documents
+            for uri in list(self.documents.keys()):
+                await self._send_notification("textDocument/didClose", {
+                    "textDocument": {"uri": uri}
+                })
+
+            # Shutdown
+            await self._send_request("shutdown", {})
+            await self._send_notification("exit", {})
+
+            # Cancel reader task
+            if self._reader_task:
+                self._reader_task.cancel()
+                try:
+                    await self._reader_task
+                except asyncio.CancelledError:
+                    pass
+
+            # Terminate process
+            if self.process:
+                self.process.terminate()
+                try:
+                    self.process.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    self.process.kill()
+
+            logger.info("LSP service shutdown complete")
+
+        except Exception as e:
+            logger.error(f"Error during LSP shutdown: {e}")

morecompute/services/pod_manager.py

@@ -96,6 +96,23 @@ class PodKernelManager:
         """
         import sys
 
+        # Check if already connected to this pod
+        if self.pod and self.pod.id == pod_id:
+            # Check if tunnel is still alive
+            if self.ssh_tunnel_proc and self.ssh_tunnel_proc.poll() is None:
+                return {
+                    "status": "ok",
+                    "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)
@@ -144,12 +161,16 @@ class PodKernelManager:
         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)
         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)
         if tunnel_result.get("status") ==  "error":
             return tunnel_result
 
@@ -159,6 +180,11 @@ class PodKernelManager:
             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}",
@@ -253,7 +279,7 @@ class PodKernelManager:
                 (
                     "cd /tmp && "
                     "tar -xzf morecompute.tar.gz && "
-                    "pip install --quiet pyzmq && "
+                    "pip install --quiet pyzmq matplotlib && "
                     "echo 'Deployment complete'"
                 )
             ])
@@ -315,19 +341,27 @@ class PodKernelManager:
 
             self.ssh_tunnel_proc = subprocess.Popen(
                 tunnel_cmd,
-                stdout=subprocess.DEVNULL,
-                stderr=subprocess.DEVNULL
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE
             )
 
             # Wait briefly for tunnel to establish
             await asyncio.sleep(2)
 
+            # Check if process is still running
+            if self.ssh_tunnel_proc is None:
+                return {
+                    "status": "error",
+                    "message": "SSH tunnel process is None"
+                }
             if self.ssh_tunnel_proc.poll() is not None:
+                # Process died, get error output
+                stdout, stderr = self.ssh_tunnel_proc.communicate()
+                error_msg = stderr.decode('utf-8') if stderr else "No error output"
                 return {
                     "status": "error",
-                    "message": "SSH tunnel failed to establish"
+                    "message": f"SSH tunnel failed to establish: {error_msg}"
                 }
-
             return {
                 "status": "ok",
                 "message": "SSH tunnel created",
@@ -335,6 +369,7 @@ 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)}"
@@ -352,7 +387,10 @@ 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
             ssh_key = self._get_ssh_key()
             worker_cmd = ["ssh", "-p", ssh_port]
 
@@ -365,14 +403,14 @@ class PodKernelManager:
                 "-o", "BatchMode=yes",
                 "-o", "ConnectTimeout=10",
                 f"root@{ssh_host}",
+                "sh", "-c",
                 (
-                    f"cd /tmp && "
-                    f"export MC_ZMQ_CMD_ADDR=tcp://0.0.0.0:{self.remote_cmd_port} && "
-                    f"export MC_ZMQ_PUB_ADDR=tcp://0.0.0.0:{self.remote_pub_port} && "
-                    f"export PYTHONPATH=/tmp:$PYTHONPATH && "
-                    f"nohup {sys.executable} -m morecompute.execution.worker "
-                    f"> /tmp/worker.log 2>&1 & "
-                    f"echo $!"
+                    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"nohup python3 /tmp/morecompute/execution/worker.py "
+                    f">/tmp/worker.log 2>&1 </dev/null & "
+                    f"echo $!'"
                 )
             ])
 
@@ -390,10 +428,13 @@ 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",
@@ -458,6 +499,64 @@ class PodKernelManager:
             "messages": messages
         }
 
+    async def execute_ssh_command(self, command: str) -> tuple[str, str, int]:
+        """
+        Execute a command on the remote pod via SSH.
+
+        args:
+            command: The command to execute
+
+        returns:
+            tuple of (stdout, stderr, return_code)
+        """
+        if not self.pod or not self.pod.sshConnection:
+            raise RuntimeError("No active pod connection")
+
+        # Parse SSH connection string to get host and port
+        ssh_parts = self.pod.sshConnection.split()
+        host_part = None
+        port = "22"  # default SSH port
+
+        for part in ssh_parts:
+            if "@" in part:
+                host_part = part
+            if part == "-p" and ssh_parts.index(part) + 1 < len(ssh_parts):
+                port = ssh_parts[ssh_parts.index(part) + 1]
+
+        if not host_part:
+            raise RuntimeError(f"Invalid SSH connection format: {self.pod.sshConnection}")
+
+        # Get SSH key
+        ssh_key = self._get_ssh_key()
+        if not ssh_key:
+            raise RuntimeError("SSH key not found. Please configure MORECOMPUTE_SSH_KEY or add key to ~/.ssh/")
+
+        # Build SSH command
+        ssh_cmd = [
+            "ssh",
+            "-i", ssh_key,
+            "-p", port,
+            "-o", "StrictHostKeyChecking=no",
+            "-o", "UserKnownHostsFile=/dev/null",
+            "-o", "LogLevel=ERROR",
+            host_part,
+            command
+        ]
+
+        # Execute command
+        proc = await asyncio.create_subprocess_exec(
+            *ssh_cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+
+        stdout, stderr = await proc.communicate()
+        return (
+            stdout.decode('utf-8', errors='replace'),
+            stderr.decode('utf-8', errors='replace'),
+            proc.returncode or 0
+        )
+
     async def get_status(self) -> dict[str, object]:
         """
         Get current connection status.
@@ -465,7 +564,10 @@ class PodKernelManager:
         returns:
             dict with status information
         """
-        if not self.pod:
+        # Cache pod reference to avoid race condition with disconnect()
+        pod = self.pod
+
+        if not pod:
             return {
                 "connected": False,
                 "pod": None
@@ -478,7 +580,7 @@ class PodKernelManager:
 
         # Get updated pod info
         try:
-            updated_pod = await self.pi_service.get_pod(self.pod.id)
+            updated_pod = await self.pi_service.get_pod(pod.id)
             pod_status = updated_pod.status
         except Exception:
             pod_status = "unknown"
@@ -486,13 +588,13 @@ class PodKernelManager:
         return {
             "connected": True,
             "pod": {
-                "id": self.pod.id,
-                "name": self.pod.name,
+                "id": pod.id,
+                "name": pod.name,
                 "status": pod_status,
-                "gpu_type": self.pod.gpuName,
-                "gpu_count": self.pod.gpuCount,
-                "price_hr": self.pod.priceHr,
-                "ssh_connection": self.pod.sshConnection
+                "gpu_type": pod.gpuName,
+                "gpu_count": pod.gpuCount,
+                "price_hr": pod.priceHr,
+                "ssh_connection": pod.sshConnection
             },
             "tunnel": {
                 "alive": tunnel_alive,