synapse-sdk 1.0.0a11__py3-none-any.whl → 2026.1.1b2__py3-none-any.whl

synapse_sdk/clients/pipeline/client.py

@@ -0,0 +1,636 @@
+"""Pipeline Service Client for communicating with dev-api.
+
+This client provides methods to interact with the pipeline orchestration
+backend for registering pipelines, creating runs, and reporting progress.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import AsyncIterator, Iterator
+from datetime import datetime
+from typing import Any
+
+import httpx
+
+from synapse_sdk.plugins.models.logger import (
+    ActionProgress,
+    LogEntry,
+    PipelineProgress,
+)
+from synapse_sdk.plugins.models.pipeline import RunStatus
+
+logger = logging.getLogger(__name__)
+
+
+class PipelineServiceClient:
+    """Client for the Pipeline Service API.
+
+    Provides methods to manage pipelines, runs, progress, checkpoints, and logs.
+
+    Attributes:
+        base_url: Base URL of the pipeline service.
+        timeout: Request timeout in seconds.
+
+    Example:
+        >>> client = PipelineServiceClient("http://localhost:8100")
+        >>> pipeline = client.create_pipeline(
+        ...     name="YOLO Training",
+        ...     actions=[{"name": "download", "entrypoint": "plugin.download.DownloadAction"}]
+        ... )
+        >>> run = client.create_run(pipeline["id"], params={"dataset_id": 123})
+        >>> client.report_progress(run["id"], current_action="download", status="running")
+    """
+
+    def __init__(
+        self,
+        base_url: str = 'http://localhost:8100',
+        timeout: float = 30.0,
+    ):
+        """Initialize the pipeline service client.
+
+        Args:
+            base_url: Base URL of the pipeline service.
+            timeout: Request timeout in seconds.
+        """
+        self.base_url = base_url.rstrip('/')
+        self.timeout = timeout
+        self._client: httpx.Client | None = None
+
+    @property
+    def client(self) -> httpx.Client:
+        """Get or create the HTTP client."""
+        if self._client is None:
+            self._client = httpx.Client(
+                base_url=self.base_url,
+                timeout=self.timeout,
+                headers={'Content-Type': 'application/json'},
+            )
+        return self._client
+
+    def close(self) -> None:
+        """Close the HTTP client."""
+        if self._client is not None:
+            self._client.close()
+            self._client = None
+
+    def __enter__(self) -> 'PipelineServiceClient':
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.close()
+
+    # -------------------------------------------------------------------------
+    # Pipeline CRUD
+    # -------------------------------------------------------------------------
+
+    def create_pipeline(
+        self,
+        name: str,
+        actions: list[dict[str, Any]],
+        description: str | None = None,
+    ) -> dict[str, Any]:
+        """Create a new pipeline definition.
+
+        Args:
+            name: Pipeline name.
+            actions: List of action definitions, each with 'name' and 'entrypoint'.
+            description: Optional pipeline description.
+
+        Returns:
+            Created pipeline data including 'id'.
+        """
+        payload = {
+            'name': name,
+            'actions': actions,
+            'description': description,
+        }
+        response = self.client.post('/api/v1/pipelines/', json=payload)
+        response.raise_for_status()
+        return response.json()
+
+    def get_pipeline(self, pipeline_id: str) -> dict[str, Any]:
+        """Get a pipeline by ID.
+
+        Args:
+            pipeline_id: Pipeline identifier.
+
+        Returns:
+            Pipeline data.
+        """
+        response = self.client.get(f'/api/v1/pipelines/{pipeline_id}')
+        response.raise_for_status()
+        return response.json()
+
+    def list_pipelines(self, skip: int = 0, limit: int = 100) -> list[dict[str, Any]]:
+        """List all pipelines.
+
+        Args:
+            skip: Number of records to skip.
+            limit: Maximum number of records to return.
+
+        Returns:
+            List of pipeline data.
+        """
+        response = self.client.get('/api/v1/pipelines/', params={'skip': skip, 'limit': limit})
+        response.raise_for_status()
+        return response.json()
+
+    def delete_pipeline(self, pipeline_id: str) -> None:
+        """Delete a pipeline.
+
+        Args:
+            pipeline_id: Pipeline identifier.
+        """
+        response = self.client.delete(f'/api/v1/pipelines/{pipeline_id}')
+        response.raise_for_status()
+
+    # -------------------------------------------------------------------------
+    # Run Management
+    # -------------------------------------------------------------------------
+
+    def create_run(
+        self,
+        pipeline_id: str,
+        params: dict[str, Any] | None = None,
+        work_dir: str | None = None,
+    ) -> dict[str, Any]:
+        """Create a new run for a pipeline.
+
+        Args:
+            pipeline_id: Pipeline to run.
+            params: Initial parameters for the run.
+            work_dir: Working directory path.
+
+        Returns:
+            Created run data including 'id' and initial progress.
+        """
+        payload: dict[str, Any] = {}
+        if params is not None:
+            payload['params'] = params
+        if work_dir is not None:
+            payload['work_dir'] = work_dir
+
+        response = self.client.post(f'/api/v1/pipelines/{pipeline_id}/runs/', json=payload)
+        response.raise_for_status()
+        return response.json()
+
+    def get_run(self, run_id: str) -> dict[str, Any]:
+        """Get a run by ID.
+
+        Args:
+            run_id: Run identifier.
+
+        Returns:
+            Run data including status and progress.
+        """
+        response = self.client.get(f'/api/v1/runs/{run_id}')
+        response.raise_for_status()
+        return response.json()
+
+    def list_runs(
+        self,
+        status: str | None = None,
+        skip: int = 0,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        """List all runs, optionally filtered by status.
+
+        Args:
+            status: Filter by run status.
+            skip: Number of records to skip.
+            limit: Maximum number of records to return.
+
+        Returns:
+            List of run data.
+        """
+        params: dict[str, Any] = {'skip': skip, 'limit': limit}
+        if status:
+            params['status_filter'] = status
+        response = self.client.get('/api/v1/runs/', params=params)
+        response.raise_for_status()
+        return response.json()
+
+    def update_run(
+        self,
+        run_id: str,
+        status: str | None = None,
+        result: dict[str, Any] | None = None,
+        error: str | None = None,
+    ) -> dict[str, Any]:
+        """Update a run's status or result.
+
+        Args:
+            run_id: Run identifier.
+            status: New status.
+            result: Final result data.
+            error: Error message if failed.
+
+        Returns:
+            Updated run data.
+        """
+        payload: dict[str, Any] = {}
+        if status is not None:
+            payload['status'] = status
+        if result is not None:
+            payload['result'] = result
+        if error is not None:
+            payload['error'] = error
+
+        response = self.client.patch(f'/api/v1/runs/{run_id}', json=payload)
+        response.raise_for_status()
+        return response.json()
+
+    def delete_run(self, run_id: str) -> None:
+        """Delete a run.
+
+        Args:
+            run_id: Run identifier.
+        """
+        response = self.client.delete(f'/api/v1/runs/{run_id}')
+        response.raise_for_status()
+
+    # -------------------------------------------------------------------------
+    # Progress Reporting
+    # -------------------------------------------------------------------------
+
+    def report_progress(
+        self,
+        run_id: str,
+        current_action: str | None = None,
+        current_action_index: int | None = None,
+        status: str | None = None,
+        action_progress: ActionProgress | dict[str, Any] | None = None,
+        error: str | None = None,
+    ) -> dict[str, Any]:
+        """Report progress update for a run.
+
+        Args:
+            run_id: Run identifier.
+            current_action: Name of current action.
+            current_action_index: Index of current action.
+            status: Overall run status.
+            action_progress: Progress for the current action.
+            error: Error message if any.
+
+        Returns:
+            Updated run data.
+        """
+        payload: dict[str, Any] = {}
+        if current_action is not None:
+            payload['current_action'] = current_action
+        if current_action_index is not None:
+            payload['current_action_index'] = current_action_index
+        if status is not None:
+            payload['status'] = status
+        if action_progress is not None:
+            if isinstance(action_progress, ActionProgress):
+                payload['action_progress'] = action_progress.to_dict()
+            else:
+                payload['action_progress'] = action_progress
+        if error is not None:
+            payload['error'] = error
+
+        response = self.client.post(f'/api/v1/runs/{run_id}/progress', json=payload)
+        response.raise_for_status()
+        return response.json()
+
+    def get_progress(self, run_id: str) -> PipelineProgress:
+        """Get current progress for a run.
+
+        Args:
+            run_id: Run identifier.
+
+        Returns:
+            PipelineProgress object with current state.
+        """
+        response = self.client.get(f'/api/v1/runs/{run_id}/progress')
+        response.raise_for_status()
+        data = response.json()
+
+        # Map API response to PipelineProgress
+        return PipelineProgress(
+            run_id=data['run_id'],
+            pipeline_id=data.get('pipeline_id', ''),
+            status=RunStatus(data.get('status', 'pending')),
+            current_action=data.get('current_action'),
+            current_action_index=data.get('current_action_index', 0),
+            actions=[ActionProgress.from_dict(a) for a in data.get('progress', [])],
+            started_at=datetime.fromisoformat(data['started_at']) if data.get('started_at') else None,
+            completed_at=datetime.fromisoformat(data['completed_at']) if data.get('completed_at') else None,
+            error=data.get('error'),
+        )
+
+    # -------------------------------------------------------------------------
+    # Checkpoints
+    # -------------------------------------------------------------------------
+
+    def create_checkpoint(
+        self,
+        run_id: str,
+        action_name: str,
+        action_index: int,
+        status: str,
+        params_snapshot: dict[str, Any] | None = None,
+        result: dict[str, Any] | None = None,
+        artifacts_path: str | None = None,
+    ) -> dict[str, Any]:
+        """Create a checkpoint for a run.
+
+        Args:
+            run_id: Run identifier.
+            action_name: Name of the action.
+            action_index: Index of the action.
+            status: Action status at checkpoint.
+            params_snapshot: Parameters at time of checkpoint.
+            result: Result from the action if completed.
+            artifacts_path: Path to saved artifacts.
+
+        Returns:
+            Created checkpoint data.
+        """
+        payload = {
+            'action_name': action_name,
+            'action_index': action_index,
+            'status': status,
+            'params_snapshot': params_snapshot,
+            'result': result,
+            'artifacts_path': artifacts_path,
+        }
+        response = self.client.post(f'/api/v1/runs/{run_id}/checkpoints/', json=payload)
+        response.raise_for_status()
+        return response.json()
+
+    def get_checkpoints(self, run_id: str) -> list[dict[str, Any]]:
+        """Get all checkpoints for a run.
+
+        Args:
+            run_id: Run identifier.
+
+        Returns:
+            List of checkpoint data.
+        """
+        response = self.client.get(f'/api/v1/runs/{run_id}/checkpoints/')
+        response.raise_for_status()
+        return response.json()
+
+    def get_latest_checkpoint(self, run_id: str) -> dict[str, Any] | None:
+        """Get the latest checkpoint for a run.
+
+        Args:
+            run_id: Run identifier.
+
+        Returns:
+            Latest checkpoint data or None if no checkpoints.
+        """
+        try:
+            response = self.client.get(f'/api/v1/runs/{run_id}/checkpoints/latest')
+            response.raise_for_status()
+            return response.json()
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                return None
+            raise
+
+    def get_checkpoint_by_action(self, run_id: str, action_name: str) -> dict[str, Any] | None:
+        """Get checkpoint for a specific action.
+
+        Args:
+            run_id: Run identifier.
+            action_name: Action name.
+
+        Returns:
+            Checkpoint data or None if not found.
+        """
+        try:
+            response = self.client.get(f'/api/v1/runs/{run_id}/checkpoints/{action_name}')
+            response.raise_for_status()
+            return response.json()
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                return None
+            raise
+
+    # -------------------------------------------------------------------------
+    # Logs
+    # -------------------------------------------------------------------------
+
+    def append_logs(
+        self,
+        run_id: str,
+        entries: list[LogEntry] | list[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
+        """Append log entries for a run.
+
+        Args:
+            run_id: Run identifier.
+            entries: List of log entries.
+
+        Returns:
+            Created log entries.
+        """
+        # Convert LogEntry objects to dicts
+        entry_dicts = []
+        for entry in entries:
+            if isinstance(entry, LogEntry):
+                entry_dicts.append({
+                    'message': entry.message,
+                    'level': entry.level.value,
+                    'action_name': entry.action_name,
+                })
+            else:
+                entry_dicts.append(entry)
+
+        payload = {'entries': entry_dicts}
+        response = self.client.post(f'/api/v1/runs/{run_id}/logs/', json=payload)
+        response.raise_for_status()
+        return response.json()
+
+    def get_logs(
+        self,
+        run_id: str,
+        action_name: str | None = None,
+        level: str | None = None,
+        since: datetime | None = None,
+        limit: int = 1000,
+    ) -> list[dict[str, Any]]:
+        """Get logs for a run with optional filters.
+
+        Args:
+            run_id: Run identifier.
+            action_name: Filter by action name.
+            level: Filter by log level.
+            since: Only return logs after this time.
+            limit: Maximum number of logs to return.
+
+        Returns:
+            List of log entries.
+        """
+        params: dict[str, Any] = {'limit': limit}
+        if action_name:
+            params['action_name'] = action_name
+        if level:
+            params['level'] = level
+        if since:
+            params['since'] = since.isoformat()
+
+        response = self.client.get(f'/api/v1/runs/{run_id}/logs/', params=params)
+        response.raise_for_status()
+        return response.json()
+
+    # -------------------------------------------------------------------------
+    # Progress Streaming
+    # -------------------------------------------------------------------------
+
+    def stream_progress(
+        self,
+        run_id: str,
+        timeout: float = 3600.0,
+    ) -> 'Iterator[PipelineProgress]':
+        """Stream progress updates via SSE.
+
+        Yields PipelineProgress objects as updates are received from the server.
+        Continues until the run completes, fails, or is cancelled.
+
+        Args:
+            run_id: Run identifier.
+            timeout: Maximum time to stream in seconds.
+
+        Yields:
+            PipelineProgress objects with current state.
+
+        Example:
+            >>> for progress in client.stream_progress(run_id):
+            ...     print(f"Status: {progress.status}, Action: {progress.current_action}")
+        """
+        url = f'{self.base_url}/api/v1/runs/{run_id}/progress/stream'
+
+        with httpx.stream('GET', url, timeout=timeout) as response:
+            response.raise_for_status()
+
+            event_type = None
+            data_buffer = []
+
+            for line in response.iter_lines():
+                line = line.strip()
+
+                if line.startswith('event:'):
+                    event_type = line[6:].strip()
+                elif line.startswith('data:'):
+                    data_buffer.append(line[5:].strip())
+                elif line == '' and data_buffer:
+                    # End of event - process it
+                    data_str = ''.join(data_buffer)
+                    data_buffer = []
+
+                    try:
+                        data = json.loads(data_str)
+                    except json.JSONDecodeError:
+                        continue
+
+                    # Convert to PipelineProgress
+                    progress = PipelineProgress(
+                        run_id=data['run_id'],
+                        pipeline_id=data.get('pipeline_id', ''),
+                        status=RunStatus(data.get('status', 'pending')),
+                        current_action=data.get('current_action'),
+                        current_action_index=data.get('current_action_index', 0),
+                        actions=[ActionProgress.from_dict(a) for a in data.get('progress', []) or []],
+                        started_at=datetime.fromisoformat(data['started_at']) if data.get('started_at') else None,
+                        completed_at=datetime.fromisoformat(data['completed_at']) if data.get('completed_at') else None,
+                        error=data.get('error'),
+                    )
+
+                    yield progress
+
+                    # Stop on terminal events
+                    if event_type in ('completed', 'failed', 'cancelled', 'error'):
+                        return
+
+    async def stream_progress_async(
+        self,
+        run_id: str,
+        timeout: float = 3600.0,
+    ) -> AsyncIterator[PipelineProgress]:
+        """Stream progress updates via SSE (async version).
+
+        Yields PipelineProgress objects as updates are received from the server.
+        Continues until the run completes, fails, or is cancelled.
+
+        Args:
+            run_id: Run identifier.
+            timeout: Maximum time to stream in seconds.
+
+        Yields:
+            PipelineProgress objects with current state.
+
+        Example:
+            >>> async for progress in client.stream_progress_async(run_id):
+            ...     print(f"Status: {progress.status}, Action: {progress.current_action}")
+        """
+        url = f'{self.base_url}/api/v1/runs/{run_id}/progress/stream'
+
+        async with httpx.AsyncClient(timeout=timeout) as async_client:
+            async with async_client.stream('GET', url) as response:
+                response.raise_for_status()
+
+                event_type = None
+                data_buffer: list[str] = []
+
+                async for line in response.aiter_lines():
+                    line = line.strip()
+
+                    if line.startswith('event:'):
+                        event_type = line[6:].strip()
+                    elif line.startswith('data:'):
+                        data_buffer.append(line[5:].strip())
+                    elif line == '' and data_buffer:
+                        # End of event - process it
+                        data_str = ''.join(data_buffer)
+                        data_buffer = []
+
+                        try:
+                            data = json.loads(data_str)
+                        except json.JSONDecodeError:
+                            continue
+
+                        # Convert to PipelineProgress
+                        started = data.get('started_at')
+                        completed = data.get('completed_at')
+                        progress = PipelineProgress(
+                            run_id=data['run_id'],
+                            pipeline_id=data.get('pipeline_id', ''),
+                            status=RunStatus(data.get('status', 'pending')),
+                            current_action=data.get('current_action'),
+                            current_action_index=data.get('current_action_index', 0),
+                            actions=[ActionProgress.from_dict(a) for a in data.get('progress', []) or []],
+                            started_at=datetime.fromisoformat(started) if started else None,
+                            completed_at=datetime.fromisoformat(completed) if completed else None,
+                            error=data.get('error'),
+                        )
+
+                        yield progress
+
+                        # Stop on terminal events
+                        if event_type in ('completed', 'failed', 'cancelled', 'error'):
+                            return
+
+    # -------------------------------------------------------------------------
+    # Health Check
+    # -------------------------------------------------------------------------
+
+    def health_check(self) -> bool:
+        """Check if the pipeline service is healthy.
+
+        Returns:
+            True if healthy, False otherwise.
+        """
+        try:
+            response = self.client.get('/health')
+            return response.status_code == 200
+        except Exception:
+            return False
+
+
+__all__ = ['PipelineServiceClient']