gmicloud 0.1.6__tar.gz → 0.1.7__tar.gz

{gmicloud-0.1.6 → gmicloud-0.1.7}/PKG-INFO

@@ -1,8 +1,8 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: gmicloud
-Version: 0.1.6
+Version: 0.1.7
 Summary: GMI Cloud Python SDK
-Author-email: GMI <support@gmicloud.ai>
+Author-email: GMI <gmi@gmitec.net>
 License: MIT
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -10,7 +10,7 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.6
 Description-Content-Type: text/markdown
 
-# GMICloud SDK (Beta)
+# GMICloud SDK
 
 ## Overview
 Before you start: Our service and GPU resource is currenly invite-only so please contact our team (getstarted@gmicloud.ai) to get invited if you don't have one yet.
@@ -45,7 +45,7 @@ There are two ways to configure the SDK:
 Set the following environment variables:
 
 ```shell
-export GMI_CLOUD_CLIENT_ID=<YOUR_CLIENT_ID>
+export GMI_CLOUD_CLIENT_ID=<YOUR_CLIENT_ID> # Pick what every ID you need.
 export GMI_CLOUD_EMAIL=<YOUR_EMAIL>
 export GMI_CLOUD_PASSWORD=<YOUR_PASSWORD>
 ```
@@ -73,7 +73,7 @@ pip install -r requirements.txt
 python -m examples.create_task_from_artifact_template.py
 ```
 
-### 2. Create an inference task from an artifact template
+### 2. Example of create an inference task from an artifact template
 
 This is the simplest example to deploy an inference task using an existing artifact template:
 
@@ -119,6 +119,97 @@ print(call_chat_completion(cli, task.task_id))
 
 ```
 
+### 3. Example of creating an inference task based on custom model with local vllm / SGLang serve command
+* Full example is available at [examples/inference_task_with_custom_model.py](https://github.com/GMISWE/python-sdk/blob/main/examples/inference_task_with_custom_model.py)
+
+1. Prepare custom model checkpoint (using a model downloaded from HF as an example)
+
+```python
+# Download model from huggingface
+from huggingface_hub import snapshot_download
+
+model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+model_checkpoint_save_dir = "files/model_garden"
+snapshot_download(repo_id=model_name, local_dir=model_checkpoint_save_dir)
+```
+
+2. Find a template of specific SGLang version
+
+```python
+# export GMI_CLOUD_CLIENT_ID=<YOUR_CLIENT_ID>
+# export GMI_CLOUD_EMAIL=<YOUR_EMAIL>
+# export GMI_CLOUD_PASSWORD=<YOUR_PASSWORD>
+cli = Client()
+
+# List templates offered by GMI cloud 
+templates = cli.artifact_manager.list_public_template_names()
+print(f"Found {len(templates)} templates: {templates}")
+```
+
+3. Pick a template (e.g. SGLang 0.4.5) and prepare a local serve command
+
+```python
+# Example for vllm server
+picked_template_name = "gmi_vllm_0.8.4"
+serve_command = "vllm serve deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --trust-remote-code --gpu-memory-utilization 0.8"
+
+# Example for sglang server
+picked_template_name = "gmi_sglang_0.4.5.post1"
+serve_command = "python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --trust-remote-code --mem-fraction-static 0.8 --tp 2"
+```
+
+4. Create an artifact and upload custom model. The artifact can be reused to create inference tasks later. Artifact also suggests recommended resources for each inference server replica
+
+```python
+artifact_id, recommended_replica_resources = cli.artifact_manager.create_artifact_from_template_name(
+    artifact_template_name=picked_template_name,
+    env_parameters={
+        "SERVER_COMMAND": serve_command,
+        "GPU_TYPE": "H100",
+    }
+)
+print(f"Created artifact {artifact_id} with recommended resources: {recommended_replica_resources}")
+
+# Upload model files to artifact
+cli.artifact_manager.upload_model_files_to_artifact(artifact_id, model_checkpoint_save_dir)
+```
+
+5. Create Inference task (defining min/max inference replica), start and wait
+
+```python
+new_task = Task(
+    config=TaskConfig(
+        ray_task_config=RayTaskConfig(
+            artifact_id=artifact_id,
+            file_path="serve",
+            deployment_name="app",
+            replica_resource=recommended_replica_resources,
+        ),
+        task_scheduling = TaskScheduling(
+            scheduling_oneoff=OneOffScheduling(
+                trigger_timestamp=int(datetime.now().timestamp()),
+                min_replicas=1,
+                max_replicas=4,
+            )
+        ),
+    ),
+)
+task = cli.task_manager.create_task(new_task)
+task_id = task.task_id
+task = cli.task_manager.get_task(task_id)
+print(f"Task created: {task.config.task_name}. You can check details at https://inference-engine.gmicloud.ai/user-console/task")
+
+# Start Task and wait for it to be ready
+cli.task_manager.start_task_and_wait(task_id)
+```
+
+6. Test with sample chat completion request
+
+```python
+print(call_chat_completion(cli, task_id))
+```
+
+
 ## API Reference
 
 ### Client
@@ -144,4 +235,3 @@ password: Optional[str] = ""
 * get_task(task_id: str): Retrieve the status and details of a specific task.
 
 ## Notes & Troubleshooting
-k

{gmicloud-0.1.6 → gmicloud-0.1.7}/README.md

@@ -1,4 +1,4 @@
-# GMICloud SDK (Beta)
+# GMICloud SDK
 
 ## Overview
 Before you start: Our service and GPU resource is currenly invite-only so please contact our team (getstarted@gmicloud.ai) to get invited if you don't have one yet.
@@ -33,7 +33,7 @@ There are two ways to configure the SDK:
 Set the following environment variables:
 
 ```shell
-export GMI_CLOUD_CLIENT_ID=<YOUR_CLIENT_ID>
+export GMI_CLOUD_CLIENT_ID=<YOUR_CLIENT_ID> # Pick what every ID you need.
 export GMI_CLOUD_EMAIL=<YOUR_EMAIL>
 export GMI_CLOUD_PASSWORD=<YOUR_PASSWORD>
 ```
@@ -61,7 +61,7 @@ pip install -r requirements.txt
 python -m examples.create_task_from_artifact_template.py
 ```
 
-### 2. Create an inference task from an artifact template
+### 2. Example of create an inference task from an artifact template
 
 This is the simplest example to deploy an inference task using an existing artifact template:
 
@@ -107,6 +107,97 @@ print(call_chat_completion(cli, task.task_id))
 
 ```
 
+### 3. Example of creating an inference task based on custom model with local vllm / SGLang serve command
+* Full example is available at [examples/inference_task_with_custom_model.py](https://github.com/GMISWE/python-sdk/blob/main/examples/inference_task_with_custom_model.py)
+
+1. Prepare custom model checkpoint (using a model downloaded from HF as an example)
+
+```python
+# Download model from huggingface
+from huggingface_hub import snapshot_download
+
+model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+model_checkpoint_save_dir = "files/model_garden"
+snapshot_download(repo_id=model_name, local_dir=model_checkpoint_save_dir)
+```
+
+2. Find a template of specific SGLang version
+
+```python
+# export GMI_CLOUD_CLIENT_ID=<YOUR_CLIENT_ID>
+# export GMI_CLOUD_EMAIL=<YOUR_EMAIL>
+# export GMI_CLOUD_PASSWORD=<YOUR_PASSWORD>
+cli = Client()
+
+# List templates offered by GMI cloud 
+templates = cli.artifact_manager.list_public_template_names()
+print(f"Found {len(templates)} templates: {templates}")
+```
+
+3. Pick a template (e.g. SGLang 0.4.5) and prepare a local serve command
+
+```python
+# Example for vllm server
+picked_template_name = "gmi_vllm_0.8.4"
+serve_command = "vllm serve deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --trust-remote-code --gpu-memory-utilization 0.8"
+
+# Example for sglang server
+picked_template_name = "gmi_sglang_0.4.5.post1"
+serve_command = "python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --trust-remote-code --mem-fraction-static 0.8 --tp 2"
+```
+
+4. Create an artifact and upload custom model. The artifact can be reused to create inference tasks later. Artifact also suggests recommended resources for each inference server replica
+
+```python
+artifact_id, recommended_replica_resources = cli.artifact_manager.create_artifact_from_template_name(
+    artifact_template_name=picked_template_name,
+    env_parameters={
+        "SERVER_COMMAND": serve_command,
+        "GPU_TYPE": "H100",
+    }
+)
+print(f"Created artifact {artifact_id} with recommended resources: {recommended_replica_resources}")
+
+# Upload model files to artifact
+cli.artifact_manager.upload_model_files_to_artifact(artifact_id, model_checkpoint_save_dir)
+```
+
+5. Create Inference task (defining min/max inference replica), start and wait
+
+```python
+new_task = Task(
+    config=TaskConfig(
+        ray_task_config=RayTaskConfig(
+            artifact_id=artifact_id,
+            file_path="serve",
+            deployment_name="app",
+            replica_resource=recommended_replica_resources,
+        ),
+        task_scheduling = TaskScheduling(
+            scheduling_oneoff=OneOffScheduling(
+                trigger_timestamp=int(datetime.now().timestamp()),
+                min_replicas=1,
+                max_replicas=4,
+            )
+        ),
+    ),
+)
+task = cli.task_manager.create_task(new_task)
+task_id = task.task_id
+task = cli.task_manager.get_task(task_id)
+print(f"Task created: {task.config.task_name}. You can check details at https://inference-engine.gmicloud.ai/user-console/task")
+
+# Start Task and wait for it to be ready
+cli.task_manager.start_task_and_wait(task_id)
+```
+
+6. Test with sample chat completion request
+
+```python
+print(call_chat_completion(cli, task_id))
+```
+
+
 ## API Reference
 
 ### Client
@@ -132,4 +223,3 @@ password: Optional[str] = ""
 * get_task(task_id: str): Retrieve the status and details of a specific task.
 
 ## Notes & Troubleshooting
-k

{gmicloud-0.1.6 → gmicloud-0.1.7}/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-0.1.6 → gmicloud-0.1.7}/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-0.1.6 → gmicloud-0.1.7}/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-0.1.7/gmicloud/_internal/_config.py

@@ -0,0 +1,9 @@
+# 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-0.1.6 → gmicloud-0.1.7}/gmicloud/_internal/_manager/_artifact_manager.py

@@ -2,11 +2,16 @@ 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
 
@@ -53,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.
@@ -69,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.
 
@@ -85,11 +100,16 @@ 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]:
         """
@@ -125,6 +145,56 @@ class ArtifactManager:
         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:
         """
@@ -211,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:
@@ -250,36 +320,64 @@ 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 wait_for_artifact_ready(self, artifact_id: str, timeout_s: int = 900) -> None:
         """
@@ -304,12 +402,12 @@ class ArtifactManager:
             time.sleep(10)
 
     
-    def get_public_templates(self) -> List[ArtifactTemplate]:
+    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()