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

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,42 +41,82 @@ 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)
         
-        # Submit a task with script
+        # Submit a task with gpu type
         result = task.submit(
-            name="my-task",
-            cluster_id=1,
-            script="#!/bin/bash\\necho 'Hello World'"
-        )
-        
-        # Or submit with command
-        result = task.submit(
-            name="my-task",
-            cluster_id=1,
-            command="echo 'Hello World'"
+            name="gpu-task-from-sdk",
+            cluster_name="slurm-cn",
+            image="/mnt/minio/images/01ai-registry.cn-shanghai.cr.aliyuncs.com+public+llamafactory+0.9.3.sqsh",
+            entry_command="llamafactory-cli train /workspace/config/test_lora.yaml",
+            resources={
+                "partition": "gpu",
+                "nodes": 2,
+                "ntasks": 2,
+                "cpus_per_task": 2,
+                "memory": "4G",
+                "time": "01:00:00",
+                "gres": "gpu:nvidia_a10:1",
+                "qos": "qos_xcloud",
+                "job_type": "batch",
+            },
+            team_id=1,
+            file_path="your file path",
         )
         
         # 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 +124,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 +134,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 +145,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,27 +154,27 @@ class Task:
 
         self._config = config
         self._client = TaskClient(config=config)
-
     def submit(
         self,
         name: str,
-        cluster_id: Optional[int] = None,
-        script: Optional[str] = None,
-        command: Optional[str] = None,
+        cluster_name: str,
+        image: str,
+        entry_command: str,
         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
-            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)
+            cluster_name: Cluster name to submit the task to
+            image: Container image reference
+            entry_command: Container entry command/script
             resources: Resource requirements dict (optional)
             team_id: Team ID (optional)
-
+            file_path: Local file path to upload (optional, support for .zip, .tar.gz, .tgz)
         Returns:
             TaskSubmitResponse containing the submitted task information
 
@@ -134,28 +182,14 @@ 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,
+            "image": image,
+            "entry_command": entry_command,
         }
-        
-        # Handle script and command (at least one is required)
-        # script is Union[Unset, str], so we need to set it or leave as UNSET
-        if script:
-            request_kwargs["script"] = script
-        # command is Union[None, Unset, str], so we can set it or leave as UNSET
-        if command:
-            request_kwargs["command"] = command
+        # Map resources dict to individual fields
+        # resources dict can contain: cpu, cpus_per_task, memory, nodes, gres, time, partition, etc.
         
         # team_id is Union[None, Unset, int]
         if team_id is not None:
@@ -165,19 +199,115 @@ class Task:
         if resources:
             if "cpu" in resources or "cpus_per_task" in resources:
                 request_kwargs["cpus_per_task"] = resources.get("cpus_per_task") or resources.get("cpu")
+            else:
+                request_kwargs["cpus_per_task"] = 1
             if "memory" in resources:
                 request_kwargs["memory"] = resources.get("memory")
+            else: 
+                request_kwargs["memory"] = "1G"
             if "nodes" in resources:
                 request_kwargs["nodes"] = resources.get("nodes")
+            else:
+                request_kwargs["nodes"] = 1
             if "gres" in resources:
                 request_kwargs["gres"] = resources.get("gres")
             if "time" in resources:
                 request_kwargs["time"] = resources.get("time")
+            else:
+                request_kwargs["time"] = "01:00:00"
             if "partition" in resources:
                 request_kwargs["partition"] = resources.get("partition")
-            if "tres" in resources:
-                request_kwargs["tres"] = resources.get("tres")
-        
+            else:
+                request_kwargs["partition"] = "debug"
+            if "qos" in resources:
+                request_kwargs["qos"] = resources.get("qos")
+            else:
+                request_kwargs["qos"] = "qos_xcloud"
+            if "ntasks" in resources:
+                request_kwargs["ntasks"] = resources.get("ntasks")
+            else:
+                request_kwargs["ntasks"] = 1
+
+        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 +360,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 +380,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 +432,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 +443,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 +459,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 +507,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 +527,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 +564,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 +584,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