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

mlops/api/client/models/task_submit_request.py

@@ -19,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.
@@ -55,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
@@ -91,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
 
@@ -269,7 +269,7 @@ class TaskSubmitRequest:
         field_dict.update(self.additional_properties)
         field_dict.update(
             {
-                "cluster_id": cluster_id,
+                "cluster_name": cluster_name,
                 "name": name,
             }
         )
@@ -340,7 +340,7 @@ class TaskSubmitRequest:
         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")
 
@@ -605,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,

mlops/connection_config.py

@@ -29,10 +29,6 @@ class ConnectionConfig:
     def _api_key():
         return os.getenv('MLOPS_API_KEY')
 
-    @staticmethod
-    def _access_token():
-        return os.getenv('MLOPS_ACCESS_TOKEN')
-
     @staticmethod
     def _api_path():
         return os.getenv('MLOPS_API_PATH', DEFAULT_API_PATH)
@@ -42,7 +38,6 @@ class ConnectionConfig:
         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,
         )

mlops/task/task.py

@@ -1,12 +1,17 @@
 """
-High-level Task SDK interface for XClient.
+High-level Task SDK interface for MLOps.
 
-This module provides a convenient interface for managing tasks through the XClient API.
-"""
+This module provides a convenient interface for managing tasks through the MLOps API.
+""" 
 
 import json
+import os
 from http import HTTPStatus
+from pathlib import Path
 from typing import Optional
+
+import httpx
+
 from ..api.client.api.tasks import (
     submit_task,
     get_task,
@@ -14,8 +19,15 @@ from ..api.client.api.tasks import (
     cancel_task,
     delete_task,
 )
+from ..api.client.api.storage import (
+    get_storage_presign_upload,
+    get_storage_presign_download,
+)
 from ..api.client.models.task import Task as TaskModel
 from ..api.client.models.task_submit_request import TaskSubmitRequest
+from ..api.client.models.task_submit_request_environment_type_0 import (
+    TaskSubmitRequestEnvironmentType0,
+)
 from ..api.client.models.task_submit_response import TaskSubmitResponse
 from ..api.client.models.task_list_response import TaskListResponse
 from ..api.client.models.task_status import TaskStatus
@@ -29,13 +41,46 @@ from ..exceptions import (
 from .client import TaskClient, handle_api_exception
 
 
+def _validate_archive_file_path(file_path: str) -> Path:
+    p = Path(os.path.expanduser(file_path)).resolve()
+    if not p.exists():
+        raise APIException(f"File not found: {p}")
+    if not p.is_file():
+        raise APIException(f"file_path must be a file: {p}")
+
+    lower = p.name.lower()
+    if not (lower.endswith(".zip") or lower.endswith(".tar.gz") or lower.endswith(".tgz")):
+        raise APIException(f"file_path must be one of .zip, .tar.gz, .tgz: {p}")
+    return p
+
+
+def _upload_file_to_presigned_url(url: str, file_path: Path, timeout: Optional[float]) -> None:
+    size = file_path.stat().st_size
+    # Use a dedicated client for S3 presigned upload (avoid leaking API auth headers).
+    with httpx.Client(timeout=timeout) as client:
+        with file_path.open("rb") as f:
+            resp = client.put(
+                url,
+                content=f,
+                headers={
+                    "Content-Length": str(size),
+                    "Content-Type": "application/octet-stream",
+                },
+            )
+    if resp.status_code < 200 or resp.status_code >= 300:
+        body = (resp.text or "")[:2048]
+        raise APIException(
+            f"Failed to upload file to presigned url: HTTP {resp.status_code}: {body}"
+        )
+
+
 class Task:
     """
     High-level interface for managing tasks.
     
     Example:
         ```python
-        from xclient import Task, ConnectionConfig
+        from mlops import Task, ConnectionConfig
         
         config = ConnectionConfig(api_key="your_api_key")
         task = Task(config=config)
@@ -43,28 +88,28 @@ class Task:
         # Submit a task with script
         result = task.submit(
             name="my-task",
-            cluster_id=1,
+            cluster_name="slurm-cn",
             script="#!/bin/bash\\necho 'Hello World'"
         )
         
         # Or submit with command
         result = task.submit(
             name="my-task",
-            cluster_id=1,
+            cluster_name="slurm-cn",
             command="echo 'Hello World'"
         )
         
         # Get task details
-        task_info = task.get(task_id=result.job_id, cluster_id=1)
+        task_info = task.get(task_id=result.job_id, cluster_name="slurm-cn")
         
         # List tasks
         tasks = task.list(status=TaskStatus.RUNNING)
         
         # Cancel a task
-        task.cancel(task_id=result.job_id, cluster_id=1)
+        task.cancel(task_id=result.job_id, cluster_name="slurm-cn")
 
         # Delete a task
-        task.delete(task_id=result.job_id, cluster_id=1)
+        task.delete(task_id=result.job_id, cluster_name="slurm-cn")
         ```
     """
 
@@ -72,7 +117,6 @@ class Task:
         self,
         config: Optional["ConnectionConfig"] = None,
         api_key: Optional[str] = None,
-        access_token: Optional[str] = None,
         domain: Optional[str] = None,
         debug: Optional[bool] = None,
         request_timeout: Optional[float] = None,
@@ -83,7 +127,6 @@ class Task:
         Args:
             config: ConnectionConfig instance. If not provided, a new one will be created.
             api_key: API key for authentication. Overrides config.api_key.
-            access_token: Access token for authentication. Overrides config.access_token.
             domain: API domain. Overrides config.domain.
             debug: Enable debug mode. Overrides config.debug.
             request_timeout: Request timeout in seconds. Overrides config.request_timeout.
@@ -95,8 +138,6 @@ class Task:
         # Override config values if provided
         if api_key is not None:
             config.api_key = api_key
-        if access_token is not None:
-            config.access_token = access_token
         if domain is not None:
             config.domain = domain
         if debug is not None:
@@ -106,22 +147,22 @@ class Task:
 
         self._config = config
         self._client = TaskClient(config=config)
-
     def submit(
         self,
         name: str,
-        cluster_id: Optional[int] = None,
+        cluster_name: str,
         script: Optional[str] = None,
         command: Optional[str] = None,
         resources: Optional[dict] = None,
         team_id: Optional[int] = None,
+        file_path: Optional[str] = None,
     ) -> TaskSubmitResponse:
         """
         Submit a new task.
 
         Args:
             name: Task name
-            cluster_id: Cluster ID to submit the task to
+            cluster_name: Cluster name to submit the task to
             script: Task script content (optional, but at least one of script or command is required)
             command: Command to execute (optional, but at least one of script or command is required)
             resources: Resource requirements dict (optional)
@@ -134,19 +175,16 @@ class Task:
             APIException: If the API returns an error
             AuthenticationException: If authentication fails
         """
-        # Validate required fields
-        if cluster_id is None:
-            raise APIException("cluster_id is required")
-        
         # At least one of script or command must be provided
         if not script and not command:
             raise APIException("At least one of 'script' or 'command' must be provided")
         
         # Map resources dict to individual fields
         # resources dict can contain: cpu, cpus_per_task, memory, nodes, gres, time, partition, etc.
+
         request_kwargs = {
             "name": name,
-            "cluster_id": cluster_id,
+            "cluster_name": cluster_name,
         }
         
         # Handle script and command (at least one is required)
@@ -177,7 +215,87 @@ class Task:
                 request_kwargs["partition"] = resources.get("partition")
             if "tres" in resources:
                 request_kwargs["tres"] = resources.get("tres")
-        
+
+        if file_path:
+            local_path = _validate_archive_file_path(file_path)
+            timeout = self._config.get_request_timeout()
+
+            # 1) Get presigned upload URL
+            presign_upload_obj = get_storage_presign_upload.sync_detailed(
+                client=self._client,
+                filename=local_path.name,
+            )
+            presign_upload = presign_upload_obj.parsed
+            if isinstance(presign_upload, ErrorResponse):
+                status_code = (
+                    presign_upload.code
+                    if presign_upload.code != UNSET and presign_upload.code != 0
+                    else presign_upload_obj.status_code.value
+                )
+                exception = handle_api_exception(
+                    Response(
+                        status_code=HTTPStatus(status_code),
+                        content=presign_upload_obj.content,
+                        headers=presign_upload_obj.headers,
+                        parsed=None,
+                    )
+                )
+                raise exception
+
+            if (
+                presign_upload is None
+                or presign_upload.url in (UNSET, None)
+                or presign_upload.key in (UNSET, None)
+            ):
+                raise APIException("Failed to get presigned upload url: empty response")
+
+            # 2) Upload file to S3 (presigned URL)
+            _upload_file_to_presigned_url(
+                url=str(presign_upload.url),
+                file_path=local_path,
+                timeout=timeout,
+            )
+
+            # 3) Get presigned download URL
+            presign_download_obj = get_storage_presign_download.sync_detailed(
+                client=self._client,
+                key=str(presign_upload.key),
+            )
+            presign_download = presign_download_obj.parsed
+            if isinstance(presign_download, ErrorResponse):
+                status_code = (
+                    presign_download.code
+                    if presign_download.code != UNSET and presign_download.code != 0
+                    else presign_download_obj.status_code.value
+                )
+                exception = handle_api_exception(
+                    Response(
+                        status_code=HTTPStatus(status_code),
+                        content=presign_download_obj.content,
+                        headers=presign_download_obj.headers,
+                        parsed=None,
+                    )
+                )
+                raise exception
+
+            if presign_download is None or presign_download.url in (UNSET, None):
+                raise APIException(
+                    "Failed to get presigned download url: empty response"
+                )
+
+            # 4) Set env var (merge if user already provided environment)
+            env: dict[str, str] = {}
+            existing_env = request_kwargs.get("environment")
+            if isinstance(existing_env, TaskSubmitRequestEnvironmentType0):
+                env.update(existing_env.additional_properties)
+            elif isinstance(existing_env, dict):
+                env.update(existing_env)
+
+            env["SYSTEM_DOWNLOAD_ARCHIVE_URL"] = str(presign_download.url)
+            request_kwargs["environment"] = TaskSubmitRequestEnvironmentType0.from_dict(
+                env
+            )
+
         request = TaskSubmitRequest(**request_kwargs)
 
         # Use sync_detailed to get full response information
@@ -230,14 +348,14 @@ class Task:
     def get(
         self,
         task_id: int,
-        cluster_id: int,
+        cluster_name: str,
     ) -> TaskModel:
         """
         Get task details by task ID.
 
         Args:
             task_id: Task ID
-            cluster_id: Cluster ID
+            cluster_name: Cluster name
 
         Returns:
             Task model with task details
@@ -250,7 +368,7 @@ class Task:
         response_obj = get_task.sync_detailed(
             id=task_id,
             client=self._client,
-            cluster_id=cluster_id,
+            cluster_name=cluster_name,
         )
         response = response_obj.parsed
 
@@ -302,7 +420,7 @@ class Task:
         status: Optional[TaskStatus] = None,
         user_id: Optional[int] = None,
         team_id: Optional[int] = None,
-        cluster_id: Optional[int] = None,
+        cluster_name: Optional[str] = None,
     ) -> TaskListResponse:
         """
         List tasks with optional filtering.
@@ -313,7 +431,7 @@ class Task:
             status: Filter by task status (optional)
             user_id: Filter by user ID (optional)
             team_id: Filter by team ID (optional)
-            cluster_id: Filter by cluster ID (optional)
+            cluster_name: Filter by cluster name (optional)
 
         Returns:
             TaskListResponse containing the list of tasks
@@ -329,7 +447,7 @@ class Task:
             status=status if status is not None else UNSET,
             user_id=user_id if user_id is not None else UNSET,
             team_id=team_id if team_id is not None else UNSET,
-            cluster_id=cluster_id if cluster_id is not None else UNSET,
+            cluster_name=cluster_name if cluster_name is not None else UNSET,
         )
         response = response_obj.parsed
 
@@ -377,14 +495,14 @@ class Task:
     def cancel(
         self,
         task_id: int,
-        cluster_id: int,
+        cluster_name: str,
     ) -> bool:
         """
         Cancel a task.
 
         Args:
             task_id: Task ID to cancel
-            cluster_id: Cluster ID where the task is running
+            cluster_name: Cluster name where the task is running
 
         Returns:
             True if the task was cancelled successfully
@@ -397,7 +515,7 @@ class Task:
         response_obj = cancel_task.sync_detailed(
             id=task_id,
             client=self._client,
-            cluster_id=cluster_id,
+            cluster_name=cluster_name,
         )
         response = response_obj.parsed
 
@@ -434,14 +552,14 @@ class Task:
     def delete(
         self,
         task_id: int,
-        cluster_id: int,
+        cluster_name: str,
     ) -> bool:
         """
         Delete a task.
 
         Args:
             task_id: Task ID to delete
-            cluster_id: Cluster ID where the task is running
+            cluster_name: Cluster name where the task is running
 
         Returns:
             True if the task was deleted successfully
@@ -454,7 +572,7 @@ class Task:
         response_obj = delete_task.sync_detailed(
             id=task_id,
             client=self._client,
-            cluster_id=cluster_id,
+            cluster_name=cluster_name,
         )
         response = response_obj.parsed
 

{mlops_python_sdk-1.0.0.dist-info → mlops_python_sdk-1.0.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: mlops-python-sdk
-Version: 1.0.0
+Version: 1.0.1
 Summary: MLOps Python SDK for XCloud Service API
 License: MIT
 Author: mlops
@@ -39,9 +39,9 @@ pip install mlops-python-sdk
 
 ### 1. Setup Authentication
 
-You can authenticate using either an API Key or an Access Token.
+You can authenticate using either an API Key.
 
-#### Option 1: API Key (Recommended for programmatic access)
+#### API Key (Recommended for programmatic access)
 
 1. Sign up at [MLOps](https://xcloud-service.com)
 2. Create an API key from [API Keys](https://xcloud-service.com/home/api-keys)
@@ -52,13 +52,6 @@ export MLOPS_API_KEY=xck_******
 export MLOPS_DOMAIN=localhost:8090  # optional, default is localhost:8090
 ```
 
-#### Option 2: Access Token (For user authentication)
-
-```bash
-export MLOPS_ACCESS_TOKEN=your_access_token
-export MLOPS_DOMAIN=localhost:8090  # optional
-```
-
 ### 2. Basic Usage
 
 ```python
@@ -127,7 +120,6 @@ task = Task()
 # With explicit configuration
 config = ConnectionConfig(
     api_key="xck_******",           # API key for authentication
-    access_token="token_******",     # Access token (alternative to API key)
     domain="localhost:8090",         # API domain
     debug=False,                      # Enable debug mode
     request_timeout=30.0              # Request timeout in seconds
@@ -280,7 +272,6 @@ TaskStatus.CREATED      # Task was created
 The SDK reads configuration from environment variables:
 
 - `MLOPS_API_KEY`: API key for authentication
-- `MLOPS_ACCESS_TOKEN`: Access token for authentication (alternative to API key)
 - `MLOPS_DOMAIN`: API domain (default: `localhost:8090`)
 - `MLOPS_DEBUG`: Enable debug mode (`true`/`false`, default: `false`)
 - `MLOPS_API_PATH`: API path prefix (default: `/api/v1`)