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

experimaestro/scheduler/remote/client.py

@@ -19,17 +19,18 @@ from concurrent.futures import Future, TimeoutError as FutureTimeoutError
 from datetime import datetime
 from importlib.metadata import version as get_package_version
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Set
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set
+
+from termcolor import colored
 
 from experimaestro.scheduler.state_provider import (
-    StateProvider,
-    StateEvent,
-    StateEventType,
+    OfflineStateProvider,
     StateListener,
     MockJob,
     MockExperiment,
     MockService,
 )
+from experimaestro.scheduler.state_status import EventBase
 from experimaestro.scheduler.interfaces import (
     BaseJob,
     BaseExperiment,
@@ -45,6 +46,9 @@ from experimaestro.scheduler.remote.protocol import (
     serialize_datetime,
 )
 
+# Type for SSH output callback
+OutputCallback = Optional["Callable[[str], None]"]
+
 if TYPE_CHECKING:
     from experimaestro.scheduler.remote.sync import RemoteFileSynchronizer
 
@@ -68,7 +72,7 @@ def _strip_dev_version(version: str) -> str:
     return re.sub(r"\.dev\d+$", "", version)
 
 
-class SSHStateProviderClient(StateProvider):
+class SSHStateProviderClient(OfflineStateProvider):
     """Client that connects to SSHStateProviderServer via SSH
 
     This client implements the StateProvider interface for remote experiment
@@ -77,7 +81,7 @@ class SSHStateProviderClient(StateProvider):
     Features:
     - JSON-RPC over SSH stdin/stdout
     - Async request/response handling with futures
-    - Server push notifications converted to StateEvents
+    - Server push notifications converted to EventBases
     - On-demand rsync for specific paths (used by services like TensorboardService)
     """
 
@@ -87,6 +91,7 @@ class SSHStateProviderClient(StateProvider):
         remote_workspace: str,
         ssh_options: Optional[List[str]] = None,
         remote_xpm_path: Optional[str] = None,
+        output_callback: Optional[Callable[[str], None]] = None,
     ):
         """Initialize the client
 
@@ -96,11 +101,18 @@ class SSHStateProviderClient(StateProvider):
             ssh_options: Additional SSH options (e.g., ["-p", "2222"])
             remote_xpm_path: Path to experimaestro executable on remote host.
                 If None, uses 'uv tool run experimaestro==<version>'.
+            output_callback: Callback for SSH process output (stderr).
+                If None, a default callback prints with colored prefix.
+                Set to False (or a no-op lambda) to disable output display.
         """
+        # Initialize base class (includes service cache)
+        super().__init__()
+
         self.host = host
         self.remote_workspace = remote_workspace
         self.ssh_options = ssh_options or []
         self.remote_xpm_path = remote_xpm_path
+        self._output_callback = output_callback
 
         # Session-specific temporary cache directory (created on connect)
         self._temp_dir: Optional[str] = None
@@ -121,19 +133,17 @@ class SSHStateProviderClient(StateProvider):
 
         self._read_thread: Optional[threading.Thread] = None
         self._notify_thread: Optional[threading.Thread] = None
+        self._stderr_thread: Optional[threading.Thread] = None
         self._running = False
         self._connected = False
 
         self._synchronizer: Optional["RemoteFileSynchronizer"] = None
 
         # Throttled notification delivery to avoid flooding UI
-        self._pending_events: List[StateEvent] = []
+        self._pending_events: List[EventBase] = []
         self._pending_events_lock = threading.Lock()
         self._notify_interval = 2.0  # Seconds between notification batches
 
-        # Service cache (from base class)
-        self._init_service_cache()
-
     def connect(self, timeout: float = 30.0):
         """Establish SSH connection and start remote server
 
@@ -161,10 +171,7 @@ class SSHStateProviderClient(StateProvider):
         # Build remote command (workdir is passed to experiments group)
         if self.remote_xpm_path:
             # Use specified path to experimaestro
-            remote_cmd = (
-                f"{self.remote_xpm_path} experiments "
-                f"--workdir {self.remote_workspace} monitor-server"
-            )
+            remote_cmd = f"{self.remote_xpm_path} experiments --workdir {self.remote_workspace} monitor-server"
         else:
             # Use uv tool run with version pinning
             try:
@@ -175,15 +182,9 @@ class SSHStateProviderClient(StateProvider):
                 xpm_version = None
 
             if xpm_version:
-                remote_cmd = (
-                    f"uv tool run experimaestro=={xpm_version} experiments "
-                    f"--workdir {self.remote_workspace} monitor-server"
-                )
+                remote_cmd = f"uv tool run experimaestro=={xpm_version} experiments --workdir {self.remote_workspace} monitor-server"
             else:
-                remote_cmd = (
-                    f"uv tool run experimaestro experiments "
-                    f"--workdir {self.remote_workspace} monitor-server"
-                )
+                remote_cmd = f"uv tool run experimaestro experiments --workdir {self.remote_workspace} monitor-server"
         cmd.append(remote_cmd)
 
         logger.info("Connecting to %s, workspace: %s", self.host, self.remote_workspace)
@@ -218,6 +219,12 @@ class SSHStateProviderClient(StateProvider):
         )
         self._notify_thread.start()
 
+        # Start stderr thread to display SSH output
+        self._stderr_thread = threading.Thread(
+            target=self._stderr_loop, daemon=True, name="SSHClient-Stderr"
+        )
+        self._stderr_thread.start()
+
         # Wait for connection to be established by sending a test request
         try:
             sync_info = self._call_sync(RPCMethod.GET_SYNC_INFO, {}, timeout=timeout)
@@ -258,6 +265,8 @@ class SSHStateProviderClient(StateProvider):
             self._read_thread.join(timeout=2.0)
         if self._notify_thread and self._notify_thread.is_alive():
             self._notify_thread.join(timeout=2.0)
+        if self._stderr_thread and self._stderr_thread.is_alive():
+            self._stderr_thread.join(timeout=2.0)
 
         # Cancel any pending requests
         with self._response_lock:
@@ -321,6 +330,33 @@ class SSHStateProviderClient(StateProvider):
             logger.warning("Connection to %s lost", self.host)
             self._connected = False
 
+    def _stderr_loop(self):
+        """Read and display SSH stderr output with colored prefix"""
+        while self._running:
+            try:
+                line = self._stderr.readline()
+                if not line:
+                    # EOF - stderr closed
+                    logger.debug("SSH stderr closed")
+                    break
+
+                line_str = line.decode("utf-8").rstrip("\n\r")
+                if not line_str:
+                    continue
+
+                # Call output callback or use default
+                if self._output_callback is not None:
+                    self._output_callback(line_str)
+                else:
+                    # Default: print with colored prefix
+                    prefix = colored("[SSH] ", "cyan", attrs=["bold"])
+                    print(f"{prefix}{line_str}")  # noqa: T201
+
+            except Exception as e:
+                if self._running:
+                    logger.debug("Error reading stderr: %s", e)
+                break
+
     def _process_message(self, line: str):
         """Process a single message from the server"""
         try:
@@ -364,7 +400,7 @@ class SSHStateProviderClient(StateProvider):
 
         logger.debug("Received notification: %s", method)
 
-        # Convert notification to StateEvent and queue for throttled delivery
+        # Convert notification to EventBase and queue for throttled delivery
         event = self._notification_to_event(method, params)
         if event:
             with self._pending_events_lock:
@@ -402,8 +438,9 @@ class SSHStateProviderClient(StateProvider):
             seen_types = set()
             unique_events = []
             for event in reversed(events):
-                if event.event_type not in seen_types:
-                    seen_types.add(event.event_type)
+                event_type = type(event)
+                if event_type not in seen_types:
+                    seen_types.add(event_type)
                     unique_events.append(event)
             unique_events.reverse()
 
@@ -411,31 +448,31 @@ class SSHStateProviderClient(StateProvider):
             for event in unique_events:
                 self._notify_listeners(event)
 
-    def _notification_to_event(self, method: str, params: Dict) -> Optional[StateEvent]:
-        """Convert a notification to a StateEvent"""
-        if method == NotificationMethod.EXPERIMENT_UPDATED.value:
-            return StateEvent(
-                event_type=StateEventType.EXPERIMENT_UPDATED,
-                data=params.get("data", params),
-            )
-        elif method == NotificationMethod.RUN_UPDATED.value:
-            return StateEvent(
-                event_type=StateEventType.RUN_UPDATED,
-                data=params.get("data", params),
-            )
-        elif method == NotificationMethod.JOB_UPDATED.value:
-            return StateEvent(
-                event_type=StateEventType.JOB_UPDATED,
-                data=params.get("data", params),
-            )
-        elif method == NotificationMethod.SERVICE_UPDATED.value:
-            return StateEvent(
-                event_type=StateEventType.SERVICE_UPDATED,
-                data=params.get("data", params),
-            )
-        return None
+    def _notification_to_event(self, method: str, params: Dict) -> Optional[EventBase]:
+        """Convert a notification to a EventBase"""
+        if method != NotificationMethod.STATE_EVENT.value:
+            # Don't warn for known control notifications (handled elsewhere)
+            if method not in (
+                NotificationMethod.SHUTDOWN.value,
+                NotificationMethod.FILE_CHANGED.value,
+            ):
+                logger.warning("Unhandled notification method: %s", method)
+            return None
+
+        event_type = params.get("event_type")
+        data = params.get("data", {})
+        event_class = EventBase.get_class(event_type)
+        if event_class is None:
+            logger.warning("Unknown event type: %s", event_type)
+            return None
 
-    def _notify_listeners(self, event: StateEvent):
+        try:
+            return event_class(**data)
+        except TypeError as e:
+            logger.warning("Error deserializing event %s: %s", event_type, e)
+            return None
+
+    def _notify_listeners(self, event: EventBase):
         """Notify all registered listeners of a state event"""
         with self._listener_lock:
             listeners = list(self._listeners)
@@ -539,7 +576,7 @@ class SSHStateProviderClient(StateProvider):
         exp = self.get_experiment(experiment_id)
         if exp is None:
             return None
-        return exp.current_run_id
+        return exp.run_id
 
     def get_jobs(
         self,
@@ -591,6 +628,32 @@ class SSHStateProviderClient(StateProvider):
         result = self._call_sync(RPCMethod.GET_ALL_JOBS, params)
         return [self._dict_to_job(d) for d in result]
 
+    def get_tags_map(
+        self,
+        experiment_id: str,
+        run_id: Optional[str] = None,
+    ) -> Dict[str, Dict[str, str]]:
+        """Get tags map for jobs in an experiment/run"""
+        params = {
+            "experiment_id": experiment_id,
+            "run_id": run_id,
+        }
+        result = self._call_sync(RPCMethod.GET_TAGS_MAP, params)
+        return result or {}
+
+    def get_dependencies_map(
+        self,
+        experiment_id: str,
+        run_id: Optional[str] = None,
+    ) -> dict[str, list[str]]:
+        """Get dependencies map for jobs in an experiment/run"""
+        params = {
+            "experiment_id": experiment_id,
+            "run_id": run_id,
+        }
+        result = self._call_sync(RPCMethod.GET_DEPENDENCIES_MAP, params)
+        return result or {}
+
     def _fetch_services_from_storage(
         self, experiment_id: Optional[str], run_id: Optional[str]
     ) -> List[BaseService]:
@@ -611,16 +674,6 @@ class SSHStateProviderClient(StateProvider):
 
         return services
 
-    def get_services_raw(
-        self, experiment_id: Optional[str] = None, run_id: Optional[str] = None
-    ) -> List[Dict]:
-        """Get raw service data as dictionaries"""
-        params = {
-            "experiment_id": experiment_id,
-            "run_id": run_id,
-        }
-        return self._call_sync(RPCMethod.GET_SERVICES, params)
-
     def kill_job(self, job: BaseJob, perform: bool = False) -> bool:
         """Kill a running job"""
         if not perform:
@@ -649,69 +702,71 @@ class SSHStateProviderClient(StateProvider):
         result = self._call_sync(RPCMethod.CLEAN_JOB, params)
         return result.get("success", False)
 
+    def get_process_info(self, job: BaseJob):
+        """Get process information for a job
+
+        Returns None if the remote server doesn't support this method.
+        """
+        from experimaestro.scheduler.state_provider import ProcessInfo
+
+        params = {
+            "job_id": job.identifier,
+            "experiment_id": getattr(job, "experiment_id", ""),
+            "run_id": getattr(job, "run_id", ""),
+        }
+
+        try:
+            result = self._call_sync(RPCMethod.GET_PROCESS_INFO, params)
+        except RuntimeError:
+            # Server doesn't support this method (older version)
+            return None
+
+        if result is None:
+            return None
+
+        return ProcessInfo(
+            pid=result["pid"],
+            type=result["type"],
+            running=result.get("running", False),
+        )
+
     # -------------------------------------------------------------------------
     # Data Conversion
     # -------------------------------------------------------------------------
 
     def _dict_to_job(self, d: Dict) -> MockJob:
-        """Convert a dictionary to a MockJob"""
-        state_str = d.get("state", "waiting")
-
-        # Map local cache path for the job
-        path = None
+        """Convert a dictionary to a MockJob using from_state_dict"""
+        # Translate remote path to local cache path
         if d.get("path"):
-            # The path from remote is absolute on remote system
-            # We map it to local cache
             remote_path = d["path"]
             if remote_path.startswith(self.remote_workspace):
                 relative = remote_path[len(self.remote_workspace) :].lstrip("/")
-                path = self.local_cache_dir / relative
+                d["path"] = self.local_cache_dir / relative
             else:
-                path = Path(remote_path)
-
-        return MockJob(
-            identifier=d["identifier"],
-            task_id=d["task_id"],
-            locator=d["locator"],
-            path=path,
-            state=state_str,
-            submittime=self._parse_datetime_to_timestamp(d.get("submittime")),
-            starttime=self._parse_datetime_to_timestamp(d.get("starttime")),
-            endtime=self._parse_datetime_to_timestamp(d.get("endtime")),
-            progress=d.get("progress", []),
-            tags=d.get("tags", {}),
-            experiment_id=d.get("experiment_id", ""),
-            run_id=d.get("run_id", ""),
-            updated_at="",
-        )
+                d["path"] = Path(remote_path)
+
+        # Convert ISO datetime strings back to timestamps (floats)
+        for key in ("submitted_time", "started_time", "ended_time"):
+            d[key] = self._parse_datetime_to_timestamp(d.get(key))
+
+        return MockJob.from_state_dict(d, self.local_cache_dir)
 
     def _dict_to_experiment(self, d: Dict) -> MockExperiment:
-        """Convert a dictionary to a MockExperiment"""
-        # Map local cache path for the experiment
-        workdir = None
+        """Convert a dictionary to a MockExperiment using from_state_dict"""
+        # Translate remote workdir to local cache path
         if d.get("workdir"):
             remote_path = d["workdir"]
             if remote_path.startswith(self.remote_workspace):
                 relative = remote_path[len(self.remote_workspace) :].lstrip("/")
-                workdir = self.local_cache_dir / relative
+                d["workdir"] = self.local_cache_dir / relative
             else:
-                workdir = Path(remote_path)
-
-        # Convert ISO datetime strings to Unix timestamps for TUI compatibility
-        started_at = self._parse_datetime_to_timestamp(d.get("started_at"))
-        ended_at = self._parse_datetime_to_timestamp(d.get("ended_at"))
-
-        return MockExperiment(
-            workdir=workdir or self.local_cache_dir / "xp" / d["experiment_id"],
-            current_run_id=d.get("current_run_id"),
-            total_jobs=d.get("total_jobs", 0),
-            finished_jobs=d.get("finished_jobs", 0),
-            failed_jobs=d.get("failed_jobs", 0),
-            updated_at=d.get("updated_at", ""),
-            started_at=started_at,
-            ended_at=ended_at,
-            hostname=d.get("hostname"),
-        )
+                d["workdir"] = Path(remote_path)
+
+        # Convert ISO datetime strings back to timestamps (floats)
+        for key in ("started_at", "ended_at"):
+            d[key] = self._parse_datetime_to_timestamp(d.get(key))
+
+        return MockExperiment.from_state_dict(d, self.local_cache_dir)
 
     def _dict_to_service(self, d: Dict) -> BaseService:
         """Convert a dictionary to a Service or MockService
@@ -720,6 +775,7 @@ class SSHStateProviderClient(StateProvider):
         Falls back to MockService with error message if module is missing.
         """
         state_dict = d.get("state_dict", {})
+        service_class = d.get("class", "")
         service_id = d.get("service_id", "")
 
         # Check for unserializable marker
@@ -729,13 +785,14 @@ class SSHStateProviderClient(StateProvider):
                 service_id=service_id,
                 description_text=f"[{reason}]",
                 state_dict_data=state_dict,
+                service_class=service_class,
                 experiment_id=d.get("experiment_id"),
                 run_id=d.get("run_id"),
                 url=d.get("url"),
             )
 
         # Try to recreate actual Service from state_dict
-        if state_dict and "__class__" in state_dict:
+        if service_class:
             try:
                 from experimaestro.scheduler.services import Service
 
@@ -751,7 +808,9 @@ class SSHStateProviderClient(StateProvider):
                         return self.local_cache_dir / relative
                     return Path(remote_path)
 
-                service = Service.from_state_dict(state_dict, path_translator)
+                service = Service.from_state_dict(
+                    service_class, state_dict, path_translator
+                )
                 service.id = service_id
                 # Copy additional attributes
                 if d.get("experiment_id"):
@@ -766,6 +825,7 @@ class SSHStateProviderClient(StateProvider):
                     service_id=service_id,
                     description_text=f"[Missing module: {missing_module}]",
                     state_dict_data=state_dict,
+                    service_class=service_class,
                     experiment_id=d.get("experiment_id"),
                     run_id=d.get("run_id"),
                     url=d.get("url"),
@@ -776,20 +836,14 @@ class SSHStateProviderClient(StateProvider):
                     service_id=service_id,
                     description_text=f"[Error: {e}]",
                     state_dict_data=state_dict,
+                    service_class=service_class,
                     experiment_id=d.get("experiment_id"),
                     run_id=d.get("run_id"),
                     url=d.get("url"),
                 )
 
-        # No state_dict or no __class__ - use MockService with original description
-        return MockService(
-            service_id=service_id,
-            description_text=d.get("description", ""),
-            state_dict_data=state_dict,
-            experiment_id=d.get("experiment_id"),
-            run_id=d.get("run_id"),
-            url=d.get("url"),
-        )
+        # No class - use MockService.from_full_state_dict
+        return MockService.from_full_state_dict(d)
 
     def _parse_datetime_to_timestamp(self, value) -> Optional[float]:
         """Convert datetime value to Unix timestamp