Update test fixtures

Author: dmchoiboi

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

@@ -2331,30 +2331,32 @@ async def create_batch_job_bundle(
         return batch_bundle
 
     async def execute(
-        self, user: User, _request: CreateBatchCompletionsRequest
+        self, user: User, request: CreateBatchCompletionsRequest
     ) -> CreateBatchCompletionsResponse:
-        hardware = infer_hardware_from_model_name(_request.model_config.model)
+        hardware = infer_hardware_from_model_name(request.model_config.model)
         # Reconcile gpus count with num_shards from request
         assert hardware.gpus is not None
-        if _request.model_config.num_shards:
-            hardware.gpus = max(hardware.gpus, _request.model_config.num_shards)
+        if request.model_config.num_shards:
+            hardware.gpus = max(hardware.gpus, request.model_config.num_shards)
 
-        request = CreateBatchCompletionsEngineRequest.from_api(_request)
-        request.model_config.num_shards = hardware.gpus
+        engine_request = CreateBatchCompletionsEngineRequest.from_api(request)
+        engine_request.model_config.num_shards = hardware.gpus
 
-        if request.tool_config and request.tool_config.name != "code_evaluator":
+        if engine_request.tool_config and engine_request.tool_config.name != "code_evaluator":
             raise ObjectHasInvalidValueException(
                 "Only code_evaluator tool is supported for batch completions."
             )
 
         additional_engine_args = infer_addition_engine_args_from_model_name(
-            request.model_config.model
+            engine_request.model_config.model
         )
 
         if additional_engine_args.gpu_memory_utilization is not None:
-            request.max_gpu_memory_utilization = additional_engine_args.gpu_memory_utilization
+            engine_request.max_gpu_memory_utilization = (
+                additional_engine_args.gpu_memory_utilization
+            )
 
-        batch_bundle = await self.create_batch_job_bundle(user, request, hardware)
+        batch_bundle = await self.create_batch_job_bundle(user, engine_request, hardware)
 
         validate_resource_requests(
             bundle=batch_bundle,
@@ -2365,21 +2367,21 @@ async def execute(
             gpu_type=hardware.gpu_type,
         )
 
-        if request.max_runtime_sec is None or request.max_runtime_sec < 1:
+        if engine_request.max_runtime_sec is None or engine_request.max_runtime_sec < 1:
             raise ObjectHasInvalidValueException("max_runtime_sec must be a positive integer.")
 
         job_id = await self.docker_image_batch_job_gateway.create_docker_image_batch_job(
             created_by=user.user_id,
             owner=user.team_id,
-            job_config=request.dict(),
+            job_config=engine_request.dict(),
             env=batch_bundle.env,
             command=batch_bundle.command,
             repo=batch_bundle.image_repository,
             tag=batch_bundle.image_tag,
             resource_requests=hardware,
-            labels=request.model_config.labels,
+            labels=engine_request.model_config.labels,
             mount_location=batch_bundle.mount_location,
-            override_job_max_runtime_s=request.max_runtime_sec,
-            num_workers=request.data_parallelism,
+            override_job_max_runtime_s=engine_request.max_runtime_sec,
+            num_workers=engine_request.data_parallelism,
         )
         return CreateBatchCompletionsResponse(job_id=job_id)

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

@@ -3,6 +3,7 @@
 import pytest
 from model_engine_server.common.dtos.llms import (
     CompletionOutput,
+    CreateBatchCompletionsEngineRequest,
     CreateBatchCompletionsModelConfig,
     CreateBatchCompletionsRequest,
     CreateBatchCompletionsRequestContent,
@@ -12,14 +13,20 @@
 
 
 @pytest.fixture
-def create_batch_completions_request():
-    return CreateBatchCompletionsRequest(
+def create_batch_completions_engine_request() -> CreateBatchCompletionsEngineRequest:
+    return CreateBatchCompletionsEngineRequest(
+        input_data_path="input_data_path",
+        output_data_path="output_data_path",
         model_config=CreateBatchCompletionsModelConfig(
-            checkpoint_path="checkpoint_path", model="model", num_shards=4, seed=123, labels={}
+            model="model",
+            checkpoint_path="checkpoint_path",
+            labels={},
+            seed=123,
+            num_shards=4,
         ),
         data_parallelism=1,
-        input_data_path="input_data_path",
-        output_data_path="output_data_path",
+        max_runtime_sec=86400,
+        max_gpu_memory_utilization=0.95,
     )
 
 

model-engine/tests/unit/inference/test_vllm_batch.py

@@ -7,7 +7,9 @@
 
 @pytest.mark.asyncio
 @patch("model_engine_server.inference.batch_inference.vllm_batch.get_vllm_engine")
-@patch("model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequest")
+@patch(
+    "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsEngineRequest"
+)
 @patch(
     "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequestContent"
 )
@@ -25,9 +27,9 @@ async def test_batch_inference(
     mock_get_s3_client,
     mock_generate_with_vllm,
     mock_create_batch_completions_request_content,
-    mock_create_batch_completions_request,
+    mock_create_batch_completions_engine_request,
     mock_vllm,
-    create_batch_completions_request,
+    create_batch_completions_engine_request,
     create_batch_completions_request_content,
     mock_s3_client,
     mock_process,
@@ -36,7 +38,9 @@ async def test_batch_inference(
     # Mock the necessary objects and data
     mock_popen.return_value = mock_process
     mock_get_s3_client.return_value = mock_s3_client
-    mock_create_batch_completions_request.parse_file.return_value = create_batch_completions_request
+    mock_create_batch_completions_engine_request.parse_file.return_value = (
+        create_batch_completions_engine_request
+    )
     mock_create_batch_completions_request_content.parse_raw.return_value = (
         create_batch_completions_request_content
     )
@@ -48,7 +52,7 @@ async def test_batch_inference(
     await batch_inference()
 
     # Assertions
-    mock_create_batch_completions_request.parse_file.assert_called_once()
+    mock_create_batch_completions_engine_request.parse_file.assert_called_once()
     mock_open_func.assert_has_calls(
         [
             call("input_data_path", "r"),
@@ -61,7 +65,9 @@ async def test_batch_inference(
 
 @pytest.mark.asyncio
 @patch("model_engine_server.inference.batch_inference.vllm_batch.get_vllm_engine")
-@patch("model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequest")
+@patch(
+    "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsEngineRequest"
+)
 @patch(
     "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequestContent"
 )
@@ -79,9 +85,9 @@ async def test_batch_inference_failed_to_download_model_but_proceed(
     mock_get_s3_client,
     mock_generate_with_vllm,
     mock_create_batch_completions_request_content,
-    mock_create_batch_completions_request,
+    mock_create_batch_completions_engine_request,
     mock_vllm,
-    create_batch_completions_request,
+    create_batch_completions_engine_request,
     create_batch_completions_request_content,
     mock_s3_client,
     mock_process,
@@ -91,7 +97,9 @@ async def test_batch_inference_failed_to_download_model_but_proceed(
     mock_process.returncode = 1  # Failed to download model
     mock_popen.return_value = mock_process
     mock_get_s3_client.return_value = mock_s3_client
-    mock_create_batch_completions_request.parse_file.return_value = create_batch_completions_request
+    mock_create_batch_completions_engine_request.parse_file.return_value = (
+        create_batch_completions_engine_request
+    )
     mock_create_batch_completions_request_content.parse_raw.return_value = (
         create_batch_completions_request_content
     )
@@ -103,7 +111,7 @@ async def test_batch_inference_failed_to_download_model_but_proceed(
     await batch_inference()
 
     # Assertions
-    mock_create_batch_completions_request.parse_file.assert_called_once()
+    mock_create_batch_completions_engine_request.parse_file.assert_called_once()
     mock_open_func.assert_has_calls(
         [
             call("input_data_path", "r"),
@@ -116,7 +124,9 @@ async def test_batch_inference_failed_to_download_model_but_proceed(
 
 @pytest.mark.asyncio
 @patch("model_engine_server.inference.batch_inference.vllm_batch.get_vllm_engine")
-@patch("model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequest")
+@patch(
+    "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsEngineRequest"
+)
 @patch(
     "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequestContent"
 )
@@ -136,9 +146,9 @@ async def test_batch_inference_two_workers(
     mock_get_s3_client,
     mock_generate_with_vllm,
     mock_create_batch_completions_request_content,
-    mock_create_batch_completions_request,
+    mock_create_batch_completions_engine_request,
     mock_vllm,
-    create_batch_completions_request,
+    create_batch_completions_engine_request,
     create_batch_completions_request_content,
     mock_s3_client,
     mock_process,
@@ -147,8 +157,10 @@ async def test_batch_inference_two_workers(
     # Mock the necessary objects and data
     mock_popen.return_value = mock_process
     mock_get_s3_client.return_value = mock_s3_client
-    create_batch_completions_request.data_parallelism = 2
-    mock_create_batch_completions_request.parse_file.return_value = create_batch_completions_request
+    create_batch_completions_engine_request.data_parallelism = 2
+    mock_create_batch_completions_engine_request.parse_file.return_value = (
+        create_batch_completions_engine_request
+    )
     mock_create_batch_completions_request_content.parse_raw.return_value = (
         create_batch_completions_request_content
     )
@@ -168,7 +180,7 @@ def side_effect(key, default):
     await batch_inference()
 
     # Assertions
-    mock_create_batch_completions_request.parse_file.assert_called_once()
+    mock_create_batch_completions_engine_request.parse_file.assert_called_once()
     mock_open_func.assert_has_calls(
         [
             call("input_data_path", "r"),
@@ -198,7 +210,9 @@ def side_effect(key, default):
 
 @pytest.mark.asyncio
 @patch("model_engine_server.inference.batch_inference.vllm_batch.get_vllm_engine")
-@patch("model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequest")
+@patch(
+    "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsEngineRequest"
+)
 @patch(
     "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequestContent"
 )
@@ -218,9 +232,9 @@ async def test_batch_inference_delete_chunks(
     mock_get_s3_client,
     mock_generate_with_vllm,
     mock_create_batch_completions_request_content,
-    mock_create_batch_completions_request,
+    mock_create_batch_completions_engine_request,
     mock_vllm,
-    create_batch_completions_request,
+    create_batch_completions_engine_request,
     create_batch_completions_request_content,
     mock_s3_client,
     mock_process,
@@ -229,9 +243,11 @@ async def test_batch_inference_delete_chunks(
     # Mock the necessary objects and data
     mock_popen.return_value = mock_process
     mock_get_s3_client.return_value = mock_s3_client
-    create_batch_completions_request.data_parallelism = 2
-    create_batch_completions_request.output_data_path = "s3://bucket/key"
-    mock_create_batch_completions_request.parse_file.return_value = create_batch_completions_request
+    create_batch_completions_engine_request.data_parallelism = 2
+    create_batch_completions_engine_request.output_data_path = "s3://bucket/key"
+    mock_create_batch_completions_engine_request.parse_file.return_value = (
+        create_batch_completions_engine_request
+    )
     mock_create_batch_completions_request_content.parse_raw.return_value = (
         create_batch_completions_request_content
     )
@@ -251,7 +267,7 @@ def side_effect(key, default):
     await batch_inference()
 
     # Assertions
-    mock_create_batch_completions_request.parse_file.assert_called_once()
+    mock_create_batch_completions_engine_request.parse_file.assert_called_once()
     mock_open_func.assert_has_calls(
         [
             call("input_data_path", "r"),
@@ -310,7 +326,9 @@ def test_file_exists_no_such_key():
 
 @pytest.mark.asyncio
 @patch("model_engine_server.inference.batch_inference.vllm_batch.get_vllm_engine")
-@patch("model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequest")
+@patch(
+    "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsEngineRequest"
+)
 @patch(
     "model_engine_server.inference.batch_inference.vllm_batch.CreateBatchCompletionsRequestContent"
 )
@@ -330,7 +348,7 @@ async def test_batch_inference_tool_completion(
     mock_get_s3_client,
     mock_generate_with_vllm,
     mock_create_batch_completions_request_content,
-    mock_create_batch_completions_request,
+    mock_create_batch_completions_engine_request,
     mock_vllm,
     create_batch_completions_tool_completion_request,
     create_batch_completions_tool_completion_request_content,
@@ -344,7 +362,7 @@ async def test_batch_inference_tool_completion(
     mock_run.return_value = mock_run_output
     mock_popen.return_value = mock_process
     mock_get_s3_client.return_value = mock_s3_client
-    mock_create_batch_completions_request.parse_file.return_value = (
+    mock_create_batch_completions_engine_request.parse_file.return_value = (
         create_batch_completions_tool_completion_request
     )
     mock_create_batch_completions_request_content.parse_raw.return_value = (
@@ -361,7 +379,7 @@ async def test_batch_inference_tool_completion(
     await batch_inference()
 
     # Assertions
-    mock_create_batch_completions_request.parse_file.assert_called_once()
+    mock_create_batch_completions_engine_request.parse_file.assert_called_once()
     mock_open_func.assert_has_calls(
         [
             call("input_data_path", "r"),