more-compute 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

morecompute/notebook.py

@@ -1,6 +1,8 @@
 import json
+from pathlib import Path
 from typing import List, Dict, Any
 from uuid import uuid4
+from .utils.py_percent_parser import parse_py_percent, generate_py_percent
 
 class Notebook:
     """Manages the state of a notebook's cells."""
@@ -71,6 +73,8 @@ class Notebook:
             if cell.get('cell_type') == 'code':
                 cell['outputs'] = []
                 cell['execution_count'] = None
+                cell['execution_time'] = None
+                cell['error'] = None
 
     def to_json(self) -> str:
         # Basic notebook format
@@ -84,10 +88,18 @@ class Notebook:
 
     def load_from_file(self, file_path: str):
         try:
-            with open(file_path, 'r') as f:
-                data = json.load(f)
+            path = Path(file_path)
+
+            # Check file extension
+            if path.suffix == '.py':
+                # Load .py file with py:percent format
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+
+                data = parse_py_percent(content)
                 loaded_cells = data.get('cells', [])
-                # Ensure stable IDs for all cells (back-compat for notebooks without IDs)
+
+                # Ensure stable IDs for all cells
                 self.cells = []
                 for cell in loaded_cells:
                     if not isinstance(cell, dict):
@@ -95,9 +107,40 @@ class Notebook:
                     if 'id' not in cell or not cell['id']:
                         cell['id'] = self._generate_cell_id()
                     self.cells.append(cell)
+
                 self.metadata = data.get('metadata', {})
                 self.file_path = file_path
-        except (FileNotFoundError, json.JSONDecodeError) as e:
+
+            elif path.suffix == '.ipynb':
+                # Block .ipynb files with helpful error
+                raise ValueError(
+                    f"MoreCompute only supports .py notebooks.\n\n"
+                    f"Convert your notebook with:\n"
+                    f"  more-compute convert {path.name} -o {path.stem}.py\n\n"
+                    f"Then open with:\n"
+                    f"  more-compute {path.stem}.py"
+                )
+
+            else:
+                raise ValueError(f"Unsupported file format: {path.suffix}. Use .py files.")
+
+        except FileNotFoundError as e:
+            print(f"Error: File not found: {e}")
+            # Initialize with a default cell if loading fails
+            self.cells = [{
+                'id': self._generate_cell_id(),
+                'cell_type': 'code',
+                'source': '',
+                'metadata': {},
+                'outputs': [],
+                'execution_count': None
+            }]
+            self.metadata = {}
+            self.file_path = file_path
+        except ValueError as e:
+            # Re-raise validation errors (like .ipynb block)
+            raise
+        except Exception as e:
             print(f"Error loading notebook: {e}")
             # Initialize with a default cell if loading fails
             self.cells = [{
@@ -116,8 +159,24 @@ class Notebook:
         if not path_to_save:
             raise ValueError("No file path specified for saving.")
 
-        with open(path_to_save, 'w') as f:
-            f.write(self.to_json())
+        path = Path(path_to_save)
+
+        # Save in appropriate format based on extension
+        if path.suffix == '.py':
+            # Save as py:percent format
+            content = generate_py_percent(self.cells)
+            with open(path_to_save, 'w', encoding='utf-8') as f:
+                f.write(content)
+        elif path.suffix == '.ipynb':
+            # Block saving as .ipynb
+            raise ValueError("MoreCompute only supports .py notebooks. Use .py extension.")
+        else:
+            # Default to .py if no extension
+            path_to_save = str(path.with_suffix('.py'))
+            content = generate_py_percent(self.cells)
+            with open(path_to_save, 'w', encoding='utf-8') as f:
+                f.write(content)
+
         self.file_path = path_to_save
 
     def _generate_cell_id(self) -> str:

morecompute/server.py

@@ -18,7 +18,7 @@ from .utils.system_environment_util import DeviceMetrics
 from .utils.error_utils import ErrorUtils
 from .utils.cache_util import make_cache_key
 from .utils.notebook_util import coerce_cell_source
-from .utils.config_util import load_api_key_from_env, save_api_key_to_env
+from .utils.config_util import load_api_key, save_api_key
 from .utils.zmq_util import reconnect_zmq_sockets, reset_to_local_zmq
 from .services.prime_intellect import PrimeIntellectService
 from .services.pod_manager import PodKernelManager
@@ -67,7 +67,7 @@ else:
 error_utils = ErrorUtils()
 executor = NextZmqExecutor(error_utils=error_utils)
 metrics = DeviceMetrics()
-prime_api_key = load_api_key_from_env("PRIME_INTELLECT_API_KEY", BASE_DIR / ".env")
+prime_api_key = load_api_key("PRIME_INTELLECT_API_KEY")
 prime_intellect = PrimeIntellectService(api_key=prime_api_key) if prime_api_key else None
 pod_manager: PodKernelManager | None = None
 data_manager = DataManager(prime_intellect=prime_intellect)
@@ -99,7 +99,23 @@ async def startup_event():
 @app.on_event("shutdown")
 async def shutdown_event():
     """Cleanup services on shutdown."""
-    global lsp_service
+    global lsp_service, executor
+
+    # Shutdown executor and worker process
+    if executor and executor.worker_proc:
+        try:
+            print("[EXECUTOR] Shutting down worker process...", file=sys.stderr, flush=True)
+            executor.worker_proc.terminate()
+            executor.worker_proc.wait(timeout=2)
+            print("[EXECUTOR] Worker process shutdown complete", file=sys.stderr, flush=True)
+        except Exception as e:
+            print(f"[EXECUTOR] Error during worker shutdown, forcing kill: {e}", file=sys.stderr, flush=True)
+            try:
+                executor.worker_proc.kill()
+            except Exception:
+                pass
+
+    # Shutdown LSP service
     if lsp_service:
         try:
             await lsp_service.shutdown()
@@ -441,12 +457,34 @@ class WebSocketManager:
 
     async def handle_message_loop(self, websocket: WebSocket):
         """Main loop to handle incoming WebSocket messages."""
+        tasks = set()
+
+        def task_done_callback(task):
+            tasks.discard(task)
+            # Check for exceptions in completed tasks
+            try:
+                exc = task.exception()
+                if exc:
+                    print(f"[SERVER] Task raised exception: {exc}", file=sys.stderr, flush=True)
+                    import traceback
+                    traceback.print_exception(type(exc), exc, exc.__traceback__)
+            except asyncio.CancelledError:
+                pass
+            except Exception as e:
+                print(f"[SERVER] Error in task_done_callback: {e}", file=sys.stderr, flush=True)
+
         while True:
             try:
                 message = await websocket.receive_json()
-                await self._handle_message(websocket, message)
+                # Process messages concurrently so interrupts can arrive during execution
+                task = asyncio.create_task(self._handle_message(websocket, message))
+                tasks.add(task)
+                task.add_done_callback(task_done_callback)
             except WebSocketDisconnect:
                 self.disconnect(websocket)
+                # Cancel all pending tasks
+                for task in tasks:
+                    task.cancel()
                 break
             except Exception as e:
                 await self._send_error(websocket, f"Unhandled error: {e}")
@@ -534,12 +572,24 @@ class WebSocketManager:
         else:
             # Normal add cell
             self.notebook.add_cell(index=index, cell_type=cell_type, source=source)
+
+        # Save the notebook after adding cell
+        try:
+            self.notebook.save_to_file()
+        except Exception as e:
+            print(f"Warning: Failed to save notebook after adding cell: {e}", file=sys.stderr)
+
         await self.broadcast_notebook_update()
 
     async def _handle_delete_cell(self, websocket: WebSocket, data: dict):
         index = data.get('cell_index')
         if index is not None:
             self.notebook.delete_cell(index)
+            # Save the notebook after deleting cell
+            try:
+                self.notebook.save_to_file()
+            except Exception as e:
+                print(f"Warning: Failed to save notebook after deleting cell: {e}", file=sys.stderr)
             await self.broadcast_notebook_update()
 
     async def _handle_update_cell(self, websocket: WebSocket, data: dict):
@@ -587,49 +637,30 @@ class WebSocketManager:
         print(f"[SERVER] Interrupt request received for cell {cell_index}", file=sys.stderr, flush=True)
 
         # Perform the interrupt (this may take up to 1 second)
+        # The execution handler will send the appropriate error and completion messages
         await self.executor.interrupt_kernel(cell_index=cell_index)
 
-        print(f"[SERVER] Interrupt completed, sending error message", file=sys.stderr, flush=True)
+        print(f"[SERVER] Interrupt completed, execution handler will send completion messages", file=sys.stderr, flush=True)
 
-        # Inform all clients that the currently running cell (if any) is interrupted
-        try:
-            await websocket.send_json({
-                "type": "execution_error",
-                "data": {
-                    "cell_index": cell_index,
-                    "error": {
-                        "output_type": "error",
-                        "ename": "KeyboardInterrupt",
-                        "evalue": "Execution interrupted by user",
-                        "traceback": ["KeyboardInterrupt: Execution was stopped by user"]
-                    }
-                }
-            })
-            await websocket.send_json({
-                "type": "execution_complete",
-                "data": {
-                    "cell_index": cell_index,
-                    "result": {
-                        "status": "error",
-                        "execution_count": None,
-                        "execution_time": "interrupted",
-                        "outputs": [],
-                        "error": {
-                            "output_type": "error",
-                            "ename": "KeyboardInterrupt",
-                            "evalue": "Execution interrupted by user",
-                            "traceback": ["KeyboardInterrupt: Execution was stopped by user"]
-                        }
-                    }
-                }
-            })
-            print(f"[SERVER] Error messages sent for cell {cell_index}", file=sys.stderr, flush=True)
-        except Exception as e:
-            print(f"[SERVER] Failed to send error messages: {e}", file=sys.stderr, flush=True)
+        # Note: We don't send completion messages here anymore because:
+        # 1. For shell commands: AsyncSpecialCommandHandler._execute_shell_command sends them
+        # 2. For Python code: The worker sends them
+        # Sending duplicate messages causes the frontend to get confused
 
     async def _handle_reset_kernel(self, websocket: WebSocket, data: dict):
+        import sys
+        print(f"[SERVER] Resetting kernel", file=sys.stderr, flush=True)
         self.executor.reset_kernel()
         self.notebook.clear_all_outputs()
+
+        # Note: We don't save the notebook here - this preserves execution times
+        # from the last session, which is useful for seeing how long things took
+
+        # Broadcast kernel restart to all clients
+        await self.broadcast_pod_update({
+            "type": "kernel_restarted",
+            "data": {}
+        })
         await self.broadcast_notebook_update()
 
     async def _send_error(self, websocket: WebSocket, error_message: str):
@@ -654,14 +685,14 @@ async def get_gpu_config() -> ConfigStatusResponse:
 
 @app.post("/api/gpu/config", response_model=ApiKeyResponse)
 async def set_gpu_config(request: ApiKeyRequest) -> ApiKeyResponse:
-    """Save Prime Intellect API key to .env file and reinitialize service."""
+    """Save Prime Intellect API key to user config (~/.morecompute/config.json) and reinitialize service."""
     global prime_intellect, pod_monitor
 
     if not request.api_key.strip():
         raise HTTPException(status_code=400, detail="API key is required")
 
     try:
-        save_api_key_to_env("PRIME_INTELLECT_API_KEY", request.api_key, BASE_DIR / ".env")
+        save_api_key("PRIME_INTELLECT_API_KEY", request.api_key)
         prime_intellect = PrimeIntellectService(api_key=request.api_key)
         if prime_intellect:
             pod_monitor = PodMonitor(
@@ -811,9 +842,10 @@ async def _connect_to_pod_background(pod_id: str):
             reconnect_zmq_sockets(
                 executor,
                 cmd_addr=addresses["cmd_addr"],
-                pub_addr=addresses["pub_addr"]
+                pub_addr=addresses["pub_addr"],
+                is_remote=True  # Critical: Tell executor this is a remote worker
             )
-            print(f"[CONNECT BACKGROUND] Successfully connected to pod {pod_id}", file=sys.stderr, flush=True)
+            print(f"[CONNECT BACKGROUND] Successfully connected to pod {pod_id}, executor.is_remote=True", file=sys.stderr, flush=True)
         else:
             # Connection failed - clean up
             print(f"[CONNECT BACKGROUND] Failed to connect: {result}", file=sys.stderr, flush=True)

morecompute/services/pod_manager.py

@@ -408,8 +408,8 @@ class PodKernelManager:
                     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"setsid python3 -u /tmp/morecompute/execution/worker.py "
+                    f"</dev/null >/tmp/worker.log 2>&1 & "
                     f"echo $!'"
                 )
             ])

morecompute/utils/cell_magics.py

@@ -248,6 +248,11 @@ class CellMagicHandlers:
                     env=env
                 )
 
+                # Track process for interrupt handling
+                if hasattr(cell_magic_handler, 'special_handler'):
+                    cell_magic_handler.special_handler.current_process_sync = process
+                    print(f"[CELL_MAGIC] Tracking sync subprocess PID={process.pid}", file=sys.stderr, flush=True)
+
                 # Read and print output line by line (real-time streaming)
                 def read_stream(stream, output_type):
                     """Read stream line by line and print immediately"""
@@ -276,12 +281,38 @@ class CellMagicHandlers:
                 stdout_thread.start()
                 stderr_thread.start()
 
-                # Wait for process to complete
-                return_code = process.wait()
+                # Wait for process to complete, checking if it was killed
+                try:
+                    # Poll with timeout to detect if process was killed externally
+                    while process.poll() is None:
+                        try:
+                            process.wait(timeout=0.1)
+                        except subprocess.TimeoutExpired:
+                            # Check if interrupted
+                            if hasattr(cell_magic_handler, 'special_handler'):
+                                if cell_magic_handler.special_handler.sync_interrupted:
+                                    # Process was killed by interrupt handler
+                                    print(f"[CELL_MAGIC] Process was interrupted, raising KeyboardInterrupt", file=sys.stderr, flush=True)
+                                    raise KeyboardInterrupt("Execution interrupted by user")
+
+                    return_code = process.returncode
+                except KeyboardInterrupt:
+                    # Kill process if KeyboardInterrupt
+                    try:
+                        process.kill()
+                        process.wait()
+                    except Exception:
+                        pass
+                    raise
 
                 # Wait for output threads to finish
-                stdout_thread.join()
-                stderr_thread.join()
+                stdout_thread.join(timeout=1)
+                stderr_thread.join(timeout=1)
+
+                # Clear process reference
+                if hasattr(cell_magic_handler, 'special_handler'):
+                    cell_magic_handler.special_handler.current_process_sync = None
+                    print(f"[CELL_MAGIC] Cleared sync subprocess reference", file=sys.stderr, flush=True)
 
                 return return_code
 

morecompute/utils/config_util.py

@@ -1,45 +1,67 @@
 """Configuration utilities for managing API keys and environment variables."""
 
 from pathlib import Path
+from typing import Optional
 import os
+import json
 
 
-def load_api_key_from_env(env_var: str, env_file_path: Path | None = None) -> str | None:
+# Global config directory in user's home
+CONFIG_DIR = Path.home() / ".morecompute"
+CONFIG_FILE = CONFIG_DIR / "config.json"
+
+
+def _ensure_config_dir() -> None:
+    """Ensure the config directory exists."""
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def _load_config() -> dict:
+    """Load config from JSON file."""
+    if not CONFIG_FILE.exists():
+        return {}
+    try:
+        with CONFIG_FILE.open("r", encoding="utf-8") as f:
+            return json.load(f)
+    except (json.JSONDecodeError, IOError):
+        return {}
+
+
+def _save_config(config: dict) -> None:
+    """Save config to JSON file."""
+    _ensure_config_dir()
+    with CONFIG_FILE.open("w", encoding="utf-8") as f:
+        json.dump(config, f, indent=2)
+
+
+def load_api_key(key_name: str) -> Optional[str]:
     """
-    Load API key from environment variable or .env file.
+    Load API key from user config directory (~/.morecompute/config.json).
+    Falls back to environment variable if not found in config.
 
     Args:
-        env_var: Environment variable name to check
-        env_file_path: Path to .env file (optional)
+        key_name: Key name (e.g., "PRIME_INTELLECT_API_KEY")
 
     Returns:
         API key string or None if not found
     """
-    api_key = os.getenv(env_var)
-    if api_key:
-        return api_key
+    # Check environment variable first
+    env_key = os.getenv(key_name)
+    if env_key:
+        return env_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
+    # Check config file
+    config = _load_config()
+    return config.get(key_name)
 
-    return None
 
-
-def save_api_key_to_env(env_var: str, api_key: str, env_file_path: Path) -> None:
+def save_api_key(key_name: str, api_key: str) -> None:
     """
-    Save API key to .env file, replacing existing value if present.
+    Save API key to user config directory (~/.morecompute/config.json).
 
     Args:
-        env_var: Environment variable name
+        key_name: Key name (e.g., "PRIME_INTELLECT_API_KEY")
         api_key: API key value to save
-        env_file_path: Path to .env file
 
     Raises:
         ValueError: If API key is empty
@@ -48,12 +70,6 @@ def save_api_key_to_env(env_var: str, api_key: str, env_file_path: Path) -> None
     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)
+    config = _load_config()
+    config[key_name] = api_key
+    _save_config(config)

morecompute/utils/notebook_converter.py

@@ -0,0 +1,129 @@
+"""Converter utilities for notebook formats."""
+
+import json
+import re
+from pathlib import Path
+from typing import List, Set
+from .py_percent_parser import generate_py_percent, parse_py_percent
+
+
+def extract_pip_dependencies(notebook_data: dict) -> Set[str]:
+    """
+    Extract package names from !pip install and %pip install commands.
+
+    Args:
+        notebook_data: Parsed notebook JSON
+
+    Returns:
+        Set of package names
+    """
+    packages = set()
+
+    for cell in notebook_data.get('cells', []):
+        if cell.get('cell_type') != 'code':
+            continue
+
+        source = cell.get('source', [])
+        if isinstance(source, list):
+            source = ''.join(source)
+
+        # Match: !pip install package1 package2
+        # Match: %pip install package1 package2
+        pip_pattern = r'[!%]pip\s+install\s+([^\n]+)'
+        matches = re.finditer(pip_pattern, source)
+
+        for match in matches:
+            install_line = match.group(1)
+            # Remove common flags
+            install_line = re.sub(r'--[^\s]+\s*', '', install_line)
+            install_line = re.sub(r'-[qU]\s*', '', install_line)
+
+            # Extract package names (handle package==version format)
+            parts = install_line.split()
+            for part in parts:
+                part = part.strip()
+                if part and not part.startswith('-'):
+                    packages.add(part)
+
+    return packages
+
+
+def convert_ipynb_to_py(ipynb_path: Path, output_path: Path, include_uv_deps: bool = True) -> None:
+    """
+    Convert .ipynb notebook to .py format with py:percent cell markers.
+
+    Args:
+        ipynb_path: Path to input .ipynb file
+        output_path: Path to output .py file
+        include_uv_deps: Whether to add UV inline script dependencies
+    """
+    # Read notebook
+    with open(ipynb_path, 'r', encoding='utf-8') as f:
+        notebook_data = json.load(f)
+
+    cells = notebook_data.get('cells', [])
+
+    # Generate UV dependencies header if requested
+    header_lines = []
+    if include_uv_deps:
+        dependencies = extract_pip_dependencies(notebook_data)
+        if dependencies:
+            header_lines.append('# /// script')
+            header_lines.append('# dependencies = [')
+            for dep in sorted(dependencies):
+                header_lines.append(f'#   "{dep}",')
+            header_lines.append('# ]')
+            header_lines.append('# ///')
+            header_lines.append('')
+
+    # Generate py:percent format
+    py_content = generate_py_percent(cells)
+
+    # Combine header and content
+    if header_lines:
+        final_content = '\n'.join(header_lines) + '\n' + py_content
+    else:
+        final_content = py_content
+
+    # Write output
+    with open(output_path, 'w', encoding='utf-8') as f:
+        f.write(final_content)
+
+    print(f"✓ Converted {ipynb_path.name} → {output_path.name}")
+
+    # Show dependencies if found
+    if include_uv_deps and dependencies:
+        print(f"  Found dependencies: {', '.join(sorted(dependencies))}")
+        print(f"  Run with: more-compute {output_path.name}")
+
+
+def convert_py_to_ipynb(py_path: Path, output_path: Path) -> None:
+    """
+    Convert .py notebook to .ipynb format.
+
+    Args:
+        py_path: Path to input .py file
+        output_path: Path to output .ipynb file
+    """
+    # Read .py file
+    with open(py_path, 'r', encoding='utf-8') as f:
+        py_content = f.read()
+
+    # Parse py:percent format to notebook structure
+    notebook_data = parse_py_percent(py_content)
+
+    # Ensure source is in list format (Jupyter notebook standard)
+    for cell in notebook_data.get('cells', []):
+        source = cell.get('source', '')
+        if isinstance(source, str):
+            # Split into lines and keep newlines (Jupyter format)
+            lines = source.split('\n')
+            # Add \n to each line except the last
+            cell['source'] = [line + '\n' for line in lines[:-1]] + ([lines[-1]] if lines[-1] else [])
+
+    # Write .ipynb file
+    with open(output_path, 'w', encoding='utf-8') as f:
+        json.dump(notebook_data, f, indent=1, ensure_ascii=False)
+
+    print(f"Converted {py_path.name} -> {output_path.name}")
+    print(f"  Upload to Google Colab or open in Jupyter")