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

morecompute/execution/worker.py

@@ -150,6 +150,74 @@ def worker_main():
                 exec_count = requested_count - 1
             command_type = msg.get('command_type')
             pub.send_json({'type': 'execution_start', 'cell_index': cell_index, 'execution_count': exec_count + 1})
+
+            # Check if this is a special command (shell command starting with ! or magic command)
+            is_special_cmd = code.strip().startswith('!') or code.strip().startswith('%')
+
+            if is_special_cmd:
+                # Handle special commands on remote worker
+                exec_count += 1
+                status = 'ok'
+                error_payload = None
+                start = time.time()
+
+                try:
+                    import subprocess
+                    import shlex
+
+                    # Strip the ! prefix for shell commands
+                    if code.strip().startswith('!'):
+                        shell_cmd = code.strip()[1:].strip()
+
+                        # Run shell command
+                        process = subprocess.Popen(
+                            ['/bin/bash', '-c', shell_cmd],
+                            stdout=subprocess.PIPE,
+                            stderr=subprocess.PIPE,
+                            text=True
+                        )
+                        stdout, stderr = process.communicate()
+
+                        # Send stdout
+                        if stdout:
+                            pub.send_json({'type': 'stream', 'name': 'stdout', 'text': stdout, 'cell_index': cell_index})
+
+                        # Send stderr
+                        if stderr:
+                            pub.send_json({'type': 'stream', 'name': 'stderr', 'text': stderr, 'cell_index': cell_index})
+
+                        # Check return code
+                        if process.returncode != 0:
+                            status = 'error'
+                            # Only set error if we don't have detailed stderr
+                            if not stderr.strip():
+                                error_payload = {
+                                    'ename': 'ShellCommandError',
+                                    'evalue': f'Command failed with return code {process.returncode}',
+                                    'traceback': [f'Shell command failed: {shell_cmd}']
+                                }
+                    else:
+                        # Magic commands not fully supported on remote yet
+                        status = 'error'
+                        error_payload = {
+                            'ename': 'NotImplementedError',
+                            'evalue': 'Magic commands (%) not yet supported on remote GPU pods',
+                            'traceback': ['Use ! for shell commands instead']
+                        }
+
+                except Exception as exc:
+                    status = 'error'
+                    error_payload = {'ename': type(exc).__name__, 'evalue': str(exc), 'traceback': traceback.format_exc().split('\n')}
+
+                duration_ms = f"{(time.time()-start)*1000:.1f}ms"
+                if error_payload:
+                    pub.send_json({'type': 'execution_error', 'cell_index': cell_index, 'error': error_payload})
+                pub.send_json({'type': 'execution_complete', 'cell_index': cell_index, 'result': {'status': status, 'execution_count': exec_count, 'execution_time': duration_ms, 'outputs': [], 'error': error_payload}})
+                rep.send_json({'ok': True, 'pid': os.getpid()})
+                current_cell = None
+                continue
+
+            # Regular Python code execution
             # Redirect streams
             sf = _StreamForwarder(pub, cell_index)
             old_out, old_err = sys.stdout, sys.stderr
@@ -164,13 +232,6 @@ def worker_main():
             error_payload = None
             start = time.time()
             try:
-                if command_type == 'special':
-                    # This path should be handled in-process; worker only handles python execution
-                    exec_count += 1
-                    pub.send_json({'type': 'execution_complete', 'cell_index': cell_index, 'result': {'status': 'ok', 'execution_count': exec_count, 'execution_time': '0.0ms', 'outputs': [], 'error': None}})
-                    rep.send_json({'ok': True})
-                    current_cell = None
-                    continue
                 compiled = compile(code, '<cell>', 'exec')
                 exec(compiled, g, l)
 

morecompute/models/__init__.py

@@ -0,0 +1,31 @@
+"""Models package for typed API data structures."""
+
+from .api_models import (
+    ApiKeyRequest,
+    ApiKeyResponse,
+    ConfigStatusResponse,
+    DatasetInfoRequest,
+    DatasetCheckRequest,
+    DatasetDiskRequest,
+    PackageInfo,
+    PackagesResponse,
+    EnvironmentInfo,
+    EnvironmentsResponse,
+    FileItem,
+    FileTreeResponse,
+)
+
+__all__ = [
+    "ApiKeyRequest",
+    "ApiKeyResponse",
+    "ConfigStatusResponse",
+    "DatasetInfoRequest",
+    "DatasetCheckRequest",
+    "DatasetDiskRequest",
+    "PackageInfo",
+    "PackagesResponse",
+    "EnvironmentInfo",
+    "EnvironmentsResponse",
+    "FileItem",
+    "FileTreeResponse",
+]

morecompute/models/api_models.py

@@ -0,0 +1,197 @@
+"""Typed models for API requests and responses."""
+
+from pydantic import BaseModel
+from datetime import datetime
+
+
+class ApiKeyRequest(BaseModel):
+    """Request model for setting API key."""
+    api_key: str
+
+
+class ApiKeyResponse(BaseModel):
+    """Response model for API key configuration."""
+    configured: bool
+    message: str | None = None
+
+
+class ConfigStatusResponse(BaseModel):
+    """Response model for configuration status."""
+    configured: bool
+
+
+class DatasetInfoRequest(BaseModel):
+    """Request model for dataset info."""
+    name: str
+    config: str | None = None
+
+
+class DatasetCheckRequest(BaseModel):
+    """Request model for checking dataset load."""
+    name: str
+    config: str | None = None
+    split: str | None = None
+    auto_stream_threshold_gb: float = 10.0
+
+
+class DatasetDiskRequest(BaseModel):
+    """Request model for creating dataset disk."""
+    pod_id: str
+    disk_name: str
+    size_gb: int
+    provider_type: str = "runpod"
+
+
+class PackageInfo(BaseModel):
+    """Model for installed package information."""
+    name: str
+    version: str
+
+
+class PackagesResponse(BaseModel):
+    """Response model for package list."""
+    packages: list[PackageInfo]
+
+
+class EnvironmentInfo(BaseModel):
+    """Model for Python environment information."""
+    name: str
+    path: str
+    version: str
+    is_current: bool
+
+
+class EnvironmentsResponse(BaseModel):
+    """Response model for environment list."""
+    status: str
+    environments: list[EnvironmentInfo]
+    current: EnvironmentInfo | None
+
+
+class FileItem(BaseModel):
+    """Model for file/directory item."""
+    name: str
+    path: str
+    type: str
+    size: int | None = None
+    modified: str | None = None
+
+
+class FileTreeResponse(BaseModel):
+    """Response model for file tree listing."""
+    root: str
+    path: str
+    items: list[FileItem]
+
+
+# ============================================================================
+# Prime Intellect GPU API Models
+# ============================================================================
+
+class EnvVar(BaseModel):
+    """Environment variable for pod configuration."""
+    key: str
+    value: str
+
+
+class PodConfig(BaseModel):
+    """Configuration for creating a GPU pod."""
+    # Required fields
+    name: str
+    cloudId: str
+    gpuType: str
+    socket: str
+    gpuCount: int = 1
+
+    # Optional fields
+    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):
+    """Cloud provider configuration."""
+    type: str = "runpod"
+
+
+class TeamConfig(BaseModel):
+    """Team configuration for shared resources."""
+    teamId: str | None = None
+
+
+class CreatePodRequest(BaseModel):
+    """Request to create a new GPU pod."""
+    pod: PodConfig
+    provider: ProviderConfig
+    team: TeamConfig | None = None
+
+
+class DiskConfig(BaseModel):
+    """Configuration for creating a persistent disk."""
+    name: str
+    size: int  # Size in GB
+    cloudId: str | None = None
+    dataCenterId: str | None = None
+    country: str | None = None
+
+
+class CreateDiskRequest(BaseModel):
+    """Request to create a new disk."""
+    disk: DiskConfig
+    provider: ProviderConfig
+    team: TeamConfig | None = None
+
+
+class DiskResponse(BaseModel):
+    """Response with disk information."""
+    id: str
+    name: str
+    remoteId: str
+    providerType: str
+    status: str
+    size: int
+    createdAt: datetime
+    updatedAt: datetime
+    terminatedAt: datetime | None
+    priceHr: float | None
+    stoppedPriceHr: float | None
+    provisioningPriceHr: float | None
+    userId: str | None
+    teamId: str | None
+    walletId: str | None
+    pods: list[str]
+    clusters: list[str]
+    info: dict[str, object] | None
+
+
+class PodResponse(BaseModel):
+    """Response with GPU pod information."""
+    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):
+    """Query parameters for GPU availability."""
+    regions: list[str] | None = None
+    gpu_count: int | None = None
+    gpu_type: str | None = None
+    security: str | None = None

morecompute/notebook.py

@@ -13,13 +13,43 @@ class Notebook:
             self.load_from_file(file_path)
         else:
             # Default empty notebook structure
-            self.cells.append({'id': self._generate_cell_id(), 'cell_type': 'code', 'source': '', 'outputs': []})
+            self.cells.append({
+                'id': self._generate_cell_id(),
+                'cell_type': 'code',
+                'source': '',
+                'metadata': {},
+                'outputs': [],
+                'execution_count': None
+            })
 
     def get_notebook_data(self) -> Dict[str, Any]:
         return {"cells": self.cells, "metadata": self.metadata, "file_path": self.file_path}
 
-    def add_cell(self, index: int, cell_type: str = 'code', source: str = ''):
-        new_cell = {'id': self._generate_cell_id(), 'cell_type': cell_type, 'source': source, 'outputs': []}
+    def add_cell(self, index: int, cell_type: str = 'code', source: str = '', full_cell: dict = None):
+        if full_cell:
+            actual_cell_type = full_cell.get('cell_type', cell_type)
+            new_cell = {
+                'id': full_cell.get('id', self._generate_cell_id()),
+                'cell_type': actual_cell_type,
+                'source': full_cell.get('source', source),
+                'metadata': full_cell.get('metadata', {}),
+            }
+            # Only add outputs and execution_count for code cells
+            if actual_cell_type == 'code':
+                new_cell['outputs'] = full_cell.get('outputs', [])
+                new_cell['execution_count'] = full_cell.get('execution_count')
+        else:
+            # Normal new cell creation
+            new_cell = {
+                'id': self._generate_cell_id(),
+                'cell_type': cell_type,
+                'source': source,
+                'metadata': {}
+            }
+            # Only add outputs for code cells
+            if cell_type == 'code':
+                new_cell['outputs'] = []
+                new_cell['execution_count'] = None
         self.cells.insert(index, new_cell)
 
     def delete_cell(self, index: int):
@@ -30,10 +60,16 @@ class Notebook:
         if 0 <= index < len(self.cells):
             self.cells[index]['source'] = source
 
+    def move_cell(self, from_index: int, to_index: int):
+        if 0 <= from_index < len(self.cells) and 0 <= to_index < len(self.cells):
+            cell = self.cells.pop(from_index)
+            self.cells.insert(to_index, cell)
+
     def clear_all_outputs(self):
         for cell in self.cells:
-            cell['outputs'] = []
-            if 'execution_count' in cell:
+            # Only clear outputs for code cells
+            if cell.get('cell_type') == 'code':
+                cell['outputs'] = []
                 cell['execution_count'] = None
 
     def to_json(self) -> str:
@@ -64,7 +100,14 @@ class Notebook:
         except (FileNotFoundError, json.JSONDecodeError) as e:
             print(f"Error loading notebook: {e}")
             # Initialize with a default cell if loading fails
-            self.cells = [{'id': self._generate_cell_id(), 'cell_type': 'code', 'source': '', 'outputs': []}]
+            self.cells = [{
+                'id': self._generate_cell_id(),
+                'cell_type': 'code',
+                'source': '',
+                'metadata': {},
+                'outputs': [],
+                'execution_count': None
+            }]
             self.metadata = {}
             self.file_path = file_path
 
@@ -72,7 +115,7 @@ class Notebook:
         path_to_save = file_path or self.file_path
         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())
         self.file_path = path_to_save