gmicloud 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

gmicloud/_internal/_client/_task_client.py

@@ -1,9 +1,14 @@
+import logging
+from requests.exceptions import RequestException
+
 from ._http_client import HTTPClient
 from ._decorator import handle_refresh_token
 from ._iam_client import IAMClient
 from .._config import TASK_SERVICE_BASE_URL
 from .._models import *
 
+logger = logging.getLogger(__name__)
+
 
 class TaskClient:
     """
@@ -21,17 +26,19 @@ class TaskClient:
         self.iam_client = iam_client
 
     @handle_refresh_token
-    def get_task(self, task_id: str) -> Task:
+    def get_task(self, task_id: str) -> Optional[Task]:
         """
         Retrieves a task from the task service using the given task ID.
 
         :param task_id: The ID of the task to be retrieved.
-        :return: An instance of Task containing the details of the retrieved task.
-        :rtype: Task
+        :return: An instance of Task containing the details of the retrieved task, or None if an error occurs.
         """
-        result = self.client.get("/get_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
-
-        return Task.model_validate(result)
+        try:
+            response = self.client.get("/get_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+            return Task.model_validate(response) if response else None
+        except (RequestException, ValueError) as e:
+            logger.error(f"Failed to retrieve task {task_id}: {e}")
+            return None
 
     @handle_refresh_token
     def get_all_tasks(self) -> GetAllTasksResponse:
@@ -39,70 +46,108 @@ class TaskClient:
         Retrieves all tasks from the task service.
 
         :return: An instance of GetAllTasksResponse containing the retrieved tasks.
-        :rtype: GetAllTasksResponse
         """
-        result = self.client.get("/get_tasks", self.iam_client.get_custom_headers())
-        if not result:
+        try:
+            response = self.client.get("/get_tasks", self.iam_client.get_custom_headers())
+            if not response:
+                logger.error("Empty response from /get_tasks")
+                return GetAllTasksResponse(tasks=[])
+            return GetAllTasksResponse.model_validate(response)
+        except (RequestException, ValueError) as e:
+            logger.error(f"Failed to retrieve all tasks: {e}")
             return GetAllTasksResponse(tasks=[])
 
-        return GetAllTasksResponse.model_validate(result)
-
     @handle_refresh_token
-    def create_task(self, task: Task) -> CreateTaskResponse:
+    def create_task(self, task: Task) -> Optional[CreateTaskResponse]:
         """
         Creates a new task using the provided task object.
 
         :param task: The Task object containing the details of the task to be created.
+        :return: The response object containing created task details, or None if an error occurs.
         """
-        result = self.client.post("/create_task", self.iam_client.get_custom_headers(), task.model_dump())
-
-        return CreateTaskResponse.model_validate(result)
+        try:
+            response = self.client.post("/create_task", self.iam_client.get_custom_headers(), task.model_dump())
+            return CreateTaskResponse.model_validate(response) if response else None
+        except (RequestException, ValueError) as e:
+            logger.error(f"Failed to create task: {e}")
+            return None
 
     @handle_refresh_token
-    def update_task_schedule(self, task: Task):
+    def update_task_schedule(self, task: Task) -> bool:
         """
         Updates the schedule of an existing task.
 
         :param task: The Task object containing the updated task details.
+        :return: True if update is successful, False otherwise.
         """
-        self.client.put("/update_schedule", self.iam_client.get_custom_headers(), task.model_dump())
+        try:
+            response = self.client.put("/update_schedule", self.iam_client.get_custom_headers(), task.model_dump())
+            return response is not None
+        except RequestException as e:
+            logger.error(f"Failed to update schedule for task {task.task_id}: {e}")
+            return False
 
     @handle_refresh_token
-    def start_task(self, task_id: str):
+    def start_task(self, task_id: str) -> bool:
         """
         Starts a task using the given task ID.
 
         :param task_id: The ID of the task to be started.
+        :return: True if start is successful, False otherwise.
         """
-        self.client.post("/start_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+        try:
+            response = self.client.post("/start_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+            return response is not None
+        except RequestException as e:
+            logger.error(f"Failed to start task {task_id}: {e}")
+            return False
 
     @handle_refresh_token
-    def stop_task(self, task_id: str):
+    def stop_task(self, task_id: str) -> bool:
         """
         Stops a running task using the given task ID.
 
         :param task_id: The ID of the task to be stopped.
+        :return: True if stop is successful, False otherwise.
         """
-        self.client.post("/stop_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+        try:
+            response = self.client.post("/stop_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+            return response is not None
+        except RequestException as e:
+            logger.error(f"Failed to stop task {task_id}: {e}")
+            return False
 
     @handle_refresh_token
-    def get_usage_data(self, start_timestamp: str, end_timestamp: str) -> GetUsageDataResponse:
+    def get_usage_data(self, start_timestamp: str, end_timestamp: str) -> Optional[GetUsageDataResponse]:
         """
         Retrieves the usage data of a task using the given task ID.
 
         :param start_timestamp: The start timestamp of the usage data.
         :param end_timestamp: The end timestamp of the usage data.
-        """
-        result = self.client.get("/get_usage_data", self.iam_client.get_custom_headers(),
-                                 {"start_timestamp": start_timestamp, "end_timestamp": end_timestamp})
-
-        return result
+        :return: An instance of GetUsageDataResponse, or None if an error occurs.
+        """
+        try:
+            response = self.client.get(
+                "/get_usage_data",
+                self.iam_client.get_custom_headers(),
+                {"start_timestamp": start_timestamp, "end_timestamp": end_timestamp}
+            )
+            return GetUsageDataResponse.model_validate(response) if response else None
+        except (RequestException, ValueError) as e:
+            logger.error(f"Failed to retrieve usage data from {start_timestamp} to {end_timestamp}: {e}")
+            return None
 
     @handle_refresh_token
-    def archive_task(self, task_id: str):
+    def archive_task(self, task_id: str) -> bool:
         """
         Archives a task using the given task ID.
 
         :param task_id: The ID of the task to be archived.
-        """
-        self.client.post("/archive_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+        :return: True if archiving is successful, False otherwise.
+        """
+        try:
+            response = self.client.post("/archive_task", self.iam_client.get_custom_headers(), {"task_id": task_id})
+            return response is not None
+        except RequestException as e:
+            logger.error(f"Failed to archive task {task_id}: {e}")
+            return False

gmicloud/_internal/_enums.py

@@ -23,3 +23,16 @@ class TaskEndpointStatus(str, Enum):
     READY = "ready"
     UNREADY = "unready"
     NEW = "new"
+
+class TaskStatus(str, Enum):
+    IDLE = "idle"
+    STARTING = "starting"
+    IN_QUEUE = "in-queue"
+    RUNNING = "running"
+    NEEDSTOP = "needstop"
+    ARCHIVED = "archived"
+
+class ModelParameterType(str, Enum):
+    NUMERIC = "numeric"
+    TEXT = "text"
+    BOOLEAN = "boolean"

gmicloud/_internal/_manager/_artifact_manager.py

@@ -1,4 +1,5 @@
 import os
+import time
 from typing import List
 import mimetypes
 
@@ -7,6 +8,9 @@ from .._client._artifact_client import ArtifactClient
 from .._client._file_upload_client import FileUploadClient
 from .._models import *
 
+import logging
+
+logger = logging.getLogger(__name__)
 
 class ArtifactManager:
     """
@@ -81,7 +85,46 @@ class ArtifactManager:
         if not artifact_template_id or not artifact_template_id.strip():
             raise ValueError("Artifact template ID is required and cannot be empty.")
 
-        return self.artifact_client.create_artifact_from_template(artifact_template_id).artifact_id
+        resp = self.artifact_client.create_artifact_from_template(artifact_template_id)
+        if not resp or not resp.artifact_id:
+            raise ValueError("Failed to create artifact from template.")
+
+        return resp.artifact_id
+    
+    def create_artifact_from_template_name(self, artifact_template_name: str) -> tuple[str, ReplicaResource]:
+        """
+        Create an artifact from a template.
+        :param artifact_template_name: The name of the template to use.
+        :return: A tuple containing the artifact ID and the recommended replica resources.
+        :rtype: tuple[str, ReplicaResource]
+        """
+
+        recommended_replica_resources = None
+        template_id = None
+        try:
+            templates = self.get_public_templates()
+        except Exception as e:
+            logger.error(f"Failed to get artifact templates, Error: {e}")
+        for template in templates:
+            if template.template_data and template.template_data.name == artifact_template_name:
+                resources_template = template.template_data.resources
+                recommended_replica_resources = ReplicaResource(
+                    cpu=resources_template.cpu,
+                    ram_gb=resources_template.memory,
+                    gpu=resources_template.gpu,
+                    gpu_name=resources_template.gpu_name,
+                )
+                template_id = template.template_id
+                break
+        if not template_id:
+            raise ValueError(f"Template with name {artifact_template_name} not found.")
+        try: 
+            artifact_id = self.create_artifact_from_template(template_id)
+            self.wait_for_artifact_ready(artifact_id)
+            return artifact_id, recommended_replica_resources
+        except Exception as e:
+            logger.error(f"Failed to create artifact from template, Error: {e}")
+            raise e
 
     def rebuild_artifact(self, artifact_id: str) -> RebuildArtifactResponse:
         """
@@ -170,7 +213,11 @@ class ArtifactManager:
 
         req = GetBigFileUploadUrlRequest(artifact_id=artifact_id, file_name=model_file_name, file_type=model_file_type)
 
-        return self.artifact_client.get_bigfile_upload_url(req).upload_link
+        resp = self.artifact_client.get_bigfile_upload_url(req)
+        if not resp or not resp.upload_link:
+            raise ValueError("Failed to get bigfile upload URL.")
+
+        return resp.upload_link
 
     def delete_bigfile(self, artifact_id: str, file_name: str) -> str:
         """
@@ -182,7 +229,11 @@ class ArtifactManager:
         self._validate_artifact_id(artifact_id)
         self._validate_file_name(file_name)
 
-        return self.artifact_client.delete_bigfile(artifact_id, file_name).status
+        resp = self.artifact_client.delete_bigfile(artifact_id, file_name)
+        if not resp or not resp.status:
+            raise ValueError("Failed to delete bigfile.")
+
+        return resp.status
 
     def upload_large_file(self, upload_link: str, file_path: str) -> None:
         """
@@ -228,15 +279,59 @@ class ArtifactManager:
             FileUploadClient.upload_large_file(bigfile_upload_url_resp.upload_link, model_file_path)
 
         return artifact_id
+    
+
+    def wait_for_artifact_ready(self, artifact_id: str, timeout_s: int = 900) -> None:
+        """
+        Wait for an artifact to be ready.
+
+        :param artifact_id: The ID of the artifact to wait for.
+        :param timeout_s: The timeout in seconds.
+        :return: None
+        """
+        start_time = time.time()
+        while True:
+            try:
+                artifact = self.get_artifact(artifact_id)
+                if artifact.build_status == BuildStatus.SUCCESS:
+                    return
+                elif artifact.build_status in [BuildStatus.FAILED, BuildStatus.TIMEOUT, BuildStatus.CANCELLED]:
+                    raise Exception(f"Artifact build failed, status: {artifact.build_status}")
+            except Exception as e:
+                logger.error(f"Failed to get artifact, Error: {e}")
+            if time.time() - start_time > timeout_s:
+                raise Exception(f"Artifact build takes more than {timeout_s // 60} minutes. Testing aborted.")
+            time.sleep(10)
 
-    def get_artifact_templates(self) -> List[ArtifactTemplate]:
+    
+    def get_public_templates(self) -> List[ArtifactTemplate]:
         """
         Fetch all artifact templates.
 
         :return: A list of ArtifactTemplate objects.
         :rtype: List[ArtifactTemplate]
         """
-        return self.artifact_client.get_artifact_templates().artifact_templates
+        return self.artifact_client.get_public_templates()
+        
+
+    def list_public_template_names(self) -> list[str]:
+        """
+        List all public templates.
+
+        :return: A list of template names.
+        :rtype: list[str]
+        """
+        template_names = []
+        try: 
+            templates = self.get_public_templates()
+            for template in templates:
+                if template.template_data and template.template_data.name:
+                    template_names.append(template.template_data.name)
+            return template_names
+        except Exception as e:
+            logger.error(f"Failed to get artifact templates, Error: {e}")
+            return []
+
 
     @staticmethod
     def _validate_file_name(file_name: str) -> None:

gmicloud/_internal/_manager/_iam_manager.py

@@ -0,0 +1,36 @@
+from datetime import datetime
+from .._client._iam_client import IAMClient
+from .._models import *
+
+
+class IAMManager:
+    """
+    IamManager handles operations related to IAM, including user authentication and authorization.
+    """
+
+    def __init__(self, iam_client: IAMClient):
+        """
+        Initialize the IAMManager instance and the associated IAMClient.
+        """
+        self.iam_client = iam_client
+
+    def create_org_api_key(self, name: str, expires_at: Optional[int] = None) -> str:
+        """
+        Creates a new API key for the current user.
+        """
+
+        if not name:
+            raise ValueError("API key name cannot be empty")
+        if not expires_at:
+            # Set the expiration date to 30 days from now
+            expires_at = int(datetime.now().timestamp()) + 30 * 24 * 60 * 60
+            print(expires_at)
+
+        return self.iam_client.create_org_api_key(
+            CreateAPIKeyRequest(name=name, type="ie_model", expiresAt=expires_at))
+
+    def get_org_api_keys(self) -> List[APIKey]:
+        """
+        Fetches all API keys for the current user.
+        """
+        return self.iam_client.get_org_api_keys().keys

gmicloud/_internal/_manager/_task_manager.py

@@ -4,6 +4,10 @@ from .._client._iam_client import IAMClient
 from .._client._task_client import TaskClient
 from .._models import *
 
+import time
+import logging
+
+logger = logging.getLogger(__name__)
 
 class TaskManager:
     """
@@ -37,7 +41,11 @@ class TaskManager:
 
         :return: A list of `Task` objects.
         """
-        return self.task_client.get_all_tasks(self.iam_client.get_user_id()).tasks
+        resp = self.task_client.get_all_tasks(self.iam_client.get_user_id())
+        if not resp or not resp.tasks:
+            return []
+
+        return resp.tasks
 
     def create_task(self, task: Task) -> Task:
         """
@@ -51,8 +59,11 @@ class TaskManager:
         self._validate_task(task)
         if not task.owner:
             task.owner = TaskOwner(user_id=self.iam_client.get_user_id())
+        resp = self.task_client.create_task(task)
+        if not resp or not resp.task:
+            raise ValueError("Failed to create task.")
 
-        return self.task_client.create_task(task).task
+        return resp.task
 
     def create_task_from_file(self, artifact_id: str, config_file_path: str, trigger_timestamp: int = None) -> Task:
         """
@@ -76,7 +87,7 @@ class TaskManager:
 
         return self.create_task(task)
 
-    def update_task_schedule(self, task: Task):
+    def update_task_schedule(self, task: Task) -> bool:
         """
         Update the schedule of an existing task.
 
@@ -87,10 +98,10 @@ class TaskManager:
         self._validate_task(task)
         self._validate_not_empty(task.task_id, "Task ID")
 
-        self.task_client.update_task_schedule(task)
+        return self.task_client.update_task_schedule(task)
 
     def update_task_schedule_from_file(self, artifact_id: str, task_id: str, config_file_path: str,
-                                       trigger_timestamp: int = None):
+                                       trigger_timestamp: int = None) -> bool:
         """
         Update the schedule of an existing task using data from a file. The file should contain a valid task definition.
 
@@ -112,9 +123,9 @@ class TaskManager:
         if trigger_timestamp:
             task.config.task_scheduling.scheduling_oneoff.trigger_timestamp = trigger_timestamp
 
-        self.update_task_schedule(task)
+        return self.update_task_schedule(task)
 
-    def start_task(self, task_id: str):
+    def start_task(self, task_id: str) -> bool:
         """
         Start a task by its ID.
 
@@ -124,9 +135,53 @@ class TaskManager:
         """
         self._validate_not_empty(task_id, "Task ID")
 
-        self.task_client.start_task(task_id)
+        return self.task_client.start_task(task_id)
+    
 
-    def stop_task(self, task_id: str):
+    def start_task_and_wait(self, task_id: str, timeout_s: int = 900) -> Task:
+        """
+        Start a task and wait for it to be ready.
+
+        :param task_id: The ID of the task to start.
+        :param timeout_s: The timeout in seconds.
+        :return: The task object.
+        :rtype: Task
+        """
+        # trigger start task
+        try:
+            self.start_task(task_id)
+            logger.info(f"Started task ID: {task_id}")
+        except Exception as e:
+            logger.error(f"Failed to start task, Error: {e}")
+            raise e
+        
+        start_time = time.time()
+        while True:
+            try:
+                task = self.get_task(task_id)
+                if task.task_status == TaskStatus.RUNNING:
+                    return task
+                elif task.task_status in [TaskStatus.NEEDSTOP, TaskStatus.ARCHIVED]:
+                    raise Exception(f"Unexpected task status after starting: {task.task_status}")
+                # Also check endpoint status. 
+                elif task.task_status == TaskStatus.RUNNING:
+                    if task.endpoint_info and task.endpoint_info.endpoint_status == TaskEndpointStatus.RUNNING:
+                        return task
+                    elif task.endpoint_info and task.endpoint_info.endpoint_status in [TaskEndpointStatus.UNKNOWN, TaskEndpointStatus.ARCHIVED]:
+                        raise Exception(f"Unexpected endpoint status after starting: {task.endpoint_info.endpoint_status}")
+                    else:
+                        logger.info(f"Pending endpoint starting. endpoint status: {task.endpoint_info.endpoint_status}")
+                else:
+                    logger.info(f"Pending task starting. Task status: {task.task_status}")
+
+            except Exception as e:
+                logger.error(f"Failed to get task, Error: {e}")
+            if time.time() - start_time > timeout_s:
+                raise Exception(f"Task creation takes more than {timeout_s // 60} minutes. Testing aborted.")
+            time.sleep(10)
+
+
+    def stop_task(self, task_id: str) -> bool:
         """
         Stop a task by its ID.
 
@@ -136,7 +191,28 @@ class TaskManager:
         """
         self._validate_not_empty(task_id, "Task ID")
 
-        self.task_client.stop_task(task_id)
+        
+    def stop_task_and_wait(self, task_id: str, timeout_s: int = 900):
+        task_manager = self.task_manager
+        try:
+            self.task_manager.stop_task(task_id)
+            logger.info(f"Stopping task ID: {task_id}")
+        except Exception as e:
+            logger.error(f"Failed to stop task, Error: {e}")
+        task_manager = self.task_manager
+        start_time = time.time()
+        while True:
+            try:
+                task = self.get_task(task_id)
+                if task.task_status == TaskStatus.IDLE:
+                    break
+            except Exception as e:
+                logger.error(f"Failed to get task, Error: {e}")
+            if time.time() - start_time > timeout_s:
+                raise Exception(f"Task stopping takes more than {timeout_s // 60} minutes. Testing aborted.")
+            time.sleep(10)
+
+        return self.task_client.stop_task(task_id)
 
     def get_usage_data(self, start_timestamp: str, end_timestamp: str) -> GetUsageDataResponse:
         """
@@ -151,7 +227,7 @@ class TaskManager:
 
         return self.task_client.get_usage_data(start_timestamp, end_timestamp)
 
-    def archive_task(self, task_id: str):
+    def archive_task(self, task_id: str) -> bool:
         """
         Archive a task by its ID.
 
@@ -161,7 +237,7 @@ class TaskManager:
         """
         self._validate_not_empty(task_id, "Task ID")
 
-        self.task_client.archive_task(task_id)
+        return self.task_client.archive_task(task_id)
 
     @staticmethod
     def _validate_not_empty(value: str, name: str):