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

synapse_sdk/clients/agent/ray.py

@@ -1,54 +1,311 @@
-from synapse_sdk.clients.base import BaseClient
-from synapse_sdk.clients.exceptions import ClientError
+from __future__ import annotations
 
+import json
+from typing import TYPE_CHECKING, Generator, Literal
 
-class RayClientMixin(BaseClient):
-    def get_job(self, pk):
-        path = f'jobs/{pk}/'
-        return self._get(path)
+from synapse_sdk.exceptions import ClientError
+from synapse_sdk.utils.network import (
+    StreamLimits,
+    http_to_websocket_url,
+    sanitize_error_message,
+    validate_resource_id,
+    validate_timeout,
+)
 
-    def list_jobs(self):
-        path = 'jobs/'
-        return self._get(path)
+if TYPE_CHECKING:
+    from synapse_sdk.clients.protocols import ClientProtocol
 
-    def list_job_logs(self, pk):
-        path = f'jobs/{pk}/logs/'
-        return self._get(path)
+StreamProtocol = Literal['websocket', 'http', 'auto']
 
-    def tail_job_logs(self, pk):
-        if self.long_poll_handler:
-            raise ClientError(400, '"tail_job_logs" does not support long polling')
 
-        path = f'jobs/{pk}/tail_logs/'
+class RayClientMixin:
+    """Mixin for Ray cluster management endpoints."""
 
-        url = self._get_url(path)
+    _stream_limits: StreamLimits | None = None
+
+    @property
+    def stream_limits(self) -> StreamLimits:
+        """Get stream limits configuration."""
+        if self._stream_limits is None:
+            self._stream_limits = StreamLimits()
+        return self._stream_limits
+
+    @stream_limits.setter
+    def stream_limits(self, value: StreamLimits) -> None:
+        """Set stream limits configuration."""
+        self._stream_limits = value
+
+    # -------------------------------------------------------------------------
+    # Jobs
+    # -------------------------------------------------------------------------
+
+    def list_jobs(self: ClientProtocol) -> list[dict]:
+        """List all Ray jobs."""
+        return self._get('ray/jobs/')
+
+    def get_job(self: ClientProtocol, job_id: str) -> dict:
+        """Get a Ray job by ID."""
+        return self._get(f'ray/jobs/{job_id}')
+
+    def get_job_logs(self: ClientProtocol, job_id: str) -> str:
+        """Get all logs for a job (non-streaming)."""
+        return self._get(f'ray/jobs/{job_id}/logs')
+
+    def stop_job(self: ClientProtocol, job_id: str) -> dict:
+        """Stop a running job."""
+        return self._post(f'ray/jobs/{job_id}/stop/')
+
+    def websocket_tail_job_logs(
+        self: ClientProtocol,
+        job_id: str,
+        timeout: float = 30.0,
+    ) -> Generator[str, None, None]:
+        """Stream job logs via WebSocket protocol.
+
+        Establishes a WebSocket connection for real-time log streaming.
+        Preferred method for low-latency log monitoring.
+
+        Args:
+            job_id: The Ray job ID to tail logs for.
+            timeout: Connection and read timeout in seconds.
+
+        Yields:
+            Log message strings.
+
+        Raises:
+            ClientError: On connection, protocol, or validation errors.
+
+        Example:
+            >>> for line in client.websocket_tail_job_logs('raysubmit_abc123'):
+            ...     print(line)
+        """
+        validated_id = validate_resource_id(job_id, 'job')
+        validated_timeout = validate_timeout(timeout)
+
+        url = self._get_url(f'ray/jobs/{validated_id}/logs/')
+        ws_url = http_to_websocket_url(url)
         headers = self._get_headers()
 
-        response = self.requests_session.get(url, headers=headers, stream=True)
-        for line in response.iter_lines(decode_unicode=True):
-            if line:
-                yield f'{line}\n'
+        try:
+            import websocket
+        except ImportError:
+            raise ClientError(500, 'websocket-client package required for WebSocket streaming')
+
+        header_list = [f'{k}: {v}' for k, v in headers.items()]
+
+        try:
+            ws = websocket.create_connection(
+                ws_url,
+                header=header_list,
+                timeout=validated_timeout,
+            )
+        except websocket.WebSocketException as e:
+            raise ClientError(503, f'WebSocket connection failed: {e}')
+        except Exception as e:
+            raise ClientError(503, sanitize_error_message(str(e), 'connection error'))
+
+        limits = self.stream_limits
+        message_count = 0
+
+        try:
+            while True:
+                try:
+                    data = ws.recv()
+                except websocket.WebSocketTimeoutException:
+                    break
+                except websocket.WebSocketConnectionClosedException:
+                    break
+
+                if not data:
+                    break
+
+                message_count += 1
+                if message_count > limits.max_messages:
+                    raise ClientError(429, 'Stream message limit exceeded')
+
+                # Skip oversized messages
+                if len(data) > limits.max_message_size:
+                    continue
+
+                # Parse JSON event format
+                try:
+                    event = json.loads(data)
+                except json.JSONDecodeError:
+                    event = {'message': data}
+
+                # Handle event types
+                event_type = event.get('type')
+                if event_type == 'error':
+                    raise ClientError(500, event.get('message', 'Unknown error'))
+                elif event_type == 'complete':
+                    return
+
+                if msg := event.get('message'):
+                    yield msg
+        finally:
+            ws.close()
+
+    def stream_tail_job_logs(
+        self: ClientProtocol,
+        job_id: str,
+        timeout: float = 30.0,
+    ) -> Generator[str, None, None]:
+        """Stream job logs via HTTP chunked transfer encoding.
+
+        Uses HTTP streaming as an alternative when WebSocket is unavailable.
+
+        Args:
+            job_id: The Ray job ID to tail logs for.
+            timeout: Connection timeout in seconds (read timeout is infinite).
+
+        Yields:
+            Log lines as strings.
+
+        Raises:
+            ClientError: On connection, protocol, or validation errors.
+
+        Example:
+            >>> for line in client.stream_tail_job_logs('raysubmit_abc123'):
+            ...     print(line)
+        """
+        validated_id = validate_resource_id(job_id, 'job')
+        validated_timeout = validate_timeout(timeout)
+
+        url = self._get_url(f'ray/jobs/{validated_id}/logs/')
+        headers = self._get_headers()
+
+        response = None
+        try:
+            response = self.requests_session.get(
+                url,
+                headers=headers,
+                stream=True,
+                timeout=(validated_timeout, None),  # No read timeout for streaming
+            )
+            response.raise_for_status()
+
+            limits = self.stream_limits
+            line_count = 0
+            total_bytes = 0
+
+            for line in response.iter_lines(decode_unicode=True):
+                if line:
+                    line_count += 1
+                    total_bytes += len(line.encode('utf-8'))
+
+                    if line_count > limits.max_lines:
+                        raise ClientError(429, 'Stream line limit exceeded')
+
+                    if total_bytes > limits.max_bytes:
+                        raise ClientError(429, 'Stream size limit exceeded')
+
+                    if len(line) > limits.max_message_size:
+                        continue
+
+                    yield line
+
+        except ClientError:
+            raise
+        except Exception as e:
+            raise ClientError(503, sanitize_error_message(str(e), 'HTTP streaming error'))
+        finally:
+            if response is not None:
+                response.close()
+
+    def tail_job_logs(
+        self: ClientProtocol,
+        job_id: str,
+        timeout: float = 30.0,
+        *,
+        protocol: StreamProtocol = 'auto',
+    ) -> Generator[str, None, None]:
+        """Stream job logs with automatic protocol selection.
+
+        Unified method that supports WebSocket, HTTP, and auto-selection.
+
+        Args:
+            job_id: The Ray job ID to tail logs for.
+            timeout: Connection timeout in seconds.
+            protocol: Protocol to use:
+                - 'websocket': Use WebSocket only
+                - 'http': Use HTTP streaming only
+                - 'auto': Try WebSocket, fall back to HTTP on connection failure
+
+        Yields:
+            Log message strings.
+
+        Raises:
+            ClientError: On connection, protocol, or validation errors.
+
+        Example:
+            >>> # Auto protocol selection (recommended)
+            >>> for line in client.tail_job_logs('raysubmit_abc123'):
+            ...     print(line)
+
+            >>> # Explicit WebSocket
+            >>> for line in client.tail_job_logs('raysubmit_abc123', protocol='websocket'):
+            ...     print(line)
+
+            >>> # Explicit HTTP streaming
+            >>> for line in client.tail_job_logs('raysubmit_abc123', protocol='http'):
+            ...     print(line)
+        """
+        # Pre-validate to fail fast
+        validate_resource_id(job_id, 'job')
+        validate_timeout(timeout)
+
+        if protocol == 'websocket':
+            yield from self.websocket_tail_job_logs(job_id, timeout)
+        elif protocol == 'http':
+            yield from self.stream_tail_job_logs(job_id, timeout)
+        elif protocol == 'auto':
+            try:
+                yield from self.websocket_tail_job_logs(job_id, timeout)
+            except ClientError as e:
+                # Fall back to HTTP on WebSocket failures
+                if e.status_code in (500, 503):
+                    yield from self.stream_tail_job_logs(job_id, timeout)
+                else:
+                    raise
+        else:
+            raise ClientError(400, f'Invalid protocol: {protocol}')
+
+    # -------------------------------------------------------------------------
+    # Nodes
+    # -------------------------------------------------------------------------
+
+    def list_nodes(self: ClientProtocol) -> list[dict]:
+        """List all Ray nodes."""
+        return self._get('ray/nodes/')
+
+    def get_node(self: ClientProtocol, node_id: str) -> dict:
+        """Get a Ray node by ID."""
+        return self._get(f'ray/nodes/{node_id}')
+
+    # -------------------------------------------------------------------------
+    # Tasks
+    # -------------------------------------------------------------------------
 
-    def get_node(self, pk):
-        path = f'nodes/{pk}/'
-        return self._get(path)
+    def list_tasks(self: ClientProtocol) -> list[dict]:
+        """List all Ray tasks."""
+        return self._get('ray/tasks/')
 
-    def list_nodes(self):
-        path = 'nodes/'
-        return self._get(path)
+    def get_task(self: ClientProtocol, task_id: str) -> dict:
+        """Get a Ray task by ID."""
+        return self._get(f'ray/tasks/{task_id}')
 
-    def get_task(self, pk):
-        path = f'tasks/{pk}/'
-        return self._get(path)
+    # -------------------------------------------------------------------------
+    # Serve Applications
+    # -------------------------------------------------------------------------
 
-    def list_tasks(self):
-        path = 'tasks/'
-        return self._get(path)
+    def list_serve_applications(self: ClientProtocol) -> list[dict]:
+        """List all Ray Serve applications."""
+        return self._get('ray/serve_applications/')
 
-    def get_serve_application(self, pk):
-        path = f'serve_application/{pk}/'
-        return self._get(path)
+    def get_serve_application(self: ClientProtocol, name: str) -> dict:
+        """Get a Ray Serve application by name."""
+        return self._get(f'ray/serve_applications/{name}')
 
-    def list_serve_applications(self):
-        path = 'serve_applications/'
-        return self._get(path)
+    def delete_serve_application(self: ClientProtocol, name: str) -> None:
+        """Delete a Ray Serve application."""
+        self._delete(f'ray/serve_applications/{name}')

synapse_sdk/clients/backend/__init__.py

@@ -1,23 +1,163 @@
+"""Backend API client for synapse-backend.
+
+This module provides the BackendClient for interacting with the synapse-backend API.
+It composes functionality from multiple mixins for different API domains.
+
+Example:
+    >>> from synapse_sdk.clients.backend import BackendClient
+    >>>
+    >>> client = BackendClient(
+    ...     'https://api.example.com',
+    ...     access_token='your_token',
+    ...     tenant='your_tenant',
+    ... )
+    >>>
+    >>> # Get a project
+    >>> project = client.get_project(123)
+    >>>
+    >>> # Upload data collection
+    >>> client.upload_data_collection(456, data, project_id=789)
+"""
+
+from __future__ import annotations
+
 from synapse_sdk.clients.backend.annotation import AnnotationClientMixin
-from synapse_sdk.clients.backend.dataset import DatasetClientMixin
+from synapse_sdk.clients.backend.bulk_upload import BulkUploadClientMixin
+from synapse_sdk.clients.backend.core import CoreClientMixin
+from synapse_sdk.clients.backend.data_collection import DataCollectionClientMixin
+from synapse_sdk.clients.backend.hitl import HITLClientMixin
 from synapse_sdk.clients.backend.integration import IntegrationClientMixin
 from synapse_sdk.clients.backend.ml import MLClientMixin
+from synapse_sdk.clients.base import BaseClient
+
+
+class BackendClient(
+    AnnotationClientMixin,
+    BulkUploadClientMixin,
+    CoreClientMixin,
+    DataCollectionClientMixin,
+    HITLClientMixin,
+    IntegrationClientMixin,
+    MLClientMixin,
+    BaseClient,
+):
+    """Synchronous client for synapse-backend API.
 
+    Composes functionality from multiple mixins:
+    - AnnotationClientMixin: Project and task operations
+    - BulkUploadClientMixin: High-performance presigned URL uploads
+    - CoreClientMixin: Chunked file upload
+    - DataCollectionClientMixin: Data collection management
+    - HITLClientMixin: Assignment operations
+    - IntegrationClientMixin: Plugin, job, and storage operations
+    - MLClientMixin: Model and ground truth operations
+
+    Args:
+        base_url: Backend API base URL.
+        access_token: API access token.
+        authorization_token: Optional authorization token (legacy).
+        tenant: Optional tenant identifier for multi-tenancy.
+        agent_token: Optional agent token for agent-initiated requests.
+        timeout: Request timeout dict with 'connect' and 'read' keys.
+
+    Example:
+        >>> client = BackendClient(
+        ...     'https://api.example.com',
+        ...     access_token='abc123',
+        ...     tenant='my-tenant',
+        ... )
+        >>> project = client.get_project(1)
+    """
 
-class BackendClient(AnnotationClientMixin, DatasetClientMixin, IntegrationClientMixin, MLClientMixin):
     name = 'Backend'
-    token = None
-    tenant = None
 
-    def __init__(self, base_url, token=None, tenant=None):
-        super().__init__(base_url)
-        self.token = token
+    def __init__(
+        self,
+        base_url: str,
+        access_token: str | None = None,
+        *,
+        authorization_token: str | None = None,
+        tenant: str | None = None,
+        agent_token: str | None = None,
+        timeout: dict[str, int] | None = None,
+    ):
+        """Initialize the backend client.
+
+        Args:
+            base_url: Backend API base URL.
+            access_token: API access token for authentication.
+            authorization_token: Legacy auth token (deprecated, use access_token).
+            tenant: Tenant code for multi-tenant deployments.
+            agent_token: Agent token for agent-initiated requests.
+            timeout: Request timeout configuration.
+        """
+        super().__init__(base_url, timeout=timeout)
+        self.access_token = access_token
+        self.authorization_token = authorization_token
         self.tenant = tenant
+        self.agent_token = agent_token
+
+    def _get_headers(self) -> dict[str, str]:
+        """Return authentication headers.
+
+        Multiple authentication methods are supported:
+        - Synapse-Access-Token: Primary authentication
+        - Authorization: Legacy token authentication
+        - Synapse-Tenant: Multi-tenant identifier
+        - SYNAPSE-Agent: Agent-initiated request identifier
+        """
+        headers: dict[str, str] = {}
+
+        if self.access_token:
+            headers['Synapse-Access-Token'] = f'Token {self.access_token}'
+
+        if self.authorization_token:
+            headers['Authorization'] = f'Token {self.authorization_token}'
 
-    def _get_headers(self):
-        headers = {}
-        if self.token:
-            headers = {'Authorization': f'Token {self.token}'}
         if self.tenant:
-            headers['SYNAPSE-Tenant'] = f'Token {self.tenant}'
+            headers['Synapse-Tenant'] = f'Token {self.tenant}'
+
+        if self.agent_token:
+            headers['SYNAPSE-Agent'] = f'Token {self.agent_token}'
+
         return headers
+
+    def close(self) -> None:
+        """Close the HTTP session.
+
+        Call this when done with the client to release resources.
+        """
+        if self._session is not None:
+            self._session.close()
+            self._session = None
+
+    def __enter__(self) -> BackendClient:
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, *args) -> None:
+        """Context manager exit - closes session."""
+        self.close()
+
+
+# Re-export models for convenience
+from synapse_sdk.clients.backend.models import (  # noqa: E402
+    Agent,
+    JobStatus,
+    Storage,
+    StorageCategory,
+    StorageProvider,
+    UpdateJobRequest,
+)
+
+__all__ = [
+    # Client
+    'BackendClient',
+    # Models
+    'Agent',
+    'JobStatus',
+    'Storage',
+    'StorageCategory',
+    'StorageProvider',
+    'UpdateJobRequest',
+]