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

experimaestro/scheduler/remote/protocol.py

@@ -12,11 +12,7 @@ import logging
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
-
-if TYPE_CHECKING:
-    from experimaestro.scheduler.interfaces import JobState
+from typing import Any, Dict, Optional, Union
 
 logger = logging.getLogger("xpm.remote.protocol")
 
@@ -40,10 +36,10 @@ TIMEOUT_ERROR = -32004
 class NotificationMethod(str, Enum):
     """Server-to-client notification methods"""
 
-    EXPERIMENT_UPDATED = "notification.experiment_updated"
-    RUN_UPDATED = "notification.run_updated"
-    JOB_UPDATED = "notification.job_updated"
-    SERVICE_UPDATED = "notification.service_updated"
+    # Generic state event notification (serialized dataclass)
+    STATE_EVENT = "notification.state_event"
+
+    # Control notifications
     FILE_CHANGED = "notification.file_changed"
     SHUTDOWN = "notification.shutdown"
 
@@ -58,9 +54,12 @@ class RPCMethod(str, Enum):
     GET_JOB = "get_job"
     GET_ALL_JOBS = "get_all_jobs"
     GET_SERVICES = "get_services"
+    GET_TAGS_MAP = "get_tags_map"
+    GET_DEPENDENCIES_MAP = "get_dependencies_map"
     KILL_JOB = "kill_job"
     CLEAN_JOB = "clean_job"
     GET_SYNC_INFO = "get_sync_info"
+    GET_PROCESS_INFO = "get_process_info"
 
 
 @dataclass
@@ -281,187 +280,3 @@ def deserialize_datetime(s: Optional[str]) -> Optional[datetime]:
     if s is None:
         return None
     return datetime.fromisoformat(s)
-
-
-def serialize_job(job) -> Dict:
-    """Serialize a job (MockJob or Job) to a dictionary for JSON-RPC"""
-    from experimaestro.scheduler.interfaces import JobState
-
-    result = {
-        "identifier": job.identifier,
-        "task_id": job.task_id,
-        "locator": job.locator,
-        "path": str(job.path) if job.path else None,
-        "state": job.state.name if isinstance(job.state, JobState) else str(job.state),
-        "submittime": serialize_datetime(job.submittime),
-        "starttime": serialize_datetime(job.starttime),
-        "endtime": serialize_datetime(job.endtime),
-        "progress": job.progress,
-        "tags": job.tags,
-        "experiment_id": getattr(job, "experiment_id", None),
-        "run_id": getattr(job, "run_id", None),
-    }
-    return result
-
-
-def deserialize_job(d: Dict) -> "MockJobData":
-    """Deserialize a dictionary to MockJobData"""
-    from experimaestro.scheduler.interfaces import JobState, STATE_NAME_TO_JOBSTATE
-    from pathlib import Path
-
-    state = STATE_NAME_TO_JOBSTATE.get(d["state"], JobState.WAITING)
-    return MockJobData(
-        identifier=d["identifier"],
-        task_id=d["task_id"],
-        locator=d["locator"],
-        path=Path(d["path"]) if d["path"] else None,
-        state=state,
-        submittime=deserialize_datetime(d.get("submittime")),
-        starttime=deserialize_datetime(d.get("starttime")),
-        endtime=deserialize_datetime(d.get("endtime")),
-        progress=d.get("progress"),
-        tags=d.get("tags", {}),
-        experiment_id=d.get("experiment_id"),
-        run_id=d.get("run_id"),
-    )
-
-
-def serialize_experiment(experiment) -> Dict:
-    """Serialize a MockExperiment to a dictionary for JSON-RPC"""
-    result = {
-        "experiment_id": experiment.experiment_id,
-        "workdir": str(experiment.workdir) if experiment.workdir else None,
-        "current_run_id": experiment.current_run_id,
-        "total_jobs": experiment.total_jobs,
-        "finished_jobs": experiment.finished_jobs,
-        "failed_jobs": experiment.failed_jobs,
-        "updated_at": serialize_datetime(experiment.updated_at),
-        "started_at": serialize_datetime(experiment.started_at),
-        "ended_at": serialize_datetime(experiment.ended_at),
-        "hostname": experiment.hostname,
-    }
-    return result
-
-
-def deserialize_experiment(d: Dict) -> "MockExperimentData":
-    """Deserialize a dictionary to MockExperimentData"""
-    from pathlib import Path
-
-    return MockExperimentData(
-        experiment_id=d["experiment_id"],
-        workdir=Path(d["workdir"]) if d["workdir"] else None,
-        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=deserialize_datetime(d.get("updated_at")),
-        started_at=deserialize_datetime(d.get("started_at")),
-        ended_at=deserialize_datetime(d.get("ended_at")),
-        hostname=d.get("hostname"),
-    )
-
-
-def serialize_service(service) -> Dict:
-    """Serialize a service to a dictionary for JSON-RPC"""
-    from experimaestro.scheduler.services import Service
-
-    # Service has id attribute, description() method, state property, state_dict() method
-    state = service.state
-    if hasattr(state, "name"):
-        state = state.name  # Convert ServiceState enum to string
-    elif hasattr(state, "value"):
-        state = state.value
-
-    # Get URL if service has it (e.g., TensorboardService)
-    url = None
-    if hasattr(service, "url"):
-        url = service.url
-    elif hasattr(service, "get_url"):
-        try:
-            url = service.get_url()
-        except Exception:
-            pass
-
-    # Get state_dict with __class__ and serialize paths
-    if hasattr(service, "_full_state_dict"):
-        state_dict = Service.serialize_state_dict(service._full_state_dict())
-    elif callable(getattr(service, "state_dict", None)):
-        # Fallback: serialize paths in the raw state_dict
-        state_dict = Service.serialize_state_dict(service.state_dict())
-    else:
-        state_dict = getattr(service, "state_dict", {})
-
-    return {
-        "service_id": getattr(service, "id", None),
-        "description": (
-            service.description()
-            if callable(service.description)
-            else service.description
-        ),
-        "state": state,
-        "state_dict": state_dict,
-        "experiment_id": getattr(service, "experiment_id", None),
-        "run_id": getattr(service, "run_id", None),
-        "url": url,
-    }
-
-
-def serialize_run(run) -> Dict:
-    """Serialize an experiment run to a dictionary for JSON-RPC
-
-    Handles both dictionary and object inputs (get_experiment_runs returns dicts).
-    """
-    if isinstance(run, dict):
-        # Already a dictionary - just ensure datetime serialization
-        return {
-            "run_id": run.get("run_id"),
-            "experiment_id": run.get("experiment_id"),
-            "hostname": run.get("hostname"),
-            "started_at": run.get("started_at"),  # Already serialized
-            "ended_at": run.get("ended_at"),  # Already serialized
-            "status": run.get("status"),
-        }
-    else:
-        # Object with attributes
-        return {
-            "run_id": run.run_id,
-            "experiment_id": run.experiment_id,
-            "hostname": getattr(run, "hostname", None),
-            "started_at": serialize_datetime(run.started_at),
-            "ended_at": serialize_datetime(run.ended_at),
-            "status": run.status,
-        }
-
-
-@dataclass
-class MockJobData:
-    """Deserialized job data from remote"""
-
-    identifier: str
-    task_id: str
-    locator: str
-    path: Optional["Path"]
-    state: "JobState"
-    submittime: Optional[datetime]
-    starttime: Optional[datetime]
-    endtime: Optional[datetime]
-    progress: Optional[float]
-    tags: Dict
-    experiment_id: Optional[str]
-    run_id: Optional[str]
-
-
-@dataclass
-class MockExperimentData:
-    """Deserialized experiment data from remote"""
-
-    experiment_id: str
-    workdir: Optional["Path"]
-    current_run_id: Optional[str]
-    total_jobs: int
-    finished_jobs: int
-    failed_jobs: int
-    updated_at: Optional[datetime]
-    started_at: Optional[datetime]
-    ended_at: Optional[datetime]
-    hostname: Optional[str]

experimaestro/scheduler/remote/server.py

@@ -13,11 +13,10 @@ 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.workspace_state_provider import WorkspaceStateProvider
+from dataclasses import asdict
+
+from experimaestro.scheduler.state_status import EventBase
 from experimaestro.scheduler.remote.protocol import (
     RPCMethod,
     NotificationMethod,
@@ -25,9 +24,6 @@ from experimaestro.scheduler.remote.protocol import (
     create_success_response,
     create_error_response,
     create_notification,
-    serialize_job,
-    serialize_experiment,
-    serialize_run,
     serialize_datetime,
     deserialize_datetime,
     PARSE_ERROR,
@@ -82,9 +78,12 @@ class SSHStateProviderServer:
             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.GET_TAGS_MAP.value: self._handle_get_tags_map,
+            RPCMethod.GET_DEPENDENCIES_MAP.value: self._handle_get_dependencies_map,
             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,
+            RPCMethod.GET_PROCESS_INFO.value: self._handle_get_process_info,
         }
 
     def start(self):
@@ -101,12 +100,10 @@ class SSHStateProviderServer:
             )
             return
 
-        # Initialize state provider in read-only mode
+        # Initialize state provider in read-only mode with event watcher
         try:
             self._state_provider = WorkspaceStateProvider.get_instance(
-                self.workspace_path,
-                read_only=True,
-                sync_on_start=True,
+                self.workspace_path
             )
         except Exception as e:
             logger.exception("Failed to initialize state provider")
@@ -231,60 +228,34 @@ class SSHStateProviderServer:
             {"reason": "error", "code": code, "message": message},
         )
 
-    def _on_state_event(self, event: StateEvent):
+    def _on_state_event(self, event: EventBase):
         """Handle state change events from the state provider
 
         Converts events to JSON-RPC notifications and sends them to the client.
+        Uses generic serialization via dataclasses.asdict.
         """
         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,
-                    },
-                )
+            # Serialize event to dict, filtering out None values and non-serializable objects
+            event_dict = {}
+            for key, value in asdict(event).items():
+                # Skip None values and complex objects (like job references)
+                if value is not None and not isinstance(value, (Path,)):
+                    # Try to serialize - skip if not JSON-serializable
+                    try:
+                        import json
+
+                        json.dumps(value)
+                        event_dict[key] = value
+                    except (TypeError, ValueError):
+                        pass
+
+            self._send_notification(
+                NotificationMethod.STATE_EVENT,
+                {
+                    "event_type": type(event).__name__,
+                    "data": event_dict,
+                },
+            )
         except Exception as e:
             logger.exception("Error sending notification: %s", e)
 
@@ -296,7 +267,7 @@ class SSHStateProviderServer:
         """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]
+        return [exp.state_dict() for exp in experiments]
 
     def _handle_get_experiment(self, params: Dict) -> Optional[Dict]:
         """Handle get_experiment request"""
@@ -307,7 +278,7 @@ class SSHStateProviderServer:
         experiment = self._state_provider.get_experiment(experiment_id)
         if experiment is None:
             return None
-        return serialize_experiment(experiment)
+        return experiment.state_dict()
 
     def _handle_get_experiment_runs(self, params: Dict) -> list:
         """Handle get_experiment_runs request"""
@@ -316,7 +287,7 @@ class SSHStateProviderServer:
             raise TypeError("experiment_id is required")
 
         runs = self._state_provider.get_experiment_runs(experiment_id)
-        return [serialize_run(run) for run in runs]
+        return [run.state_dict() for run in runs]
 
     def _handle_get_jobs(self, params: Dict) -> list:
         """Handle get_jobs request"""
@@ -329,7 +300,7 @@ class SSHStateProviderServer:
             tags=params.get("tags"),
             since=since,
         )
-        return [serialize_job(job) for job in jobs]
+        return [job.state_dict() for job in jobs]
 
     def _handle_get_job(self, params: Dict) -> Optional[Dict]:
         """Handle get_job request"""
@@ -345,7 +316,7 @@ class SSHStateProviderServer:
         )
         if job is None:
             return None
-        return serialize_job(job)
+        return job.state_dict()
 
     def _handle_get_all_jobs(self, params: Dict) -> list:
         """Handle get_all_jobs request"""
@@ -355,19 +326,46 @@ class SSHStateProviderServer:
             tags=params.get("tags"),
             since=since,
         )
-        return [serialize_job(job) for job in jobs]
+        return [job.state_dict() 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.
+        Returns serialized service data using full_state_dict().
         """
-        return self._state_provider.get_services_raw(
+        services = self._state_provider.get_services(
             experiment_id=params.get("experiment_id"),
             run_id=params.get("run_id"),
         )
+        return [svc.full_state_dict() for svc in services]
+
+    def _handle_get_tags_map(self, params: Dict) -> Dict[str, Dict[str, str]]:
+        """Handle get_tags_map request
+
+        Returns tags map for jobs in an experiment/run.
+        """
+        experiment_id = params.get("experiment_id")
+        if not experiment_id:
+            raise TypeError("experiment_id is required")
+
+        return self._state_provider.get_tags_map(
+            experiment_id=experiment_id,
+            run_id=params.get("run_id"),
+        )
+
+    def _handle_get_dependencies_map(self, params: Dict) -> dict[str, list[str]]:
+        """Handle get_dependencies_map request
+
+        Returns dependencies map for jobs in an experiment/run.
+        """
+        experiment_id = params.get("experiment_id")
+        if not experiment_id:
+            raise TypeError("experiment_id is required")
+
+        return self._state_provider.get_dependencies_map(
+            experiment_id=experiment_id,
+            run_id=params.get("run_id"),
+        )
 
     def _handle_kill_job(self, params: Dict) -> Dict:
         """Handle kill_job request"""
@@ -421,3 +419,29 @@ class SSHStateProviderServer:
                 else None
             ),
         }
+
+    def _handle_get_process_info(self, params: Dict) -> Optional[Dict]:
+        """Handle get_process_info 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:
+            raise TypeError("job_id and experiment_id are required")
+
+        # Get the job first
+        job = self._state_provider.get_job(job_id, experiment_id, run_id)
+        if job is None:
+            return None
+
+        # Get process info
+        pinfo = self._state_provider.get_process_info(job)
+        if pinfo is None:
+            return None
+
+        # Serialize ProcessInfo to dict
+        return {
+            "pid": pinfo.pid,
+            "type": pinfo.type,
+            "running": pinfo.running,
+        }

experimaestro/scheduler/services.py

@@ -16,7 +16,7 @@ logger = logging.getLogger(__name__)
 class ServiceListener:
     """A service listener"""
 
-    def service_state_changed(service):
+    def service_state_changed(self, service):
         pass
 
 
@@ -78,15 +78,21 @@ class Service(BaseService):
                 }
 
         Returns:
-            Dict with constructor kwargs (no need to include __class__).
+            Dict with constructor kwargs.
         """
         return {}
 
-    def _full_state_dict(self) -> dict:
-        """Get complete state_dict including __class__ for serialization."""
-        d = self.state_dict()
-        d["__class__"] = f"{self.__class__.__module__}.{self.__class__.__name__}"
-        return d
+    def full_state_dict(self) -> dict:
+        """Serialize service to dictionary for JSON serialization.
+
+        Overrides BaseService.full_state_dict() to properly serialize Path objects.
+        """
+        return {
+            "service_id": self.id,
+            "description": self.description(),
+            "class": f"{self.__class__.__module__}.{self.__class__.__name__}",
+            "state_dict": self.serialize_state_dict(self.state_dict()),
+        }
 
     @staticmethod
     def serialize_state_dict(data: dict) -> dict:
@@ -101,21 +107,29 @@ class Service(BaseService):
         Returns:
             Serializable dictionary with paths converted
         """
-        result = {}
-        for k, v in data.items():
+
+        def serialize_value(v):
             if isinstance(v, Path):
-                result[k] = {"__path__": str(v)}
+                return {"__path__": str(v)}
+            elif isinstance(v, dict):
+                return {k: serialize_value(val) for k, val in v.items()}
+            elif isinstance(v, (list, tuple)):
+                return [serialize_value(item) for item in v]
             else:
-                result[k] = v
-        return result
+                return v
+
+        return {k: serialize_value(v) for k, v in data.items()}
 
     @staticmethod
     def from_state_dict(
-        data: dict, path_translator: Optional[Callable[[str], Path]] = None
+        service_class: str,
+        data: dict,
+        path_translator: Optional[Callable[[str], Path]] = None,
     ) -> "Service":
         """Recreate a service from a state dictionary.
 
         Args:
+            service_class: Fully qualified class name (e.g., "module.ClassName")
             data: Dictionary from :meth:`state_dict` (may be serialized)
             path_translator: Optional function to translate remote paths to local.
                 Used by remote clients to map paths to local cache.
@@ -124,7 +138,7 @@ class Service(BaseService):
             A new Service instance, or raises if the class cannot be loaded.
 
         Raises:
-            ValueError: If __unserializable__ is True or __class__ is missing
+            ValueError: If __unserializable__ is True or class cannot be loaded
         """
         import importlib
 
@@ -134,28 +148,39 @@ class Service(BaseService):
                 f"Service cannot be recreated: {data.get('__reason__', 'unknown reason')}"
             )
 
-        class_path = data.get("__class__")
-        if not class_path:
-            raise ValueError("Missing '__class__' in service state_dict")
+        if not service_class:
+            raise ValueError("Missing service_class")
 
-        module_name, class_name = class_path.rsplit(".", 1)
+        module_name, class_name = service_class.rsplit(".", 1)
         module = importlib.import_module(module_name)
         cls = getattr(module, class_name)
 
-        # Build kwargs, detecting and translating paths automatically
+        # Build kwargs, detecting and translating paths automatically (handles nested)
+        def deserialize_value(v):
+            if isinstance(v, dict):
+                if "__path__" in v:
+                    # Serialized path - deserialize with optional translation
+                    path_str = v["__path__"]
+                    if path_translator:
+                        return path_translator(path_str)
+                    else:
+                        return Path(path_str)
+                else:
+                    return {
+                        k: deserialize_value(val)
+                        for k, val in v.items()
+                        if not k.startswith("__")
+                    }
+            elif isinstance(v, list):
+                return [deserialize_value(item) for item in v]
+            else:
+                return v
+
         kwargs = {}
         for k, v in data.items():
             if k.startswith("__"):
                 continue  # Skip special keys
-            if isinstance(v, dict) and "__path__" in v:
-                # Serialized path - deserialize with optional translation
-                path_str = v["__path__"]
-                if path_translator:
-                    kwargs[k] = path_translator(path_str)
-                else:
-                    kwargs[k] = Path(path_str)
-            else:
-                kwargs[k] = v
+            kwargs[k] = deserialize_value(v)
 
         logger.debug("Creating %s with kwargs: %s", cls.__name__, kwargs)
         return cls(**kwargs)