[bug-fix] LLM Artifact Gateway .list_files() (#416)

* fix test to catch use case

* fix parsing for prefixes

Author: saiatmakuri

model-engine/model_engine_server/core/utils/url.py

@@ -32,7 +32,7 @@ class InvalidAttachmentUrl(ValueError):
     pass
 
 
-def parse_attachment_url(url: str) -> ParsedURL:
+def parse_attachment_url(url: str, clean_key: bool = True) -> ParsedURL:
     """Extracts protocol, bucket, region, and key from the :param:`url`.
 
     :raises: InvalidAttachmentUrl Iff the input `url` is not a valid AWS S3 or GCS url.
@@ -102,5 +102,5 @@ def clean(v):
         protocol=clean(protocol),
         bucket=clean(bucket),
         region=clean(region),
-        key=clean(key),
+        key=clean(key) if clean_key else key,
     )

model-engine/model_engine_server/infra/gateways/s3_llm_artifact_gateway.py

@@ -23,7 +23,7 @@ def _get_s3_resource(self, kwargs):
 
     def list_files(self, path: str, **kwargs) -> List[str]:
         s3 = self._get_s3_resource(kwargs)
-        parsed_remote = parse_attachment_url(path)
+        parsed_remote = parse_attachment_url(path, clean_key=False)
         bucket = parsed_remote.bucket
         key = parsed_remote.key
 
@@ -33,7 +33,7 @@ def list_files(self, path: str, **kwargs) -> List[str]:
 
     def download_files(self, path: str, target_path: str, overwrite=False, **kwargs) -> List[str]:
         s3 = self._get_s3_resource(kwargs)
-        parsed_remote = parse_attachment_url(path)
+        parsed_remote = parse_attachment_url(path, clean_key=False)
         bucket = parsed_remote.bucket
         key = parsed_remote.key
 
@@ -58,7 +58,9 @@ def download_files(self, path: str, target_path: str, overwrite=False, **kwargs)
 
     def get_model_weights_urls(self, owner: str, model_name: str, **kwargs) -> List[str]:
         s3 = self._get_s3_resource(kwargs)
-        parsed_remote = parse_attachment_url(hmi_config.hf_user_fine_tuned_weights_prefix)
+        parsed_remote = parse_attachment_url(
+            hmi_config.hf_user_fine_tuned_weights_prefix, clean_key=False
+        )
         bucket = parsed_remote.bucket
         fine_tuned_weights_prefix = parsed_remote.key
 

model-engine/tests/unit/infra/gateways/test_s3_llm_artifact_gateway.py

@@ -14,7 +14,7 @@ def llm_artifact_gateway():
 
 @pytest.fixture
 def fake_files():
-    return ["fake-prefix/fake1", "fake-prefix/fake2", "fake-prefix/fake3"]
+    return ["fake-prefix/fake1", "fake-prefix/fake2", "fake-prefix/fake3", "fake-prefix-ext/fake1"]
 
 
 def mock_boto3_session(fake_files: List[str]):
@@ -39,11 +39,13 @@ def filter_files(*args, **kwargs):
     lambda *args, **kwargs: None,  # noqa
 )
 def test_s3_llm_artifact_gateway_download_folder(llm_artifact_gateway, fake_files):
-    prefix = "/".join(fake_files[0].split("/")[:-1])
+    prefix = "/".join(fake_files[0].split("/")[:-1]) + "/"
     uri_prefix = f"s3://fake-bucket/{prefix}"
     target_dir = "fake-target"
 
-    expected_files = [f"{target_dir}/{file.split('/')[-1]}" for file in fake_files]
+    expected_files = [
+        f"{target_dir}/{file.split('/')[-1]}" for file in fake_files if file.startswith(prefix)
+    ]
     with mock.patch(
         "model_engine_server.infra.gateways.s3_llm_artifact_gateway.boto3.Session",
         mock_boto3_session(fake_files),