mlops-python-sdk 0.0.1__py3-none-any.whl

mlops/api/client/models/task_submit_response.py

@@ -0,0 +1,78 @@
+from collections.abc import Mapping
+from typing import Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="TaskSubmitResponse")
+
+
+@_attrs_define
+class TaskSubmitResponse:
+    """Task submission response
+
+    Attributes:
+        id (Union[Unset, int]): Database ID (slurm_jobs.id) Example: 1.
+        job_id (Union[Unset, str]): Slurm job ID (assigned by scheduler) Example: 12345.
+        message (Union[Unset, str]):  Example: Task submitted successfully.
+    """
+
+    id: Union[Unset, int] = UNSET
+    job_id: Union[Unset, str] = UNSET
+    message: Union[Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        id = self.id
+
+        job_id = self.job_id
+
+        message = self.message
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if id is not UNSET:
+            field_dict["id"] = id
+        if job_id is not UNSET:
+            field_dict["job_id"] = job_id
+        if message is not UNSET:
+            field_dict["message"] = message
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        id = d.pop("id", UNSET)
+
+        job_id = d.pop("job_id", UNSET)
+
+        message = d.pop("message", UNSET)
+
+        task_submit_response = cls(
+            id=id,
+            job_id=job_id,
+            message=message,
+        )
+
+        task_submit_response.additional_properties = d
+        return task_submit_response
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

mlops/api/client/models/task_tres_type_0.py

@@ -0,0 +1,49 @@
+from collections.abc import Mapping
+from typing import Any, TypeVar
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+T = TypeVar("T", bound="TaskTresType0")
+
+
+@_attrs_define
+class TaskTresType0:
+    """Trackable Resources
+
+    Example:
+        {'cpu': 4, 'mem': 8589934592}
+
+    """
+
+    additional_properties: dict[str, int] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        task_tres_type_0 = cls()
+
+        task_tres_type_0.additional_properties = d
+        return task_tres_type_0
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> int:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: int) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

mlops/api/client/models/task_tres_used_type_0.py

@@ -0,0 +1,49 @@
+from collections.abc import Mapping
+from typing import Any, TypeVar
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+T = TypeVar("T", bound="TaskTresUsedType0")
+
+
+@_attrs_define
+class TaskTresUsedType0:
+    """Trackable Resources Used
+
+    Example:
+        {'cpu': 4, 'mem': 4294967296}
+
+    """
+
+    additional_properties: dict[str, int] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        task_tres_used_type_0 = cls()
+
+        task_tres_used_type_0.additional_properties = d
+        return task_tres_used_type_0
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> int:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: int) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

mlops/api/client/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561

mlops/api/client/types.py

@@ -0,0 +1,54 @@
+"""Contains some shared types for properties"""
+
+from collections.abc import Mapping, MutableMapping
+from http import HTTPStatus
+from typing import IO, BinaryIO, Generic, Literal, Optional, TypeVar, Union
+
+from attrs import define
+
+
+class Unset:
+    def __bool__(self) -> Literal[False]:
+        return False
+
+
+UNSET: Unset = Unset()
+
+# The types that `httpx.Client(files=)` can accept, copied from that library.
+FileContent = Union[IO[bytes], bytes, str]
+FileTypes = Union[
+    # (filename, file (or bytes), content_type)
+    tuple[Optional[str], FileContent, Optional[str]],
+    # (filename, file (or bytes), content_type, headers)
+    tuple[Optional[str], FileContent, Optional[str], Mapping[str, str]],
+]
+RequestFiles = list[tuple[str, FileTypes]]
+
+
+@define
+class File:
+    """Contains information for file uploads"""
+
+    payload: BinaryIO
+    file_name: Optional[str] = None
+    mime_type: Optional[str] = None
+
+    def to_tuple(self) -> FileTypes:
+        """Return a tuple representation that httpx will accept for multipart/form-data"""
+        return self.file_name, self.payload, self.mime_type
+
+
+T = TypeVar("T")
+
+
+@define
+class Response(Generic[T]):
+    """A response from an endpoint"""
+
+    status_code: HTTPStatus
+    content: bytes
+    headers: MutableMapping[str, str]
+    parsed: Optional[T]
+
+
+__all__ = ["UNSET", "File", "FileTypes", "RequestFiles", "Response", "Unset"]

mlops/connection_config.py

@@ -0,0 +1,106 @@
+import os
+
+from typing import Literal, Optional, Dict
+from httpx._types import ProxyTypes
+
+REQUEST_TIMEOUT: float = 30.0  # 30 seconds
+
+KEEPALIVE_PING_INTERVAL_SEC = 50  # 50 seconds
+KEEPALIVE_PING_HEADER = "Keepalive-Ping-Interval"
+
+# API path prefix
+DEFAULT_API_PATH: str = "/api/v1"
+
+
+class ConnectionConfig:
+    """
+    Configuration for the connection to the API.
+    """
+
+    @staticmethod
+    def _domain():
+        return os.getenv("XCLIENT_DOMAIN", "localhost:8090")
+
+    @staticmethod
+    def _debug():
+        return os.getenv("XCLIENT_DEBUG", "false").lower() == "true"
+
+    @staticmethod
+    def _api_key():
+        return os.getenv("XCLIENT_API_KEY")
+
+    @staticmethod
+    def _access_token():
+        return os.getenv("XCLIENT_ACCESS_TOKEN")
+
+    @staticmethod
+    def _api_path():
+        return os.getenv("XCLIENT_API_PATH", DEFAULT_API_PATH)
+
+    def __init__(
+        self,
+        domain: Optional[str] = None,
+        debug: Optional[bool] = None,
+        api_key: Optional[str] = None,
+        access_token: Optional[str] = None,
+        request_timeout: Optional[float] = None,
+        headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
+        api_path: Optional[str] = None,
+    ):
+        self.domain = domain or ConnectionConfig._domain()
+        self.debug = debug or ConnectionConfig._debug()
+        self.api_key = api_key or ConnectionConfig._api_key()
+        self.access_token = access_token or ConnectionConfig._access_token()
+        self.headers = headers or {}
+        self.proxy = proxy
+        self.api_path = api_path or ConnectionConfig._api_path()
+
+        self.request_timeout = ConnectionConfig._get_request_timeout(
+            REQUEST_TIMEOUT,
+            request_timeout,
+        )
+
+        if request_timeout == 0:
+            self.request_timeout = None
+        elif request_timeout is not None:
+            self.request_timeout = request_timeout
+        else:
+            self.request_timeout = REQUEST_TIMEOUT
+
+        # Ensure api_path starts with /
+        if not self.api_path.startswith("/"):
+            self.api_path = "/" + self.api_path
+
+        # Build API URL
+        if self.debug:
+            base_url = "http://localhost:8090"
+        else:
+            # If domain already includes protocol, use it as-is
+            # Otherwise, default to http:// for backward compatibility
+            if self.domain.startswith(("http://", "https://")):
+                base_url = self.domain
+            else:
+                base_url = f"http://{self.domain}"
+        
+        self.api_url = f"{base_url}{self.api_path}"
+
+    @staticmethod
+    def _get_request_timeout(
+        default_timeout: Optional[float],
+        request_timeout: Optional[float],
+    ):
+        if request_timeout == 0:
+            return None
+        elif request_timeout is not None:
+            return request_timeout
+        else:
+            return default_timeout
+
+    def get_request_timeout(self, request_timeout: Optional[float] = None):
+        return self._get_request_timeout(self.request_timeout, request_timeout)
+
+
+# Re-export ProxyTypes for convenience
+__all__ = ["ConnectionConfig", "ProxyTypes"]
+

mlops/exceptions.py

@@ -0,0 +1,82 @@
+def format_request_timeout_error() -> Exception:
+    return TimeoutException(
+        "Request timed out — the 'request_timeout' option can be used to increase this timeout",
+    )
+
+
+def format_execution_timeout_error() -> Exception:
+    return TimeoutException(
+        "Execution timed out — the 'timeout' option can be used to increase this timeout",
+    )
+
+
+class XClientException(Exception):
+    """
+    Base class for all XClient errors.
+
+    Raised when a general XClient exception occurs.
+    """
+
+    pass
+
+
+class TimeoutException(XClientException):
+    """
+    Raised when a timeout occurs.
+
+    The `unavailable` exception type is caused by service timeout.\n
+    The `canceled` exception type is caused by exceeding request timeout.\n
+    The `deadline_exceeded` exception type is caused by exceeding the timeout for process, watch, etc.\n
+    The `unknown` exception type is sometimes caused by the service timeout when the request is not processed correctly.\n
+    """
+
+    pass
+
+
+class InvalidArgumentException(XClientException):
+    """
+    Raised when an invalid argument is provided.
+    """
+
+    pass
+
+
+class NotEnoughSpaceException(XClientException):
+    """
+    Raised when there is not enough disk space.
+    """
+
+    pass
+
+
+class NotFoundException(XClientException):
+    """
+    Raised when a resource is not found.
+    """
+
+    pass
+
+
+class AuthenticationException(XClientException):
+    """
+    Raised when authentication fails.
+    """
+
+    pass
+
+
+class RateLimitException(XClientException):
+    """
+    Raised when the API rate limit is exceeded.
+    """
+
+    pass
+
+
+class APIException(XClientException):
+    """
+    Raised when an API error occurs.
+    """
+
+    pass
+

mlops/task/__init__.py

@@ -0,0 +1,10 @@
+"""Task SDK module for XClient"""
+
+from .client import TaskClient
+from .task import Task
+
+__all__ = [
+    "Task",
+    "TaskClient",
+]
+

mlops/task/client.py

@@ -0,0 +1,146 @@
+import json
+import logging
+from typing import Optional
+from httpx import Limits
+
+from ..api.client.client import AuthenticatedClient
+from ..connection_config import ConnectionConfig
+from ..exceptions import (
+    AuthenticationException,
+    RateLimitException,
+    NotFoundException,
+    APIException,
+)
+from ..api.client.types import Response
+
+logger = logging.getLogger(__name__)
+
+
+def handle_api_exception(e: Response):
+    """Handle API exceptions and convert them to appropriate XClient exceptions."""
+    try:
+        body = json.loads(e.content) if e.content else {}
+    except json.JSONDecodeError:
+        body = {}
+
+    if e.status_code == 401:
+        return AuthenticationException(
+            f"Authentication failed: {body.get('error', 'Invalid credentials')}"
+        )
+    
+    if e.status_code == 404:
+        return NotFoundException(
+            f"Resource not found: {body.get('error', 'The requested resource was not found')}"
+        )
+
+    if e.status_code == 429:
+        return RateLimitException(
+            f"{e.status_code}: Rate limit exceeded, please try again later."
+        )
+
+    if "error" in body:
+        return APIException(f"{e.status_code}: {body['error']}")
+    
+    if "message" in body:
+        return APIException(f"{e.status_code}: {body['message']}")
+    
+    return APIException(f"{e.status_code}: {e.content}")
+
+
+class TaskClient(AuthenticatedClient):
+    """
+    The client for interacting with the XClient Task API.
+    """
+
+    def __init__(
+        self,
+        config: ConnectionConfig,
+        require_api_key: bool = True,
+        require_access_token: bool = False,
+        limits: Optional[Limits] = None,
+        *args,
+        **kwargs,
+    ):
+        if require_api_key and require_access_token:
+            raise AuthenticationException(
+                "Only one of api_key or access_token can be required, not both",
+            )
+
+        if not require_api_key and not require_access_token:
+            raise AuthenticationException(
+                "Either api_key or access_token is required",
+            )
+
+        token = None
+        if require_api_key:
+            if config.api_key is None:
+                raise AuthenticationException(
+                    "API key is required. "
+                    "You can either set the environment variable `XCLIENT_API_KEY` "
+                    'or pass it directly like TaskClient(api_key="xclient_...")',
+                )
+            token = config.api_key
+
+        if require_access_token:
+            if config.access_token is None:
+                raise AuthenticationException(
+                    "Access token is required. "
+                    "You can set the environment variable `XCLIENT_ACCESS_TOKEN` "
+                    "or pass the `access_token` in options.",
+                )
+            token = config.access_token
+
+        # API Key header: X-API-Key (per OpenAPI spec)
+        # JWT header: Authorization: Bearer <token>
+        auth_header_name = "X-API-Key" if require_api_key else "Authorization"
+        prefix = "" if require_api_key else "Bearer"
+
+        headers = {
+            **(config.headers or {}),
+        }
+
+        httpx_args = {
+            "event_hooks": {
+                "request": [self._log_request],
+                "response": [self._log_response],
+            },
+        }
+        
+        if config.proxy is not None:
+            httpx_args["proxy"] = config.proxy
+        
+        if limits is not None:
+            httpx_args["limits"] = limits
+
+        super().__init__(
+            base_url=config.api_url,
+            httpx_args=httpx_args,
+            headers=headers,
+            token=token,
+            auth_header_name=auth_header_name,
+            prefix=prefix,
+            *args,
+            **kwargs,
+        )
+
+    def _log_request(self, request):
+        logger.info(f"Request {request.method} {request.url}")
+
+    def _log_response(self, response: Response):
+        if response.status_code >= 400:
+            logger.error(f"Response {response.status_code}")
+        else:
+            logger.info(f"Response {response.status_code}")
+
+
+# We need to override the logging hooks for the async usage
+class AsyncTaskClient(TaskClient):
+    async def _log_request(self, request):
+        logger.info(f"Request {request.method} {request.url}")
+
+    async def _log_response(self, response: Response):
+        if response.status_code >= 400:
+            logger.error(f"Response {response.status_code}")
+        else:
+            logger.info(f"Response {response.status_code}")
+