enable sensitive log mode (#415)

Enable sensitive log mode

Author: song-william

charts/model-engine/values_circleci.yaml

@@ -146,6 +146,7 @@ config:
       s3_file_llm_fine_tune_repository: "s3://$CIRCLECI_AWS_S3_BUCKET/fine_tune_repository"
       dd_trace_enabled: false
       istio_enabled: true
+      sensitive_log_mode: false
       tgi_repository: "text-generation-inference"
       vllm_repository: "vllm"
       lightllm_repository: "lightllm"

charts/model-engine/values_sample.yaml

@@ -163,6 +163,7 @@ config:
       # dd_trace_enabled specifies whether to enable datadog tracing, datadog must be installed in the cluster
       dd_trace_enabled: false
       istio_enabled: true
+      sensitive_log_mode: false
 
       # Asynchronous endpoints configs (coming soon)
       sqs_profile: default

model-engine/model_engine_server/api/llms_v1.py

@@ -12,6 +12,7 @@
     get_external_interfaces_read_only,
     verify_authentication,
 )
+from model_engine_server.common.config import hmi_config
 from model_engine_server.common.dtos.llms import (
     CancelFineTuneResponse,
     CompletionStreamV1Request,
@@ -307,9 +308,10 @@ async def create_completion_sync_task(
     """
     Runs a sync prompt completion on an LLM.
     """
-    logger.info(
-        f"POST /completion_sync with {request} to endpoint {model_endpoint_name} for {auth}"
-    )
+    if not hmi_config.sensitive_log_mode:
+        logger.info(
+            f"POST /completion_sync with {request} to endpoint {model_endpoint_name} for {auth}"
+        )
     try:
         use_case = CompletionSyncV1UseCase(
             model_endpoint_service=external_interfaces.model_endpoint_service,
@@ -369,9 +371,10 @@ async def create_completion_stream_task(
     """
     Runs a stream prompt completion on an LLM.
     """
-    logger.info(
-        f"POST /completion_stream with {request} to endpoint {model_endpoint_name} for {auth}"
-    )
+    if not hmi_config.sensitive_log_mode:  # pragma: no cover
+        logger.info(
+            f"POST /completion_stream with {request} to endpoint {model_endpoint_name} for {auth}"
+        )
     use_case = CompletionStreamV1UseCase(
         model_endpoint_service=external_interfaces.model_endpoint_service,
         llm_model_endpoint_service=external_interfaces.llm_model_endpoint_service,

model-engine/model_engine_server/common/config.py

@@ -62,6 +62,7 @@ class HostedModelInferenceServiceConfig:
     user_inference_pytorch_repository: str
     user_inference_tensorflow_repository: str
     docker_image_layer_cache_repository: str
+    sensitive_log_mode: bool
 
     @classmethod
     def from_yaml(cls, yaml_path):

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

@@ -682,6 +682,9 @@ async def create_vllm_bundle(
             else:
                 raise InvalidRequestException(f"Quantization {quantize} is not supported by vLLM.")
 
+        if hmi_config.sensitive_log_mode:  # pragma: no cover
+            subcommands[-1] = subcommands[-1] + " --disable-log-requests"
+
         command = [
             "/bin/bash",
             "-c",

model-engine/service_configs/service_config_circleci.yaml

@@ -54,6 +54,7 @@ s3_file_llm_fine_tune_repository: "s3://model-engine-integration-tests/fine_tune
 
 dd_trace_enabled: false
 istio_enabled: true
+sensitive_log_mode: false
 tgi_repository: "text-generation-inference"
 vllm_repository: "vllm"
 lightllm_repository: "lightllm"

model-engine/tests/unit/api/test_llms.py

@@ -1,6 +1,6 @@
 import json
-import re
 from typing import Any, Dict, Tuple
+from unittest import mock
 
 import pytest
 from model_engine_server.common.dtos.llms import GetLLMModelEndpointV1Response
@@ -156,12 +156,14 @@ def test_completion_sync_endpoint_not_found_returns_404(
     assert response_1.status_code == 404
 
 
+# When enabling this test, other tests fail with "RunTumeError got Future <Future pending> attached to a different loop"
+# https://github.com/encode/starlette/issues/1315#issuecomment-980784457
 @pytest.mark.skip(reason="Need to figure out FastAPI test client asyncio funkiness")
 def test_completion_stream_success(
     llm_model_endpoint_streaming: ModelEndpoint,
     completion_stream_request: Dict[str, Any],
     get_test_client_wrapper,
-):
+):  # pragma: no cover
     client = get_test_client_wrapper(
         fake_docker_repository_image_always_exists=True,
         fake_model_bundle_repository_contents={},
@@ -175,19 +177,28 @@ def test_completion_stream_success(
         fake_batch_job_progress_gateway_contents={},
         fake_docker_image_batch_job_bundle_repository_contents={},
     )
-    response_1 = client.post(
-        f"/v1/llm/completions-stream?model_endpoint_name={llm_model_endpoint_streaming.record.name}",
-        auth=("no_user", ""),
-        json=completion_stream_request,
-        stream=True,
-    )
+    with mock.patch(
+        "model_engine_server.domain.use_cases.llm_model_endpoint_use_cases.count_tokens",
+        return_value=5,
+    ):
+        response_1 = client.post(
+            f"/v1/llm/completions-stream?model_endpoint_name={llm_model_endpoint_streaming.record.name}",
+            auth=("no_user", ""),
+            json=completion_stream_request,
+            stream=True,
+        )
     assert response_1.status_code == 200
     count = 0
     for message in response_1:
-        assert re.fullmatch(
-            'data: {"request_id"}: ".*", "output": null}\r\n\r\n',
-            message.decode("utf-8"),
-        )
+        decoded_message = message.decode("utf-8")
+        assert decoded_message.startswith("data: "), "SSE does not start with 'data: '"
+
+        # strip 'data: ' prefix from  Server-sent events format
+        json_str = decoded_message[len("data: ") :]
+        parsed_data = json.loads(json_str.strip())
+        assert parsed_data["request_id"] is not None
+        assert parsed_data["output"] is None
+        assert parsed_data["error"] is None
         count += 1
     assert count == 1