gmicloud 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

gmicloud/__init__.py

@@ -15,7 +15,7 @@ from ._internal._models import (
     OneOffScheduling,
     DailyScheduling,
     DailyTrigger,
-    ArtifactTemplate,
+    Template,
 )
 from ._internal._enums import (
     BuildStatus,
@@ -39,7 +39,7 @@ __all__ = [
     "OneOffScheduling",
     "DailyScheduling",
     "DailyTrigger",
-    "ArtifactTemplate",
+    "Template",
     "BuildStatus",
     "TaskEndpointStatus",
 ]

gmicloud/_internal/_client/_artifact_client.py

@@ -1,7 +1,7 @@
 from typing import List
 import logging
 from requests.exceptions import RequestException
-
+import json
 from ._http_client import HTTPClient
 from ._iam_client import IAMClient
 from ._decorator import handle_refresh_token
@@ -120,6 +120,39 @@ class ArtifactClient:
             logger.error(f"Failed to rebuild artifact {artifact_id}: {e}")
             return None
 
+    @handle_refresh_token
+    def add_env_parameters_to_artifact(self, artifact_id: str, env_parameters: dict[str, str]) -> None:
+        """
+        Updates an artifact by its ID.
+
+        :param artifact_id: The ID of the artifact to update.
+        :param request: The request object containing the updated artifact details.
+        """
+        try:
+            old_artifact = self.get_artifact(artifact_id)
+            if not old_artifact:
+                logger.error(f"Artifact {artifact_id} not found")
+                return
+            request = UpdateArtifactRequestBody(
+                artifact_description=old_artifact.artifact_metadata.artifact_description,
+                artifact_name=old_artifact.artifact_metadata.artifact_name,
+                artifact_tags=old_artifact.artifact_metadata.artifact_tags,
+                env_parameters=old_artifact.artifact_parameters.env_parameters,
+                model_parameters=old_artifact.artifact_parameters.model_parameters
+            )
+            new_env_parameters = [EnvParameter(key=k, value=v) for k, v in env_parameters.items()]
+            if not request.env_parameters:
+                request.env_parameters = []
+            request.env_parameters.extend(new_env_parameters)
+            response = self.client.put(
+                f"/update_artifact?artifact_id={artifact_id}",
+                self.iam_client.get_custom_headers(),
+                request.model_dump()
+            )
+        except (RequestException, ValueError) as e:
+            logger.error(f"Failed to add env parameters to artifact {artifact_id}: {e}")
+            return 
+
     @handle_refresh_token
     def delete_artifact(self, artifact_id: str) -> Optional[DeleteArtifactResponse]:
         """
@@ -140,7 +173,7 @@ class ArtifactClient:
             return None
 
     @handle_refresh_token
-    def get_bigfile_upload_url(self, request: GetBigFileUploadUrlRequest) -> Optional[GetBigFileUploadUrlResponse]:
+    def get_bigfile_upload_url(self, request: ResumableUploadLinkRequest) -> Optional[ResumableUploadLinkResponse]:
         """
         Generates a pre-signed URL for uploading a large file.
 
@@ -156,7 +189,7 @@ class ArtifactClient:
                 logger.error("Empty response from /get_bigfile_upload_url")
                 return None
 
-            return GetBigFileUploadUrlResponse.model_validate(response)
+            return ResumableUploadLinkResponse.model_validate(response)
 
         except (RequestException, ValueError) as e:
             logger.error(f"Failed to generate upload URL: {e}")
@@ -186,12 +219,12 @@ class ArtifactClient:
             return None
 
     @handle_refresh_token
-    def get_public_templates(self) -> List[ArtifactTemplate]:
+    def get_public_templates(self) -> List[Template]:
         """
         Fetches all artifact templates.
 
-        :return: A list of ArtifactTemplate objects.
-        :rtype: List[ArtifactTemplate]
+        :return: A list of Template objects.
+        :rtype: List[Template]
         """
         try:
             response = self.client.get("/get_public_templates", self.iam_client.get_custom_headers())
@@ -201,7 +234,7 @@ class ArtifactClient:
                 return []
 
             try:
-                result = GetPublicTemplatesResponse.model_validate(response)
+                result = GetTemplatesResponse.model_validate(response)
                 return result.artifact_templates
             except ValueError as ve:
                 logger.error(f"Failed to validate response data: {ve}")

gmicloud/_internal/_client/_file_upload_client.py

@@ -1,8 +1,10 @@
 import os
 import requests
+import logging
 
 from .._exceptions import UploadFileError
 
+logger = logging.getLogger()
 
 class FileUploadClient:
     CHUNK_SIZE = 10 * 1024 * 1024  # 10MB Default Chunk Size
@@ -45,13 +47,13 @@ class FileUploadClient:
         """
         try:
             file_size = os.path.getsize(file_path)
-            print(f"File Size: {file_size} bytes")
+            logger.info(f"File {file_path} size: {file_size} bytes")
 
             start_byte = 0
             uploaded_range = FileUploadClient._check_file_status(upload_url, file_size)
             if uploaded_range:
                 start_byte = int(uploaded_range.split("-")[1]) + 1
-                print(f"Resuming upload from {start_byte} bytes")
+                logger.info(f"Resuming uploading {file_path} from {start_byte} bytes")
 
             with open(file_path, "rb") as file:
                 while start_byte < file_size:
@@ -74,14 +76,15 @@ class FileUploadClient:
                     # Ensure upload is successful for this chunk
                     if resp.status_code not in (200, 201, 308):
                         raise UploadFileError(
-                            f"Failed to upload file, code:{resp.status_code} ,message: {resp.text}")
+                            f"Failed to upload file {file_path}, code:{resp.status_code} ,message: {resp.text}")
 
                     start_byte = end_byte + 1
-                    print(f"Uploaded {end_byte + 1}/{file_size} bytes")
+                    percentage = (start_byte / file_size) * 100
+                    logger.info(f"File {file_path} uploaded {end_byte + 1:,}/{file_size:,} bytes ({percentage:.2f}%)")
 
-                print("Upload completed successfully.")
+                logger.info(f"File {file_path} uploaded successfully.")
         except Exception as e:
-            raise UploadFileError(f"Failed to upload file: {str(e)}")
+            raise UploadFileError(f"Failed to upload file {file_path}, got error: {str(e)}")
 
     @staticmethod
     def _check_file_status(upload_url: str, file_size: int) -> str:
@@ -104,7 +107,7 @@ class FileUploadClient:
             if resp.status_code == 308:
                 range_header = resp.headers.get("Range")
                 if range_header:
-                    print(f"Server reports partial upload range: {range_header}")
+                    logger.info(f"Server reports partial upload range: {range_header}")
                 return range_header
 
             if resp.status_code in (200, 201):

gmicloud/_internal/_config.py

@@ -1,3 +1,9 @@
-ARTIFACT_SERVICE_BASE_URL = "https://ce-tot.gmicloud-dev.com/api/v1/ie/artifact"
-TASK_SERVICE_BASE_URL = "https://ce-tot.gmicloud-dev.com/api/v1/ie/task"
-IAM_SERVICE_BASE_URL = "https://ce-tot.gmicloud-dev.com/api/v1"
+# Dev environment
+# ARTIFACT_SERVICE_BASE_URL = "https://ce-tot.gmicloud-dev.com/api/v1/ie/artifact"
+# TASK_SERVICE_BASE_URL = "https://ce-tot.gmicloud-dev.com/api/v1/ie/task"
+# IAM_SERVICE_BASE_URL = "https://ce-tot.gmicloud-dev.com/api/v1"
+
+# Prod environment
+ARTIFACT_SERVICE_BASE_URL = "https://inference-engine.gmicloud.ai/api/v1/ie/artifact"
+TASK_SERVICE_BASE_URL = "https://inference-engine.gmicloud.ai/api/v1/ie/task"
+IAM_SERVICE_BASE_URL = "https://inference-engine.gmicloud.ai/api/v1"

gmicloud/_internal/_enums.py

@@ -31,3 +31,8 @@ class TaskStatus(str, Enum):
     RUNNING = "running"
     NEEDSTOP = "needstop"
     ARCHIVED = "archived"
+
+class ModelParameterType(str, Enum):
+    NUMERIC = "numeric"
+    TEXT = "text"
+    BOOLEAN = "boolean"

gmicloud/_internal/_manager/_artifact_manager.py

@@ -1,12 +1,21 @@
 import os
+import time
 from typing import List
 import mimetypes
+import concurrent.futures
+import re
+from tqdm import tqdm
+from tqdm.contrib.logging import logging_redirect_tqdm
 
 from .._client._iam_client import IAMClient
 from .._client._artifact_client import ArtifactClient
 from .._client._file_upload_client import FileUploadClient
 from .._models import *
+from .._manager.serve_command_utils import parse_server_command, extract_gpu_num_from_serve_command
 
+import logging
+
+logger = logging.getLogger(__name__)
 
 class ArtifactManager:
     """
@@ -49,7 +58,12 @@ class ArtifactManager:
             self,
             artifact_name: str,
             description: Optional[str] = "",
-            tags: Optional[List[str]] = None
+            tags: Optional[List[str]] = None,
+            deployment_type: Optional[str] = "",
+            template_id: Optional[str] = "",
+            env_parameters: Optional[List["EnvParameter"]] = None,
+            model_description: Optional[str] = "",
+            model_parameters: Optional[List["ModelParameter"]] = None,
     ) -> CreateArtifactResponse:
         """
         Create a new artifact for a user.
@@ -65,11 +79,16 @@ class ArtifactManager:
 
         req = CreateArtifactRequest(artifact_name=artifact_name,
                                     artifact_description=description,
-                                    artifact_tags=tags, )
+                                    artifact_tags=tags,
+                                    deployment_type=deployment_type,
+                                    template_id=template_id,
+                                    env_parameters=env_parameters,
+                                    model_description=model_description,
+                                    model_parameters=model_parameters)
 
         return self.artifact_client.create_artifact(req)
 
-    def create_artifact_from_template(self, artifact_template_id: str) -> str:
+    def create_artifact_from_template(self, artifact_template_id: str, env_parameters: Optional[dict[str, str]] = None) -> str:
         """
         Create a new artifact for a user using a template.
 
@@ -81,12 +100,102 @@ class ArtifactManager:
         if not artifact_template_id or not artifact_template_id.strip():
             raise ValueError("Artifact template ID is required and cannot be empty.")
 
+    
         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.")
 
+        if env_parameters:
+            self.artifact_client.add_env_parameters_to_artifact(resp.artifact_id, env_parameters)
+
         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 create_artifact_for_serve_command_and_custom_model(self, template_name: str, artifact_name: str, serve_command: str, gpu_type: str, artifact_description: str = "") -> tuple[str, ReplicaResource]:
+        """
+        Create an artifact from a template and support custom model.
+        :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
+        picked_template = 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 == template_name:
+                picked_template = template
+                break
+        if not picked_template:
+            raise ValueError(f"Template with name {template_name} not found.")
+        
+        try:
+            if gpu_type not in ["H100", "H200"]:
+                raise ValueError("Only support A100 and H100 for now")
+            
+            type, env_vars, serve_args_dict = parse_server_command(serve_command)
+            if type.lower() not in template_name.lower():
+                raise ValueError(f"Template {template_name} does not support inference with {type}.")
+            num_gpus = extract_gpu_num_from_serve_command(serve_args_dict)
+            recommended_replica_resources = ReplicaResource(
+                cpu=num_gpus * 16,
+                ram_gb=num_gpus * 100,
+                gpu=num_gpus,
+                gpu_name=gpu_type,
+            )
+        except Exception as e:
+            raise ValueError(f"Failed to parse serve command, Error: {e}")
+
+        try:
+            env_vars = [
+                EnvParameter(key="SERVE_COMMAND", value=serve_command),
+                EnvParameter(key="GPU_TYPE", value=gpu_type),
+            ]
+            resp = self.create_artifact(artifact_name, artifact_description, deployment_type="template", template_id=picked_template.template_id, env_parameters=env_vars)
+            # Assume Artifact is already with BuildStatus.SUCCESS status
+            return resp.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:
         """
         Rebuild an existing artifact.
@@ -172,7 +281,7 @@ class ArtifactManager:
         model_file_name = os.path.basename(model_file_path)
         model_file_type = mimetypes.guess_type(model_file_path)[0]
 
-        req = GetBigFileUploadUrlRequest(artifact_id=artifact_id, file_name=model_file_name, file_type=model_file_type)
+        req = ResumableUploadLinkRequest(artifact_id=artifact_id, file_name=model_file_name, file_type=model_file_type)
 
         resp = self.artifact_client.get_bigfile_upload_url(req)
         if not resp or not resp.upload_link:
@@ -211,44 +320,116 @@ class ArtifactManager:
 
         FileUploadClient.upload_large_file(upload_link, file_path)
 
+
+    def upload_model_files_to_artifact(self, artifact_id: str, model_directory: str) -> None:
+        """
+        Upload model files to an existing artifact.
+
+        :param artifact_id: The ID of the artifact to upload the model files to.
+        :param model_directory: The path to the model directory.
+        """
+
+        # List all files in the model directory recursively
+        model_file_paths = []
+        for root, _, files in os.walk(model_directory):
+            for file in files:
+                model_file_paths.append(os.path.join(root, file))
+
+        def upload_file(model_file_path):
+            self._validate_file_path(model_file_path)
+            bigfile_upload_url_resp = self.artifact_client.get_bigfile_upload_url(
+                ResumableUploadLinkRequest(artifact_id=artifact_id, file_name=os.path.basename(model_file_path))
+            )
+            FileUploadClient.upload_large_file(bigfile_upload_url_resp.upload_link, model_file_path)
+
+        # Upload files in parallel with progress bar
+        with tqdm(total=len(model_file_paths), desc="Uploading model files") as progress_bar:
+            with logging_redirect_tqdm():
+                with concurrent.futures.ThreadPoolExecutor() as executor:
+                    futures = {executor.submit(upload_file, path): path for path in model_file_paths}
+                    for future in concurrent.futures.as_completed(futures):
+                        try:
+                            future.result()
+                        except Exception as e:
+                            logger.error(f"Failed to upload file {futures[future]}, Error: {e}")
+                        progress_bar.update(1)
+
     def create_artifact_with_model_files(
             self,
             artifact_name: str,
             artifact_file_path: str,
-            model_file_paths: List[str],
+            model_directory: str,
             description: Optional[str] = "",
             tags: Optional[str] = None
     ) -> str:
         """
         Create a new artifact for a user and upload model files associated with the artifact.
-
         :param artifact_name: The name of the artifact.
         :param artifact_file_path: The path to the artifact file(Dockerfile+serve.py).
-        :param model_file_paths: The paths to the model files.
+        :param model_directory: The path to the model directory.
         :param description: An optional description for the artifact.
         :param tags: Optional tags associated with the artifact, as a comma-separated string.
         :return: The `artifact_id` of the created artifact.
-        :raises FileNotFoundError: If the provided `file_path` does not exist.
         """
         artifact_id = self.create_artifact_with_file(artifact_name, artifact_file_path, description, tags)
+        logger.info(f"Artifact created: {artifact_id}")
 
-        for model_file_path in model_file_paths:
-            self._validate_file_path(model_file_path)
-            bigfile_upload_url_resp = self.artifact_client.get_bigfile_upload_url(
-                GetBigFileUploadUrlRequest(artifact_id=artifact_id, model_file_path=model_file_path)
-            )
-            FileUploadClient.upload_large_file(bigfile_upload_url_resp.upload_link, model_file_path)
+        self.upload_model_files_to_artifact(artifact_id, model_directory)
 
         return artifact_id
 
-    def get_public_templates(self) -> List[ArtifactTemplate]:
+
+    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_public_templates(self) -> List[Template]:
         """
         Fetch all artifact templates.
 
-        :return: A list of ArtifactTemplate objects.
-        :rtype: List[ArtifactTemplate]
+        :return: A list of Template objects.
+        :rtype: List[Template]
         """
         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/_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,7 @@ class TaskManager:
 
         :return: A list of `Task` objects.
         """
-        resp = self.task_client.get_all_tasks(self.iam_client.get_user_id())
+        resp = self.task_client.get_all_tasks()
         if not resp or not resp.tasks:
             return []
 
@@ -59,6 +63,7 @@ class TaskManager:
         if not resp or not resp.task:
             raise ValueError("Failed to create task.")
 
+        logger.info(f"Task created: {resp.task.task_id}")
         return resp.task
 
     def create_task_from_file(self, artifact_id: str, config_file_path: str, trigger_timestamp: int = None) -> Task:
@@ -132,6 +137,56 @@ class TaskManager:
         self._validate_not_empty(task_id, "Task ID")
 
         return self.task_client.start_task(task_id)
+    
+
+    def wait_for_task(self, task_id: str, timeout_s: int = 900) -> Task:
+        """
+        Wait for a task to reach the RUNNING state or raise an exception if it fails.
+
+        :param task_id: The ID of the task to wait for.
+        :param timeout_s: The timeout in seconds.
+        :return: The task object.
+        :rtype: Task
+        """
+        start_time = time.time()
+        while True:
+            try:
+                task = self.get_task(task_id)
+                if task.task_status == TaskStatus.RUNNING:
+                    if task.endpoint_info is not None and task.endpoint_info.endpoint_status == TaskEndpointStatus.RUNNING:
+                        return task
+                    else:
+                        if task.cluster_endpoints:
+                            for ce in task.cluster_endpoints:
+                                if ce.endpoint_status == TaskEndpointStatus.RUNNING:
+                                    return task
+                if task.task_status in [TaskStatus.NEEDSTOP, TaskStatus.ARCHIVED]:
+                    raise Exception(f"Unexpected task status after starting: {task.task_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 start_task_and_wait(self, task_id: str, timeout_s: int = 3600) -> 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
+        """
+        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
+
+        return self.wait_for_task(task_id, timeout_s)
 
     def stop_task(self, task_id: str) -> bool:
         """
@@ -142,9 +197,28 @@ class TaskManager:
         :raises ValueError: If `task_id` is invalid (None or empty string).
         """
         self._validate_not_empty(task_id, "Task ID")
-
         return self.task_client.stop_task(task_id)
 
+        
+    def stop_task_and_wait(self, task_id: str, timeout_s: int = 3600):
+        try:
+            self.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}")
+        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)
+
+
     def get_usage_data(self, start_timestamp: str, end_timestamp: str) -> GetUsageDataResponse:
         """
         Retrieve the usage data of a task within a given time range.