More rigorous endpoint update handling (#558)

* Fix metadata update

* Update tests

Author: dmchoiboi

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

@@ -111,6 +111,8 @@
 
 logger = make_logger(logger_name())
 
+LLM_METADATA_KEY = "_llm"
+RESERVED_METADATA_KEYS = [LLM_METADATA_KEY, CONVERTED_FROM_ARTIFACT_LIKE_KEY]
 
 INFERENCE_FRAMEWORK_REPOSITORY: Dict[LLMInferenceFramework, str] = {
     LLMInferenceFramework.DEEPSPEED: "instant-llm",
@@ -279,11 +281,14 @@ async def _get_recommended_hardware_config_map() -> Dict[str, Any]:
 def _model_endpoint_entity_to_get_llm_model_endpoint_response(
     model_endpoint: ModelEndpoint,
 ) -> GetLLMModelEndpointV1Response:
-    if model_endpoint.record.metadata is None or "_llm" not in model_endpoint.record.metadata:
+    if (
+        model_endpoint.record.metadata is None
+        or LLM_METADATA_KEY not in model_endpoint.record.metadata
+    ):
         raise ObjectHasInvalidValueException(
             f"Can't translate model entity to response, endpoint {model_endpoint.record.id} does not have LLM metadata."
         )
-    llm_metadata = model_endpoint.record.metadata.get("_llm", {})
+    llm_metadata = model_endpoint.record.metadata.get(LLM_METADATA_KEY, {})
     response = GetLLMModelEndpointV1Response(
         id=model_endpoint.record.id,
         name=model_endpoint.record.name,
@@ -962,7 +967,7 @@ async def execute(
         aws_role = self.authz_module.get_aws_role_for_user(user)
         results_s3_bucket = self.authz_module.get_s3_bucket_for_user(user)
 
-        request.metadata["_llm"] = asdict(
+        request.metadata[LLM_METADATA_KEY] = asdict(
             LLMMetadata(
                 model_name=request.model_name,
                 source=request.source,
@@ -1088,6 +1093,16 @@ async def execute(self, user: User, model_endpoint_name: str) -> GetLLMModelEndp
         return _model_endpoint_entity_to_get_llm_model_endpoint_response(model_endpoint)
 
 
+def merge_metadata(
+    request: Optional[Dict[str, Any]], record: Optional[Dict[str, Any]]
+) -> Optional[Dict[str, Any]]:
+    if request is None:
+        return record
+    if record is None:
+        return request
+    return {**record, **request}
+
+
 class UpdateLLMModelEndpointV1UseCase:
     def __init__(
         self,
@@ -1131,6 +1146,7 @@ async def execute(
             raise EndpointInfraStateNotFound(error_msg)
 
         infra_state = model_endpoint.infra_state
+        metadata: Optional[Dict[str, Any]]
 
         if (
             request.model_name
@@ -1140,7 +1156,7 @@ async def execute(
             or request.quantize
             or request.checkpoint_path
         ):
-            llm_metadata = (model_endpoint.record.metadata or {}).get("_llm", {})
+            llm_metadata = (model_endpoint.record.metadata or {}).get(LLM_METADATA_KEY, {})
             inference_framework = llm_metadata["inference_framework"]
 
             if request.inference_framework_image_tag == "latest":
@@ -1177,7 +1193,7 @@ async def execute(
             )
 
             metadata = endpoint_record.metadata or {}
-            metadata["_llm"] = asdict(
+            metadata[LLM_METADATA_KEY] = asdict(
                 LLMMetadata(
                     model_name=model_name,
                     source=source,
@@ -1188,7 +1204,7 @@ async def execute(
                     checkpoint_path=checkpoint_path,
                 )
             )
-            request.metadata = metadata
+            endpoint_record.metadata = metadata
 
         # For resources that are not specified in the update endpoint request, pass in resource from
         # infra_state to make sure that after the update, all resources are valid and in sync.
@@ -1209,15 +1225,20 @@ async def execute(
             endpoint_type=endpoint_record.endpoint_type,
         )
 
-        if request.metadata is not None and CONVERTED_FROM_ARTIFACT_LIKE_KEY in request.metadata:
-            raise ObjectHasInvalidValueException(
-                f"{CONVERTED_FROM_ARTIFACT_LIKE_KEY} is a reserved metadata key and cannot be used by user."
-            )
+        if request.metadata is not None:
+            # If reserved metadata key is provided, throw ObjectHasInvalidValueException
+            for key in RESERVED_METADATA_KEYS:
+                if key in request.metadata:
+                    raise ObjectHasInvalidValueException(
+                        f"{key} is a reserved metadata key and cannot be used by user."
+                    )
+
+        metadata = merge_metadata(request.metadata, endpoint_record.metadata)
 
         updated_endpoint_record = await self.model_endpoint_service.update_model_endpoint(
             model_endpoint_id=model_endpoint_id,
             model_bundle_id=bundle.id,
-            metadata=request.metadata,
+            metadata=metadata,
             post_inference_hooks=request.post_inference_hooks,
             cpus=request.cpus,
             gpus=request.gpus,

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

@@ -31,6 +31,9 @@
     Quantization,
     StreamingEnhancedRunnableImageFlavor,
 )
+from model_engine_server.domain.use_cases.model_endpoint_use_cases import (
+    CONVERTED_FROM_ARTIFACT_LIKE_KEY,
+)
 
 
 @pytest.fixture
@@ -265,6 +268,19 @@ def update_llm_model_endpoint_request() -> UpdateLLMModelEndpointV1Request:
     )
 
 
+@pytest.fixture
+def update_llm_model_endpoint_request_only_workers() -> UpdateLLMModelEndpointV1Request:
+    return UpdateLLMModelEndpointV1Request(
+        min_workers=5,
+        max_workers=10,
+    )
+
+
+@pytest.fixture
+def update_llm_model_endpoint_request_bad_metadata() -> UpdateLLMModelEndpointV1Request:
+    return UpdateLLMModelEndpointV1Request(metadata={CONVERTED_FROM_ARTIFACT_LIKE_KEY: {}})
+
+
 @pytest.fixture
 def create_llm_model_endpoint_request_llama_2() -> CreateLLMModelEndpointV1Request:
     return CreateLLMModelEndpointV1Request(

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

@@ -51,6 +51,7 @@
     UpdateLLMModelEndpointV1UseCase,
     _fill_hardware_info,
     _infer_hardware,
+    merge_metadata,
     validate_and_update_completion_params,
     validate_checkpoint_files,
 )
@@ -614,6 +615,7 @@ async def test_update_model_endpoint_use_case_success(
     fake_llm_model_endpoint_service,
     create_llm_model_endpoint_request_streaming: CreateLLMModelEndpointV1Request,
     update_llm_model_endpoint_request: UpdateLLMModelEndpointV1Request,
+    update_llm_model_endpoint_request_only_workers: UpdateLLMModelEndpointV1Request,
 ):
     fake_model_endpoint_service.model_bundle_repository = fake_model_bundle_repository
     bundle_use_case = CreateModelBundleV2UseCase(
@@ -687,6 +689,102 @@ async def test_update_model_endpoint_use_case_success(
         == update_llm_model_endpoint_request.max_workers
     )
 
+    update_response2 = await update_use_case.execute(
+        user=user,
+        model_endpoint_name=create_llm_model_endpoint_request_streaming.name,
+        request=update_llm_model_endpoint_request_only_workers,
+    )
+    assert update_response2.endpoint_creation_task_id
+
+    endpoint = (
+        await fake_model_endpoint_service.list_model_endpoints(
+            owner=None,
+            name=create_llm_model_endpoint_request_streaming.name,
+            order_by=None,
+        )
+    )[0]
+    assert endpoint.record.metadata == {
+        "_llm": {
+            "model_name": create_llm_model_endpoint_request_streaming.model_name,
+            "source": create_llm_model_endpoint_request_streaming.source,
+            "inference_framework": create_llm_model_endpoint_request_streaming.inference_framework,
+            "inference_framework_image_tag": "fake_docker_repository_latest_image_tag",
+            "num_shards": create_llm_model_endpoint_request_streaming.num_shards,
+            "quantize": None,
+            "checkpoint_path": update_llm_model_endpoint_request.checkpoint_path,
+        }
+    }
+    assert endpoint.infra_state.resource_state.memory == update_llm_model_endpoint_request.memory
+    assert (
+        endpoint.infra_state.deployment_state.min_workers
+        == update_llm_model_endpoint_request_only_workers.min_workers
+    )
+    assert (
+        endpoint.infra_state.deployment_state.max_workers
+        == update_llm_model_endpoint_request_only_workers.max_workers
+    )
+
+
+@pytest.mark.asyncio
+@mock.patch(
+    "model_engine_server.domain.use_cases.llm_model_endpoint_use_cases._get_latest_tag",
+    mocked__get_latest_tag(),
+)
+async def test_update_model_endpoint_use_case_failure(
+    test_api_key: str,
+    fake_model_bundle_repository,
+    fake_model_endpoint_service,
+    fake_docker_repository_image_always_exists,
+    fake_model_primitive_gateway,
+    fake_llm_artifact_gateway,
+    fake_llm_model_endpoint_service,
+    create_llm_model_endpoint_request_streaming: CreateLLMModelEndpointV1Request,
+    update_llm_model_endpoint_request_bad_metadata: UpdateLLMModelEndpointV1Request,
+):
+    fake_model_endpoint_service.model_bundle_repository = fake_model_bundle_repository
+    bundle_use_case = CreateModelBundleV2UseCase(
+        model_bundle_repository=fake_model_bundle_repository,
+        docker_repository=fake_docker_repository_image_always_exists,
+        model_primitive_gateway=fake_model_primitive_gateway,
+    )
+    llm_bundle_use_case = CreateLLMModelBundleV1UseCase(
+        create_model_bundle_use_case=bundle_use_case,
+        model_bundle_repository=fake_model_bundle_repository,
+        llm_artifact_gateway=fake_llm_artifact_gateway,
+        docker_repository=fake_docker_repository_image_always_exists,
+    )
+    create_use_case = CreateLLMModelEndpointV1UseCase(
+        create_llm_model_bundle_use_case=llm_bundle_use_case,
+        model_endpoint_service=fake_model_endpoint_service,
+        docker_repository=fake_docker_repository_image_always_exists,
+        llm_artifact_gateway=fake_llm_artifact_gateway,
+    )
+    update_use_case = UpdateLLMModelEndpointV1UseCase(
+        create_llm_model_bundle_use_case=llm_bundle_use_case,
+        model_endpoint_service=fake_model_endpoint_service,
+        llm_model_endpoint_service=fake_llm_model_endpoint_service,
+        docker_repository=fake_docker_repository_image_always_exists,
+    )
+
+    user = User(user_id=test_api_key, team_id=test_api_key, is_privileged_user=True)
+
+    await create_use_case.execute(user=user, request=create_llm_model_endpoint_request_streaming)
+    endpoint = (
+        await fake_model_endpoint_service.list_model_endpoints(
+            owner=None,
+            name=create_llm_model_endpoint_request_streaming.name,
+            order_by=None,
+        )
+    )[0]
+    fake_llm_model_endpoint_service.add_model_endpoint(endpoint)
+
+    with pytest.raises(ObjectHasInvalidValueException):
+        await update_use_case.execute(
+            user=user,
+            model_endpoint_name=create_llm_model_endpoint_request_streaming.name,
+            request=update_llm_model_endpoint_request_bad_metadata,
+        )
+
 
 def mocked_auto_tokenizer_from_pretrained(*args, **kwargs):  # noqa
     class mocked_encode:
@@ -2241,3 +2339,23 @@ async def test_create_batch_completions(
         "-c",
         "ddtrace-run python vllm_batch.py",
     ]
+
+
+def test_merge_metadata():
+    request_metadata = {
+        "key1": "value1",
+        "key2": "value2",
+    }
+
+    endpoint_metadata = {
+        "key1": "value0",
+        "key3": "value3",
+    }
+
+    assert merge_metadata(request_metadata, None) == request_metadata
+    assert merge_metadata(None, endpoint_metadata) == endpoint_metadata
+    assert merge_metadata(request_metadata, endpoint_metadata) == {
+        "key1": "value1",
+        "key2": "value2",
+        "key3": "value3",
+    }