Smartly check safetensors vs. bin (#296)

* Smartly check safetensors vs. bin

* Fix formatting

* Add unit test

* Add unit test

* heh hope this works.

* refactoring

* adding new utils file, removing test

* adding in unit test, refactoring again

* adding artifact gateway to use case

* renaming gateway function

* whoops

* cleanup

---------

Co-authored-by: Ian Macleod <[email protected]>
Co-authored-by: Ian Macleod <[email protected]>

Author: 3 people

model-engine/model_engine_server/api/llms_v1.py

@@ -92,6 +92,7 @@ async def create_model_endpoint(
             create_model_bundle_use_case=create_model_bundle_use_case,
             model_bundle_repository=external_interfaces.model_bundle_repository,
             model_endpoint_service=external_interfaces.model_endpoint_service,
+            llm_artifact_gateway=external_interfaces.llm_artifact_gateway,
         )
         return await use_case.execute(user=auth, request=request)
     except ObjectAlreadyExistsException as exc:

model-engine/model_engine_server/core/aws/storage_client.py

@@ -20,7 +20,7 @@
 
 
 def sync_storage_client(**kwargs) -> BaseClient:
-    return session(infra_config().profile_ml_worker).client("s3", **kwargs)
+    return session(infra_config().profile_ml_worker).client("s3", **kwargs)  # type: ignore
 
 
 def open(uri: str, mode: str = "rt", **kwargs) -> IO:  # pylint: disable=redefined-builtin

model-engine/model_engine_server/domain/gateways/llm_artifact_gateway.py

@@ -7,6 +7,13 @@ class LLMArtifactGateway(ABC):
     Abstract Base Class for interacting with llm artifacts.
     """
 
+    @abstractmethod
+    def list_files(self, path: str, **kwargs) -> List[str]:
+        """
+        Gets a list of files from a given path.
+        """
+        pass
+
     @abstractmethod
     def get_model_weights_urls(self, owner: str, model_name: str, **kwargs) -> List[str]:
         """

model-engine/model_engine_server/domain/use_cases/llm_model_endpoint_use_cases.py

@@ -134,6 +134,25 @@
 DOWNSTREAM_REQUEST_TIMEOUT_SECONDS = 5 * 60  # 5 minutes
 
 
+def _exclude_safetensors_or_bin(model_files: List[str]) -> Optional[str]:
+    """
+    This function is used to determine whether to exclude "*.safetensors" or "*.bin" files
+    based on which file type is present more often in the checkpoint folder. The less
+    frequently present file type is excluded.
+    If both files are equally present, no exclusion string is returned.
+    """
+    exclude_str = None
+    if len([f for f in model_files if f.endswith(".safetensors")]) > len(
+        [f for f in model_files if f.endswith(".bin")]
+    ):
+        exclude_str = "*.bin"
+    elif len([f for f in model_files if f.endswith(".safetensors")]) < len(
+        [f for f in model_files if f.endswith(".bin")]
+    ):
+        exclude_str = "*.safetensors"
+    return exclude_str
+
+
 def _model_endpoint_entity_to_get_llm_model_endpoint_response(
     model_endpoint: ModelEndpoint,
 ) -> GetLLMModelEndpointV1Response:
@@ -182,11 +201,13 @@ def __init__(
         create_model_bundle_use_case: CreateModelBundleV2UseCase,
         model_bundle_repository: ModelBundleRepository,
         model_endpoint_service: ModelEndpointService,
+        llm_artifact_gateway: LLMArtifactGateway,
     ):
         self.authz_module = LiveAuthorizationModule()
         self.create_model_bundle_use_case = create_model_bundle_use_case
         self.model_bundle_repository = model_bundle_repository
         self.model_endpoint_service = model_endpoint_service
+        self.llm_artifact_gateway = llm_artifact_gateway
 
     async def create_model_bundle(
         self,
@@ -358,14 +379,21 @@ def load_model_weights_sub_commands(
                 ]
             )
         else:
-            if framework == LLMInferenceFramework.TEXT_GENERATION_INFERENCE:
+            # Let's check whether to exclude "*.safetensors" or "*.bin" files
+            checkpoint_files = self.llm_artifact_gateway.list_files(checkpoint_path)
+            model_files = [f for f in checkpoint_files if "model" in f]
+
+            exclude_str = _exclude_safetensors_or_bin(model_files)
+
+            if exclude_str is None:
                 subcommands.append(
                     f"{s5cmd} --numworkers 512 cp --concurrency 10 {os.path.join(checkpoint_path, '*')} {final_weights_folder}"
                 )
             else:
                 subcommands.append(
-                    f"{s5cmd} --numworkers 512 cp --concurrency 10 --exclude '*.safetensors'  {os.path.join(checkpoint_path, '*')} {final_weights_folder}"
+                    f"{s5cmd} --numworkers 512 cp --concurrency 10 --exclude '{exclude_str}' {os.path.join(checkpoint_path, '*')} {final_weights_folder}"
                 )
+
         return subcommands
 
     async def create_deepspeed_bundle(

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

@@ -3,6 +3,7 @@
 
 import boto3
 from model_engine_server.common.config import get_model_cache_directory_name, hmi_config
+from model_engine_server.core.utils.url import parse_attachment_url
 from model_engine_server.domain.gateways import LLMArtifactGateway
 
 
@@ -17,6 +18,18 @@ def _get_s3_resource(self, kwargs):
         resource = session.resource("s3")
         return resource
 
+    def list_files(self, path: str, **kwargs) -> List[str]:
+        s3 = self._get_s3_resource(kwargs)
+        parsed_remote = parse_attachment_url(path)
+        bucket = parsed_remote.bucket
+        key = parsed_remote.key
+        # From here: https://dev.to/aws-builders/how-to-list-contents-of-s3-bucket-using-boto3-python-47mm
+        files = [
+            bucket_object["Key"]
+            for bucket_object in s3.list_objects_v2(Bucket=bucket, Prefix=key)["Contents"]
+        ]
+        return files
+
     def get_model_weights_urls(self, owner: str, model_name: str, **kwargs) -> List[str]:
         s3 = self._get_s3_resource(kwargs)
         # parsing prefix to get S3 bucket name

model-engine/tests/unit/conftest.py

@@ -748,11 +748,16 @@ async def initialize_events(self, user_id: str, model_endpoint_name: str):
 class FakeLLMArtifactGateway(LLMArtifactGateway):
     def __init__(self):
         self.existing_models = []
+        self.s3_bucket = {"fake-checkpoint": ["fake.bin, fake2.bin", "fake3.safetensors"]}
         self.urls = {"filename": "https://test-bucket.s3.amazonaws.com/llm/llm-1.0.0.tar.gz"}
 
     def _add_model(self, owner: str, model_name: str):
         self.existing_models.append((owner, model_name))
 
+    def list_files(self, path: str, **kwargs) -> List[str]:
+        if path in self.s3_bucket:
+            return self.s3_bucket[path]
+
     def get_model_weights_urls(self, owner: str, model_name: str):
         if (owner, model_name) in self.existing_models:
             return self.urls

model-engine/tests/unit/domain/test_llm_use_cases.py

@@ -36,6 +36,7 @@
     DeleteLLMEndpointByNameUseCase,
     GetLLMModelEndpointByNameV1UseCase,
     ModelDownloadV1UseCase,
+    _exclude_safetensors_or_bin,
 )
 from model_engine_server.domain.use_cases.model_bundle_use_cases import CreateModelBundleV2UseCase
 
@@ -47,6 +48,7 @@ async def test_create_model_endpoint_use_case_success(
     fake_model_endpoint_service,
     fake_docker_repository_image_always_exists,
     fake_model_primitive_gateway,
+    fake_llm_artifact_gateway,
     create_llm_model_endpoint_request_async: CreateLLMModelEndpointV1Request,
     create_llm_model_endpoint_request_sync: CreateLLMModelEndpointV1Request,
     create_llm_model_endpoint_request_streaming: CreateLLMModelEndpointV1Request,
@@ -62,6 +64,7 @@ async def test_create_model_endpoint_use_case_success(
         create_model_bundle_use_case=bundle_use_case,
         model_bundle_repository=fake_model_bundle_repository,
         model_endpoint_service=fake_model_endpoint_service,
+        llm_artifact_gateway=fake_llm_artifact_gateway,
     )
     user = User(user_id=test_api_key, team_id=test_api_key, is_privileged_user=True)
     response_1 = await use_case.execute(user=user, request=create_llm_model_endpoint_request_async)
@@ -150,6 +153,7 @@ async def test_create_model_endpoint_text_generation_inference_use_case_success(
     fake_model_endpoint_service,
     fake_docker_repository_image_always_exists,
     fake_model_primitive_gateway,
+    fake_llm_artifact_gateway,
     create_llm_model_endpoint_text_generation_inference_request_async: CreateLLMModelEndpointV1Request,
     create_llm_model_endpoint_text_generation_inference_request_streaming: CreateLLMModelEndpointV1Request,
 ):
@@ -163,6 +167,7 @@ async def test_create_model_endpoint_text_generation_inference_use_case_success(
         create_model_bundle_use_case=bundle_use_case,
         model_bundle_repository=fake_model_bundle_repository,
         model_endpoint_service=fake_model_endpoint_service,
+        llm_artifact_gateway=fake_llm_artifact_gateway,
     )
     user = User(user_id=test_api_key, team_id=test_api_key, is_privileged_user=True)
     response_1 = await use_case.execute(
@@ -202,6 +207,7 @@ async def test_create_llm_model_endpoint_use_case_raises_invalid_value_exception
     fake_model_endpoint_service,
     fake_docker_repository_image_always_exists,
     fake_model_primitive_gateway,
+    fake_llm_artifact_gateway,
     create_llm_model_endpoint_request_invalid_model_name: CreateLLMModelEndpointV1Request,
 ):
     fake_model_endpoint_service.model_bundle_repository = fake_model_bundle_repository
@@ -214,6 +220,7 @@ async def test_create_llm_model_endpoint_use_case_raises_invalid_value_exception
         create_model_bundle_use_case=bundle_use_case,
         model_bundle_repository=fake_model_bundle_repository,
         model_endpoint_service=fake_model_endpoint_service,
+        llm_artifact_gateway=fake_llm_artifact_gateway,
     )
     user = User(user_id=test_api_key, team_id=test_api_key, is_privileged_user=True)
     with pytest.raises(ObjectHasInvalidValueException):
@@ -953,3 +960,34 @@ async def test_delete_public_inference_model_raises_not_authorized(
         await use_case.execute(
             user=user, model_endpoint_name=llm_model_endpoint_sync[0].record.name
         )
+
+
+@pytest.mark.asyncio
+async def test_exclude_safetensors_or_bin_majority_bin_returns_exclude_safetensors():
+    fake_model_files = ["fake.bin", "fake2.bin", "fake3.safetensors", "model.json", "optimizer.pt"]
+    assert _exclude_safetensors_or_bin(fake_model_files) == "*.safetensors"
+
+
+@pytest.mark.asyncio
+async def test_exclude_safetensors_or_bin_majority_safetensors_returns_exclude_bin():
+    fake_model_files = [
+        "fake.bin",
+        "fake2.safetensors",
+        "fake3.safetensors",
+        "model.json",
+        "optimizer.pt",
+    ]
+    assert _exclude_safetensors_or_bin(fake_model_files) == "*.bin"
+
+
+@pytest.mark.asyncio
+async def test_exclude_safetensors_or_bin_equal_bins_and_safetensors_returns_none():
+    fake_model_files = [
+        "fake.bin",
+        "fake2.safetensors",
+        "fake3.safetensors",
+        "fake4.bin",
+        "model.json",
+        "optimizer.pt",
+    ]
+    assert _exclude_safetensors_or_bin(fake_model_files) is None