[Bug] Refactor downloader and artifact_service to be non-blocking (#1895)

* Downloader minor refactor and gcp_test add timeout setting
* Refactor downloaders and artifact_service to be non-blocking
* Add unit tests for downloader
* Remove unused libs
* Move test_downloaders.py
* Fix minor typos in downloader tests
* Recover KV Cache status func

---------

Signed-off-by: Liguang Xie <[email protected]>
Co-authored-by: Liguang Xie <[email protected]>

Author: xieus

deployment/terraform/tests/gcp_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 	"net"
 	"net/http"
+	"time"
 
 	"github.com/gruntwork-io/terratest/modules/terraform"
 	"github.com/stretchr/testify/assert"
@@ -75,8 +76,11 @@ func TestAIBrixGCPDeployment(t *testing.T) {
 		option.WithMaxRetries(0),
 	)
 
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
 	// Run a chat completion against the model endpoint
-	chatCompletion, err := client.Chat.Completions.New(context.TODO(), openai.ChatCompletionNewParams{
+	chatCompletion, err := client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{
 		Messages: openai.F([]openai.ChatCompletionMessageParamUnion{
 			openai.UserMessage("What can you tell me about San Francisco?"),
 		}),

python/aibrix/aibrix/runtime/artifact_service.py

@@ -14,6 +14,7 @@
 
 """Artifact delegation service for LoRA adapters."""
 
+import asyncio
 import os
 import shutil
 from pathlib import Path
@@ -313,7 +314,8 @@ async def unload_adapter(
                 local_path = self._get_local_path_for_adapter(lora_name)
                 if os.path.exists(local_path):
                     try:
-                        shutil.rmtree(local_path)
+                        loop = asyncio.get_running_loop()
+                        await loop.run_in_executor(None, shutil.rmtree, local_path)
                         logger.info(f"Cleaned up local artifacts for {lora_name}")
                     except Exception as e:
                         logger.warning(
@@ -337,7 +339,8 @@ async def cleanup_artifact(self, lora_name: str) -> None:
 
         if os.path.exists(local_path):
             try:
-                shutil.rmtree(local_path)
+                loop = asyncio.get_running_loop()
+                await loop.run_in_executor(None, shutil.rmtree, local_path)
                 logger.info(f"Cleaned up artifacts for {lora_name}")
             except Exception as e:
                 logger.warning(f"Failed to clean up artifacts for {lora_name}: {e}")

python/aibrix/aibrix/runtime/downloaders.py

@@ -14,6 +14,8 @@
 
 """Artifact downloaders for different storage backends."""
 
+import asyncio
+import functools
 import os
 from abc import ABC, abstractmethod
 from pathlib import Path
@@ -58,6 +60,18 @@ class S3ArtifactDownloader(ArtifactDownloader):
 
     async def download(
         self, source_url: str, local_path: str, credentials: Optional[Dict] = None
+    ) -> str:
+        """
+        Download from S3 (Async wrapper).
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            None,
+            functools.partial(self._download_sync, source_url, local_path, credentials),
+        )
+
+    def _download_sync(
+        self, source_url: str, local_path: str, credentials: Optional[Dict] = None
     ) -> str:
         """
         Download from S3.
@@ -154,7 +168,7 @@ def _download_s3_directory(
                 local_file_path = os.path.join(local_path, relative_path)
 
                 # Ensure directory exists
-                os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+                self._ensure_directory(os.path.dirname(local_file_path))
 
                 # Download file
                 s3_client.download_file(bucket_name, key, local_file_path)
@@ -168,6 +182,18 @@ class GCSArtifactDownloader(ArtifactDownloader):
 
     async def download(
         self, source_url: str, local_path: str, credentials: Optional[Dict] = None
+    ) -> str:
+        """
+        Download from GCS (Async wrapper).
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            None,
+            functools.partial(self._download_sync, source_url, local_path, credentials),
+        )
+
+    def _download_sync(
+        self, source_url: str, local_path: str, credentials: Optional[Dict] = None
     ) -> str:
         """
         Download from GCS.
@@ -227,7 +253,7 @@ async def download(
                     local_file_path = os.path.join(local_path, relative_path)
 
                     # Ensure directory exists
-                    os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+                    self._ensure_directory(os.path.dirname(local_file_path))
 
                     # Download file
                     blob.download_to_filename(local_file_path)
@@ -257,6 +283,18 @@ class HuggingFaceArtifactDownloader(ArtifactDownloader):
 
     async def download(
         self, source_url: str, local_path: str, credentials: Optional[Dict] = None
+    ) -> str:
+        """
+        Download from HuggingFace Hub (Async wrapper).
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            None,
+            functools.partial(self._download_sync, source_url, local_path, credentials),
+        )
+
+    def _download_sync(
+        self, source_url: str, local_path: str, credentials: Optional[Dict] = None
     ) -> str:
         """
         Download from HuggingFace Hub.

python/aibrix/tests/downloader/test_downloader_hf.py

@@ -29,7 +29,7 @@ def test_get_downloader_hf():
 
 def test_get_downloader_hf_not_exist():
     with pytest.raises(ModelNotFoundError) as exception:
-        get_downloader("not_exsit_path/model")
+        get_downloader("not_exist_path/model")
     assert "Model not found" in str(exception.value)
 
 
@@ -39,5 +39,5 @@ def test_get_downloader_hf_invalid_uri():
     assert "not in the expected format: repo/name" in str(exception.value)
 
     with pytest.raises(ArgNotFormatError) as exception:
-        get_downloader("multi/filed/repo")
+        get_downloader("multi/field/repo")
     assert "not in the expected format: repo/name" in str(exception.value)

python/aibrix/tests/downloader/test_downloader_tos.py

@@ -63,7 +63,7 @@ def test_get_downloader_tos_path_not_exist(mock_boto3):
     mock_not_exsit_tos(mock_boto3)
 
     with pytest.raises(ModelNotFoundError) as exception:
-        get_downloader("tos://bucket/not_exsit_path")
+        get_downloader("tos://bucket/not_exist_path")
     assert "Model not found" in str(exception.value)
 
 

python/aibrix/tests/downloader/test_downloader_tos_v1.py

@@ -61,7 +61,7 @@ def test_get_downloader_tos_path_not_exist(mock_tos):
     mock_not_exsit_tos(mock_tos)
 
     with pytest.raises(ModelNotFoundError) as exception:
-        get_downloader("tos://bucket/not_exsit_path")
+        get_downloader("tos://bucket/not_exist_path")
     assert "Model not found" in str(exception.value)
 
 

python/aibrix_kvcache/aibrix_kvcache/status.py

@@ -74,6 +74,27 @@ def __init__(  # type: ignore[misc]
     def is_ok(self) -> bool:
         return self.error_code == StatusCodes.OK
 
+    def is_error(self) -> bool:
+        return self.error_code == StatusCodes.ERROR
+
+    def is_not_found(self) -> bool:
+        return self.error_code == StatusCodes.NOT_FOUND
+
+    def is_invalid(self) -> bool:
+        return self.error_code == StatusCodes.INVALID
+
+    def is_out_of_memory(self) -> bool:
+        return self.error_code == StatusCodes.OUT_OF_MEMORY
+
+    def is_timeout(self) -> bool:
+        return self.error_code == StatusCodes.TIMEOUT
+
+    def is_denied(self) -> bool:
+        return self.error_code == StatusCodes.DENIED
+
+    def is_cancelled(self) -> bool:
+        return self.error_code == StatusCodes.CANCELLED
+
     def get(self, default=None) -> T:
         """Returns the value if successful, otherwise returns default."""
         if not self.is_ok():