experimaestro 2.0.0a8__py3-none-any.whl → 2.0.0b8__py3-none-any.whl

experimaestro/scheduler/remote/server.py

@@ -0,0 +1,423 @@
+"""SSH State Provider Server
+
+JSON-RPC server that wraps WorkspaceStateProvider and communicates via stdio.
+Designed to be run over SSH for remote experiment monitoring.
+
+Usage:
+    experimaestro experiments monitor-server --workdir /path/to/workspace
+"""
+
+import logging
+import sys
+import threading
+from pathlib import Path
+from typing import IO, Callable, Dict, Optional
+
+from experimaestro.scheduler.state_provider import (
+    WorkspaceStateProvider,
+    StateEvent,
+    StateEventType,
+)
+from experimaestro.scheduler.remote.protocol import (
+    RPCMethod,
+    NotificationMethod,
+    parse_message,
+    create_success_response,
+    create_error_response,
+    create_notification,
+    serialize_job,
+    serialize_experiment,
+    serialize_run,
+    serialize_datetime,
+    deserialize_datetime,
+    PARSE_ERROR,
+    METHOD_NOT_FOUND,
+    INVALID_PARAMS,
+    INTERNAL_ERROR,
+    WORKSPACE_NOT_FOUND,
+)
+
+logger = logging.getLogger("xpm.remote.server")
+
+
+class SSHStateProviderServer:
+    """JSON-RPC server that wraps WorkspaceStateProvider for SSH-based monitoring
+
+    This server reads JSON-RPC requests from stdin and writes responses to stdout.
+    It registers as a listener with the WorkspaceStateProvider to push notifications
+    when state changes occur.
+
+    Thread safety:
+    - Writes to stdout are serialized with a lock
+    - The main read loop runs in the calling thread
+    - Event notifications may come from the state provider's change detector thread
+    """
+
+    def __init__(
+        self,
+        workspace_path: Path,
+        stdin: IO[bytes] = None,
+        stdout: IO[bytes] = None,
+    ):
+        """Initialize the server
+
+        Args:
+            workspace_path: Path to the workspace directory
+            stdin: Input stream for reading requests (default: sys.stdin.buffer)
+            stdout: Output stream for writing responses (default: sys.stdout.buffer)
+        """
+        self.workspace_path = workspace_path
+        self.stdin = stdin if stdin is not None else sys.stdin.buffer
+        self.stdout = stdout if stdout is not None else sys.stdout.buffer
+        self._state_provider: Optional[WorkspaceStateProvider] = None
+        self._running = False
+        self._write_lock = threading.Lock()
+
+        # Map of method names to handler functions
+        self._handlers: Dict[str, Callable] = {
+            RPCMethod.GET_EXPERIMENTS.value: self._handle_get_experiments,
+            RPCMethod.GET_EXPERIMENT.value: self._handle_get_experiment,
+            RPCMethod.GET_EXPERIMENT_RUNS.value: self._handle_get_experiment_runs,
+            RPCMethod.GET_JOBS.value: self._handle_get_jobs,
+            RPCMethod.GET_JOB.value: self._handle_get_job,
+            RPCMethod.GET_ALL_JOBS.value: self._handle_get_all_jobs,
+            RPCMethod.GET_SERVICES.value: self._handle_get_services,
+            RPCMethod.KILL_JOB.value: self._handle_kill_job,
+            RPCMethod.CLEAN_JOB.value: self._handle_clean_job,
+            RPCMethod.GET_SYNC_INFO.value: self._handle_get_sync_info,
+        }
+
+    def start(self):
+        """Start the server and begin processing requests
+
+        This method blocks until the server is stopped or stdin is closed.
+        """
+        # Verify workspace exists
+        if not self.workspace_path.exists():
+            logger.error("Workspace path does not exist: %s", self.workspace_path)
+            self._send_error_and_exit(
+                WORKSPACE_NOT_FOUND,
+                f"Workspace path does not exist: {self.workspace_path}",
+            )
+            return
+
+        # Initialize state provider in read-only mode
+        try:
+            self._state_provider = WorkspaceStateProvider.get_instance(
+                self.workspace_path,
+                read_only=True,
+                sync_on_start=True,
+            )
+        except Exception as e:
+            logger.exception("Failed to initialize state provider")
+            self._send_error_and_exit(INTERNAL_ERROR, f"Failed to initialize: {e}")
+            return
+
+        # Register as listener for state changes
+        self._state_provider.add_listener(self._on_state_event)
+
+        self._running = True
+        logger.info("SSH State Provider Server started for %s", self.workspace_path)
+
+        try:
+            self._read_loop()
+        finally:
+            self.stop()
+
+    def stop(self):
+        """Stop the server and clean up resources"""
+        self._running = False
+
+        # Unregister listener
+        if self._state_provider is not None:
+            try:
+                self._state_provider.remove_listener(self._on_state_event)
+            except Exception:
+                pass
+
+        # Send shutdown notification
+        try:
+            self._send_notification(
+                NotificationMethod.SHUTDOWN, {"reason": "server_shutdown"}
+            )
+        except Exception:
+            pass
+
+        logger.info("SSH State Provider Server stopped")
+
+    def _read_loop(self):
+        """Main loop: read JSON-RPC requests from stdin and process them"""
+        while self._running:
+            try:
+                line = self.stdin.readline()
+                if not line:
+                    # EOF - stdin closed
+                    logger.debug("stdin closed, stopping server")
+                    break
+
+                line_str = line.decode("utf-8").strip()
+                if not line_str:
+                    continue
+
+                self._process_request(line_str)
+
+            except Exception as e:
+                logger.exception("Error in read loop: %s", e)
+                # Continue processing - don't crash on individual request errors
+
+    def _process_request(self, line: str):
+        """Process a single JSON-RPC request"""
+        try:
+            msg = parse_message(line)
+        except ValueError as e:
+            self._send_response(create_error_response(0, PARSE_ERROR, str(e)))
+            return
+
+        # We only handle requests (with id), not responses or notifications
+        from experimaestro.scheduler.remote.protocol import RPCRequest
+
+        if not isinstance(msg, RPCRequest):
+            logger.warning("Received non-request message: %s", type(msg).__name__)
+            return
+
+        request = msg
+        method = request.method
+        params = request.params
+        request_id = request.id
+
+        if request_id is None:
+            # Notification from client - we don't handle these currently
+            logger.debug("Received notification: %s", method)
+            return
+
+        # Dispatch to handler
+        handler = self._handlers.get(method)
+        if handler is None:
+            self._send_response(
+                create_error_response(
+                    request_id, METHOD_NOT_FOUND, f"Unknown method: {method}"
+                )
+            )
+            return
+
+        try:
+            result = handler(params)
+            self._send_response(create_success_response(request_id, result))
+        except TypeError as e:
+            self._send_response(
+                create_error_response(request_id, INVALID_PARAMS, str(e))
+            )
+        except Exception as e:
+            logger.exception("Error handling %s", method)
+            self._send_response(
+                create_error_response(request_id, INTERNAL_ERROR, str(e))
+            )
+
+    def _send_response(self, response: str):
+        """Send a JSON-RPC response (thread-safe)"""
+        with self._write_lock:
+            self.stdout.write((response + "\n").encode("utf-8"))
+            self.stdout.flush()
+
+    def _send_notification(self, method: NotificationMethod, params: Dict):
+        """Send a JSON-RPC notification (thread-safe)"""
+        notification = create_notification(method, params)
+        self._send_response(notification)
+
+    def _send_error_and_exit(self, code: int, message: str):
+        """Send an error notification and exit"""
+        self._send_notification(
+            NotificationMethod.SHUTDOWN,
+            {"reason": "error", "code": code, "message": message},
+        )
+
+    def _on_state_event(self, event: StateEvent):
+        """Handle state change events from the state provider
+
+        Converts events to JSON-RPC notifications and sends them to the client.
+        """
+        try:
+            if event.event_type == StateEventType.EXPERIMENT_UPDATED:
+                self._send_notification(
+                    NotificationMethod.EXPERIMENT_UPDATED,
+                    {
+                        "experiment_id": event.data.get("experiment_id"),
+                        "data": event.data,
+                    },
+                )
+            elif event.event_type == StateEventType.RUN_UPDATED:
+                self._send_notification(
+                    NotificationMethod.RUN_UPDATED,
+                    {
+                        "experiment_id": event.data.get("experiment_id"),
+                        "run_id": event.data.get("run_id"),
+                        "data": event.data,
+                    },
+                )
+            elif event.event_type == StateEventType.JOB_UPDATED:
+                self._send_notification(
+                    NotificationMethod.JOB_UPDATED,
+                    {
+                        "job_id": event.data.get("job_id"),
+                        "experiment_id": event.data.get("experiment_id"),
+                        "run_id": event.data.get("run_id"),
+                        "state": event.data.get("state"),
+                        "data": event.data,
+                    },
+                )
+                # Also send file_changed notification for job metadata
+                if "path" in event.data and event.data["path"]:
+                    self._send_notification(
+                        NotificationMethod.FILE_CHANGED,
+                        {
+                            "path": f"{event.data['path']}/.experimaestro/",
+                            "change_type": "modified",
+                        },
+                    )
+            elif event.event_type == StateEventType.SERVICE_UPDATED:
+                self._send_notification(
+                    NotificationMethod.SERVICE_UPDATED,
+                    {
+                        "service_id": event.data.get("service_id"),
+                        "experiment_id": event.data.get("experiment_id"),
+                        "run_id": event.data.get("run_id"),
+                        "state": event.data.get("state"),
+                        "data": event.data,
+                    },
+                )
+        except Exception as e:
+            logger.exception("Error sending notification: %s", e)
+
+    # -------------------------------------------------------------------------
+    # Request Handlers
+    # -------------------------------------------------------------------------
+
+    def _handle_get_experiments(self, params: Dict) -> list:
+        """Handle get_experiments request"""
+        since = deserialize_datetime(params.get("since"))
+        experiments = self._state_provider.get_experiments(since=since)
+        return [serialize_experiment(exp) for exp in experiments]
+
+    def _handle_get_experiment(self, params: Dict) -> Optional[Dict]:
+        """Handle get_experiment request"""
+        experiment_id = params.get("experiment_id")
+        if not experiment_id:
+            raise TypeError("experiment_id is required")
+
+        experiment = self._state_provider.get_experiment(experiment_id)
+        if experiment is None:
+            return None
+        return serialize_experiment(experiment)
+
+    def _handle_get_experiment_runs(self, params: Dict) -> list:
+        """Handle get_experiment_runs request"""
+        experiment_id = params.get("experiment_id")
+        if not experiment_id:
+            raise TypeError("experiment_id is required")
+
+        runs = self._state_provider.get_experiment_runs(experiment_id)
+        return [serialize_run(run) for run in runs]
+
+    def _handle_get_jobs(self, params: Dict) -> list:
+        """Handle get_jobs request"""
+        since = deserialize_datetime(params.get("since"))
+        jobs = self._state_provider.get_jobs(
+            experiment_id=params.get("experiment_id"),
+            run_id=params.get("run_id"),
+            task_id=params.get("task_id"),
+            state=params.get("state"),
+            tags=params.get("tags"),
+            since=since,
+        )
+        return [serialize_job(job) for job in jobs]
+
+    def _handle_get_job(self, params: Dict) -> Optional[Dict]:
+        """Handle get_job request"""
+        job_id = params.get("job_id")
+        experiment_id = params.get("experiment_id")
+        if not job_id or not experiment_id:
+            raise TypeError("job_id and experiment_id are required")
+
+        job = self._state_provider.get_job(
+            job_id=job_id,
+            experiment_id=experiment_id,
+            run_id=params.get("run_id"),
+        )
+        if job is None:
+            return None
+        return serialize_job(job)
+
+    def _handle_get_all_jobs(self, params: Dict) -> list:
+        """Handle get_all_jobs request"""
+        since = deserialize_datetime(params.get("since"))
+        jobs = self._state_provider.get_all_jobs(
+            state=params.get("state"),
+            tags=params.get("tags"),
+            since=since,
+        )
+        return [serialize_job(job) for job in jobs]
+
+    def _handle_get_services(self, params: Dict) -> list:
+        """Handle get_services request
+
+        Uses get_services_raw to return raw service data without trying to
+        recreate Service objects. This allows the client to handle module
+        loading and show appropriate error messages.
+        """
+        return self._state_provider.get_services_raw(
+            experiment_id=params.get("experiment_id"),
+            run_id=params.get("run_id"),
+        )
+
+    def _handle_kill_job(self, params: Dict) -> Dict:
+        """Handle kill_job request"""
+        job_id = params.get("job_id")
+        experiment_id = params.get("experiment_id")
+        run_id = params.get("run_id")
+
+        if not job_id or not experiment_id or not run_id:
+            raise TypeError("job_id, experiment_id, and run_id are required")
+
+        # Get the job first
+        job = self._state_provider.get_job(job_id, experiment_id, run_id)
+        if job is None:
+            return {"success": False, "error": "Job not found"}
+
+        # Kill the job
+        try:
+            result = self._state_provider.kill_job(job, perform=True)
+            return {"success": result}
+        except Exception as e:
+            return {"success": False, "error": str(e)}
+
+    def _handle_clean_job(self, params: Dict) -> Dict:
+        """Handle clean_job request"""
+        job_id = params.get("job_id")
+        experiment_id = params.get("experiment_id")
+        run_id = params.get("run_id")
+
+        if not job_id or not experiment_id or not run_id:
+            raise TypeError("job_id, experiment_id, and run_id are required")
+
+        # Get the job first
+        job = self._state_provider.get_job(job_id, experiment_id, run_id)
+        if job is None:
+            return {"success": False, "error": "Job not found"}
+
+        # Clean the job
+        try:
+            result = self._state_provider.clean_job(job, perform=True)
+            return {"success": result}
+        except Exception as e:
+            return {"success": False, "error": str(e)}
+
+    def _handle_get_sync_info(self, params: Dict) -> Dict:
+        """Handle get_sync_info request"""
+        return {
+            "workspace_path": str(self.workspace_path),
+            "last_sync_time": (
+                serialize_datetime(self._state_provider.get_last_sync_time())
+                if hasattr(self._state_provider, "get_last_sync_time")
+                else None
+            ),
+        }

experimaestro/scheduler/remote/sync.py

@@ -0,0 +1,144 @@
+"""Remote File Synchronizer
+
+Handles rsync-based file synchronization between remote and local workspaces
+for SSH-based experiment monitoring. Only syncs specific paths on-demand
+when services need them (e.g., TensorboardService).
+"""
+
+import logging
+import subprocess
+from pathlib import Path
+from typing import List, Optional
+
+logger = logging.getLogger("xpm.remote.sync")
+
+
+class RemoteFileSynchronizer:
+    """Handles rsync-based file synchronization for remote monitoring
+
+    Syncs specific paths on-demand from a remote host to a local cache
+    directory. Used when services need access to remote files.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        remote_workspace: Path,
+        local_cache: Path,
+        ssh_options: Optional[List[str]] = None,
+    ):
+        """Initialize the synchronizer
+
+        Args:
+            host: SSH host (user@host or just host)
+            remote_workspace: Path to workspace on the remote host
+            local_cache: Local directory to sync files to
+            ssh_options: Additional SSH options (e.g., ["-p", "2222"])
+        """
+        self.host = host
+        self.remote_workspace = remote_workspace
+        self.local_cache = local_cache
+        self.ssh_options = ssh_options or []
+
+    def sync_path(self, remote_path: str) -> Path:
+        """Sync a specific path from remote
+
+        Args:
+            remote_path: Absolute path on remote or path relative to workspace
+
+        Returns:
+            Local path where the files were synced to
+        """
+        # Normalize the path - get relative path within workspace
+        if remote_path.startswith(str(self.remote_workspace)):
+            relative_path = remote_path[len(str(self.remote_workspace)) :].lstrip("/")
+        else:
+            relative_path = remote_path.lstrip("/")
+
+        if not relative_path:
+            raise ValueError("Cannot sync empty path")
+
+        logger.info("Syncing path: %s", relative_path)
+
+        # Build source and destination
+        source = f"{self.host}:{self.remote_workspace}/{relative_path}/"
+        local_path = self.local_cache / relative_path
+        local_path.mkdir(parents=True, exist_ok=True)
+        dest = f"{local_path}/"
+
+        self._rsync(source, dest)
+
+        return local_path
+
+    def _rsync(self, source: str, dest: str):
+        """Execute rsync command
+
+        Args:
+            source: Remote source path (host:path/)
+            dest: Local destination path
+        """
+        cmd = [
+            "rsync",
+            "--inplace",  # Update destination files in-place
+            "--delete",  # Delete extraneous files from destination
+            "-L",  # Transform symlinks into referent file/dir
+            "-a",  # Archive mode (preserves permissions, times, etc.)
+            "-z",  # Compress during transfer
+            "-v",  # Verbose
+        ]
+
+        # SSH options
+        if self.ssh_options:
+            ssh_cmd = "ssh " + " ".join(self.ssh_options)
+            cmd.extend(["-e", ssh_cmd])
+
+        cmd.extend([source, dest])
+
+        logger.debug("Running rsync: %s", " ".join(cmd))
+
+        try:
+            result = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True,
+                timeout=300,  # 5 minute timeout
+            )
+
+            if result.returncode != 0:
+                # rsync returns non-zero for some warnings
+                if result.returncode == 23:
+                    # Partial transfer due to error - some files may be missing
+                    logger.warning("Rsync partial transfer: %s", result.stderr)
+                elif result.returncode == 24:
+                    # Partial transfer due to vanished source files
+                    logger.debug("Rsync: some source files vanished")
+                else:
+                    logger.error(
+                        "Rsync failed (code %d): %s",
+                        result.returncode,
+                        result.stderr,
+                    )
+                    raise RuntimeError(f"Rsync failed: {result.stderr}")
+            else:
+                logger.debug("Rsync completed successfully")
+
+        except subprocess.TimeoutExpired:
+            logger.error("Rsync timed out")
+            raise
+        except FileNotFoundError:
+            logger.error("rsync command not found - please install rsync")
+            raise RuntimeError("rsync command not found")
+
+    def get_local_path(self, remote_path: str) -> Path:
+        """Get the local cache path for a remote path
+
+        Args:
+            remote_path: Absolute path on the remote system
+
+        Returns:
+            Corresponding path in the local cache
+        """
+        if remote_path.startswith(str(self.remote_workspace)):
+            relative = remote_path[len(str(self.remote_workspace)) :].lstrip("/")
+            return self.local_cache / relative
+        return Path(remote_path)