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

mlops/api/client/models/log_pagination.py

@@ -0,0 +1,90 @@
+from collections.abc import Mapping
+from typing import Any, TypeVar, Union, cast
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="LogPagination")
+
+
+@_attrs_define
+class LogPagination:
+    """Pagination information for logs
+
+    Attributes:
+        has_more (Union[Unset, bool]): Whether there are more logs available Example: True.
+        next_cursor (Union[None, Unset, str]): Cursor for the next page (timestamp in nanoseconds) Example:
+            1706612400123456789.
+        total_fetched (Union[Unset, int]): Number of log entries fetched in this request Example: 1000.
+    """
+
+    has_more: Union[Unset, bool] = UNSET
+    next_cursor: Union[None, Unset, str] = UNSET
+    total_fetched: Union[Unset, int] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        has_more = self.has_more
+
+        next_cursor: Union[None, Unset, str]
+        if isinstance(self.next_cursor, Unset):
+            next_cursor = UNSET
+        else:
+            next_cursor = self.next_cursor
+
+        total_fetched = self.total_fetched
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if has_more is not UNSET:
+            field_dict["has_more"] = has_more
+        if next_cursor is not UNSET:
+            field_dict["next_cursor"] = next_cursor
+        if total_fetched is not UNSET:
+            field_dict["total_fetched"] = total_fetched
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        has_more = d.pop("has_more", UNSET)
+
+        def _parse_next_cursor(data: object) -> Union[None, Unset, str]:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(Union[None, Unset, str], data)
+
+        next_cursor = _parse_next_cursor(d.pop("next_cursor", UNSET))
+
+        total_fetched = d.pop("total_fetched", UNSET)
+
+        log_pagination = cls(
+            has_more=has_more,
+            next_cursor=next_cursor,
+            total_fetched=total_fetched,
+        )
+
+        log_pagination.additional_properties = d
+        return log_pagination
+
+    @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_log_entry.py

@@ -0,0 +1,105 @@
+import datetime
+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 dateutil.parser import isoparse
+
+from ..models.task_log_entry_log_type import TaskLogEntryLogType
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="TaskLogEntry")
+
+
+@_attrs_define
+class TaskLogEntry:
+    """A single log entry
+
+    Attributes:
+        content (Union[Unset, str]): Log line content Example: Training epoch 1/100....
+        log_type (Union[Unset, TaskLogEntryLogType]): Log type (stdout or stderr) Example: stdout.
+        node (Union[Unset, str]): Node hostname where the log was generated Example: compute-node-01.
+        timestamp (Union[Unset, datetime.datetime]): Log timestamp in RFC3339 format Example:
+            2024-01-30T10:30:00.123456789Z.
+    """
+
+    content: Union[Unset, str] = UNSET
+    log_type: Union[Unset, TaskLogEntryLogType] = UNSET
+    node: Union[Unset, str] = UNSET
+    timestamp: Union[Unset, datetime.datetime] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        content = self.content
+
+        log_type: Union[Unset, str] = UNSET
+        if not isinstance(self.log_type, Unset):
+            log_type = self.log_type.value
+
+        node = self.node
+
+        timestamp: Union[Unset, str] = UNSET
+        if not isinstance(self.timestamp, Unset):
+            timestamp = self.timestamp.isoformat()
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if content is not UNSET:
+            field_dict["content"] = content
+        if log_type is not UNSET:
+            field_dict["log_type"] = log_type
+        if node is not UNSET:
+            field_dict["node"] = node
+        if timestamp is not UNSET:
+            field_dict["timestamp"] = timestamp
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        content = d.pop("content", UNSET)
+
+        _log_type = d.pop("log_type", UNSET)
+        log_type: Union[Unset, TaskLogEntryLogType]
+        if isinstance(_log_type, Unset):
+            log_type = UNSET
+        else:
+            log_type = TaskLogEntryLogType(_log_type)
+
+        node = d.pop("node", UNSET)
+
+        _timestamp = d.pop("timestamp", UNSET)
+        timestamp: Union[Unset, datetime.datetime]
+        if isinstance(_timestamp, Unset):
+            timestamp = UNSET
+        else:
+            timestamp = isoparse(_timestamp)
+
+        task_log_entry = cls(
+            content=content,
+            log_type=log_type,
+            node=node,
+            timestamp=timestamp,
+        )
+
+        task_log_entry.additional_properties = d
+        return task_log_entry
+
+    @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_log_entry_log_type.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class TaskLogEntryLogType(str, Enum):
+    STDERR = "stderr"
+    STDOUT = "stdout"
+
+    def __str__(self) -> str:
+        return str(self.value)

mlops/api/client/models/task_logs_response.py

@@ -0,0 +1,112 @@
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+if TYPE_CHECKING:
+    from ..models.log_pagination import LogPagination
+    from ..models.task_log_entry import TaskLogEntry
+
+
+T = TypeVar("T", bound="TaskLogsResponse")
+
+
+@_attrs_define
+class TaskLogsResponse:
+    """Task logs response with pagination support
+
+    Attributes:
+        cluster_id (Union[Unset, int]): Cluster ID Example: 1.
+        job_id (Union[Unset, int]): Slurm Job ID Example: 12345.
+        logs (Union[Unset, list['TaskLogEntry']]): List of log entries
+        pagination (Union[Unset, LogPagination]): Pagination information for logs
+    """
+
+    cluster_id: Union[Unset, int] = UNSET
+    job_id: Union[Unset, int] = UNSET
+    logs: Union[Unset, list["TaskLogEntry"]] = UNSET
+    pagination: Union[Unset, "LogPagination"] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        cluster_id = self.cluster_id
+
+        job_id = self.job_id
+
+        logs: Union[Unset, list[dict[str, Any]]] = UNSET
+        if not isinstance(self.logs, Unset):
+            logs = []
+            for logs_item_data in self.logs:
+                logs_item = logs_item_data.to_dict()
+                logs.append(logs_item)
+
+        pagination: Union[Unset, dict[str, Any]] = UNSET
+        if not isinstance(self.pagination, Unset):
+            pagination = self.pagination.to_dict()
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if cluster_id is not UNSET:
+            field_dict["cluster_id"] = cluster_id
+        if job_id is not UNSET:
+            field_dict["job_id"] = job_id
+        if logs is not UNSET:
+            field_dict["logs"] = logs
+        if pagination is not UNSET:
+            field_dict["pagination"] = pagination
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.log_pagination import LogPagination
+        from ..models.task_log_entry import TaskLogEntry
+
+        d = dict(src_dict)
+        cluster_id = d.pop("cluster_id", UNSET)
+
+        job_id = d.pop("job_id", UNSET)
+
+        logs = []
+        _logs = d.pop("logs", UNSET)
+        for logs_item_data in _logs or []:
+            logs_item = TaskLogEntry.from_dict(logs_item_data)
+
+            logs.append(logs_item)
+
+        _pagination = d.pop("pagination", UNSET)
+        pagination: Union[Unset, LogPagination]
+        if isinstance(_pagination, Unset):
+            pagination = UNSET
+        else:
+            pagination = LogPagination.from_dict(_pagination)
+
+        task_logs_response = cls(
+            cluster_id=cluster_id,
+            job_id=job_id,
+            logs=logs,
+            pagination=pagination,
+        )
+
+        task_logs_response.additional_properties = d
+        return task_logs_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_submit_request.py

@@ -7,6 +7,7 @@ from attrs import field as _attrs_field
 from ..types import UNSET, Unset
 
 if TYPE_CHECKING:
+    from ..models.job_spec import JobSpec
     from ..models.task_submit_request_environment_type_0 import TaskSubmitRequestEnvironmentType0
 
 
@@ -18,7 +19,7 @@ class TaskSubmitRequest:
     """Task submission request
 
     Attributes:
-        cluster_id (int): Slurm cluster ID to submit task to Example: 1.
+        cluster_name (str): Slurm cluster name to submit task to Example: slurm-prod.
         name (str): Task name Example: training-job.
         account (Union[None, Unset, str]): Account Example: research.
         command (Union[None, Unset, str]): Command to execute (alternative to script) Example: python train.py.
@@ -35,6 +36,7 @@ class TaskSubmitRequest:
         export (Union[None, Unset, str]): Environment export Example: ALL.
         gres (Union[None, Unset, str]): Generic resources (e.g., "gpu:1", "gpu:tesla:2") Example: gpu:1.
         input_ (Union[None, Unset, str]): Standard input file
+        job_spec (Union[Unset, JobSpec]): Domain-specific job specification (rendered into slurm script)
         mem_bind (Union[None, Unset, str]): Memory binding
         memory (Union[None, Unset, str]): Memory requirement (e.g., "8G", "4096M") Example: 8G.
         nice (Union[None, Unset, int]): Nice value
@@ -53,7 +55,7 @@ class TaskSubmitRequest:
         tres (Union[None, Unset, str]): Trackable resources string Example: cpu=4,mem=8G.
     """
 
-    cluster_id: int
+    cluster_name: str
     name: str
     account: Union[None, Unset, str] = UNSET
     command: Union[None, Unset, str] = UNSET
@@ -69,6 +71,7 @@ class TaskSubmitRequest:
     export: Union[None, Unset, str] = UNSET
     gres: Union[None, Unset, str] = UNSET
     input_: Union[None, Unset, str] = UNSET
+    job_spec: Union[Unset, "JobSpec"] = UNSET
     mem_bind: Union[None, Unset, str] = UNSET
     memory: Union[None, Unset, str] = UNSET
     nice: Union[None, Unset, int] = UNSET
@@ -88,7 +91,7 @@ class TaskSubmitRequest:
     def to_dict(self) -> dict[str, Any]:
         from ..models.task_submit_request_environment_type_0 import TaskSubmitRequestEnvironmentType0
 
-        cluster_id = self.cluster_id
+        cluster_name = self.cluster_name
 
         name = self.name
 
@@ -178,6 +181,10 @@ class TaskSubmitRequest:
         else:
             input_ = self.input_
 
+        job_spec: Union[Unset, dict[str, Any]] = UNSET
+        if not isinstance(self.job_spec, Unset):
+            job_spec = self.job_spec.to_dict()
+
         mem_bind: Union[None, Unset, str]
         if isinstance(self.mem_bind, Unset):
             mem_bind = UNSET
@@ -262,7 +269,7 @@ class TaskSubmitRequest:
         field_dict.update(self.additional_properties)
         field_dict.update(
             {
-                "cluster_id": cluster_id,
+                "cluster_name": cluster_name,
                 "name": name,
             }
         )
@@ -294,6 +301,8 @@ class TaskSubmitRequest:
             field_dict["gres"] = gres
         if input_ is not UNSET:
             field_dict["input"] = input_
+        if job_spec is not UNSET:
+            field_dict["job_spec"] = job_spec
         if mem_bind is not UNSET:
             field_dict["mem_bind"] = mem_bind
         if memory is not UNSET:
@@ -327,10 +336,11 @@ class TaskSubmitRequest:
 
     @classmethod
     def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.job_spec import JobSpec
         from ..models.task_submit_request_environment_type_0 import TaskSubmitRequestEnvironmentType0
 
         d = dict(src_dict)
-        cluster_id = d.pop("cluster_id")
+        cluster_name = d.pop("cluster_name")
 
         name = d.pop("name")
 
@@ -468,6 +478,13 @@ class TaskSubmitRequest:
 
         input_ = _parse_input_(d.pop("input", UNSET))
 
+        _job_spec = d.pop("job_spec", UNSET)
+        job_spec: Union[Unset, JobSpec]
+        if isinstance(_job_spec, Unset):
+            job_spec = UNSET
+        else:
+            job_spec = JobSpec.from_dict(_job_spec)
+
         def _parse_mem_bind(data: object) -> Union[None, Unset, str]:
             if data is None:
                 return data
@@ -588,7 +605,7 @@ class TaskSubmitRequest:
         tres = _parse_tres(d.pop("tres", UNSET))
 
         task_submit_request = cls(
-            cluster_id=cluster_id,
+            cluster_name=cluster_name,
             name=name,
             account=account,
             command=command,
@@ -604,6 +621,7 @@ class TaskSubmitRequest:
             export=export,
             gres=gres,
             input_=input_,
+            job_spec=job_spec,
             mem_bind=mem_bind,
             memory=memory,
             nice=nice,

mlops/connection_config.py

@@ -19,30 +19,25 @@ class ConnectionConfig:
 
     @staticmethod
     def _domain():
-        return os.getenv("XCLIENT_DOMAIN", "localhost:8090")
+        return os.getenv('MLOPS_DOMAIN', "localhost:8090")
 
     @staticmethod
     def _debug():
-        return os.getenv("XCLIENT_DEBUG", "false").lower() == "true"
+        return os.getenv('MLOPS_DEBUG', "false").lower() == "true"
 
     @staticmethod
     def _api_key():
-        return os.getenv("XCLIENT_API_KEY")
-
-    @staticmethod
-    def _access_token():
-        return os.getenv("XCLIENT_ACCESS_TOKEN")
+        return os.getenv('MLOPS_API_KEY')
 
     @staticmethod
     def _api_path():
-        return os.getenv("XCLIENT_API_PATH", DEFAULT_API_PATH)
+        return os.getenv('MLOPS_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,
@@ -51,11 +46,9 @@ class ConnectionConfig:
         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,

mlops/exceptions.py

@@ -10,17 +10,17 @@ def format_execution_timeout_error() -> Exception:
     )
 
 
-class XClientException(Exception):
+class MLOpsException(Exception):
     """
-    Base class for all XClient errors.
+    Base class for all MLOps errors.
 
-    Raised when a general XClient exception occurs.
+    Raised when a general MLOps exception occurs.
     """
 
     pass
 
 
-class TimeoutException(XClientException):
+class TimeoutException(MLOpsException):
     """
     Raised when a timeout occurs.
 
@@ -33,7 +33,7 @@ class TimeoutException(XClientException):
     pass
 
 
-class InvalidArgumentException(XClientException):
+class InvalidArgumentException(MLOpsException):
     """
     Raised when an invalid argument is provided.
     """
@@ -41,7 +41,7 @@ class InvalidArgumentException(XClientException):
     pass
 
 
-class NotEnoughSpaceException(XClientException):
+class NotEnoughSpaceException(MLOpsException):
     """
     Raised when there is not enough disk space.
     """
@@ -49,7 +49,7 @@ class NotEnoughSpaceException(XClientException):
     pass
 
 
-class NotFoundException(XClientException):
+class NotFoundException(MLOpsException):
     """
     Raised when a resource is not found.
     """
@@ -57,7 +57,7 @@ class NotFoundException(XClientException):
     pass
 
 
-class AuthenticationException(XClientException):
+class AuthenticationException(MLOpsException):
     """
     Raised when authentication fails.
     """
@@ -65,7 +65,7 @@ class AuthenticationException(XClientException):
     pass
 
 
-class RateLimitException(XClientException):
+class RateLimitException(MLOpsException):
     """
     Raised when the API rate limit is exceeded.
     """
@@ -73,7 +73,7 @@ class RateLimitException(XClientException):
     pass
 
 
-class APIException(XClientException):
+class APIException(MLOpsException):
     """
     Raised when an API error occurs.
     """

mlops/task/__init__.py

@@ -1,4 +1,4 @@
-"""Task SDK module for XClient"""
+"""Task SDK module for MLOps"""
 
 from .client import TaskClient
 from .task import Task

mlops/task/client.py

@@ -17,7 +17,7 @@ logger = logging.getLogger(__name__)
 
 
 def handle_api_exception(e: Response):
-    """Handle API exceptions and convert them to appropriate XClient exceptions."""
+    """Handle API exceptions and convert them to appropriate MLOps exceptions."""
     try:
         body = json.loads(e.content) if e.content else {}
     except json.JSONDecodeError:
@@ -49,51 +49,27 @@ def handle_api_exception(e: Response):
 
 class TaskClient(AuthenticatedClient):
     """
-    The client for interacting with the XClient Task API.
+    The client for interacting with the MLOps 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:
+        # NOTE: This SDK client only supports API key authentication.
+        # Header: X-API-Key (per OpenAPI spec)
+        if config.api_key is None:
             raise AuthenticationException(
-                "Only one of api_key or access_token can be required, not both",
+                "API key is required. "
+                "You can set the environment variable `MLOPS_API_KEY`"
+                ' to pass it directly like ConnectionConfig(api_key="mlops_...")',
             )
 
-        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"
+        auth_header_name = "X-API-Key"
+        prefix = ""
 
         headers = {
             **(config.headers or {}),
@@ -116,9 +92,9 @@ class TaskClient(AuthenticatedClient):
             base_url=config.api_url,
             httpx_args=httpx_args,
             headers=headers,
-            token=token,
             auth_header_name=auth_header_name,
             prefix=prefix,
+            token=config.api_key,
             *args,
             **kwargs,
         )