diff --git a/tests/entrypoints/unit_tests/test_remote_vllm_server.py b/tests/entrypoints/unit_tests/test_remote_vllm_server.py index 7f22746c88..c86b07a5cf 100644 --- a/tests/entrypoints/unit_tests/test_remote_vllm_server.py +++ b/tests/entrypoints/unit_tests/test_remote_vllm_server.py @@ -6,7 +6,69 @@ from unittest.mock import Mock import pytest import tests.utils as test_utils -from tests.utils import RemoteOpenAIServer +from tests.utils import RemoteLaunchRenderServer, RemoteOpenAIServer + + +@pytest.mark.parametrize( + ("args", "expected"), + [ + (["--api-key", "key1", "key2"], ["--api-key", "***", "***"]), + (["--api_key", "key1", "key2"], ["--api_key", "***", "***"]), + (["--api-key=key1", "key2"], ["--api-key=***", "***"]), + (["--api_key=key1", "key2"], ["--api_key=***", "***"]), + (["--hf-token", "token"], ["--hf-token", "***"]), + (["--hf_token", "token"], ["--hf_token", "***"]), + (["--hf-token=token"], ["--hf-token=***"]), + (["--hf_token=token"], ["--hf_token=***"]), + ], +) +def test_redact_sensitive_cli_arg_variants(args: list[str], expected: list[str]): + assert test_utils._redact_sensitive_cli_args(args) == expected + + +@pytest.mark.parametrize( + ("server_cls", "command"), + [ + (RemoteOpenAIServer, ["vllm", "serve"]), + (RemoteLaunchRenderServer, ["vllm", "launch", "render"]), + ], +) +def test_server_redacts_sensitive_values( + server_cls, + command: list[str], + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +): + server = object.__new__(server_cls) + popen = Mock() + monkeypatch.setattr(test_utils.subprocess, "Popen", popen) + monkeypatch.setenv("INHERITED_SECRET", "inherited-secret") + serve_args = [ + "--api_key", + "api-secret-1", + "api-secret-2", + "--hf-token=hf-secret", + "--max-num-seqs", + "2", + ] + + server._start_server( + "test-model", serve_args, {"OVERRIDE_SECRET": "override-secret"} + ) + + stdout = capsys.readouterr().out + assert "api-secret-1" not in stdout + assert "api-secret-2" not in stdout + assert "hf-secret" not in stdout + assert "inherited-secret" not in stdout + assert "override-secret" not in stdout + assert "--api_key *** ***" in stdout + assert "--hf-token=***" in stdout + assert "--max-num-seqs 2" in stdout + assert popen.call_args.args[0] == [*command, "test-model", *serve_args] + child_env = popen.call_args.kwargs["env"] + assert child_env["INHERITED_SECRET"] == "inherited-secret" + assert child_env["OVERRIDE_SECRET"] == "override-secret" def test_openai_server_shutdown_wait_covers_engine_cleanup( diff --git a/tests/utils.py b/tests/utils.py index 0dd86a75af..0c2cfeb6a0 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -147,6 +147,28 @@ ROCM_ENGINE_KWARGS: dict = ( _TILELANG_TVM_PYTHONPATH_FRAGMENT = os.path.join( "tilelang", "3rdparty", "tvm", "python" ) +_SENSITIVE_CLI_ARG_NARGS = {"--api-key": "+", "--hf-token": "?"} + + +def _redact_sensitive_cli_args(args: Sequence[str]) -> list[str]: + redacted_args = list(args) + index = 0 + while index < len(args): + name, separator, _ = args[index].partition("=") + nargs = _SENSITIVE_CLI_ARG_NARGS.get(name.replace("_", "-")) + if nargs is None: + index += 1 + continue + if separator: + redacted_args[index] = f"{name}=***" + index += 1 + if not separator or nargs == "+": + while index < len(args) and not args[index].startswith("-"): + redacted_args[index] = "***" + index += 1 + if nargs == "?": + break + return redacted_args def _sanitize_pythonpath_value(pythonpath: str | None) -> str: @@ -784,8 +806,8 @@ class RemoteOpenAIServer(RemoteVLLMServer): env.update(env_dict) _sanitize_pythonpath_env(env) serve_cmd = ["vllm", "serve", model, *vllm_serve_args] - print(f"Launching RemoteOpenAIServer with: {' '.join(serve_cmd)}") - print(f"Environment variables: {env}") + redacted_serve_cmd = _redact_sensitive_cli_args(serve_cmd) + print(f"Launching RemoteOpenAIServer with: {' '.join(redacted_serve_cmd)}") self.proc: subprocess.Popen = subprocess.Popen( serve_cmd, env=env, @@ -812,7 +834,10 @@ class RemoteLaunchRenderServer(RemoteVLLMServer): env.update(env_dict) _sanitize_pythonpath_env(env) serve_cmd = ["vllm", "launch", "render", model, *vllm_serve_args] - print(f"Launching RemoteLaunchRenderServer with: {' '.join(serve_cmd)}") + redacted_serve_cmd = _redact_sensitive_cli_args(serve_cmd) + print( + f"Launching RemoteLaunchRenderServer with: {' '.join(redacted_serve_cmd)}" + ) self.proc: subprocess.Popen = subprocess.Popen( serve_cmd, env=env,