diff --git a/.coderabbit.yaml b/.coderabbit.yaml index 056b573734..f8025b9ab3 100644 --- a/.coderabbit.yaml +++ b/.coderabbit.yaml @@ -138,6 +138,20 @@ reviews: outputs, owns family validation policy, or becomes a dependency of core or family implementation. + - path: "server/**" + instructions: | + Treat the local server as an optional application over public runtime + contracts. Preserve the one-way server-to-library dependency, fixed + bounded execution lanes, sanitized public errors, and source-build and + installed-wheel parity. Flag model-specific policy, reverse dependencies + from core or families, hidden queues, or claims of distributed scaling, + automatic placement, or worker self-healing that are not implemented. + + - path: "server/tests/**" + instructions: | + Do not suggest weakening assertions, lifecycle and security boundaries, + expected values, or acceptance criteria merely to make tests pass. + - path: "apps/benchmark/**" instructions: | Check semantic equivalence of timed regions, synchronization, warmup, diff --git a/CMakeLists.txt b/CMakeLists.txt index 1098879680..9619cfb3b8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -263,6 +263,8 @@ foreach(_trtmc_runtime_cmake IN LISTS _trtmc_family_runtime_cmake) ) endforeach() +add_subdirectory(server) + add_library(trtmc_cli STATIC apps/cli/cli.cpp apps/cli/io.cpp @@ -288,7 +290,8 @@ set_source_files_properties(apps/cli/io.cpp PROPERTIES add_executable(trtmc apps/cli/main.cpp) target_include_directories(trtmc PRIVATE ${PROJECT_SOURCE_DIR}/apps) -target_link_libraries(trtmc PRIVATE trtmc_cli) +target_link_libraries(trtmc PRIVATE trtmc_cli trtmc_server_native) +add_dependencies(trtmc trtmc_server_python) target_compile_options(trtmc PRIVATE -Wall -Wextra -Wpedantic) set_target_properties(trtmc PROPERTIES BUILD_RPATH "\$ORIGIN" diff --git a/Dockerfile b/Dockerfile index eabbfd4af0..e4eae01bc8 100644 --- a/Dockerfile +++ b/Dockerfile @@ -78,6 +78,16 @@ RUN python3.12 -m venv "$VIRTUAL_ENV" \ --index-url https://download.pytorch.org/whl/cu130 \ && pip install "setuptools>=80,<82" +# Server control-plane and CPU contract-test dependencies. Keep these bounds +# aligned with the serve and test extras in pyproject.toml. +RUN pip install --target /opt/trtmc-server-test-deps \ + "fastapi>=0.115,<0.142" \ + "httpx>=0.27,<0.29" \ + "pydantic>=2.11,<3" \ + "python-multipart>=0.0.9,<1" \ + "uvicorn>=0.30,<0.53" \ + "websockets>=13,<17" + ENV TRT_LIB_DIR=/opt/venv/lib/python3.12/site-packages/tensorrt_libs ENV NCCL_LIB_DIR=/opt/venv/lib/python3.12/site-packages/nvidia/nccl/lib ENV TVM_FFI_LIB_DIR=/opt/venv/lib/python3.12/site-packages/tvm_ffi/lib diff --git a/apps/cli/cli.cpp b/apps/cli/cli.cpp index 6aee9dd1df..8b70b642ad 100644 --- a/apps/cli/cli.cpp +++ b/apps/cli/cli.cpp @@ -1280,6 +1280,7 @@ void print_usage(std::ostream& output) { output << "Usage:\n" " trtmc version\n" " trtmc inspect BUNDLE\n" + " trtmc serve --runtime-root DIR [SERVER OPTIONS]\n" " trtmc COMMAND BUNDLE --runtime-root DIR [OPTIONS]\n\n" "Execution commands:\n" " run, encode, embed, rerank, classify, detect, extract-features,\n" diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index dbb5c93bb8..082bfa4e39 100644 --- a/apps/cli/main.cpp +++ b/apps/cli/main.cpp @@ -4,9 +4,15 @@ */ #include "cli/cli.h" +#include "native/entrypoint.h" #include +#include int main(int argc, char** argv) { + if (argc >= 2 && std::string(argv[1]) == "serve") + return trtmc::server::run_server_frontend(argc - 2, argv + 2); + if (argc >= 2 && std::string(argv[1]) == "_serve-worker") + return trtmc::server::run_native_worker(argc - 2, argv + 2); return trtmc::cli::run(argc, argv, std::cout, std::cerr); } diff --git a/apps/cli/tests/test_cli.cpp b/apps/cli/tests/test_cli.cpp index 2023b7ce0d..e3f15cd522 100644 --- a/apps/cli/tests/test_cli.cpp +++ b/apps/cli/tests/test_cli.cpp @@ -732,7 +732,8 @@ int main() { std::ostringstream usage; trtmc::cli::print_usage(usage); - check(usage.str().find("--source-language-token-id") != std::string::npos && + check(usage.str().find("trtmc serve --runtime-root DIR") != std::string::npos && + usage.str().find("--source-language-token-id") != std::string::npos && usage.str().find("--segment-overlap-seconds") != std::string::npos && usage.str().find("--runtime-cache") != std::string::npos && usage.str().find("--cuda-graphs") != std::string::npos, diff --git a/pyproject.toml b/pyproject.toml index c7435458d9..5439e9c421 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,24 @@ dependencies = [ trtmc-bench = "trtmc_benchmark.cli:main" [project.optional-dependencies] -test = ["pytest>=7.0", "torch>=2.0", "jsonschema>=4.23,<5"] +serve = [ + "fastapi>=0.115,<0.142", + "pydantic>=2.11,<3", + "python-multipart>=0.0.9,<1", + "uvicorn>=0.30,<0.53", + "websockets>=13,<17", +] +test = [ + "pytest>=7.0", + "torch>=2.0", + "jsonschema>=4.23,<5", + "fastapi>=0.115,<0.142", + "httpx>=0.27,<0.29", + "pydantic>=2.11,<3", + "python-multipart>=0.0.9,<1", + "uvicorn>=0.30,<0.53", + "websockets>=13,<17", +] cutedsl = ["nvidia-cutlass-dsl==4.7.1"] [tool.conan-py-build.wheel] @@ -43,6 +60,7 @@ packages = [ "core/builder/tensorrt_model_connect", "apps/benchmark/trtmc_benchmark", "families", + "server/python/trtmc_server", ] [tool.conan-py-build.sdist] @@ -57,6 +75,7 @@ include = [ "apps", "families", "requirements", + "server", "third_party", "examples", "tools/perf_matrix.py", @@ -66,11 +85,13 @@ exclude = [ "*.pyc", "core/builder/tensorrt_model_connect/build", "core/builder/tensorrt_model_connect/*.egg-info", + "server/python/trtmc_server/*.egg-info", ] [tool.pytest.ini_options] enable_assertion_pass_hook = true testpaths = [ + "server/tests", "families", "core/builder/tests", "apps/benchmark/trtmc_benchmark/tests", diff --git a/requirements/community-ci.txt b/requirements/community-ci.txt index a2b01a3c96..d93f798fd6 100644 --- a/requirements/community-ci.txt +++ b/requirements/community-ci.txt @@ -2,9 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 # Keep local source checks and the public Community CPU workflow on the same -# tool versions. The development images consume this file as well. +# tool and server-test dependencies. The development images consume this file +# as well. pre-commit==4.6.2 apache-tvm-ffi==0.1.12 +fastapi>=0.115,<0.142 +httpx>=0.27,<0.29 +pydantic>=2.11,<3 +python-multipart>=0.0.9,<1 ruff==0.16.4 clang-format==22.1.8 lizard==1.21.2 @@ -13,3 +18,5 @@ pytest-xdist==3.8.0 jsonschema==4.26.0 Pillow==12.2.0 PyYAML==6.0.2 +uvicorn>=0.30,<0.53 +websockets>=13,<17 diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt new file mode 100644 index 0000000000..3a51e548c3 --- /dev/null +++ b/server/CMakeLists.txt @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +add_library(trtmc_server_native STATIC + native/entrypoint.cpp + native/worker.cpp +) +target_include_directories(trtmc_server_native + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} +) +target_link_libraries(trtmc_server_native + PRIVATE + trtmc_runtime + nlohmann_json::nlohmann_json +) +target_compile_options(trtmc_server_native PRIVATE -Wall -Wextra -Wpedantic) + +# Keep source-build CLI smoke tests independent of the caller's working directory +# without embedding checkout paths in the native executable. +add_custom_target(trtmc_server_python ALL + COMMAND ${CMAKE_COMMAND} -E remove_directory + "${PROJECT_BINARY_DIR}/server/python/trtmc_server" + COMMAND ${CMAKE_COMMAND} -E copy_directory + "${CMAKE_CURRENT_SOURCE_DIR}/python/trtmc_server" + "${PROJECT_BINARY_DIR}/server/python/trtmc_server" + COMMENT "Copying the local server Python module" +) + +if(TRTMC_BUILD_TESTS) + add_executable(test_serve_worker tests/test_serve_worker.cpp) + target_include_directories(test_serve_worker PRIVATE + ${PROJECT_SOURCE_DIR}/core/runtime/include + ) + target_link_libraries(test_serve_worker PRIVATE + trtmc_server_native + nlohmann_json::nlohmann_json + ) + target_compile_options(test_serve_worker PRIVATE -Wall -Wextra -Wpedantic) + add_test(NAME serve_worker COMMAND test_serve_worker) + set_tests_properties(serve_worker PROPERTIES LABELS cpu) + + set(_trtmc_serve_shadow_dir "${PROJECT_BINARY_DIR}/server/tests/cwd-shadow") + file(MAKE_DIRECTORY "${_trtmc_serve_shadow_dir}/trtmc_server") + file(WRITE "${_trtmc_serve_shadow_dir}/trtmc_server/__init__.py" "") + file(WRITE "${_trtmc_serve_shadow_dir}/trtmc_server/__main__.py" + "print('TRTMC_CWD_SHADOW_EXECUTED')\n" + ) + add_test( + NAME serve_cli_ignores_cwd_shadow + COMMAND $ serve --help + ) + set_tests_properties(serve_cli_ignores_cwd_shadow PROPERTIES + WORKING_DIRECTORY "${_trtmc_serve_shadow_dir}" + PASS_REGULAR_EXPRESSION "Serve TensorRT-Model-Connect bundles" + FAIL_REGULAR_EXPRESSION "TRTMC_CWD_SHADOW_EXECUTED" + LABELS cpu + ) +endif() diff --git a/server/README.md b/server/README.md new file mode 100644 index 0000000000..6e64d85d03 --- /dev/null +++ b/server/README.md @@ -0,0 +1,31 @@ + + +# Local server + +This directory owns the optional process behind `trtmc serve`: + +- `python/trtmc_server/` provides the HTTP and WebSocket control plane. +- `native/` adapts the public `load_task()` and `ITask` contracts to a private + JSONL worker process. +- `tests/` owns server API, process-lifecycle, protocol, and dependency checks. + +The dependency is one-way: the server may use public library contracts, while +core and model families never depend on server implementation. Applications +consume the `trtmc serve` process and its HTTP/WebSocket APIs; they do not +import `trtmc_server`. + +Concurrency is a fixed set of serial worker lanes configured at startup. The +server has no waiting queue, dynamic placement, continuous batching, cluster +scheduler, or worker restart. Saturation fails immediately, failed lanes leave +the model degraded while another lane remains healthy, and recovery belongs to +an external supervisor. One server process is one local placement domain, not +a generic distributed serving system. + +Replica counts above one apply only to independently loadable single-process +bundles. MPI/NCCL distributed bundles are not supported by `trtmc serve`. +For multiple GPUs, run independent single-process server instances and pin +each instance with `CUDA_VISIBLE_DEVICES`; external routing remains the +supervisor's responsibility. diff --git a/server/native/entrypoint.cpp b/server/native/entrypoint.cpp new file mode 100644 index 0000000000..d61504429e --- /dev/null +++ b/server/native/entrypoint.cpp @@ -0,0 +1,217 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "native/entrypoint.h" + +#include "native/worker.h" +#include "trtmc/runtime/family_loader.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc::server { +namespace { + +struct ByteSizeParts { + std::string number; + std::uint64_t multiplier{1}; +}; + +ByteSizeParts split_byte_size(const std::string& text) { + if (text.size() > 3 && text.compare(text.size() - 3, 3, "GiB") == 0) { + return {text.substr(0, text.size() - 3), 1024ULL * 1024ULL * 1024ULL}; + } + if (text.size() > 2 && text.compare(text.size() - 2, 2, "GB") == 0) + return {text.substr(0, text.size() - 2), 1000ULL * 1000ULL * 1000ULL}; + return {text, 1}; +} + +std::uint64_t parse_byte_size(const std::string& text) { + const ByteSizeParts parts = split_byte_size(text); + if (parts.number.empty()) + throw std::invalid_argument( + "--kv-cache-size must be integer bytes or a value like 1GB or 1GiB"); + if (!std::all_of(parts.number.begin(), parts.number.end(), + [](unsigned char value) { return value >= '0' && value <= '9'; })) { + throw std::invalid_argument( + "--kv-cache-size must be integer bytes or a value like 1GB or 1GiB"); + } + std::size_t consumed = 0; + std::uint64_t value = 0; + try { + value = std::stoull(parts.number, &consumed); + } catch (const std::exception&) { + throw std::invalid_argument("--kv-cache-size is outside its valid range"); + } + if (value == 0 || consumed != parts.number.size()) + throw std::invalid_argument("--kv-cache-size is outside its valid range"); + if (value > std::numeric_limits::max() / parts.multiplier) + throw std::invalid_argument("--kv-cache-size is outside its valid range"); + return value * parts.multiplier; +} + +std::string take_value(int argc, char** argv, int& index, const std::string& option) { + if (index + 1 >= argc) + throw std::invalid_argument(option + " requires a value"); + ++index; + const std::string value = argv[index]; + if (value.empty()) + throw std::invalid_argument(option + " requires a non-empty value"); + return value; +} + +NativeWorkerOptions parse_options(int argc, char** argv) { + if (argc < 1 || argv[0] == nullptr || argv[0][0] == '\0') + throw std::invalid_argument("_serve-worker requires a .bundle artifact file"); + + NativeWorkerOptions options; + options.bundle_path = argv[0]; + for (int index = 1; index < argc; ++index) { + const std::string option = argv[index]; + if (option == "--runtime-root") { + options.runtime_root = take_value(argc, argv, index, option); + } else if (option == "--kv-cache-size") { + options.kv_cache_size_bytes = parse_byte_size(take_value(argc, argv, index, option)); + } else if (option == "--runtime-cache") { + options.runtime_cache_path = take_value(argc, argv, index, option); + } else if (option == "--cuda-graphs") { + options.cuda_graphs = true; + } else { + throw std::invalid_argument("unknown _serve-worker option: " + option); + } + } + if (options.runtime_root.empty()) + throw std::invalid_argument("_serve-worker requires --runtime-root DIR"); + return options; +} + +std::filesystem::path current_executable_path() { + std::array buffer{}; + const ssize_t length = readlink("/proc/self/exe", buffer.data(), buffer.size() - 1U); + if (length <= 0) + return {}; + buffer[static_cast(length)] = '\0'; + return std::filesystem::path(buffer.data()); +} + +std::string python_executable(const std::filesystem::path& executable) { + if (!executable.empty()) { + for (const char* name : {"python3", "python"}) { + const auto candidate = executable.parent_path() / name; + if (access(candidate.c_str(), X_OK) == 0) + return candidate.string(); + } + } + return "python3"; +} + +std::filesystem::path source_server_python(const std::filesystem::path& executable) { + if (executable.empty()) + return {}; + const auto server_python = executable.parent_path() / "server" / "python"; + std::error_code directory_error; + if (!std::filesystem::is_directory(server_python / "trtmc_server", directory_error)) + return {}; + return server_python; +} + +void configure_source_pythonpath(const std::filesystem::path& server_python) { + std::string pythonpath = server_python.string(); + if (const char* existing = std::getenv("PYTHONPATH"); + existing != nullptr && existing[0] != '\0') + pythonpath += ":" + std::string(existing); + if (setenv("PYTHONPATH", pythonpath.c_str(), 1) != 0) + throw std::runtime_error("failed to configure the source-build Python module path"); +} + +} // namespace + +int run_server_frontend(int argc, char** argv) { + const auto executable = current_executable_path(); + if (executable.empty()) { + std::cerr << "Error: cannot resolve the current trtmc binary for serve\n"; + return EXIT_FAILURE; + } + + const auto source_python = source_server_python(executable); + const bool source_mode = !source_python.empty(); + std::vector command{python_executable(executable), source_mode ? "-P" : "-I", "-m", + "trtmc_server"}; + bool binary_provided = false; + for (int index = 0; index < argc; ++index) { + std::string argument = argv[index]; + if (argument == "--trtmc-binary" || argument.rfind("--trtmc-binary=", 0) == 0) + binary_provided = true; + command.push_back(std::move(argument)); + } + if (!binary_provided) { + command.emplace_back("--trtmc-binary"); + command.push_back(executable.string()); + } + + std::vector exec_arguments; + exec_arguments.reserve(command.size() + 1U); + for (auto& argument : command) + exec_arguments.push_back(argument.data()); + exec_arguments.push_back(nullptr); + try { + if (source_mode) + configure_source_pythonpath(source_python); + } catch (const std::exception& error) { + std::cerr << "Error: " << error.what() << '\n'; + return EXIT_FAILURE; + } + execvp(exec_arguments[0], exec_arguments.data()); + std::cerr << "Error: failed to execute Python serving module: " << std::strerror(errno) << '\n'; + return 127; +} + +int run_native_worker(const NativeWorkerOptions& options) { + if (options.bundle_path.empty()) { + std::cerr << "Error: _serve-worker requires a .bundle artifact file\n"; + return EXIT_FAILURE; + } + if (options.runtime_root.empty()) { + std::cerr << "Error: _serve-worker requires --runtime-root DIR\n"; + return EXIT_FAILURE; + } + + try { + auto task = + load_task(options.bundle_path, options.runtime_root, options.kv_cache_size_bytes, + options.runtime_cache_path, options.cuda_graphs); + if (!task) + throw std::runtime_error("native worker task is unavailable"); + return serve::run_worker_protocol(*task, std::cin, std::cout); + } catch (const std::exception& error) { + // Stderr is the private diagnostic channel; JSONL stdout stays redacted. + std::cerr << "Error: native worker failed: " << error.what() << '\n'; + return EXIT_FAILURE; + } +} + +int run_native_worker(int argc, char** argv) { + try { + return run_native_worker(parse_options(argc, argv)); + } catch (const std::exception& error) { + std::cerr << "Error: " << error.what() << '\n'; + return EXIT_FAILURE; + } +} + +} // namespace trtmc::server diff --git a/server/native/entrypoint.h b/server/native/entrypoint.h new file mode 100644 index 0000000000..09f7650ec4 --- /dev/null +++ b/server/native/entrypoint.h @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +namespace trtmc::server { + +struct NativeWorkerOptions { + std::string bundle_path; + std::string runtime_root; + std::uint64_t kv_cache_size_bytes{0}; + std::string runtime_cache_path; + bool cuda_graphs{false}; +}; + +// Private native data-plane entry points used by the `trtmc serve` facade. +int run_server_frontend(int argc, char** argv); +int run_native_worker(const NativeWorkerOptions& options); +int run_native_worker(int argc, char** argv); + +} // namespace trtmc::server diff --git a/server/native/worker.cpp b/server/native/worker.cpp new file mode 100644 index 0000000000..2da87146a5 --- /dev/null +++ b/server/native/worker.cpp @@ -0,0 +1,753 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "native/worker.h" + +#include "trtmc/task.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc::serve { +namespace { + +using Json = nlohmann::json; + +constexpr std::size_t kMaxRequestLineBytes = 16U * 1024U * 1024U; +constexpr std::size_t kWavReadBufferBytes = 64U * 1024U; +constexpr const char* kRuntimeErrorMessage = "native worker operation failed"; + +class ProtocolError final : public std::runtime_error { + public: + using std::runtime_error::runtime_error; +}; + +class WavFormatError final : public std::runtime_error { + public: + using std::runtime_error::runtime_error; +}; + +class UnsupportedMediaTypeError final : public std::runtime_error { + public: + using std::runtime_error::runtime_error; +}; + +struct DispatchResult { + Json result; + bool shutdown{false}; +}; + +template +void assign_if_present(const Json& object, const char* name, T& destination) { + const auto value = object.find(name); + if (value != object.end()) + destination = value->get(); +} + +Json request_config(const Json& request) { + const auto value = request.find("config"); + if (value == request.end()) + return Json::object(); + if (!value->is_object()) + throw ProtocolError("config must be a JSON object"); + return *value; +} + +void require_config_fields(const Json& config, std::initializer_list allowed) { + for (auto field = config.begin(); field != config.end(); ++field) { + const bool supported = std::any_of(allowed.begin(), allowed.end(), + [&](const char* name) { return field.key() == name; }); + if (!supported) + throw ProtocolError("config." + field.key() + " is unsupported"); + } +} + +void require_finite(float value, const char* name) { + if (!std::isfinite(value)) + throw ProtocolError(std::string(name) + " must be finite"); +} + +void require_non_negative(float value, const char* name) { + require_finite(value, name); + if (value < 0.0F) + throw ProtocolError(std::string(name) + " must be non-negative"); +} + +void validate_generate_config(const TextGenerationConfig& config) { + if (config.max_new_tokens <= 0) + throw ProtocolError("config.max_new_tokens must be positive"); + if (config.top_k < 0) + throw ProtocolError("config.top_k must be non-negative"); + require_non_negative(config.temperature, "config.temperature"); + require_finite(config.top_p, "config.top_p"); + require_finite(config.min_p, "config.min_p"); + if (config.top_p < 0.0F || config.top_p > 1.0F) + throw ProtocolError("config.top_p must be in [0, 1]"); + if (config.min_p < 0.0F || config.min_p > 1.0F) + throw ProtocolError("config.min_p must be in [0, 1]"); +} + +TextGenerationConfig parse_generate_config(const Json& request, std::int32_t default_tokens) { + const Json config = request_config(request); + require_config_fields(config, {"max_new_tokens", "temperature", "top_p", "min_p", "top_k", + "seed", "use_chat_template", "enable_thinking"}); + TextGenerationConfig result; + result.max_new_tokens = default_tokens > 0 ? default_tokens : 128; + try { + assign_if_present(config, "max_new_tokens", result.max_new_tokens); + assign_if_present(config, "temperature", result.temperature); + assign_if_present(config, "top_p", result.top_p); + assign_if_present(config, "min_p", result.min_p); + assign_if_present(config, "top_k", result.top_k); + assign_if_present(config, "seed", result.seed); + assign_if_present(config, "use_chat_template", result.use_chat_template); + assign_if_present(config, "enable_thinking", result.enable_thinking); + } catch (const nlohmann::json::exception&) { + throw ProtocolError("generation config contains an invalid value"); + } + validate_generate_config(result); + return result; +} + +TranscriptionConfig parse_transcription_config(const Json& request, std::int32_t sample_rate) { + const Json config = request_config(request); + require_config_fields(config, {"language"}); + TranscriptionConfig result; + result.input_sample_rate = sample_rate; + try { + assign_if_present(config, "language", result.source_language); + } catch (const nlohmann::json::exception&) { + throw ProtocolError("transcription config contains an invalid value"); + } + return result; +} + +TranscriptionStreamConfig parse_stream_config(const Json& request) { + const Json config = request_config(request); + require_config_fields(config, {"sample_rate_hz", "channels", "audio_format", "language"}); + TranscriptionStreamConfig result; + std::int32_t channels = 1; + std::string audio_format{"pcm16le"}; + try { + assign_if_present(config, "sample_rate_hz", result.input_sample_rate); + assign_if_present(config, "language", result.language); + channels = config.value("channels", 1); + audio_format = config.value("audio_format", std::string{"pcm16le"}); + } catch (const nlohmann::json::exception&) { + throw ProtocolError("stream config contains an invalid value"); + } + if (channels != 1) + throw ProtocolError("config.channels must be 1 because streaming input is mono"); + if (audio_format != "pcm16le") + throw ProtocolError("config.audio_format must be 'pcm16le'"); + if (result.input_sample_rate <= 0) + throw ProtocolError("config.sample_rate_hz must be positive"); + return result; +} + +Json transcription_segments_json(const std::vector& segments) { + Json result = Json::array(); + for (const auto& segment : segments) { + result.push_back({ + {"start_seconds", segment.start_seconds}, + {"end_seconds", segment.end_seconds}, + {"text", segment.text}, + {"token_ids", segment.token_ids}, + }); + } + return result; +} + +Json text_result_json(const TextResult& result) { + return { + {"text", result.text}, + {"token_ids", result.token_ids}, + {"completion_tokens", result.token_ids.size()}, + {"segments", transcription_segments_json(result.segments)}, + {"setup_ms", result.setup_ms}, + {"prefill_ms", result.prefill_ms}, + {"decode_ms", result.decode_ms}, + }; +} + +Json stream_result_json(const TranscriptionStreamResult& result) { + return {{"text", result.text}}; +} + +std::uint16_t read_u16_le(const char* data) { + const auto* bytes = reinterpret_cast(data); + return static_cast(bytes[0]) | (static_cast(bytes[1]) << 8U); +} + +std::uint32_t read_u32_le(const char* data) { + const auto* bytes = reinterpret_cast(data); + return static_cast(bytes[0]) | (static_cast(bytes[1]) << 8U) | + (static_cast(bytes[2]) << 16U) | + (static_cast(bytes[3]) << 24U); +} + +void read_exact(std::ifstream& input, char* destination, std::size_t size, const char* message) { + if (size > static_cast(std::numeric_limits::max())) + throw WavFormatError(message); + input.read(destination, static_cast(size)); + if (input.gcount() != static_cast(size)) + throw WavFormatError(message); +} + +void seek_exact(std::ifstream& input, std::uint64_t position) { + if (position > static_cast(std::numeric_limits::max())) + throw WavFormatError("WAV file is too large"); + input.clear(); + input.seekg(static_cast(position), std::ios::beg); + if (!input) + throw WavFormatError("WAV seek failed"); +} + +struct WavLayout { + std::uint16_t format{0}; + std::uint16_t channels{0}; + std::uint32_t sample_rate{0}; + std::uint16_t bits_per_sample{0}; + std::uint64_t data_offset{0}; + std::uint32_t data_size{0}; + bool have_format{false}; + bool have_data{false}; +}; + +std::uint64_t read_wav_container_end(std::ifstream& input, std::uint64_t file_size) { + if (file_size < 12U) + throw WavFormatError("WAV file is too small"); + std::array header{}; + seek_exact(input, 0); + read_exact(input, header.data(), header.size(), "WAV header is truncated"); + if (std::memcmp(header.data(), "RIFF", 4) != 0) + throw WavFormatError("file is not a RIFF container"); + if (std::memcmp(header.data() + 8, "WAVE", 4) != 0) + throw WavFormatError("RIFF container is not WAVE audio"); + const std::uint32_t declared_size = read_u32_le(header.data() + 4); + if (declared_size < 4U) + throw WavFormatError("WAV RIFF chunk is too small"); + const std::uint64_t container_end = 8U + static_cast(declared_size); + if (container_end > file_size) + throw WavFormatError("WAV contains a truncated RIFF chunk"); + return container_end; +} + +struct WavChunk { + std::array id{}; + std::uint32_t size{0}; + std::uint64_t data_offset{0}; + std::uint64_t padded_size{0}; +}; + +WavChunk read_wav_chunk(std::ifstream& input, std::uint64_t position, std::uint64_t container_end) { + if (container_end - position < 8U) + throw WavFormatError("WAV contains a truncated chunk header"); + std::array header{}; + seek_exact(input, position); + read_exact(input, header.data(), header.size(), "WAV chunk header is truncated"); + + WavChunk chunk; + std::copy_n(header.data(), chunk.id.size(), chunk.id.data()); + chunk.size = read_u32_le(header.data() + 4); + chunk.data_offset = position + 8U; + chunk.padded_size = + static_cast(chunk.size) + static_cast(chunk.size & 1U); + if (chunk.padded_size > container_end - chunk.data_offset) + throw WavFormatError("WAV contains a truncated chunk"); + return chunk; +} + +void read_wav_format(std::ifstream& input, const WavChunk& chunk, WavLayout& layout) { + if (chunk.size < 16U) + throw WavFormatError("WAV fmt chunk is too small"); + std::array format{}; + seek_exact(input, chunk.data_offset); + read_exact(input, format.data(), format.size(), "WAV fmt chunk is truncated"); + layout.format = read_u16_le(format.data()); + layout.channels = read_u16_le(format.data() + 2); + layout.sample_rate = read_u32_le(format.data() + 4); + layout.bits_per_sample = read_u16_le(format.data() + 14); + layout.have_format = true; +} + +void apply_wav_chunk(std::ifstream& input, const WavChunk& chunk, WavLayout& layout) { + if (std::memcmp(chunk.id.data(), "fmt ", 4) == 0) { + read_wav_format(input, chunk, layout); + return; + } + if (std::memcmp(chunk.id.data(), "data", 4) != 0 || chunk.size == 0U || layout.have_data) + return; + layout.data_offset = chunk.data_offset; + layout.data_size = chunk.size; + layout.have_data = true; +} + +void validate_wav_layout(const WavLayout& layout) { + if (!layout.have_format || !layout.have_data) + throw WavFormatError("WAV must contain non-empty fmt and data chunks"); + if (layout.channels == 0U || layout.sample_rate == 0U) + throw WavFormatError("WAV channels and sample rate must be positive"); + if (layout.sample_rate > static_cast(std::numeric_limits::max())) + throw WavFormatError("WAV sample rate exceeds the supported range"); + const bool pcm16 = layout.format == 1U && layout.bits_per_sample == 16U; + const bool float32 = layout.format == 3U && layout.bits_per_sample == 32U; + if (!pcm16 && !float32) + throw WavFormatError("WAV samples must be PCM16 or IEEE float32"); +} + +WavLayout parse_wav_layout(std::ifstream& input, std::uint64_t file_size) { + const std::uint64_t container_end = read_wav_container_end(input, file_size); + + WavLayout layout; + std::uint64_t position = 12U; + while (position < container_end) { + const WavChunk chunk = read_wav_chunk(input, position, container_end); + apply_wav_chunk(input, chunk, layout); + position = chunk.data_offset + chunk.padded_size; + } + validate_wav_layout(layout); + return layout; +} + +float decode_wav_sample(const char* data, std::uint16_t format) { + if (format == 3U) { + const std::uint32_t bits = read_u32_le(data); + float value = 0.0F; + std::memcpy(&value, &bits, sizeof(value)); + if (!std::isfinite(value)) + throw WavFormatError("WAV contains a non-finite float sample"); + return value; + } + const std::uint16_t raw = read_u16_le(data); + const std::int32_t value = + raw >= 0x8000U ? static_cast(raw) - 0x10000 : static_cast(raw); + return static_cast(value) / 32768.0F; +} + +struct DecodedAudio { + std::vector samples; + std::int32_t sample_rate{0}; +}; + +DecodedAudio read_wav(const std::string& path) { + std::ifstream input(path, std::ios::binary | std::ios::ate); + if (!input) + throw std::runtime_error("cannot open uploaded audio file"); + const std::streampos end = input.tellg(); + if (end < 0) + throw std::runtime_error("cannot determine uploaded audio file size"); + const auto file_size = static_cast(end); + const WavLayout layout = parse_wav_layout(input, file_size); + const std::size_t sample_width = layout.bits_per_sample / 8U; + const std::size_t frame_width = sample_width * static_cast(layout.channels); + if (frame_width == 0U || layout.data_size % frame_width != 0U) + throw WavFormatError("WAV data does not contain complete audio frames"); + const std::size_t frame_count = layout.data_size / frame_width; + if (frame_count > static_cast(std::numeric_limits::max())) + throw WavFormatError("WAV contains too many audio frames"); + + DecodedAudio result; + result.samples.resize(frame_count); + result.sample_rate = static_cast(layout.sample_rate); + seek_exact(input, layout.data_offset); + const std::size_t frames_per_read = + std::max(1U, kWavReadBufferBytes / frame_width); + std::vector buffer(frames_per_read * frame_width); + for (std::size_t first = 0; first < frame_count; first += frames_per_read) { + const std::size_t count = std::min(frames_per_read, frame_count - first); + const std::size_t bytes = count * frame_width; + read_exact(input, buffer.data(), bytes, "WAV audio data is truncated"); + for (std::size_t frame = 0; frame < count; ++frame) { + float sum = 0.0F; + for (std::uint16_t channel = 0; channel < layout.channels; ++channel) { + const std::size_t offset = + frame * frame_width + static_cast(channel) * sample_width; + sum += decode_wav_sample(buffer.data() + offset, layout.format); + } + result.samples[first + frame] = sum / static_cast(layout.channels); + } + } + return result; +} + +int decode_base64_character(unsigned char value) { + if (value >= 'A' && value <= 'Z') + return value - 'A'; + if (value >= 'a' && value <= 'z') + return value - 'a' + 26; + if (value >= '0' && value <= '9') + return value - '0' + 52; + if (value == '+') + return 62; + if (value == '/') + return 63; + return -1; +} + +int decode_required_base64_character(char value) { + const int decoded = decode_base64_character(static_cast(value)); + if (decoded < 0) + throw ProtocolError("audio is not valid base64"); + return decoded; +} + +int decode_optional_base64_character(char value) { + return value == '=' ? 0 : decode_required_base64_character(value); +} + +void validate_base64_padding(char third, char fourth, bool last) { + if ((third == '=' && fourth != '=') || (!last && (third == '=' || fourth == '='))) + throw ProtocolError("audio is not valid base64"); +} + +std::vector decode_base64(const std::string& encoded) { + if (encoded.empty()) + return {}; + if (encoded.size() % 4U != 0U) + throw ProtocolError("audio has invalid base64 length"); + + std::vector decoded; + decoded.reserve(encoded.size() / 4U * 3U); + for (std::size_t offset = 0; offset < encoded.size(); offset += 4U) { + const bool last = offset + 4U == encoded.size(); + const char third = encoded[offset + 2U]; + const char fourth = encoded[offset + 3U]; + validate_base64_padding(third, fourth, last); + const std::uint32_t bits = + (static_cast(decode_required_base64_character(encoded[offset])) << 18U) | + (static_cast(decode_required_base64_character(encoded[offset + 1U])) + << 12U) | + (static_cast(decode_optional_base64_character(third)) << 6U) | + static_cast(decode_optional_base64_character(fourth)); + decoded.push_back(static_cast((bits >> 16U) & 0xFFU)); + if (third != '=') + decoded.push_back(static_cast((bits >> 8U) & 0xFFU)); + if (fourth != '=') + decoded.push_back(static_cast(bits & 0xFFU)); + } + return decoded; +} + +std::vector decode_pcm16_base64(const std::string& encoded) { + const auto bytes = decode_base64(encoded); + if (bytes.size() % 2U != 0U) + throw ProtocolError("audio must contain complete little-endian int16 samples"); + if (bytes.size() / 2U > static_cast(std::numeric_limits::max())) + throw ProtocolError("PCM chunk contains too many samples"); + + std::vector samples(bytes.size() / 2U); + for (std::size_t index = 0; index < samples.size(); ++index) { + const std::uint16_t raw = static_cast(bytes[index * 2U]) | + (static_cast(bytes[index * 2U + 1U]) << 8U); + const std::int32_t value = raw >= 0x8000U ? static_cast(raw) - 0x10000 + : static_cast(raw); + samples[index] = static_cast(value) / 32768.0F; + } + return samples; +} + +std::string required_string(const Json& request, const char* field, bool allow_empty = false) { + const auto value = request.find(field); + if (value == request.end() || !value->is_string()) + throw ProtocolError(std::string(field) + " must be a string"); + const std::string result = value->get(); + if (result.empty() && !allow_empty) + throw ProtocolError(std::string(field) + " must not be empty"); + return result; +} + +class Worker final { + public: + explicit Worker(ITask& task) + : text_(dynamic_cast(&task)), + transcription_(dynamic_cast(&task)), + streaming_(dynamic_cast(&task)) {} + + Json ready_event() const { + Json capabilities = Json::array(); + if (text_ != nullptr) + capabilities.push_back(ITextGeneration::kTask); + if (transcription_ != nullptr) + capabilities.push_back(ITranscription::kTask); + if (streaming_ != nullptr) + capabilities.push_back(IStreamingTranscription::kTask); + Json result = { + {"event", "ready"}, + {"protocol_version", 3}, + {"capabilities", std::move(capabilities)}, + }; + if (text_ != nullptr) + result["default_max_new_tokens"] = text_->default_max_new_tokens(); + return result; + } + + DispatchResult dispatch(const Json& request) { + if (!request.is_object()) + throw ProtocolError("request must be a JSON object"); + const std::string operation = required_string(request, "op"); + if (operation == "shutdown") + return {{{"status", "shutting_down"}}, true}; + if (operation == "generate") + return {generate(request), false}; + if (operation == "transcribe") + return {transcribe(request), false}; + if (operation == "probe_transcription_stream") + return {probe_transcription_stream(request), false}; + if (operation == "stream_start") + return {stream_start(request), false}; + if (operation == "stream_chunk") + return {stream_chunk(request), false}; + if (operation == "stream_finish") + return {stream_finish(), false}; + if (operation == "stream_reset") + return {stream_reset(), false}; + throw ProtocolError("unknown operation: " + operation); + } + + private: + void require_no_active_stream(const char* operation) const { + if (active_stream_) + throw ProtocolError(std::string(operation) + + " is unavailable while a transcription stream is active"); + } + + ITextGeneration& text() const { + if (text_ == nullptr) + throw ProtocolError("loaded task does not support text generation"); + return *text_; + } + + ITranscription& transcription() const { + if (transcription_ == nullptr) + throw ProtocolError("loaded task does not support transcription"); + return *transcription_; + } + + IStreamingTranscription& streaming() const { + if (streaming_ == nullptr) + throw ProtocolError("loaded task does not support streaming transcription"); + return *streaming_; + } + + Json generate(const Json& request) { + require_no_active_stream("generate"); + auto& interface = text(); + const std::string prompt = required_string(request, "prompt", true); + const auto config = parse_generate_config(request, interface.default_max_new_tokens()); + return text_result_json(interface.generate(prompt, config)); + } + + Json transcribe(const Json& request) { + require_no_active_stream("transcribe"); + const std::string audio_path = required_string(request, "audio_path"); + DecodedAudio audio; + try { + audio = read_wav(audio_path); + } catch (const WavFormatError&) { + throw UnsupportedMediaTypeError( + "uploaded audio must be a supported PCM16 or IEEE float32 WAV file"); + } + const auto config = parse_transcription_config(request, audio.sample_rate); + return text_result_json(transcription().transcribe( + audio.samples.data(), static_cast(audio.samples.size()), config)); + } + + Json probe_transcription_stream(const Json& request) { + require_no_active_stream("probe_transcription_stream"); + auto stream = streaming().create_transcription_stream(parse_stream_config(request)); + if (!stream) + throw std::runtime_error("streaming transcription task returned a null stream"); + stream->reset(); + return {{"supported", true}}; + } + + Json stream_start(const Json& request) { + if (active_stream_) + throw ProtocolError("this worker already has an active transcription stream"); + auto stream = streaming().create_transcription_stream(parse_stream_config(request)); + if (!stream) + throw std::runtime_error("streaming transcription task returned a null stream"); + active_stream_ = std::move(stream); + return Json::object(); + } + + ITranscriptionStream& active_stream() { + if (!active_stream_) + throw ProtocolError("no active transcription stream"); + return *active_stream_; + } + + std::unique_ptr take_stream() { + if (!active_stream_) + throw ProtocolError("no active transcription stream"); + return std::move(active_stream_); + } + + Json stream_chunk(const Json& request) { + const auto encoded = request.find("audio"); + if (encoded == request.end() || !encoded->is_string()) + throw ProtocolError("audio must be a base64 PCM16 string"); + auto samples = decode_pcm16_base64(encoded->get()); + const auto result = + active_stream().accept_audio(samples.empty() ? nullptr : samples.data(), + static_cast(samples.size()), false); + return stream_result_json(result); + } + + Json stream_finish() { + auto stream = take_stream(); + return stream_result_json(stream->finish()); + } + + Json stream_reset() { + auto stream = take_stream(); + stream->reset(); + return Json::object(); + } + + ITextGeneration* text_{nullptr}; + ITranscription* transcription_{nullptr}; + IStreamingTranscription* streaming_{nullptr}; + std::unique_ptr active_stream_; +}; + +bool valid_request_id(const Json& id) { + return id.is_string() && !id.get_ref().empty(); +} + +Json success_response(const Json& id, Json result) { + return {{"id", id}, {"ok", true}, {"result", std::move(result)}}; +} + +Json error_response(const Json& id, const char* type, const std::string& message, + const char* code = nullptr, const char* param = nullptr) { + Json error = {{"type", type}, {"message", message}}; + if (code != nullptr) + error["code"] = code; + if (param != nullptr) + error["param"] = param; + return {{"id", id}, {"ok", false}, {"error", std::move(error)}}; +} + +bool write_message(std::ostream& output, const Json& message) { + output << message.dump(-1, ' ', false, Json::error_handler_t::replace) << '\n'; + output.flush(); + return static_cast(output); +} + +struct ProcessedRequest { + Json response; + bool shutdown{false}; +}; + +enum class RequestLineRead { kRecord, kEnd, kError }; + +RequestLineRead read_request_line(std::istream& input, std::vector& buffer, + std::size_t& line_size) { + input.getline(buffer.data(), static_cast(buffer.size())); + const std::streamsize extracted = input.gcount(); + if (input.bad()) + return RequestLineRead::kError; + if (input.eof() && extracted == 0) + return RequestLineRead::kEnd; + if (input.fail()) { + input.clear(input.rdstate() & ~std::ios::failbit); + input.ignore(std::numeric_limits::max(), '\n'); + if (input.bad()) + return RequestLineRead::kError; + line_size = kMaxRequestLineBytes + 1U; + return RequestLineRead::kRecord; + } + line_size = static_cast(extracted); + if (!input.eof()) + --line_size; + return RequestLineRead::kRecord; +} + +Json extract_request_id(const Json& request) { + if (!request.is_object()) + return nullptr; + const auto id = request.find("id"); + if (id == request.end() || !valid_request_id(*id)) + return nullptr; + return *id; +} + +ProcessedRequest process_request_line(Worker& worker, std::string_view line) { + Json request_id = nullptr; + try { + if (line.size() > kMaxRequestLineBytes) + throw ProtocolError("request exceeds the 16 MiB JSONL limit"); + Json request; + try { + request = Json::parse(line.begin(), line.end()); + } catch (const nlohmann::json::parse_error&) { + throw ProtocolError("request is not valid JSON"); + } + request_id = extract_request_id(request); + if (request_id.is_null()) + throw ProtocolError("id must be a non-empty string"); + auto dispatched = worker.dispatch(request); + return {success_response(request_id, std::move(dispatched.result)), dispatched.shutdown}; + } catch (const UnsupportedMediaTypeError& error) { + return {error_response(request_id, "invalid_request_error", error.what(), + "unsupported_media_type", "file"), + false}; + } catch (const std::exception& error) { + if (dynamic_cast(&error) != nullptr) + return {error_response(request_id, "invalid_request_error", error.what()), false}; + std::cerr << "[trtmc.serve.worker] " << error.what() << '\n'; + return {error_response(request_id, "runtime_error", kRuntimeErrorMessage), false}; + } catch (...) { + std::cerr << "[trtmc.serve.worker] unknown native worker error\n"; + return {error_response(request_id, "runtime_error", kRuntimeErrorMessage), false}; + } +} + +} // namespace + +int run_worker_protocol(ITask& task, std::istream& input, std::ostream& output) { + Worker worker(task); + if (!write_message(output, worker.ready_event())) + return 2; + + std::vector line_buffer(kMaxRequestLineBytes + 2U); + while (true) { + std::size_t line_size = 0; + const RequestLineRead read = read_request_line(input, line_buffer, line_size); + if (read == RequestLineRead::kEnd) + return 0; + if (read == RequestLineRead::kError) + return 2; + if (line_size == 0) + continue; + auto processed = + process_request_line(worker, std::string_view(line_buffer.data(), line_size)); + if (!write_message(output, processed.response)) + return 2; + if (processed.shutdown) + return 0; + } +} + +} // namespace trtmc::serve diff --git a/server/native/worker.h b/server/native/worker.h new file mode 100644 index 0000000000..a9705285ab --- /dev/null +++ b/server/native/worker.h @@ -0,0 +1,19 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +namespace trtmc { +class ITask; +} + +namespace trtmc::serve { + +// Run the private, serialized JSONL data-plane protocol for one resident task. +int run_worker_protocol(ITask& task, std::istream& input, std::ostream& output); + +} // namespace trtmc::serve diff --git a/server/python/trtmc_server/__init__.py b/server/python/trtmc_server/__init__.py new file mode 100644 index 0000000000..20dd54acbb --- /dev/null +++ b/server/python/trtmc_server/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Private implementation package for the ``trtmc serve`` process.""" diff --git a/server/python/trtmc_server/__main__.py b/server/python/trtmc_server/__main__.py new file mode 100644 index 0000000000..e68077045e --- /dev/null +++ b/server/python/trtmc_server/__main__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the private control plane behind ``trtmc serve``.""" + +from .cli import main + + +raise SystemExit(main()) diff --git a/server/python/trtmc_server/app.py b/server/python/trtmc_server/app.py new file mode 100644 index 0000000000..37bd0edc90 --- /dev/null +++ b/server/python/trtmc_server/app.py @@ -0,0 +1,704 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FastAPI control plane for persistent TensorRT-Model-Connect workers.""" + +from __future__ import annotations + +import asyncio +import hmac +import logging +import tempfile +import time +import uuid +from collections.abc import AsyncIterator, Mapping +from contextlib import asynccontextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from urllib.parse import urlsplit + +from fastapi import FastAPI, File, Form, Request, UploadFile, WebSocket +from fastapi.exceptions import RequestValidationError +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse, PlainTextResponse +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from .errors import ( + ModelCapabilityError, + ModelNotFoundError, + WorkerCrashedError, + WorkerProtocolError, + WorkerRemoteError, + WorkerRequestTooLargeError, + WorkerSaturatedError, + WorkerTimeoutError, +) +from .protocol import ( + extract_transcription_segments, + extract_text, + extract_usage, + invalid_request_message, + is_text_only_content, + prepare_chat_prompt, + public_worker_error_message, +) +from .realtime import RealtimeTranscriptionConnection +from .registry import ModelRegistry +from .schemas import ChatCompletionRequest, model_to_dict +from .worker import WorkerSession + + +_LOGGER = logging.getLogger(__name__) +_HTTP_ENVELOPE_OVERHEAD_BYTES = 64 * 1024 + + +class _RequestBodyTooLarge(Exception): + pass + + +class _BodyLimitMiddleware: + """Reject oversized HTTP bodies before framework parsing or spooling.""" + + def __init__(self, app: ASGIApp, *, limits: Mapping[str, int]) -> None: + self.app = app + self.limits = dict(limits) + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http" or (limit := self.limits.get(scope["path"])) is None: + await self.app(scope, receive, send) + return + + content_length = _content_length(scope) + if content_length is not None and content_length > limit: + await _request_body_too_large()(scope, receive, send) + return + + received = 0 + response_started = False + + async def limited_receive() -> Message: + nonlocal received + message = await receive() + if message["type"] == "http.request": + received += len(message.get("body", b"")) + if received > limit: + raise _RequestBodyTooLarge + return message + + async def tracked_send(message: Message) -> None: + nonlocal response_started + if message["type"] == "http.response.start": + response_started = True + await send(message) + + try: + await self.app(scope, limited_receive, tracked_send) + except _RequestBodyTooLarge: + if response_started: + raise + await _request_body_too_large()(scope, receive, send) + + +def _content_length(scope: Scope) -> int | None: + for name, raw_value in scope.get("headers", ()): + if name.lower() != b"content-length": + continue + try: + value = int(raw_value) + except ValueError: + return None + return value if value >= 0 else None + return None + + +def _request_body_too_large() -> JSONResponse: + return _error_response( + 413, + "request_body_too_large", + "Request body exceeds the server ingress limit", + headers={"Connection": "close"}, + ) + + +@dataclass(frozen=True) +class ServerConfig: + """HTTP-layer policy independent of model process configuration.""" + + api_key: str | None = None + max_generation_tokens: int = 4096 + max_prompt_bytes: int = 8 * 1024 * 1024 + max_upload_bytes: int = 512 * 1024 * 1024 + max_realtime_chunk_bytes: int = 1024 * 1024 + max_realtime_session_bytes: int = 512 * 1024 * 1024 + realtime_idle_timeout_seconds: float = 30.0 + realtime_max_session_seconds: float = 4 * 60 * 60 + + def __post_init__(self) -> None: + if self.api_key is not None and not self.api_key: + raise ValueError("api_key cannot be empty") + if ( + self.max_prompt_bytes <= 0 + or self.max_generation_tokens <= 0 + or self.max_upload_bytes <= 0 + or self.max_realtime_chunk_bytes <= 0 + or self.max_realtime_session_bytes <= 0 + or self.realtime_idle_timeout_seconds <= 0 + or self.realtime_max_session_seconds <= 0 + ): + raise ValueError("server limits must be positive") + + +def create_app( + registry: ModelRegistry, + *, + config: ServerConfig | None = None, +) -> FastAPI: + """Create an application whose lifespan owns the supplied registry.""" + + server_config = config or ServerConfig() + + @asynccontextmanager + async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + await asyncio.to_thread(registry.start) + try: + yield + finally: + await asyncio.to_thread(registry.close) + + app = FastAPI( + title="TensorRT-Model-Connect Serve", + version="0.1.0", + lifespan=lifespan, + ) + app.add_middleware( + _BodyLimitMiddleware, + limits={ + "/v1/chat/completions": ( + server_config.max_prompt_bytes + _HTTP_ENVELOPE_OVERHEAD_BYTES + ), + "/v1/audio/transcriptions": ( + server_config.max_upload_bytes + _HTTP_ENVELOPE_OVERHEAD_BYTES + ), + }, + ) + app.add_middleware( + CORSMiddleware, + allow_origins=["null", "http://localhost", "http://127.0.0.1"], + allow_origin_regex=r"^http://(?:localhost|127\.0\.0\.1)(?::[0-9]{1,5})?$", + allow_methods=["GET", "POST", "OPTIONS"], + allow_headers=["Authorization", "Content-Type"], + allow_credentials=False, + ) + + @app.middleware("http") + async def authenticate(request: Request, call_next: Any) -> Any: + if ( + server_config.api_key is not None + and request.url.path != "/healthz" + and request.method != "OPTIONS" + and not _valid_authorization( + request.headers.get("authorization"), server_config.api_key + ) + ): + return _error_response( + 401, + "invalid_api_key", + "Missing or invalid bearer token", + headers={"WWW-Authenticate": "Bearer"}, + ) + return await call_next(request) + + @app.exception_handler(ModelNotFoundError) + async def model_not_found(_request: Request, exc: ModelNotFoundError) -> JSONResponse: + return _error_response(404, exc.code, str(exc), param="model") + + @app.exception_handler(RequestValidationError) + async def invalid_request(_request: Request, exc: RequestValidationError) -> JSONResponse: + errors = exc.errors() + first = errors[0] if errors else {} + location = first.get("loc", ()) + param = str(location[-1]) if location else None + message = str(first.get("msg") or "request validation failed") + return _error_response(422, "invalid_request", message, param=param) + + @app.exception_handler(ModelCapabilityError) + async def model_capability(_request: Request, exc: ModelCapabilityError) -> JSONResponse: + return _error_response(400, exc.code, str(exc), param="model") + + @app.exception_handler(WorkerTimeoutError) + async def worker_timeout(_request: Request, exc: WorkerTimeoutError) -> JSONResponse: + _log_worker_failure() + return _error_response( + 504, + exc.code, + public_worker_error_message(exc), + error_type="server_error", + ) + + @app.exception_handler(WorkerCrashedError) + async def worker_crashed(_request: Request, exc: WorkerCrashedError) -> JSONResponse: + _log_worker_failure() + return _error_response( + 503, + exc.code, + public_worker_error_message(exc), + error_type="server_error", + ) + + @app.exception_handler(WorkerRemoteError) + async def worker_remote(_request: Request, exc: WorkerRemoteError) -> JSONResponse: + message = invalid_request_message(exc) + if message is not None: + details = exc.details + if ( + isinstance(details, Mapping) + and details.get("code") == "unsupported_media_type" + and details.get("param") == "file" + ): + return _error_response( + 415, + "unsupported_media_type", + message, + param="file", + ) + return _error_response( + 400, + "invalid_request", + message, + ) + _log_worker_failure() + return _error_response( + 502, + exc.code, + public_worker_error_message(exc), + error_type="server_error", + ) + + @app.exception_handler(WorkerRequestTooLargeError) + async def worker_request_too_large( + _request: Request, exc: WorkerRequestTooLargeError + ) -> JSONResponse: + return _error_response(413, exc.code, public_worker_error_message(exc)) + + @app.exception_handler(WorkerSaturatedError) + async def worker_saturated(_request: Request, exc: WorkerSaturatedError) -> JSONResponse: + return _error_response( + 429, + exc.code, + public_worker_error_message(exc), + error_type="rate_limit_error", + headers={"Retry-After": "1"}, + ) + + @app.exception_handler(WorkerProtocolError) + async def worker_protocol(_request: Request, exc: WorkerProtocolError) -> JSONResponse: + _log_worker_failure() + return _error_response( + 502, + exc.code, + public_worker_error_message(exc), + error_type="server_error", + ) + + @app.get("/healthz") + async def healthz() -> JSONResponse: + available = registry.has_healthy_worker + degraded = registry.public_status()["degraded"] + return JSONResponse( + status_code=200 if available else 503, + content={ + "status": "ok" if available else "unavailable", + "degraded": degraded, + }, + ) + + @app.get("/readyz") + async def readyz() -> JSONResponse: + status = registry.public_status() + code = 200 if status["ready"] else 503 + content: dict[str, Any] = { + "status": "ready" if code == 200 else "not_ready", + "ready": status["ready"], + "degraded": status["degraded"], + } + if server_config.api_key is not None: + content.update(status) + return JSONResponse( + status_code=code, + content=content, + ) + + @app.get("/v1/models") + async def models() -> dict[str, Any]: + return {"object": "list", "data": registry.list_models()} + + @app.post("/v1/chat/completions") + async def chat_completions(request: ChatCompletionRequest) -> Any: + if request.stream: + return _error_response( + 400, + "streaming_not_supported", + "text token streaming is not available", + param="stream", + ) + raw_messages = [model_to_dict(message) for message in request.messages] + unsupported = _unsupported_chat_parameter(request, raw_messages) + if unsupported is not None: + return _error_response( + 400, + "unsupported_parameter", + f"{unsupported} is not supported", + param=unsupported, + ) + prompt = prepare_chat_prompt(raw_messages) + if len(prompt.encode("utf-8")) > server_config.max_prompt_bytes: + return _error_response( + 413, + "prompt_too_large", + f"rendered messages exceed {server_config.max_prompt_bytes} UTF-8 bytes", + param="messages", + ) + max_tokens = request.max_completion_tokens or request.max_tokens + if max_tokens is not None and max_tokens > server_config.max_generation_tokens: + return _error_response( + 400, + "max_tokens_exceeded", + "requested generation exceeds the server hard cap of " + f"{server_config.max_generation_tokens} tokens", + param="max_tokens", + ) + spec, session = registry.acquire_session("chat", request.model) + max_tokens = _effective_generation_tokens( + registry, + spec.name, + max_tokens, + server_config.max_generation_tokens, + ) + result = await _worker_request( + session, + "generate", + { + "prompt": prompt, + "config": _generation_config( + request, + max_tokens=max_tokens, + ), + }, + ) + text = extract_text(result, operation="generate") + response_id = f"chatcmpl-{uuid.uuid4().hex}" + created = int(time.time()) + return { + "id": response_id, + "object": "chat.completion", + "created": created, + "model": spec.name, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": text}, + "logprobs": None, + "finish_reason": "stop", + } + ], + "usage": extract_usage(result), + "trtmc": { + "effective_max_tokens": max_tokens, + }, + } + + @app.post("/v1/audio/transcriptions") + async def audio_transcriptions( + file: UploadFile = File(...), + model: str | None = Form(default=None), + language: str | None = Form(default=None), + response_format: str = Form(default="json"), + ) -> Any: + spec = registry.resolve_model("transcription", model) + if response_format not in {"json", "text", "verbose_json"}: + return _error_response( + 400, + "unsupported_response_format", + "response_format must be json, verbose_json, or text", + param="response_format", + ) + + suffix = Path(file.filename or "audio.bin").suffix[:16] + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + prefix="trtmc-serve-audio-", suffix=suffix, delete=False + ) as temporary: + temporary_path = Path(temporary.name) + total = 0 + while chunk := await file.read(1024 * 1024): + total += len(chunk) + if total > server_config.max_upload_bytes: + return _error_response( + 413, + "audio_file_too_large", + f"audio upload exceeds {server_config.max_upload_bytes} bytes", + param="file", + ) + temporary.write(chunk) + if total == 0: + return _error_response( + 400, "empty_audio", "uploaded audio file is empty", param="file" + ) + transcription_config: dict[str, Any] = {} + if language: + transcription_config["language"] = language + _spec, session = registry.acquire_session("transcription", spec.name) + result = await _worker_request( + session, + "transcribe", + { + "audio_path": str(temporary_path), + "config": transcription_config, + }, + ) + finally: + await file.close() + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + + text = extract_text(result, operation="transcribe") + if response_format == "text": + return PlainTextResponse(text) + if response_format == "verbose_json" and isinstance(result, Mapping): + return { + "text": text, + "model": spec.name, + "segments": extract_transcription_segments(result), + } + return {"text": text} + + @app.websocket("/v1/realtime") + async def realtime(websocket: WebSocket) -> None: + if not _valid_websocket_origin(websocket.headers.get("origin")): + await websocket.close(code=4403, reason="websocket origin is not allowed") + return + if websocket.query_params.get("intent") != "transcription": + await websocket.close(code=4400, reason="intent must be transcription") + return + if server_config.api_key is not None and not _valid_websocket_token( + websocket, server_config.api_key + ): + await websocket.close(code=4401, reason="missing or invalid access token") + return + await websocket.accept() + connection = RealtimeTranscriptionConnection( + websocket, + registry, + max_audio_chunk_bytes=server_config.max_realtime_chunk_bytes, + max_session_audio_bytes=server_config.max_realtime_session_bytes, + idle_timeout_seconds=server_config.realtime_idle_timeout_seconds, + max_session_seconds=server_config.realtime_max_session_seconds, + ) + await connection.run() + + return app + + +async def _worker_request( + session: WorkerSession, + op: str, + payload: Mapping[str, Any], +) -> Any: + try: + task = asyncio.wrap_future(session.submit(op, payload)) + except BaseException: + session.close() + raise + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + + def release_when_done(completed: asyncio.Future[Any]) -> None: + session.close() + try: + completed.exception() + except BaseException: + pass + + task.add_done_callback(release_when_done) + raise + finally: + if task.done(): + session.close() + + +def _generation_config( + request: ChatCompletionRequest, + *, + max_tokens: int | None, +) -> dict[str, Any]: + config: dict[str, Any] = {"use_chat_template": True} + if max_tokens is not None: + config["max_new_tokens"] = max_tokens + for name in ( + "temperature", + "top_p", + "min_p", + "top_k", + "seed", + "enable_thinking", + ): + value = getattr(request, name, None) + if value is not None: + config[name] = value + return config + + +def _effective_generation_tokens( + registry: ModelRegistry, + model_name: str, + requested: int | None, + hard_cap: int, +) -> int: + if requested is not None: + return requested + default = registry.metadata_for(model_name).get("default_max_new_tokens", 128) + if isinstance(default, bool) or not isinstance(default, int) or default <= 0: + default = 128 + return min(default, hard_cap) + + +def _unsupported_chat_parameter( + request: ChatCompletionRequest, + messages: list[Mapping[str, Any]], +) -> str | None: + """Reject meaningful OpenAI options that this server cannot honor.""" + + payload = model_to_dict(request, exclude_none=True) + supported = { + "model", + "messages", + "max_tokens", + "max_completion_tokens", + "temperature", + "top_p", + "min_p", + "top_k", + "seed", + "enable_thinking", + "stream", + } + no_op = { + "n", + "best_of", + "logprobs", + "top_logprobs", + "frequency_penalty", + "presence_penalty", + "logit_bias", + "ignore_eos", + "tools", + "tool_choice", + "parallel_tool_calls", + "response_format", + "stream_options", + "stop", + } + metadata = {"user", "metadata"} + unknown = sorted(set(payload) - supported - no_op - metadata) + if unknown: + return unknown[0] + for name in sorted(no_op): + if name in payload and not _is_no_op_chat_parameter(name, payload[name]): + return name + if len(messages) != 1: + return "messages" + message = messages[0] + if ( + set(message) != {"role", "content"} + or message["role"] != "user" + or not is_text_only_content(message["content"]) + ): + return "messages" + return None + + +def _is_no_op_chat_parameter(name: str, value: Any) -> bool: + if name in {"n", "best_of"}: + return isinstance(value, int) and not isinstance(value, bool) and value == 1 + if name == "top_logprobs": + return isinstance(value, int) and not isinstance(value, bool) and value == 0 + if name in {"frequency_penalty", "presence_penalty"}: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value == 0 + if name in {"logprobs", "ignore_eos", "parallel_tool_calls"}: + return value is False + if name == "logit_bias": + return isinstance(value, Mapping) and not value + if name == "tools": + return isinstance(value, list) and not value + if name == "tool_choice": + return value == "none" + if name == "response_format": + return value == {"type": "text"} + if name == "stream_options": + return isinstance(value, Mapping) and not value + return False + + +def _valid_authorization(header: str | None, expected: str) -> bool: + if header is None: + return False + scheme, separator, token = header.partition(" ") + return bool(separator) and scheme.lower() == "bearer" and _tokens_match(token, expected) + + +def _tokens_match(candidate: str, expected: str) -> bool: + """Compare arbitrary text tokens without leaking timing or raising on Unicode.""" + + try: + return hmac.compare_digest(candidate.encode("utf-8"), expected.encode("utf-8")) + except UnicodeEncodeError: + return False + + +def _valid_websocket_token(websocket: WebSocket, expected: str) -> bool: + query_token = websocket.query_params.get("access_token") + if query_token is not None and _tokens_match(query_token, expected): + return True + return _valid_authorization(websocket.headers.get("authorization"), expected) + + +def _valid_websocket_origin(origin: str | None) -> bool: + """Apply the HTTP CORS boundary to browser WebSocket handshakes.""" + + if origin is None or origin == "null": + return True + try: + parsed = urlsplit(origin) + return parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1"} + except ValueError: + return False + + +def _error_response( + status_code: int, + code: str, + message: str, + *, + error_type: str = "invalid_request_error", + param: str | None = None, + headers: Mapping[str, str] | None = None, +) -> JSONResponse: + return JSONResponse( + status_code=status_code, + content={ + "error": { + "message": message, + "type": error_type, + "param": param, + "code": code, + } + }, + headers=dict(headers or {}), + ) + + +def _log_worker_failure() -> None: + _LOGGER.error("Model worker request failed") diff --git a/server/python/trtmc_server/cli.py b/server/python/trtmc_server/cli.py new file mode 100644 index 0000000000..9280baf518 --- /dev/null +++ b/server/python/trtmc_server/cli.py @@ -0,0 +1,465 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Command-line entry point for the Python serving control plane.""" + +from __future__ import annotations + +import argparse +import ipaddress +import json +import logging +import math +import os +import re +import shutil +import socket +import sys +import threading +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from .registry import ModelRegistry, ModelSpec +from .worker import WorkerLoadOptions + + +_REDACTED = "" +_SENSITIVE_LOG_KEYS = frozenset({"access_token", "authorization", "cookie"}) +_QUERY_STRING_PATTERN = re.compile(r"\?[^\s\"']+") +_SECRET_PATTERNS = ( + re.compile(r"([?&]access_token=)[^&\s\"']+", re.IGNORECASE), + re.compile( + r"(\baccess_token\b\s*(?:=|:)\s*(?:\"|')?)[^&,\s\"'}]+", + re.IGNORECASE, + ), + re.compile( + r"(\bauthorization\b\s*(?:=|:)\s*(?:\"|')?)(?:bearer\s+)?[^,\r\n\"'}]+", + re.IGNORECASE, + ), + re.compile( + r"(\bcookie\b\s*(?:=|:)\s*(?:\"|')?)[^,\r\n\"'}]+", + re.IGNORECASE, + ), +) + + +def _redact_text(value: str) -> str: + value = _QUERY_STRING_PATTERN.sub(f"?{_REDACTED}", value) + for pattern in _SECRET_PATTERNS: + value = pattern.sub(rf"\1{_REDACTED}", value) + return value + + +def _is_sensitive_log_key(value: object) -> bool: + if isinstance(value, bytes): + normalized = value.decode("ascii", errors="ignore") + elif isinstance(value, str): + normalized = value + else: + return False + return normalized.strip().lower().replace("-", "_") in _SENSITIVE_LOG_KEYS + + +def _redacted_like(value: object) -> str | bytes | bytearray: + if isinstance(value, bytes): + return _REDACTED.encode("ascii") + if isinstance(value, bytearray): + return bytearray(_REDACTED, "ascii") + return _REDACTED + + +def _redact_log_value(value: Any) -> Any: + if isinstance(value, str): + return _redact_text(value) + if isinstance(value, bytes): + return _redact_text(value.decode("utf-8", errors="replace")).encode("utf-8") + if isinstance(value, bytearray): + redacted = _redact_text(bytes(value).decode("utf-8", errors="replace")) + return bytearray(redacted, "utf-8") + if isinstance(value, Mapping): + return { + key: _redacted_like(item) if _is_sensitive_log_key(key) else _redact_log_value(item) + for key, item in value.items() + } + if isinstance(value, Sequence): + items = list(value) + if len(items) == 2 and _is_sensitive_log_key(items[0]): + items[1] = _redacted_like(items[1]) + else: + items = [_redact_log_value(item) for item in items] + if isinstance(value, tuple): + return tuple(items) + if isinstance(value, list): + return items + return items + return value + + +class _RedactAccessToken(logging.Filter): + """Recursively remove transport credentials from structured log records.""" + + def filter(self, record: logging.LogRecord) -> bool: + record.msg = _redact_log_value(record.msg) + record.args = _redact_log_value(record.args) + return True + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="trtmc serve", + description="Serve TensorRT-Model-Connect bundles over local HTTP and Realtime APIs", + ) + parser.add_argument( + "--chat-model", + action="append", + default=[], + metavar="NAME=PATH", + help="Register a text-generation bundle (repeatable)", + ) + parser.add_argument( + "--transcription-model", + action="append", + default=[], + metavar="NAME=PATH", + help="Register an offline/realtime transcription bundle (repeatable)", + ) + parser.add_argument( + "--default-chat-model", + metavar="NAME", + help="Default model for the chat endpoint (first registered chat model otherwise)", + ) + parser.add_argument( + "--default-transcription-model", + metavar="NAME", + help="Default model for audio endpoints (first registered transcription model otherwise)", + ) + parser.add_argument( + "--trtmc-binary", + metavar="PATH", + help="Native trtmc executable (injected by `trtmc serve`; PATH lookup otherwise)", + ) + parser.add_argument( + "--host", + default="127.0.0.1", + help="Bind loopback IP literal (default: 127.0.0.1)", + ) + parser.add_argument( + "--port", + type=_port, + default=8000, + help="Bind port; 0 selects a free port (default: 8000)", + ) + parser.add_argument( + "--api-key", + help="Bearer token; falls back to TRTMC_SERVE_TOKEN", + ) + parser.add_argument( + "--require-streaming-transcription", + action="append", + default=[], + metavar="MODEL", + help="Fail startup unless every MODEL replica proves native streaming support (repeatable)", + ) + parser.add_argument( + "--model-replicas", + action="append", + default=[], + metavar="NAME=N", + help="Run N native worker replicas for MODEL (repeatable; default: 1)", + ) + parser.add_argument( + "--runtime-root", + metavar="DIR", + help="Directory containing the exact family and backend runtime DSOs", + ) + parser.add_argument("--kv-cache-size", type=_positive_int, metavar="BYTES") + parser.add_argument("--runtime-cache", metavar="PATH") + parser.add_argument("--cuda-graphs", action="store_true") + parser.add_argument("--startup-timeout", type=_positive_float, default=120.0) + parser.add_argument("--request-timeout", type=_positive_float, default=120.0) + parser.add_argument("--max-generation-tokens", type=_positive_int, default=4096) + parser.add_argument("--realtime-idle-timeout", type=_positive_float, default=30.0) + parser.add_argument("--realtime-max-session-seconds", type=_positive_float, default=4 * 60 * 60) + parser.add_argument("--realtime-max-audio-bytes", type=_positive_int, default=512 * 1024 * 1024) + parser.add_argument( + "--parent-liveness-stdin", + action="store_true", + help="Gracefully stop when the parent-owned stdin pipe reaches EOF", + ) + parser.add_argument( + "--log-level", + choices=("critical", "error", "warning", "info"), + default="info", + ) + parser.add_argument( + "--access-log", + action="store_true", + help="Enable HTTP access logs (disabled by default to avoid logging WS tokens)", + ) + return parser + + +def parse_model_assignment(value: str, *, kind: str) -> ModelSpec: + name, separator, raw_path = value.partition("=") + if not separator or not name or not raw_path: + raise ValueError(f"--{kind}-model must use NAME=PATH") + bundle = Path(raw_path).expanduser().resolve() + if not bundle.is_file(): + raise ValueError(f"bundle for model {name!r} does not exist or is not a file") + model_kind = "chat" if kind == "chat" else "transcription" + return ModelSpec(name=name, bundle=bundle, kind=model_kind) # type: ignore[arg-type] + + +def parse_replica_assignment(value: str) -> tuple[str, int]: + name, separator, raw_replicas = value.partition("=") + if not separator or not name or not raw_replicas: + raise ValueError("--model-replicas must use NAME=N") + try: + replicas = int(raw_replicas) + except ValueError as exc: + raise ValueError(f"replicas for model {name!r} must be a positive integer") from exc + if replicas <= 0: + raise ValueError(f"replicas for model {name!r} must be a positive integer") + return name, replicas + + +def validate_bind_policy(host: str) -> None: + if not is_loopback_host(host): + raise ValueError(f"refusing host {host!r}; --host must be a loopback IP literal") + + +def is_loopback_host(host: str) -> bool: + normalized = host.strip().lower() + try: + return ipaddress.ip_address(normalized).is_loopback + except ValueError: + return False + + +def resolve_trtmc_binary(value: str | None) -> Path: + if value: + candidate = Path(value).expanduser() + if candidate.is_file(): + return candidate.resolve() + discovered = shutil.which(value) + if discovered: + return Path(discovered).resolve() + raise ValueError("trtmc executable does not exist or is not a file") + discovered = shutil.which("trtmc") + if discovered: + return Path(discovered).resolve() + raise ValueError("cannot find trtmc on PATH; native forwarding must pass --trtmc-binary") + + +def resolve_runtime_root(value: str | None) -> Path: + if value: + candidate = Path(value).expanduser() + if candidate.is_dir(): + return candidate.resolve() + raise ValueError("--runtime-root must name an existing directory") + + +def bind_socket(host: str, port: int) -> socket.socket: + """Pre-bind the listening socket so port 0 can be reported without a race.""" + + errors: list[OSError] = [] + addresses = socket.getaddrinfo( + host, + port, + family=socket.AF_UNSPEC, + type=socket.SOCK_STREAM, + flags=socket.AI_PASSIVE, + ) + for family, socktype, protocol, _canonical_name, sockaddr in addresses: + listener = socket.socket(family, socktype, protocol) + try: + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind(sockaddr) + listener.listen(2048) + listener.setblocking(False) + return listener + except OSError as exc: + errors.append(exc) + listener.close() + detail = "; ".join(str(error) for error in errors) or "no addresses resolved" + raise OSError(f"cannot bind {host}:{port}: {detail}") + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + + try: + specs = [ + *(parse_model_assignment(value, kind="chat") for value in args.chat_model), + *( + parse_model_assignment(value, kind="transcription") + for value in args.transcription_model + ), + ] + if not specs: + raise ValueError("at least one --chat-model or --transcription-model is required") + + replica_assignments = [parse_replica_assignment(value) for value in args.model_replicas] + replica_names = [name for name, _replicas in replica_assignments] + if len(replica_names) != len(set(replica_names)): + raise ValueError("--model-replicas may be specified only once per model") + model_replicas = dict(replica_assignments) + + trtmc_binary = resolve_trtmc_binary(args.trtmc_binary) + runtime_root = resolve_runtime_root(args.runtime_root) + api_key = args.api_key + if api_key is None: + api_key = os.environ.get("TRTMC_SERVE_TOKEN") + if api_key is not None: + api_key = api_key.strip() + if not api_key: + api_key = None + if api_key is None: + raise ValueError("authentication requires --api-key or TRTMC_SERVE_TOKEN") + validate_bind_policy(args.host) + load_options = WorkerLoadOptions( + runtime_root=str(runtime_root), + kv_cache_size_bytes=args.kv_cache_size, + runtime_cache=args.runtime_cache, + cuda_graphs=args.cuda_graphs, + ) + registry = ModelRegistry( + specs, + trtmc_binary=trtmc_binary, + default_chat_model=args.default_chat_model, + default_transcription_model=args.default_transcription_model, + startup_timeout=args.startup_timeout, + request_timeout=args.request_timeout, + load_options=load_options, + model_replicas=model_replicas, + required_streaming_transcription=args.require_streaming_transcription, + ) + except ValueError as exc: + parser.error(str(exc)) + + try: + import uvicorn + + from .app import ServerConfig, create_app + except ModuleNotFoundError as exc: + if exc.name in {"fastapi", "multipart", "uvicorn"}: + parser.error("serve dependencies are missing; install 'tensorrt-model-connect[serve]'") + raise + try: + app = create_app( + registry, + config=ServerConfig( + api_key=api_key, + max_generation_tokens=args.max_generation_tokens, + max_realtime_session_bytes=args.realtime_max_audio_bytes, + realtime_idle_timeout_seconds=args.realtime_idle_timeout, + realtime_max_session_seconds=args.realtime_max_session_seconds, + ), + ) + except RuntimeError as exc: + if "python-multipart" in str(exc): + parser.error("serve dependencies are missing; install 'tensorrt-model-connect[serve]'") + raise + try: + listener = bind_socket(args.host, args.port) + except OSError as exc: + parser.error(str(exc)) + actual_port = int(listener.getsockname()[1]) + + try: + + class ReadyServer(uvicorn.Server): + _ready_emitted = False + + async def startup(self, sockets: list[socket.socket] | None = None) -> None: + await super().startup(sockets=sockets) + if self.started and not self._ready_emitted: + self._ready_emitted = True + ready = { + "event": "ready", + "host": args.host, + "port": actual_port, + "models": registry.names, + } + print( + json.dumps(ready, separators=(",", ":"), ensure_ascii=False), + file=sys.stdout, + flush=True, + ) + + uvicorn_config = uvicorn.Config( + app, + host=args.host, + port=actual_port, + log_level=args.log_level, + access_log=args.access_log, + ) + _install_secret_log_filter() + server = ReadyServer(uvicorn_config) + if args.parent_liveness_stdin: + threading.Thread( + target=_watch_parent_stdin, + args=(server,), + name="trtmc-parent-liveness", + daemon=True, + ).start() + server.run(sockets=[listener]) + return 0 if server.started else 1 + finally: + listener.close() + + +def _port(value: str) -> int: + try: + port = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("port must be an integer") from exc + if not 0 <= port <= 65535: + raise argparse.ArgumentTypeError("port must be from 0 to 65535") + return port + + +def _positive_float(value: str) -> float: + try: + parsed = float(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be a number") from exc + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError("must be a finite positive number") + return parsed + + +def _positive_int(value: str) -> int: + try: + parsed = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be an integer") from exc + if parsed <= 0: + raise argparse.ArgumentTypeError("must be positive") + return parsed + + +def _watch_parent_stdin(server: Any) -> None: + try: + while sys.stdin.buffer.read(8192): + pass + except (OSError, ValueError): + pass + server.should_exit = True + + +def _install_secret_log_filter() -> None: + redactor = _RedactAccessToken() + for logger_name in ("uvicorn", "uvicorn.error", "uvicorn.access"): + logger = logging.getLogger(logger_name) + logger.addFilter(redactor) + for handler in logger.handlers: + handler.addFilter(redactor) + if logger_name == "uvicorn.access" and isinstance(handler, logging.StreamHandler): + handler.setStream(sys.stderr) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/server/python/trtmc_server/errors.py b/server/python/trtmc_server/errors.py new file mode 100644 index 0000000000..d3d636e83f --- /dev/null +++ b/server/python/trtmc_server/errors.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Errors raised by the TensorRT-Model-Connect serving control plane.""" + +from __future__ import annotations + +from typing import Any + + +class ServeError(RuntimeError): + """Base class for expected serving failures.""" + + code = "serve_error" + + +class WorkerError(ServeError): + """Base class for native worker failures.""" + + code = "worker_error" + + +class WorkerStartupError(WorkerError): + """A worker did not complete its ready handshake.""" + + code = "worker_startup_failed" + + +class WorkerCrashedError(WorkerError): + """A worker exited or lost its protocol stream.""" + + code = "worker_crashed" + + +class WorkerTimeoutError(WorkerError): + """A worker operation exceeded its deadline.""" + + code = "worker_timeout" + + +class WorkerSaturatedError(WorkerError): + """No native worker replica is immediately available.""" + + code = "server_busy" + + +class WorkerProtocolError(WorkerError): + """A worker emitted malformed or unexpected JSONL.""" + + code = "worker_protocol_error" + + +class WorkerRequestTooLargeError(WorkerError): + """A serialized request would exceed the native JSONL line limit.""" + + code = "request_too_large" + + +class WorkerRemoteError(WorkerError): + """A worker returned a structured operation error.""" + + code = "worker_operation_failed" + + def __init__(self, message: str, *, details: Any = None) -> None: + super().__init__(message) + self.details = details + + +class ModelNotFoundError(ServeError): + """The requested model name is not registered.""" + + code = "model_not_found" + + +class ModelCapabilityError(ServeError): + """The requested model does not implement the required API capability.""" + + code = "model_capability_mismatch" diff --git a/server/python/trtmc_server/protocol.py b/server/python/trtmc_server/protocol.py new file mode 100644 index 0000000000..55f59e6372 --- /dev/null +++ b/server/python/trtmc_server/protocol.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Translation helpers between API envelopes and the native worker protocol.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from typing import Any + +from .errors import ( + WorkerCrashedError, + WorkerError, + WorkerProtocolError, + WorkerRemoteError, + WorkerRequestTooLargeError, + WorkerSaturatedError, + WorkerTimeoutError, +) + + +def invalid_request_message(error: WorkerRemoteError) -> str | None: + """Return the native message only for client-caused worker failures.""" + + details = error.details + if not isinstance(details, Mapping) or details.get("type") != "invalid_request_error": + return None + message = details.get("message") + if isinstance(message, str) and message: + return message + return "The model worker rejected the request" + + +def public_worker_error_message(error: WorkerError) -> str: + """Return a stable public message without worker diagnostics.""" + + if isinstance(error, WorkerTimeoutError): + return "The model worker timed out" + if isinstance(error, WorkerCrashedError): + return "The model worker is unavailable" + if isinstance(error, WorkerProtocolError): + return "The model worker returned an invalid response" + if isinstance(error, WorkerRequestTooLargeError): + return "The request exceeds the model worker transport limit" + if isinstance(error, WorkerSaturatedError): + return "All model worker replicas are busy" + return "The model worker operation failed" + + +def extract_text(result: Any, *, operation: str) -> str: + """Extract text from the private v3 worker result.""" + + if isinstance(result, Mapping) and isinstance(result.get("text"), str): + return str(result["text"]) + raise WorkerProtocolError(f"worker {operation!r} result did not contain a string text field") + + +def extract_transcription_segments(result: Any) -> list[dict[str, float | str]]: + """Copy only the public fields from native transcription segments.""" + + if not isinstance(result, Mapping): + raise WorkerProtocolError("worker transcription result was not a JSON object") + raw_segments = result.get("segments", []) + if not isinstance(raw_segments, list): + raise WorkerProtocolError("worker transcription segments were not a JSON array") + + segments: list[dict[str, float | str]] = [] + for index, raw_segment in enumerate(raw_segments): + if not isinstance(raw_segment, Mapping): + raise WorkerProtocolError(f"worker transcription segment {index} was not an object") + start = raw_segment.get("start_seconds") + end = raw_segment.get("end_seconds") + text = raw_segment.get("text") + if ( + isinstance(start, bool) + or not isinstance(start, (int, float)) + or not math.isfinite(start) + or isinstance(end, bool) + or not isinstance(end, (int, float)) + or not math.isfinite(end) + or not isinstance(text, str) + ): + raise WorkerProtocolError( + f"worker transcription segment {index} has invalid public fields" + ) + segments.append( + { + "start_seconds": float(start), + "end_seconds": float(end), + "text": text, + } + ) + return segments + + +def extract_usage(result: Any) -> dict[str, int]: + if not isinstance(result, Mapping): + return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + usage = result.get("usage") + source = usage if isinstance(usage, Mapping) else result + prompt = _non_negative_int(source.get("prompt_tokens")) + completion = _non_negative_int(source.get("completion_tokens", source.get("generated_tokens"))) + total = _non_negative_int(source.get("total_tokens")) or prompt + completion + return { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": total, + } + + +def prepare_chat_prompt( + messages: list[Mapping[str, Any]], +) -> str: + """Return the one chat shape the native runtime can honor exactly.""" + + if ( + len(messages) != 1 + or messages[0].get("role") != "user" + or not is_text_only_content(messages[0].get("content")) + ): + raise ValueError("messages must contain exactly one text-only user message") + return _render_content(messages[0].get("content")) + + +def is_text_only_content(content: Any) -> bool: + if isinstance(content, str): + return True + if not isinstance(content, list): + return False + return all( + isinstance(part, str) + or ( + isinstance(part, Mapping) + and set(part) <= {"text", "type"} + and isinstance(part.get("type", "text"), str) + and part.get("type", "text") in {"text", "input_text"} + and isinstance(part.get("text"), str) + ) + for part in content + ) + + +def _render_content(content: Any) -> str: + if isinstance(content, str): + return content + if not isinstance(content, list): + raise ValueError("message content must be text-only") + return "".join(part if isinstance(part, str) else str(part["text"]) for part in content) + + +def _non_negative_int(value: Any) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, int) and value >= 0: + return value + if isinstance(value, float) and math.isfinite(value) and value >= 0: + return int(value) + return 0 diff --git a/server/python/trtmc_server/realtime.py b/server/python/trtmc_server/realtime.py new file mode 100644 index 0000000000..01e877992e --- /dev/null +++ b/server/python/trtmc_server/realtime.py @@ -0,0 +1,523 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OpenAI Realtime-style transcription WebSocket session handling.""" + +from __future__ import annotations + +import asyncio +import base64 +import binascii +import logging +import traceback +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from fastapi import WebSocket, WebSocketDisconnect + +from .errors import ( + ModelNotFoundError, + ServeError, + WorkerError, + WorkerRemoteError, + WorkerRequestTooLargeError, + WorkerSaturatedError, +) +from .protocol import extract_text, invalid_request_message, public_worker_error_message +from .registry import ModelRegistry +from .worker import WorkerSession + + +_LOGGER = logging.getLogger(__name__) +_TIMEOUT_ERRORS = (TimeoutError, asyncio.TimeoutError) + + +class RealtimeTranscriptionConnection: + """Translate browser realtime events into one exclusive native stream.""" + + def __init__( + self, + websocket: WebSocket, + registry: ModelRegistry, + *, + max_audio_chunk_bytes: int = 1024 * 1024, + max_session_audio_bytes: int = 512 * 1024 * 1024, + idle_timeout_seconds: float = 30.0, + max_session_seconds: float = 4 * 60 * 60, + ) -> None: + self.websocket = websocket + self.registry = registry + self.max_audio_chunk_bytes = max_audio_chunk_bytes + self.max_session_audio_bytes = max_session_audio_bytes + self.idle_timeout_seconds = idle_timeout_seconds + self.max_session_seconds = max_session_seconds + self.connection_id = f"sess_{uuid.uuid4().hex}" + self.item_id = f"item_{uuid.uuid4().hex}" + self.model = registry.default_transcription_model + self.language: str | None = None + self.sample_rate_hz = 24000 + self.channels = 1 + self.transcript = "" + self.total_audio_bytes = 0 + self._lease: WorkerSession | None = None + self._stop_requested = False + + async def run(self) -> None: + await self._send( + { + "type": "session.created", + "event_id": self._event_id(), + "session": self._session_payload(), + } + ) + loop = asyncio.get_running_loop() + started_at = loop.time() + try: + while True: + elapsed = loop.time() - started_at + remaining = self.max_session_seconds - elapsed + if remaining <= 0: + await self._send_failure( + "session_duration_exceeded", + f"realtime session exceeded {self.max_session_seconds:g}s", + error_type="invalid_request_error", + ) + break + try: + event = await asyncio.wait_for( + self.websocket.receive_json(), + timeout=min(self.idle_timeout_seconds, remaining), + ) + except _TIMEOUT_ERRORS: + elapsed = loop.time() - started_at + if elapsed >= self.max_session_seconds: + code = "session_duration_exceeded" + message = f"realtime session exceeded {self.max_session_seconds:g}s" + else: + code = "session_idle_timeout" + message = f"no realtime event received for {self.idle_timeout_seconds:g}s" + await self._send_failure( + code, + message, + error_type="invalid_request_error", + ) + break + if not isinstance(event, Mapping): + await self._send_error("invalid_event", "event must be a JSON object") + continue + await self._dispatch(event) + if self._stop_requested: + break + except WebSocketDisconnect: + pass + except (ValueError, TypeError) as exc: + await self._send_error("invalid_event", str(exc)) + except WorkerError as exc: + await self._send_failed(exc) + except ServeError as exc: + await self._send_failure(exc.code, str(exc)) + except Exception as exc: + frames = traceback.extract_tb(exc.__traceback__, limit=-8) + trace = " <- ".join( + f"{Path(frame.filename).name}:{frame.lineno}:{frame.name}" for frame in frames + ) + _LOGGER.error( + "Unexpected realtime session failure session=%s type=%s trace=%s", + self.connection_id, + type(exc).__name__, + trace, + ) + await self._send_failure("internal_error", "realtime transcription session failed") + finally: + try: + await self._release_stream(reset=True) + except asyncio.CancelledError: + # Client shutdown may cancel the endpoint while the native + # reset is being joined. Cleanup is complete at this point, so + # let the WebSocket handler return normally. + pass + + async def _dispatch(self, event: Mapping[str, Any]) -> None: + event_type = event.get("type") + if event_type == "session.update": + await self._update_session(event) + elif event_type == "input_audio_buffer.append": + await self._append(event) + elif event_type == "input_audio_buffer.commit": + await self._commit() + elif event_type == "input_audio_buffer.clear": + await self._clear() + else: + await self._send_error( + "unsupported_event", + f"unsupported realtime event type: {event_type!r}", + ) + + async def _update_session(self, event: Mapping[str, Any]) -> None: + raw_session = event.get("session") + if not isinstance(raw_session, Mapping): + await self._send_error("invalid_session", "session.update requires session object") + return + + transcription = raw_session.get("input_audio_transcription") + transcription_options = transcription if isinstance(transcription, Mapping) else {} + requested_model = raw_session.get("model", transcription_options.get("model")) + if requested_model is not None and not isinstance(requested_model, str): + await self._send_error("invalid_model", "session model must be a string") + return + next_model = requested_model or self.model + try: + spec = self.registry.resolve_model("transcription", next_model) + except ServeError as exc: + await self._send_error(exc.code, str(exc)) + return + + audio_format = raw_session.get("input_audio_format", "pcm16") + if isinstance(audio_format, Mapping): + format_name = audio_format.get("type", "pcm16") + format_rate = audio_format.get("rate") + else: + format_name = audio_format + format_rate = None + if str(format_name).lower() not in {"pcm16", "audio/pcm", "audio/pcm16"}: + await self._send_error( + "unsupported_audio_format", "only little-endian PCM16 audio is supported" + ) + return + + trtmc_options = raw_session.get("trtmc") + if not isinstance(trtmc_options, Mapping): + trtmc_options = {} + sample_rate = raw_session.get( + "sample_rate_hz", + trtmc_options.get("sample_rate_hz", format_rate or self.sample_rate_hz), + ) + if not isinstance(sample_rate, int) or not 8000 <= sample_rate <= 192000: + await self._send_error( + "invalid_sample_rate", "sample_rate_hz must be an integer from 8000 to 192000" + ) + return + language = raw_session.get("language", transcription_options.get("language")) + if language is not None and not isinstance(language, str): + await self._send_error("invalid_language", "language must be a string") + return + if isinstance(language, str) and len(language.encode("utf-8")) > 128: + await self._send_error("invalid_language", "language exceeds 128 UTF-8 bytes") + return + + if self._lease is not None and ( + spec.name != self.model + or sample_rate != self.sample_rate_hz + or language != self.language + ): + await self._release_stream(reset=True) + self.model = spec.name + self.sample_rate_hz = sample_rate + self.language = language + await self._send( + { + "type": "session.updated", + "event_id": self._event_id(), + "session": self._session_payload(), + } + ) + + async def _append(self, event: Mapping[str, Any]) -> None: + encoded = event.get("audio") + if not isinstance(encoded, str) or not encoded: + await self._send_error("invalid_audio", "append requires non-empty base64 audio") + return + try: + audio = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError): + await self._send_error("invalid_audio", "audio is not valid base64") + return + if len(audio) > self.max_audio_chunk_bytes: + await self._send_error( + "audio_chunk_too_large", + f"decoded audio chunk exceeds {self.max_audio_chunk_bytes} bytes", + ) + return + if len(audio) % 2: + await self._send_error( + "invalid_audio", "PCM16 audio must contain an even number of bytes" + ) + return + if self.total_audio_bytes + len(audio) > self.max_session_audio_bytes: + await self._send_failure( + "audio_session_too_large", + f"realtime session audio exceeds {self.max_session_audio_bytes} bytes", + error_type="invalid_request_error", + ) + await self._release_stream(reset=True) + await self.websocket.close(code=1009, reason="audio session limit exceeded") + self._stop_requested = True + return + self.total_audio_bytes += len(audio) + + try: + lease = await self._ensure_stream() + result = await _lease_request( + lease, + "stream_chunk", + {"audio": encoded}, + ) + cumulative = extract_text(result, operation="stream_chunk") + except WorkerRequestTooLargeError as exc: + await self._send_error(exc.code, public_worker_error_message(exc)) + await self._release_stream(reset=True) + return + except WorkerSaturatedError as exc: + await self._send_error( + exc.code, + public_worker_error_message(exc), + error_type="rate_limit_error", + ) + await self._release_stream(reset=True) + return + except WorkerRemoteError as exc: + message = invalid_request_message(exc) + if message is not None: + await self._send_error("invalid_request", message) + else: + await self._send_failed(exc) + await self._release_stream(reset=True) + return + except WorkerError as exc: + await self._send_failed(exc) + await self._release_stream(reset=True) + return + except ServeError as exc: + await self._send_error(exc.code, str(exc)) + return + + delta = ( + cumulative[len(self.transcript) :] + if cumulative.startswith(self.transcript) + else cumulative + ) + self.transcript = cumulative + await self._send( + { + "type": "conversation.item.input_audio_transcription.delta", + "event_id": self._event_id(), + "item_id": self.item_id, + "content_index": 0, + "delta": delta, + "transcript": self.transcript, + } + ) + + async def _commit(self) -> None: + try: + lease = await self._ensure_stream() + result = await _lease_request(lease, "stream_finish") + self.transcript = extract_text(result, operation="stream_finish") + except WorkerRequestTooLargeError as exc: + await self._send_error(exc.code, public_worker_error_message(exc)) + await self._release_stream(reset=False) + return + except WorkerSaturatedError as exc: + await self._send_error( + exc.code, + public_worker_error_message(exc), + error_type="rate_limit_error", + ) + await self._release_stream(reset=False) + return + except WorkerRemoteError as exc: + message = invalid_request_message(exc) + if message is not None: + await self._send_error("invalid_request", message) + else: + await self._send_failed(exc) + await self._release_stream(reset=False) + return + except WorkerError as exc: + await self._send_failed(exc) + await self._release_stream(reset=False) + return + except ServeError as exc: + await self._send_error(exc.code, str(exc)) + return + + await self._send( + { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": self._event_id(), + "item_id": self.item_id, + "content_index": 0, + "transcript": self.transcript, + } + ) + await self._release_stream(reset=False) + self.transcript = "" + self.item_id = f"item_{uuid.uuid4().hex}" + + async def _clear(self) -> None: + await self._release_stream(reset=True) + await self._send( + { + "type": "input_audio_buffer.cleared", + "event_id": self._event_id(), + } + ) + + async def _ensure_stream(self) -> WorkerSession: + if self._lease is not None: + return self._lease + if self.model is None: + raise ModelNotFoundError("no default transcription model is configured") + _spec, lease = self.registry.acquire_session("transcription", self.model) + config: dict[str, Any] = { + "sample_rate_hz": self.sample_rate_hz, + "channels": self.channels, + "audio_format": "pcm16le", + } + if self.language: + config["language"] = self.language + try: + await _lease_request(lease, "stream_start", {"config": config}) + except BaseException: + try: + await _lease_request( + lease, + "stream_reset", + timeout=min(2.0, self.registry.request_timeout), + ) + except BaseException: + pass + finally: + lease.close() + raise + self._lease = lease + return lease + + async def _release_stream(self, *, reset: bool) -> None: + lease = self._lease + self._lease = None + if reset: + self.transcript = "" + self.item_id = f"item_{uuid.uuid4().hex}" + if lease is None: + return + try: + if reset: + try: + await _lease_request( + lease, + "stream_reset", + timeout=min(2.0, self.registry.request_timeout), + ) + except WorkerError: + pass + finally: + lease.close() + + async def _send_failed(self, error: WorkerError) -> None: + _log_worker_failure() + await self._send_failure(error.code, public_worker_error_message(error)) + + async def _send_failure( + self, + code: str, + message: str, + *, + error_type: str = "server_error", + ) -> None: + await self._send( + { + "type": "conversation.item.input_audio_transcription.failed", + "event_id": self._event_id(), + "item_id": self.item_id, + "content_index": 0, + "transcript": self.transcript, + "error": { + "type": error_type, + "code": code, + "message": message, + }, + } + ) + + async def _send_error( + self, + code: str, + message: str, + *, + error_type: str = "invalid_request_error", + ) -> None: + await self._send( + { + "type": "error", + "event_id": self._event_id(), + "error": { + "type": error_type, + "code": code, + "message": message, + }, + } + ) + + async def _send(self, payload: dict[str, Any]) -> None: + try: + await self.websocket.send_json(payload) + except (RuntimeError, WebSocketDisconnect): + pass + + def _session_payload(self) -> dict[str, Any]: + return { + "id": self.connection_id, + "object": "realtime.transcription_session", + "model": self.model, + "input_audio_format": "pcm16", + "input_audio_transcription": { + "model": self.model, + "language": self.language, + }, + "trtmc": { + "sample_rate_hz": self.sample_rate_hz, + "channels": self.channels, + "audio_bytes_received": self.total_audio_bytes, + "max_session_audio_bytes": self.max_session_audio_bytes, + "idle_timeout_seconds": self.idle_timeout_seconds, + "max_session_seconds": self.max_session_seconds, + }, + } + + @staticmethod + def _event_id() -> str: + return f"event_{uuid.uuid4().hex}" + + +async def _lease_request( + lease: WorkerSession, + operation: str, + payload: Mapping[str, Any] | None = None, + *, + timeout: float | None = None, +) -> Any: + """Join a native operation before propagating client cancellation.""" + + task = asyncio.wrap_future(lease.submit(operation, payload, timeout=timeout)) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError: + continue + except BaseException: + break + try: + task.result() + except BaseException: + pass + raise + + +def _log_worker_failure() -> None: + _LOGGER.error("Realtime model worker request failed") diff --git a/server/python/trtmc_server/registry.py b/server/python/trtmc_server/registry.py new file mode 100644 index 0000000000..4ae60c6a00 --- /dev/null +++ b/server/python/trtmc_server/registry.py @@ -0,0 +1,433 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model registration and native worker lifecycle management.""" + +from __future__ import annotations + +import re +import threading +import time +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +from .errors import ModelCapabilityError, ModelNotFoundError, WorkerProtocolError +from .worker import WorkerGroup, WorkerLoadOptions, WorkerProcess, WorkerSession + + +ModelKind = Literal["chat", "transcription"] +_WORKER_PROTOCOL_VERSION = 3 +_MODEL_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_PUBLIC_MODEL_METADATA_FIELDS = ( + "default_max_new_tokens", + "streaming_transcription", +) +_PUBLIC_MODEL_STATUS_FIELDS = ( + "ready", + "degraded", + "replicas", + "ready_replicas", + "idle_replicas", + "busy", +) + + +@dataclass(frozen=True) +class ModelSpec: + """One API-visible name backed by a fixed native worker group.""" + + name: str + bundle: Path + kind: ModelKind + + def __post_init__(self) -> None: + if _MODEL_NAME.fullmatch(self.name) is None: + raise ValueError( + "model name must start with an alphanumeric character and contain only " + "letters, digits, '.', '_', or '-' (maximum 128 characters)" + ) + if self.kind not in {"chat", "transcription"}: + raise ValueError(f"unsupported model kind: {self.kind!r}") + object.__setattr__(self, "bundle", Path(self.bundle)) + + +WorkerFactory = Callable[[ModelSpec], WorkerProcess] + + +class ModelRegistry: + """Own all model workers exposed by one server process.""" + + def __init__( + self, + specs: Iterable[ModelSpec], + *, + trtmc_binary: str | Path, + default_chat_model: str | None = None, + default_transcription_model: str | None = None, + startup_timeout: float = 120.0, + request_timeout: float = 120.0, + load_options: WorkerLoadOptions | None = None, + model_replicas: Mapping[str, int] | None = None, + required_streaming_transcription: Iterable[str] = (), + worker_factory: WorkerFactory | None = None, + ) -> None: + ordered_specs = list(specs) + names = [spec.name for spec in ordered_specs] + if len(names) != len(set(names)): + duplicates = sorted({name for name in names if names.count(name) > 1}) + raise ValueError(f"duplicate model name(s): {', '.join(duplicates)}") + + self._specs = {spec.name: spec for spec in ordered_specs} + self._order = names + self._trtmc_binary = Path(trtmc_binary) + self._startup_timeout = float(startup_timeout) + self._request_timeout = float(request_timeout) + if load_options is None and worker_factory is None: + raise ValueError("native workers require explicit runtime load options") + self._load_options = load_options + self._worker_factory = worker_factory + configured_replicas = dict(model_replicas or {}) + unknown_replicas = sorted(set(configured_replicas) - set(self._specs)) + if unknown_replicas: + raise ValueError( + "replicas configured for unknown model(s): " + ", ".join(unknown_replicas) + ) + for name, replicas in configured_replicas.items(): + if isinstance(replicas, bool) or not isinstance(replicas, int) or replicas <= 0: + raise ValueError(f"replicas for model {name!r} must be a positive integer") + self._replicas = {name: configured_replicas.get(name, 1) for name in names} + self._groups: dict[str, WorkerGroup] = {} + self._metadata: dict[str, dict[str, Any]] = {} + self._started = False + self._started_at: int | None = None + self._lock = threading.RLock() + self._required_streaming = set(required_streaming_transcription) + + for name in sorted(self._required_streaming): + spec = self._specs.get(name) + if spec is None: + raise ValueError( + f"required streaming transcription model {name!r} is not registered" + ) + if spec.kind != "transcription": + raise ValueError( + f"required streaming transcription model {name!r} is not transcription" + ) + + self.default_chat_model = self._resolve_default("chat", default_chat_model) + self.default_transcription_model = self._resolve_default( + "transcription", default_transcription_model + ) + + @property + def request_timeout(self) -> float: + return self._request_timeout + + @property + def names(self) -> list[str]: + return list(self._order) + + @property + def ready(self) -> bool: + with self._lock: + return ( + self._started + and bool(self._groups) + and len(self._groups) == len(self._specs) + and all(group.ready for group in self._groups.values()) + ) + + @property + def has_healthy_worker(self) -> bool: + """Return whether this process can still execute any model request.""" + + with self._lock: + return self._started and any(group.ready for group in self._groups.values()) + + def start(self) -> None: + """Start every configured worker replica.""" + + with self._lock: + if self._started: + return + started_workers: list[WorkerProcess] = [] + try: + for name in self._order: + spec = self._specs[name] + workers: list[WorkerProcess] = [] + ready_payload: dict[str, Any] | None = None + for replica in range(self._replicas[name]): + worker = self._make_worker(spec, replica) + worker.start() + workers.append(worker) + started_workers.append(worker) + payload = worker.ready_payload + self._validate_metadata(name, spec.kind, payload) + if ready_payload is None: + ready_payload = payload + elif payload != ready_payload: + raise WorkerProtocolError( + f"worker replicas for model {name!r} reported inconsistent metadata" + ) + + assert ready_payload is not None + capabilities = list(ready_payload["capabilities"]) + metadata = { + key: ready_payload[key] + for key in _PUBLIC_MODEL_METADATA_FIELDS + if key in ready_payload + } + metadata["capabilities"] = capabilities + metadata["streaming_transcription"] = ( + "transcription_streaming" in capabilities + ) + if ( + name in self._required_streaming + and not metadata["streaming_transcription"] + ): + raise WorkerProtocolError( + f"worker {name!r} does not advertise transcription_streaming" + ) + if name in self._required_streaming: + for worker in workers: + probe = worker.request( + "probe_transcription_stream", + { + "config": { + "sample_rate_hz": 16000, + "channels": 1, + "audio_format": "pcm16le", + } + }, + ) + if ( + not isinstance(probe, Mapping) + or set(probe) != {"supported"} + or probe["supported"] is not True + ): + raise WorkerProtocolError( + f"worker {worker.name!r} returned an invalid streaming probe" + ) + self._groups[name] = WorkerGroup(name, workers) + self._metadata[name] = metadata + self._started = True + self._started_at = int(time.time()) + except BaseException: + for worker in reversed(started_workers): + worker.close() + self._groups.clear() + self._metadata.clear() + self._started = False + raise + + def close(self) -> None: + with self._lock: + self._started = False + for group in reversed(self._groups.values()): + group.close() + self._groups.clear() + + def resolve_model( + self, + kind: ModelKind, + requested_name: str | None, + ) -> ModelSpec: + """Resolve an API-visible model without exposing worker transport state.""" + + return self._resolve(kind, requested_name)[0] + + def _resolve( + self, + kind: ModelKind, + requested_name: str | None, + ) -> tuple[ModelSpec, WorkerGroup]: + if requested_name is None: + requested_name = ( + self.default_chat_model if kind == "chat" else self.default_transcription_model + ) + if requested_name is None: + raise ModelNotFoundError(f"no default {kind} model is configured") + + spec = self._specs.get(requested_name) + if spec is None: + raise ModelNotFoundError(f"model {requested_name!r} is not registered") + if spec.kind != kind: + raise ModelCapabilityError( + f"model {requested_name!r} is registered for {spec.kind}, not {kind}" + ) + with self._lock: + group = self._groups.get(requested_name) + if group is None: + raise ModelNotFoundError(f"model {requested_name!r} is not ready") + return spec, group + + def acquire_session( + self, + kind: ModelKind, + requested_name: str | None, + ) -> tuple[ModelSpec, WorkerSession]: + """Resolve a model and lease one fixed worker lane atomically.""" + + with self._lock: + spec, group = self._resolve(kind, requested_name) + return spec, group.acquire_session() + + def list_models(self) -> list[dict[str, Any]]: + with self._lock: + created = self._started_at or int(time.time()) + return [ + { + "id": name, + "object": "model", + "created": created, + "owned_by": "trtmc", + "capabilities": list( + self._metadata.get(name, {}).get("capabilities", ()) + ), + "metadata": { + key: self._metadata[name][key] + for key in _PUBLIC_MODEL_METADATA_FIELDS + if key in self._metadata.get(name, {}) + }, + } + for name in self._order + ] + + def status(self) -> dict[str, Any]: + with self._lock: + statuses = { + name: ( + self._groups[name].status() + if name in self._groups + else { + "state": "not_started", + "ready": False, + "pid": None, + "pids": [], + "returncode": None, + "error": None, + "replicas": self._replicas[name], + "ready_replicas": 0, + "idle_replicas": 0, + "busy": False, + } + ) + for name in self._order + } + ready = ( + self._started + and bool(self._groups) + and len(self._groups) == len(self._specs) + and all(model["ready"] for model in statuses.values()) + ) + available = self._started and any(model["ready"] for model in statuses.values()) + degraded = available and any( + model["ready_replicas"] < model["replicas"] + for model in statuses.values() + ) + return {"ready": ready, "degraded": degraded, "models": statuses} + + def public_status(self) -> dict[str, Any]: + """Return the minimal process-independent readiness contract.""" + + status = self.status() + return { + "ready": status["ready"], + "degraded": status["degraded"], + "models": { + name: { + key: model_status[key] + for key in _PUBLIC_MODEL_STATUS_FIELDS + if key in model_status + } + for name, model_status in status["models"].items() + }, + } + + def metadata_for(self, name: str) -> dict[str, Any]: + if name not in self._specs: + raise ModelNotFoundError(f"model {name!r} is not registered") + with self._lock: + return dict(self._metadata.get(name, {})) + + def _make_worker(self, spec: ModelSpec, replica: int) -> WorkerProcess: + if self._worker_factory is not None: + return self._worker_factory(spec) + assert self._load_options is not None + worker_name = spec.name if self._replicas[spec.name] == 1 else f"{spec.name}-{replica + 1}" + return WorkerProcess( + name=worker_name, + bundle=spec.bundle, + trtmc_binary=self._trtmc_binary, + startup_timeout=self._startup_timeout, + request_timeout=self._request_timeout, + load_options=self._load_options, + ) + + @staticmethod + def _validate_metadata( + name: str, kind: ModelKind, metadata: Mapping[str, Any] + ) -> None: + allowed = { + "event", + "protocol_version", + "capabilities", + "default_max_new_tokens", + } + unknown = sorted(set(metadata) - allowed) + if unknown: + raise WorkerProtocolError( + f"worker {name!r} metadata has unsupported field(s): {', '.join(unknown)}" + ) + if metadata.get("event") != "ready": + raise WorkerProtocolError(f"worker {name!r} metadata.event must be 'ready'") + if metadata.get("protocol_version") != _WORKER_PROTOCOL_VERSION: + raise WorkerProtocolError(f"worker {name!r} uses an unsupported protocol version") + capabilities = metadata.get("capabilities") + if ( + not isinstance(capabilities, list) + or not capabilities + or any(not isinstance(value, str) or not value for value in capabilities) + or len(capabilities) != len(set(capabilities)) + ): + raise WorkerProtocolError( + f"worker {name!r} metadata.capabilities must be unique non-empty strings" + ) + required = "text_generation" if kind == "chat" else "transcription" + if required not in capabilities: + raise WorkerProtocolError( + f"worker {name!r} does not advertise required capability {required!r}" + ) + default_max_new_tokens = metadata.get("default_max_new_tokens") + if "default_max_new_tokens" in metadata and ( + isinstance(default_max_new_tokens, bool) + or not isinstance(default_max_new_tokens, int) + or default_max_new_tokens < 0 + ): + raise WorkerProtocolError( + f"worker {name!r} metadata.default_max_new_tokens must be a non-negative integer" + ) + + def _resolve_default( + self, + kind: ModelKind, + requested_name: str | None, + ) -> str | None: + available = [spec.name for spec in self._specs.values() if spec.kind == kind] + if requested_name is None: + return available[0] if available else None + spec = self._specs.get(requested_name) + if spec is None: + raise ValueError(f"default {kind} model {requested_name!r} is not registered") + if spec.kind != kind: + raise ValueError(f"default {kind} model {requested_name!r} has kind {spec.kind!r}") + return requested_name + + def __enter__(self) -> "ModelRegistry": + self.start() + return self + + def __exit__(self, _exc_type: Any, _exc: Any, _traceback: Any) -> None: + self.close() diff --git a/server/python/trtmc_server/schemas.py b/server/python/trtmc_server/schemas.py new file mode 100644 index 0000000000..cc03f5547b --- /dev/null +++ b/server/python/trtmc_server/schemas.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small request schemas used by the serving facade.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + + +class _OpenAIRequest(BaseModel): + model_config = ConfigDict(extra="allow") + + +class ChatMessage(_OpenAIRequest): + role: str + content: Any + + +class ChatCompletionRequest(_OpenAIRequest): + model: str | None = None + messages: list[ChatMessage] = Field(min_length=1) + max_tokens: int | None = Field(default=None, ge=1) + max_completion_tokens: int | None = Field(default=None, ge=1) + temperature: float | None = Field(default=None, ge=0) + top_p: float | None = Field(default=None, ge=0, le=1) + min_p: float | None = Field(default=None, ge=0, le=1) + top_k: int | None = Field(default=None, ge=0) + seed: int | None = None + enable_thinking: bool | None = None + stop: str | list[str] | None = None + stream: bool = False + + +def model_to_dict(model: BaseModel, *, exclude_none: bool = False) -> dict[str, Any]: + return dict(model.model_dump(exclude_none=exclude_none)) diff --git a/server/python/trtmc_server/worker.py b/server/python/trtmc_server/worker.py new file mode 100644 index 0000000000..27b3054172 --- /dev/null +++ b/server/python/trtmc_server/worker.py @@ -0,0 +1,751 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Long-lived JSONL subprocess transport for ``trtmc _serve-worker``.""" + +from __future__ import annotations + +import json +import os +import queue +import subprocess +import threading +import uuid +from collections import deque +from collections.abc import Callable, Iterable, Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Any, TextIO + +from .errors import ( + WorkerCrashedError, + WorkerError, + WorkerProtocolError, + WorkerRequestTooLargeError, + WorkerSaturatedError, + WorkerRemoteError, + WorkerStartupError, + WorkerTimeoutError, +) + + +_Response = dict[str, Any] | WorkerError +_MAX_STDERR_LINE_CHARS = 8192 +_MAX_REQUEST_LINE_BYTES = 16 * 1024 * 1024 +# This local server owns only independently loadable, single-process lanes. +# MPI/NCCL rank and rendezvous state must not cross into its child workers. +_WORKER_ENVIRONMENT_ALLOWLIST = frozenset( + { + "CONDA_PREFIX", + "CUDA_CACHE_DISABLE", + "CUDA_CACHE_MAXSIZE", + "CUDA_CACHE_PATH", + "CUDA_DEVICE_ORDER", + "CUDA_MODULE_LOADING", + "CUDA_VISIBLE_DEVICES", + "HOME", + "LANG", + "LC_ALL", + "LC_CTYPE", + "LD_LIBRARY_PATH", + "LOGNAME", + "NVIDIA_DRIVER_CAPABILITIES", + "NVIDIA_VISIBLE_DEVICES", + "OMP_NUM_THREADS", + "PATH", + "TMPDIR", + "TZ", + "USER", + "VIRTUAL_ENV", + "XDG_CACHE_HOME", + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", + } +) +_WINDOWS_WORKER_ENVIRONMENT_ALLOWLIST = frozenset( + { + "COMSPEC", + "PATHEXT", + "SYSTEMROOT", + "TEMP", + "TMP", + "USERPROFILE", + } +) + + +def _worker_environment() -> dict[str, str]: + """Copy only runtime variables that a native worker is allowed to observe.""" + + allowed = _WORKER_ENVIRONMENT_ALLOWLIST + if os.name == "nt": + allowed = allowed | _WINDOWS_WORKER_ENVIRONMENT_ALLOWLIST + return {name: os.environ[name] for name in allowed if name in os.environ} + + +@dataclass(frozen=True) +class WorkerLoadOptions: + """Native load-time options forwarded to every ``_serve-worker``.""" + + runtime_root: str + kv_cache_size_bytes: int | None = None + runtime_cache: str | None = None + cuda_graphs: bool = False + + def __post_init__(self) -> None: + if not self.runtime_root: + raise ValueError("runtime_root must be non-empty") + if ( + self.kv_cache_size_bytes is not None + and ( + isinstance(self.kv_cache_size_bytes, bool) + or self.kv_cache_size_bytes <= 0 + ) + ): + raise ValueError("kv_cache_size_bytes must be positive") + + def argv(self) -> list[str]: + result = ["--runtime-root", self.runtime_root] + if self.kv_cache_size_bytes is not None: + result.extend(("--kv-cache-size", str(self.kv_cache_size_bytes))) + if self.runtime_cache: + result.extend(("--runtime-cache", self.runtime_cache)) + if self.cuda_graphs: + result.append("--cuda-graphs") + return result + + +class WorkerSession: + """Exclusive sequence of operations against one stateful worker stream.""" + + def __init__(self, worker: "WorkerProcess", release: Callable[[], None]) -> None: + self._worker = worker + self._release = release + self._closed = False + self._lock = threading.Lock() + + def request( + self, + op: str, + payload: Mapping[str, Any] | None = None, + *, + timeout: float | None = None, + ) -> Any: + if self._closed: + raise WorkerCrashedError(f"worker session for {self._worker.name!r} is closed") + return self._worker._request_locked(op, payload, timeout=timeout) # noqa: SLF001 + + def submit( + self, + op: str, + payload: Mapping[str, Any] | None = None, + *, + timeout: float | None = None, + ) -> Future[Any]: + if self._closed: + raise WorkerCrashedError(f"worker session for {self._worker.name!r} is closed") + return self._worker._executor.submit( # noqa: SLF001 + self.request, op, payload, timeout=timeout + ) + + def close(self) -> None: + with self._lock: + if self._closed: + return + self._closed = True + self._release() + + def __enter__(self) -> "WorkerSession": + return self + + def __exit__(self, _exc_type: Any, _exc: Any, _traceback: Any) -> None: + self.close() + + +class WorkerProcess: + """Own one persistent native model worker. + + Stdout is reserved for one JSON object per line. Stderr is drained on a + dedicated thread and retained as a bounded diagnostic tail. Requests are + assigned opaque IDs and serialized because the native worker owns mutable + pipeline state. ``acquire_session`` holds that serialization lock across a + realtime start/chunk/finish sequence. + """ + + def __init__( + self, + *, + name: str, + bundle: str | os.PathLike[str], + trtmc_binary: str | os.PathLike[str], + startup_timeout: float = 120.0, + request_timeout: float = 120.0, + stderr_lines: int = 100, + max_request_line_bytes: int = _MAX_REQUEST_LINE_BYTES, + load_options: WorkerLoadOptions, + ) -> None: + if startup_timeout <= 0 or request_timeout <= 0: + raise ValueError("worker timeouts must be positive") + if stderr_lines <= 0: + raise ValueError("stderr_lines must be positive") + if max_request_line_bytes <= 0: + raise ValueError("max_request_line_bytes must be positive") + self.name = name + self.bundle = Path(bundle) + self.trtmc_binary = Path(trtmc_binary) + self.startup_timeout = float(startup_timeout) + self.request_timeout = float(request_timeout) + self.max_request_line_bytes = int(max_request_line_bytes) + self.load_options = load_options + + self._state = "new" + self._state_lock = threading.RLock() + self._operation_lock = threading.Lock() + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix=f"trtmc-{name}") + self._response: queue.Queue[_Response] = queue.Queue(maxsize=1) + self._expected_request_id: str | None = None + self._process: subprocess.Popen[str] | None = None + self._stdout_thread: threading.Thread | None = None + self._stderr_thread: threading.Thread | None = None + self._ready_event = threading.Event() + self._ready_payload: dict[str, Any] = {} + self._last_error: WorkerError | None = None + self._stderr_tail: deque[str] = deque(maxlen=stderr_lines) + + @property + def state(self) -> str: + with self._state_lock: + return self._state + + @property + def ready(self) -> bool: + with self._state_lock: + return ( + self._state == "ready" + and self._process is not None + and self._process.poll() is None + ) + + @property + def pid(self) -> int | None: + with self._state_lock: + return self._process.pid if self._process is not None else None + + @property + def ready_payload(self) -> dict[str, Any]: + with self._state_lock: + return dict(self._ready_payload) + + @property + def stderr_tail(self) -> list[str]: + with self._state_lock: + return list(self._stderr_tail) + + def start(self) -> None: + """Launch the worker and wait for its JSON ready record.""" + + with self._state_lock: + if self._state == "ready": + return + if self._state != "new": + raise WorkerStartupError( + f"worker {self.name!r} cannot start from state {self._state!r}" + ) + self._state = "starting" + self._ready_event.clear() + + command = [ + str(self.trtmc_binary), + "_serve-worker", + str(self.bundle), + *self.load_options.argv(), + ] + worker_environment = _worker_environment() + try: + process = subprocess.Popen( + command, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + encoding="utf-8", + errors="replace", + bufsize=1, + start_new_session=(os.name != "nt"), + env=worker_environment, + ) + except OSError as exc: + error = WorkerStartupError(f"failed to launch worker {self.name!r}") + self._mark_failed(error) + raise error from exc + + assert process.stdout is not None + assert process.stderr is not None + with self._state_lock: + self._process = process + self._stdout_thread = threading.Thread( + target=self._read_stdout, + args=(process.stdout,), + name=f"trtmc-{self.name}-stdout", + daemon=True, + ) + self._stderr_thread = threading.Thread( + target=self._read_stderr, + args=(process.stderr,), + name=f"trtmc-{self.name}-stderr", + daemon=True, + ) + self._stdout_thread.start() + self._stderr_thread.start() + + if not self._ready_event.wait(self.startup_timeout): + error = WorkerStartupError( + f"worker {self.name!r} did not become ready within {self.startup_timeout:g}s" + ) + self._mark_failed(error) + self._terminate_process() + self._finalize_io() + raise error + + with self._state_lock: + ready = self._state == "ready" + error = self._last_error or WorkerStartupError( + f"worker {self.name!r} failed before readiness" + ) + if not ready: + self._terminate_process() + self._finalize_io() + if isinstance(error, WorkerStartupError): + raise error + raise WorkerStartupError(str(error)) from error + + def request( + self, + op: str, + payload: Mapping[str, Any] | None = None, + *, + timeout: float | None = None, + ) -> Any: + """Run one serialized worker operation and return its unwrapped result.""" + + deadline = self.request_timeout if timeout is None else float(timeout) + if deadline <= 0: + raise ValueError("request timeout must be positive") + with self.acquire_session(timeout=deadline) as session: + return session.request(op, payload, timeout=deadline) + + def acquire_session( + self, + *, + timeout: float | None = None, + on_close: Callable[["WorkerProcess"], None] | None = None, + ) -> WorkerSession: + """Acquire exclusive worker ownership for a stateful request sequence.""" + + deadline = self.request_timeout if timeout is None else float(timeout) + if deadline < 0: + raise ValueError("session timeout must be non-negative") + if not self._operation_lock.acquire(timeout=deadline): + raise WorkerTimeoutError(f"worker {self.name!r} remained busy for {deadline:g}s") + try: + self._assert_usable() + except BaseException: + self._operation_lock.release() + raise + + def release() -> None: + self._operation_lock.release() + if on_close is not None: + on_close(self) + + return WorkerSession(self, release) + + def _request_locked( + self, + op: str, + payload: Mapping[str, Any] | None, + *, + timeout: float | None, + allow_closing: bool = False, + ) -> Any: + deadline = self.request_timeout if timeout is None else float(timeout) + if not op or not isinstance(op, str): + raise ValueError("worker operation must be a non-empty string") + if payload: + reserved = {"id", "op"}.intersection(payload) + if reserved: + field = sorted(reserved)[0] + raise ValueError(f"worker payload cannot override {field!r}") + + with self._state_lock: + self._assert_usable_locked(allow_closing=allow_closing) + process = self._process + assert process is not None + request_id = uuid.uuid4().hex + if self._expected_request_id is not None: + raise WorkerProtocolError(f"worker {self.name!r} already has an active request") + self._expected_request_id = request_id + + message: dict[str, Any] = {"id": request_id, "op": op} + if payload: + message.update(payload) + + try: + encoded = json.dumps(message, separators=(",", ":"), ensure_ascii=False) + except (TypeError, ValueError) as exc: + with self._state_lock: + self._expected_request_id = None + raise WorkerProtocolError(f"worker request is not JSON serializable: {exc}") from exc + encoded_bytes = len(encoded.encode("utf-8")) + if encoded_bytes > self.max_request_line_bytes: + with self._state_lock: + self._expected_request_id = None + raise WorkerRequestTooLargeError( + f"worker request for {op!r} is {encoded_bytes} bytes; maximum is " + f"{self.max_request_line_bytes} bytes" + ) + + try: + assert process.stdin is not None + process.stdin.write(encoded + "\n") + process.stdin.flush() + except (BrokenPipeError, OSError, ValueError) as exc: + error = WorkerCrashedError(f"worker {self.name!r} transport failed during {op!r}") + self._mark_failed(error) + self._terminate_process() + self._finalize_io() + raise error from exc + + try: + response = self._response.get(timeout=deadline) + except queue.Empty as exc: + error = WorkerTimeoutError( + f"worker {self.name!r} timed out after {deadline:g}s during {op!r}" + ) + self._mark_failed(error) + self._terminate_process() + self._finalize_io() + raise error from exc + finally: + with self._state_lock: + self._expected_request_id = None + + if isinstance(response, WorkerError): + self._terminate_process() + self._finalize_io() + raise response + return self._unwrap_response(op, response) + + def close(self, *, grace_period: float = 1.0) -> None: + """Best-effort protocol shutdown followed by deterministic cleanup.""" + + with self._state_lock: + if self._state == "closed": + return + was_ready = self._state == "ready" + + acquired = self._operation_lock.acquire(timeout=max(0.0, grace_period)) + try: + with self._state_lock: + if self._state == "closed": + return + self._state = "closing" + if acquired and was_ready: + try: + self._request_locked( + "shutdown", + None, + timeout=max(0.05, grace_period), + allow_closing=True, + ) + except WorkerError: + pass + finally: + if acquired: + self._operation_lock.release() + + self._terminate_process(grace_period=grace_period) + self._finalize_io() + self._executor.shutdown(wait=True, cancel_futures=True) + error = WorkerCrashedError(f"worker {self.name!r} was closed") + self._wake_request(error) + with self._state_lock: + self._state = "closed" + + @property + def busy(self) -> bool: + return self._operation_lock.locked() + + def _assert_usable(self) -> None: + with self._state_lock: + self._assert_usable_locked() + + def _assert_usable_locked(self, *, allow_closing: bool = False) -> None: + allowed = {"ready"} + if allow_closing: + allowed.add("closing") + process = self._process + if self._state not in allowed or process is None or process.poll() is not None: + if self._last_error is not None: + raise WorkerCrashedError(str(self._last_error)) from self._last_error + raise WorkerCrashedError(f"worker {self.name!r} is unavailable (state={self._state!r})") + + def _read_stdout(self, stream: TextIO) -> None: + try: + for raw_line in stream: + line = raw_line.strip() + if not line: + continue + try: + message = json.loads(line) + except json.JSONDecodeError as exc: + self._fail_protocol(f"worker {self.name!r} emitted invalid JSONL: {exc.msg}") + return + if not isinstance(message, dict): + self._fail_protocol(f"worker {self.name!r} emitted a non-object JSON record") + return + + if message.get("event") == "ready": + with self._state_lock: + if self._state != "starting": + self._fail_protocol( + f"worker {self.name!r} emitted an unexpected ready record" + ) + return + self._ready_payload = dict(message) + self._state = "ready" + self._ready_event.set() + continue + + request_id = message.get("id") + if not isinstance(request_id, str) or not request_id: + self._fail_protocol( + f"worker {self.name!r} response is missing a string request id" + ) + return + with self._state_lock: + expected = self._expected_request_id + terminal = self._state in {"failed", "closing", "closed"} + if request_id != expected: + # A timed-out request may race with termination. Ignore only + # after failure/close; otherwise the protocol is desynchronized. + if terminal: + continue + self._fail_protocol( + f"worker {self.name!r} responded with unknown id {request_id!r}" + ) + return + ok = message.get("ok") + if not isinstance(ok, bool): + self._fail_protocol( + f"worker {self.name!r} response is missing a boolean ok field" + ) + return + if ok and "result" not in message: + self._fail_protocol(f"worker {self.name!r} success response is missing result") + return + if not ok and not isinstance(message.get("error"), Mapping): + self._fail_protocol( + f"worker {self.name!r} failure response is missing an error object" + ) + return + try: + self._response.put_nowait(message) + except queue.Full: + self._fail_protocol(f"worker {self.name!r} emitted more than one response") + return + except (OSError, ValueError): + self._mark_failed(WorkerCrashedError(f"worker {self.name!r} stdout failed")) + return + finally: + with self._state_lock: + process = self._process + state = self._state + returncode = process.poll() if process is not None else None + error = WorkerCrashedError( + f"worker {self.name!r} exited" + + (f" with code {returncode}" if returncode is not None else "") + ) + self._wake_request(error) + if state not in {"failed", "closing", "closed"}: + self._mark_failed(error) + try: + stream.close() + except OSError: + pass + if process is not None and process.stdin is not None: + try: + process.stdin.close() + except OSError: + pass + + def _read_stderr(self, stream: TextIO) -> None: + try: + for raw_line in stream: + line = raw_line.rstrip("\r\n") + if line: + with self._state_lock: + self._stderr_tail.append(line[-_MAX_STDERR_LINE_CHARS:]) + except (OSError, ValueError): + return + finally: + try: + stream.close() + except OSError: + pass + + def _unwrap_response(self, op: str, response: dict[str, Any]) -> Any: + if not response["ok"]: + details = response["error"] + message = str(details.get("message") or details.get("code") or details) + raise WorkerRemoteError( + f"worker {self.name!r} failed {op!r}: {message}", details=details + ) + return response["result"] + + def _fail_protocol(self, message: str) -> None: + self._mark_failed(WorkerProtocolError(message)) + self._terminate_process() + + def _mark_failed(self, error: WorkerError) -> None: + with self._state_lock: + if self._state in {"closing", "closed"}: + self._wake_request(error) + self._ready_event.set() + return + self._state = "failed" + self._last_error = error + self._ready_event.set() + self._wake_request(error) + + def _wake_request(self, error: WorkerError) -> None: + with self._state_lock: + pending = self._expected_request_id is not None + if pending: + try: + self._response.put_nowait(error) + except queue.Full: + pass + + def _terminate_process(self, *, grace_period: float = 0.5) -> None: + with self._state_lock: + process = self._process + if process is None or process.poll() is not None: + return + try: + process.terminate() + process.wait(timeout=max(0.05, grace_period)) + except subprocess.TimeoutExpired: + process.kill() + try: + process.wait(timeout=1.0) + except subprocess.TimeoutExpired: + pass + except OSError: + pass + + def _finalize_io(self) -> None: + with self._state_lock: + process = self._process + threads = (self._stdout_thread, self._stderr_thread) + if process is not None: + for stream in (process.stdin, process.stdout, process.stderr): + if stream is not None and not stream.closed: + try: + stream.close() + except OSError: + pass + current = threading.current_thread() + for thread in threads: + if thread is not None and thread is not current and thread.is_alive(): + thread.join(timeout=0.5) + + def __enter__(self) -> "WorkerProcess": + self.start() + return self + + def __exit__(self, _exc_type: Any, _exc: Any, _traceback: Any) -> None: + self.close() + + +class WorkerGroup: + """Route one logical model to a fixed set of native worker replicas.""" + + def __init__(self, name: str, workers: Iterable[WorkerProcess]) -> None: + self.name = name + self._workers = tuple(workers) + if not self._workers: + raise ValueError("worker group must contain at least one replica") + self._idle: queue.Queue[WorkerProcess] = queue.Queue(maxsize=len(self._workers)) + for worker in self._workers: + if not worker.ready: + raise ValueError(f"worker {worker.name!r} is not ready") + self._idle.put_nowait(worker) + self._closed = False + self._lock = threading.Lock() + + @property + def replicas(self) -> int: + return len(self._workers) + + @property + def ready(self) -> bool: + with self._lock: + return not self._closed and any(worker.ready for worker in self._workers) + + def acquire_session(self) -> WorkerSession: + """Lease one idle replica, or fail immediately when all are busy.""" + + with self._lock: + if self._closed: + raise WorkerCrashedError(f"worker group {self.name!r} is closed") + + while True: + try: + worker = self._idle.get_nowait() + except queue.Empty as exc: + if not self.ready: + raise WorkerCrashedError( + f"worker group {self.name!r} has no healthy replicas" + ) from exc + raise WorkerSaturatedError( + f"all {self.replicas} replicas for model {self.name!r} are busy" + ) from exc + + if not worker.ready: + continue + try: + return worker.acquire_session(timeout=0, on_close=self._release) + except WorkerCrashedError: + continue + + def close(self) -> None: + with self._lock: + if self._closed: + return + self._closed = True + for worker in reversed(self._workers): + worker.close() + + def status(self) -> dict[str, Any]: + with self._lock: + closed = self._closed + ready_replicas = sum(worker.ready for worker in self._workers) + idle_replicas = sum(worker.ready and not worker.busy for worker in self._workers) + ready = not closed and ready_replicas > 0 + return { + "state": "closed" if closed else "ready" if ready else "failed", + "ready": ready, + "degraded": not closed and 0 < ready_replicas < self.replicas, + "replicas": self.replicas, + "ready_replicas": ready_replicas, + "idle_replicas": idle_replicas, + "busy": idle_replicas < ready_replicas, + } + + def _release(self, worker: WorkerProcess) -> None: + with self._lock: + if self._closed or not worker.ready: + return + self._idle.put_nowait(worker) diff --git a/server/tests/conftest.py b/server/tests/conftest.py new file mode 100644 index 0000000000..2bab45774f --- /dev/null +++ b/server/tests/conftest.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Import-origin guard for the physically isolated server package.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + + +_SERVER_SOURCE_ROOT = Path(__file__).resolve().parents[1] / "python" +_TEST_INSTALLED_WHEEL = os.environ.get("TRTMC_TEST_INSTALLED_WHEEL") == "1" + +if not _TEST_INSTALLED_WHEEL and str(_SERVER_SOURCE_ROOT) not in sys.path: + sys.path.insert(0, str(_SERVER_SOURCE_ROOT)) + +if _TEST_INSTALLED_WHEEL: + import trtmc_server as _installed_server + + installed_path = Path(_installed_server.__file__).resolve() + if installed_path.is_relative_to(_SERVER_SOURCE_ROOT): + raise RuntimeError( + "TRTMC_TEST_INSTALLED_WHEEL=1 imported trtmc_server " + f"from the source checkout: {installed_path}" + ) diff --git a/server/tests/fake_serve_worker.py b/server/tests/fake_serve_worker.py new file mode 100755 index 0000000000..6146fdfef8 --- /dev/null +++ b/server/tests/fake_serve_worker.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only executable fixture for the trtmc serve JSONL protocol.""" + +from __future__ import annotations + +import base64 +import json +import os +import sys +import time +from pathlib import Path +from typing import Any + + +def emit(payload: dict[str, Any]) -> None: + print(json.dumps(payload, separators=(",", ":")), flush=True) + + +def supports_transcription_wav(path: Path) -> bool: + """Mirror the native PCM16/IEEE-float32 WAV format boundary.""" + + payload = path.read_bytes() + if len(payload) < 12 or payload[:4] != b"RIFF" or payload[8:12] != b"WAVE": + return False + container_end = 8 + int.from_bytes(payload[4:8], "little") + if container_end < 12 or container_end > len(payload): + return False + + audio_format = channels = sample_rate = bits_per_sample = data_size = 0 + position = 12 + while position + 8 <= container_end: + chunk_id = payload[position : position + 4] + chunk_size = int.from_bytes(payload[position + 4 : position + 8], "little") + data_offset = position + 8 + next_position = data_offset + chunk_size + (chunk_size & 1) + if next_position > container_end: + return False + if chunk_id == b"fmt ": + if chunk_size < 16: + return False + audio_format = int.from_bytes(payload[data_offset : data_offset + 2], "little") + channels = int.from_bytes(payload[data_offset + 2 : data_offset + 4], "little") + sample_rate = int.from_bytes(payload[data_offset + 4 : data_offset + 8], "little") + bits_per_sample = int.from_bytes(payload[data_offset + 14 : data_offset + 16], "little") + elif chunk_id == b"data" and chunk_size > 0: + data_size = chunk_size + position = next_position + + supported_format = (audio_format, bits_per_sample) in {(1, 16), (3, 32)} + frame_width = channels * (bits_per_sample // 8) + return ( + position == container_end + and supported_format + and channels > 0 + and 0 < sample_rate <= 0x7FFF_FFFF + and data_size > 0 + and frame_width > 0 + and data_size % frame_width == 0 + ) + + +def main() -> int: + if len(sys.argv) < 3 or sys.argv[1] != "_serve-worker": + print("expected _serve-worker BUNDLE", file=sys.stderr, flush=True) + return 2 + bundle = Path(sys.argv[2]) + worker_args = sys.argv[3:] + mode = bundle.stem + if "startup-timeout" in mode: + print(f"startup detail for {bundle.resolve()}", file=sys.stderr, flush=True) + time.sleep(60) + return 1 + if "bad-ready" in mode: + print("not-json", flush=True) + return 1 + + kind = ( + "transcription" + if any(token in mode for token in ("asr", "transcription", "stream")) + else "chat" + ) + capabilities = ( + ["transcription", "transcription_streaming"] + if kind == "transcription" + else ["text_generation"] + ) + if "no-stream" in mode: + capabilities.remove("transcription_streaming") + ready = { + "event": "ready", + "protocol_version": 2 if "protocol-v2" in mode else 3, + "capabilities": capabilities, + } + if kind == "chat": + ready["default_max_new_tokens"] = 64 + if "wrong-capability" in mode: + ready["capabilities"] = ["transcription"] + if "duplicate-capability" in mode: + ready["capabilities"] = [capabilities[0], capabilities[0]] + if "unknown-ready-field" in mode: + ready["private_detail"] = "must be rejected" + if "inspect-environment" in mode: + ready.update( + { + "serve_token_present": "TRTMC_SERVE_TOKEN" in os.environ, + "allowed_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "secret_environment_present": any( + name in os.environ + for name in ( + "ACCESS_TOKEN", + "AUTHORIZATION", + "AWS_SECRET_ACCESS_KEY", + "COOKIE", + "GITHUB_TOKEN", + "HF_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "NVIDIA_API_KEY", + "TRTMC_SERVE_TOKEN", + "TRTMC_MODEL_PLUGIN_DIR", + "TRTMC_MODEL_PLUGIN_STRICT", + "TRTMC_TRT_LIBRARY_DIR", + "UNLISTED_ENVIRONMENT", + ) + ), + "distributed_environment_present": any( + name in os.environ + for name in ( + "OMPI_COMM_WORLD_JOBID", + "OMPI_COMM_WORLD_LOCAL_RANK", + "OMPI_COMM_WORLD_RANK", + "OMPI_COMM_WORLD_SIZE", + "PMI_RANK", + "PMI_SIZE", + "PMIX_NAMESPACE", + "RANK", + "SLURM_PROCID", + "TRTMC_NCCL_RENDEZVOUS", + "TRTMC_NCCL_SKIP_DESTROY", + "WORLD_SIZE", + ) + ), + } + ) + if "inspect-argv" in mode: + ready["worker_args"] = worker_args + if "legacy-ready" in mode: + ready.pop("event") + ready["ready"] = True + emit(ready) + print(f"fake worker ready: {mode}", file=sys.stderr, flush=True) + + active_samples: int | None = None + chunk_failed = False + for line in sys.stdin: + request = json.loads(line) + request_id = request["id"] + op = request["op"] + if "slow-request" in mode: + if op == "generate": + time.sleep(2) + elif "saturation" in mode and op == "generate": + time.sleep(0.5) + elif "slow" in mode and op != "shutdown": + time.sleep(2) + if "delayed-start" in mode and op == "stream_start": + time.sleep(0.25) + if "delayed-reset" in mode and op == "stream_reset": + time.sleep(0.25) + if "crash-request" in mode and op == "generate": + print( + f"intentional fake crash at {bundle.resolve()} with access_token=worker-secret", + file=sys.stderr, + flush=True, + ) + return 17 + if "crash" in mode and "crash-request" not in mode and op != "shutdown": + print( + f"intentional fake crash at {bundle.resolve()} with access_token=worker-secret", + file=sys.stderr, + flush=True, + ) + return 17 + if "bad-response" in mode: + emit({"id": "not-the-request-id", "ok": True, "result": {}}) + continue + if "non-bool-ok" in mode: + emit({"id": request_id, "ok": "true", "result": {}}) + continue + if "missing-result" in mode: + emit({"id": request_id, "ok": True}) + continue + + try: + if op == "generate": + if "invalid-request" in mode: + raise ValueError("invalid fake generation request") + prompt = request.get("prompt", "") + thinking = request.get("config", {}).get("enable_thinking") + text = f"generated:{prompt}" + if isinstance(thinking, bool): + text = f"enable_thinking={str(thinking).lower()}:{text}" + result = { + "text": text, + "usage": { + "prompt_tokens": len(str(prompt).split()), + "completion_tokens": 1, + }, + } + elif op == "transcribe": + audio_path = Path(request["audio_path"]) + try: + supported = supports_transcription_wav(audio_path) + except FileNotFoundError as exc: + raise RuntimeError("read_wav: cannot open input file") from exc + if not supported: + raise ValueError("read_wav: WAV samples must be PCM16 or IEEE float32") + result = {"text": f"transcribed {audio_path.stat().st_size} bytes"} + elif op == "stream_start": + if active_samples is not None: + raise ValueError("another stream is already active") + active_samples = 0 + result = {} + elif op == "probe_transcription_stream": + if request.get("config") != { + "sample_rate_hz": 16000, + "channels": 1, + "audio_format": "pcm16le", + }: + raise ValueError("invalid streaming probe config") + if ( + kind != "transcription" + or "no-stream" in mode + or "probe-fails" in mode + ): + raise RuntimeError("streaming transcription is unavailable") + result = {"supported": True} + elif op == "stream_chunk": + if "chunk-error" in mode and not chunk_failed: + chunk_failed = True + raise ValueError("invalid fake audio chunk") + if active_samples is None: + raise ValueError("no active transcription stream") + audio = base64.b64decode(request["audio"], validate=True) + active_samples += len(audio) // 2 + result = { + "text": ( + request["audio"] + if "echo-wire" in mode + else f"{active_samples} samples" + ) + } + elif op == "stream_finish": + if active_samples is None: + raise ValueError("no active transcription stream") + result = {"text": f"{active_samples} samples"} + active_samples = None + elif op == "stream_reset": + if active_samples is None: + raise ValueError("no active transcription stream") + active_samples = None + result = {} + elif op == "shutdown": + emit({"id": request_id, "ok": True, "result": {"status": "shutting_down"}}) + return 0 + else: + raise ValueError(f"unsupported op: {op}") + emit({"id": request_id, "ok": True, "result": result}) + except Exception as exc: # noqa: BLE001 - fixture reports protocol errors + error_type = ( + "runtime_error" if isinstance(exc, RuntimeError) else "invalid_request_error" + ) + error = {"type": error_type, "message": str(exc)} + if op == "transcribe" and error_type == "invalid_request_error": + error.update(code="unsupported_media_type", param="file") + emit( + { + "id": request_id, + "ok": False, + "error": error, + } + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/server/tests/test_dependency_direction.py b/server/tests/test_dependency_direction.py new file mode 100644 index 0000000000..0794d0a532 --- /dev/null +++ b/server/tests/test_dependency_direction.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import ast +import re +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +CMAKE_ROOT = REPO_ROOT / "CMakeLists.txt" +SERVER_CMAKE = REPO_ROOT / "server" / "CMakeLists.txt" +SERVER_ENTRYPOINT = REPO_ROOT / "server" / "native" / "entrypoint.cpp" +CLI_MAIN = REPO_ROOT / "apps" / "cli" / "main.cpp" +LIBRARY_ROOTS = ( + REPO_ROOT / "core" / "runtime", + REPO_ROOT / "families", +) +SERVER_ROOT = REPO_ROOT / "server" / "native" +SERVER_PYTHON_ROOT = REPO_ROOT / "server" / "python" / "trtmc_server" +PYTHON_LIBRARY_ROOTS = ( + REPO_ROOT / "core", + REPO_ROOT / "families", +) +CPP_SUFFIXES = {".c", ".cc", ".cpp", ".cu", ".cuh", ".cxx", ".h", ".hpp"} +INCLUDE = re.compile(r'^\s*#\s*include\s*([<"])([^>"]+)[>"]', re.MULTILINE) +LIBRARY_PYTHON_IMPORT = re.compile( + r"^\s*(?:from\s+tensorrt_model_connect(?:[.\s])|" + r"import\s+tensorrt_model_connect(?:[.\s]))", + re.MULTILINE, +) +PRIVATE_LIBRARY_PREFIXES = ( + "../", + "apps/", + "core/", + "families/", +) + + +def _cpp_files(root: Path) -> list[Path]: + assert root.is_dir(), f"dependency boundary root is missing: {root.relative_to(REPO_ROOT)}" + files = sorted(path for path in root.rglob("*") if path.suffix in CPP_SUFFIXES) + assert files, f"dependency boundary root has no C++ sources: {root.relative_to(REPO_ROOT)}" + return files + + +def _includes(text: str) -> list[tuple[str, str]]: + return INCLUDE.findall(text) + + +def _server_import_lines(path: Path) -> list[int]: + tree = ast.parse(path.read_text(encoding="utf-8", errors="strict"), filename=str(path)) + violations: list[int] = [] + for node in ast.walk(tree): + modules: list[str] = [] + if isinstance(node, ast.Import): + modules.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0: + if node.module: + modules.append(node.module) + if node.module == "tensorrt_model_connect" and any( + alias.name == "serve" for alias in node.names + ): + modules.append("tensorrt_model_connect.serve") + if any( + module == "trtmc_server" + or module.startswith("trtmc_server.") + or module == "tensorrt_model_connect.serve" + or module.startswith("tensorrt_model_connect.serve.") + for module in modules + ): + violations.append(node.lineno) + return violations + + +def test_library_roots_do_not_reference_server_headers_or_target() -> None: + violations: list[str] = [] + for root in LIBRARY_ROOTS: + for path in _cpp_files(root): + relative = path.relative_to(REPO_ROOT) + contents = path.read_text(encoding="utf-8", errors="strict") + for _, include in _includes(contents): + if include.startswith("native/"): + violations.append(f'{relative}: includes private Server header "{include}"') + for reference in ( + "trtmc::serve", + "trtmc::server", + "trtmc_server_native", + "server/native/", + ): + if reference in contents: + violations.append(f'{relative}: references Server boundary "{reference}"') + + assert violations == [], "\n".join(violations) + + +def test_python_library_does_not_import_optional_server() -> None: + files = sorted(path for root in PYTHON_LIBRARY_ROOTS for path in root.rglob("*.py")) + assert files, "Python Library roots have no source files" + violations = [ + f"{path.relative_to(REPO_ROOT)}:{line}" + for path in files + for line in _server_import_lines(path) + ] + assert violations == [], "Python Library imports optional Server: " + ", ".join(violations) + + +def test_python_server_does_not_import_library_implementation() -> None: + files = sorted(SERVER_PYTHON_ROOT.rglob("*.py")) + assert files, "Python Server root has no source files" + violations = [ + str(path.relative_to(REPO_ROOT)) + for path in files + if LIBRARY_PYTHON_IMPORT.search(path.read_text(encoding="utf-8", errors="strict")) + ] + assert violations == [], "Python Server imports Library implementation: " + ", ".join( + violations + ) + + +def test_cmake_keeps_server_downstream_of_core() -> None: + cmake = CMAKE_ROOT.read_text(encoding="utf-8") + server_cmake = SERVER_CMAKE.read_text(encoding="utf-8") + core_start = cmake.index("add_library(trtmc_core SHARED") + server_start = cmake.index("add_subdirectory(server)") + core_region = cmake[core_start:server_start] + assert "target_link_libraries(trtmc_core" in core_region + assert "server/native/" not in core_region + assert "trtmc_server_native" not in core_region + assert "add_library(trtmc_server_native STATIC" in server_cmake + assert re.search( + r"target_link_libraries\(trtmc_server_native\s+PRIVATE\s+" + r"trtmc_runtime\b", + server_cmake, + ) + for library_target in ("trtmc_core", "trtmc_runtime"): + reverse_dependency = ( + r"(?m)^\s*(?:target_sources|target_link_libraries|add_dependencies)" + rf"\s*\(\s*{library_target}\b" + ) + assert not re.search(reverse_dependency, server_cmake) + assert not re.search(reverse_dependency, cmake[server_start:]) + assert "install(" not in server_cmake + assert not re.search( + r"install\s*\([^)]*\btrtmc_server_native\b", + cmake + "\n" + server_cmake, + re.DOTALL, + ) + + +def test_server_uses_only_local_or_public_library_headers() -> None: + violations: list[str] = [] + for path in _cpp_files(SERVER_ROOT): + relative = path.relative_to(REPO_ROOT) + contents = path.read_text(encoding="utf-8", errors="strict") + for delimiter, include in _includes(contents): + if include.startswith("native/") or include.startswith("trtmc/"): + continue + if delimiter == '"' or include.startswith(PRIVATE_LIBRARY_PREFIXES): + violations.append( + f'{relative}: Server must not include private Library header "{include}"' + ) + + assert violations == [], "\n".join(violations) + + +def test_source_build_python_copy_removes_stale_modules_first() -> None: + server_cmake = SERVER_CMAKE.read_text(encoding="utf-8") + copy_target = server_cmake.split("add_custom_target(trtmc_server_python ALL", 1)[1].split( + "\n)", 1 + )[0] + generated_module = '"${PROJECT_BINARY_DIR}/server/python/trtmc_server"' + assert copy_target.count(generated_module) == 2 + assert copy_target.index("-E remove_directory") < copy_target.index("-E copy_directory") + + +def test_server_frontend_isolates_python_module_search() -> None: + source = SERVER_ENTRYPOINT.read_text(encoding="utf-8") + assert 'source_mode ? "-P" : "-I"' in source + assert 'std::string pythonpath = server_python.string();' in source + assert 'pythonpath += ":" + std::string(existing);' in source + + cmake = SERVER_CMAKE.read_text(encoding="utf-8") + shadow_test = cmake.split("add_test(\n NAME serve_cli_ignores_cwd_shadow", 1)[1] + assert "TRTMC_CWD_SHADOW_EXECUTED" in shadow_test + assert "PASS_REGULAR_EXPRESSION" in shadow_test + assert "FAIL_REGULAR_EXPRESSION" in shadow_test + + +def test_cli_main_keeps_only_thin_server_dispatch() -> None: + source = CLI_MAIN.read_text(encoding="utf-8") + assert '#include "native/entrypoint.h"' in source + assert source.count("trtmc::server::run_server_frontend") == 1 + assert source.count("trtmc::server::run_native_worker") == 1 + for implementation_detail in ( + "execvp", + "PYTHONPATH", + "--runtime-root", + "--kv-cache-size", + "load_task", + ): + assert implementation_detail not in source diff --git a/server/tests/test_serve_api.py b/server/tests/test_serve_api.py new file mode 100644 index 0000000000..aabf9c3f4d --- /dev/null +++ b/server/tests/test_serve_api.py @@ -0,0 +1,1522 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import base64 +import io +import logging +import struct +import threading +import time +import wave +from concurrent.futures import Future, ThreadPoolExecutor +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient +from httpx import Response +from starlette.websockets import WebSocketDisconnect + +from trtmc_server import realtime as realtime_module +from trtmc_server.app import ( + ServerConfig, + _BodyLimitMiddleware, + _HTTP_ENVELOPE_OVERHEAD_BYTES, + _worker_request, + create_app, +) +from trtmc_server.errors import WorkerProtocolError, WorkerRemoteError +from trtmc_server.protocol import ( + extract_usage, + extract_text, + invalid_request_message, + prepare_chat_prompt, +) +from trtmc_server.registry import ModelRegistry, ModelSpec +from trtmc_server.realtime import RealtimeTranscriptionConnection +from trtmc_server.worker import WorkerLoadOptions, WorkerProcess, WorkerSession + + +FAKE_TRTMC = Path(__file__).with_name("fake_serve_worker.py") + + +def _http_scope(*, content_length: int | None = None) -> dict[str, object]: + headers = [] if content_length is None else [(b"content-length", str(content_length).encode())] + return { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/limited", + "raw_path": b"/limited", + "query_string": b"", + "headers": headers, + "client": ("127.0.0.1", 1234), + "server": ("127.0.0.1", 8000), + "root_path": "", + } + + +def test_body_limit_rejects_declared_and_chunked_ingress_before_response() -> None: + async def exercise_declared() -> tuple[bool, list[dict[str, object]]]: + called = False + sent: list[dict[str, object]] = [] + + async def inner(_scope: object, _receive: object, _send: object) -> None: + nonlocal called + called = True + + async def receive() -> dict[str, object]: + raise AssertionError("oversized Content-Length must not read the body") + + async def send(message: dict[str, object]) -> None: + sent.append(message) + + middleware = _BodyLimitMiddleware(inner, limits={"/limited": 4}) + await middleware(_http_scope(content_length=5), receive, send) # type: ignore[arg-type] + return called, sent + + async def exercise_chunked() -> tuple[bool, list[dict[str, object]]]: + completed = False + sent: list[dict[str, object]] = [] + chunks = iter( + ( + {"type": "http.request", "body": b"abc", "more_body": True}, + {"type": "http.request", "body": b"de", "more_body": False}, + ) + ) + + async def inner(_scope: object, receive: object, _send: object) -> None: + nonlocal completed + while (await receive())["more_body"]: # type: ignore[operator] + pass + completed = True + + async def receive() -> dict[str, object]: + return next(chunks) + + async def send(message: dict[str, object]) -> None: + sent.append(message) + + middleware = _BodyLimitMiddleware(inner, limits={"/limited": 4}) + await middleware(_http_scope(), receive, send) # type: ignore[arg-type] + return completed, sent + + declared_called, declared_messages = asyncio.run(exercise_declared()) + chunked_completed, chunked_messages = asyncio.run(exercise_chunked()) + for messages in (declared_messages, chunked_messages): + assert messages[0]["status"] == 413 + assert b"request_body_too_large" in messages[-1]["body"] + assert declared_called is False + assert chunked_completed is False + + +def test_private_v3_text_and_invalid_request_shapes_are_strict() -> None: + assert extract_text({"text": "canonical"}, operation="generate") == "canonical" + for result in ("legacy", {"transcript": "legacy"}, {"output_text": "legacy"}): + with pytest.raises(WorkerProtocolError, match="string text field"): + extract_text(result, operation="generate") + + def remote(error_type: str) -> WorkerRemoteError: + return WorkerRemoteError("bad", details={"type": error_type, "message": "detail"}) + + assert invalid_request_message(remote("invalid_request_error")) == "detail" + assert invalid_request_message(remote("invalid_request")) is None + missing_message = WorkerRemoteError( + "private provider diagnostic", + details={"type": "invalid_request_error", "provider_path": "/private/model"}, + ) + assert invalid_request_message(missing_message) == "The model worker rejected the request" + + +def test_usage_rejects_non_finite_numbers() -> None: + assert extract_usage( + { + "usage": { + "prompt_tokens": float("inf"), + "completion_tokens": float("nan"), + "total_tokens": float("-inf"), + } + } + ) == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + +def test_registry_start_is_one_transaction(tmp_path: Path) -> None: + bundle = tmp_path / "chat.bundle" + bundle.write_bytes(b"chat") + created = 0 + counter_lock = threading.Lock() + + def make_worker(spec: ModelSpec) -> WorkerProcess: + nonlocal created + time.sleep(0.05) + with counter_lock: + created += 1 + replica = created + return WorkerProcess( + name=f"{spec.name}-{replica}", + bundle=spec.bundle, + trtmc_binary=FAKE_TRTMC, + startup_timeout=1, + request_timeout=1, + load_options=WorkerLoadOptions(runtime_root=str(tmp_path)), + ) + + registry = ModelRegistry( + [ModelSpec("chat", bundle, "chat")], + trtmc_binary=FAKE_TRTMC, + model_replicas={"chat": 2}, + startup_timeout=1, + request_timeout=1, + worker_factory=make_worker, + ) + try: + with ThreadPoolExecutor(max_workers=2) as executor: + starts = [executor.submit(registry.start) for _ in range(2)] + for start in starts: + start.result() + assert created == 2 + assert registry.status()["models"]["chat"]["ready_replicas"] == 2 + finally: + registry.close() + + +def test_registry_rejects_mismatched_replicas_and_rolls_back(tmp_path: Path) -> None: + canonical = tmp_path / "asr.bundle" + incompatible = tmp_path / "protocol-v2-asr.bundle" + canonical.write_bytes(b"asr") + incompatible.write_bytes(b"asr") + workers: list[WorkerProcess] = [] + + def make_worker(spec: ModelSpec) -> WorkerProcess: + bundle = canonical if not workers else incompatible + worker = WorkerProcess( + name=f"{spec.name}-{len(workers) + 1}", + bundle=bundle, + trtmc_binary=FAKE_TRTMC, + startup_timeout=1, + request_timeout=1, + load_options=WorkerLoadOptions(runtime_root=str(tmp_path)), + ) + workers.append(worker) + return worker + + registry = ModelRegistry( + [ModelSpec("asr", canonical, "transcription")], + trtmc_binary=FAKE_TRTMC, + load_options=WorkerLoadOptions(runtime_root=str(tmp_path)), + model_replicas={"asr": 2}, + startup_timeout=1, + request_timeout=1, + worker_factory=make_worker, + ) + with pytest.raises(WorkerProtocolError, match="unsupported protocol version"): + registry.start() + assert not registry.ready + assert len(workers) == 2 + assert all(worker.state == "closed" for worker in workers) + + +@pytest.mark.parametrize( + ("bundle_stem", "message"), + ( + ("wrong-capability-chat", "required capability"), + ("duplicate-capability-chat", "unique non-empty strings"), + ("unknown-ready-field-chat", "unsupported field"), + ), +) +def test_registry_rejects_noncanonical_v3_ready_metadata( + tmp_path: Path, bundle_stem: str, message: str +) -> None: + registry = make_single_chat_registry( + tmp_path, + bundle_stem, + request_timeout=1, + ) + with pytest.raises(WorkerProtocolError, match=message): + registry.start() + assert not registry.ready + + +def wav_fixture(samples: bytes = b"\x01\x00\x02\x00", *, sample_width: int = 2) -> bytes: + output = io.BytesIO() + with wave.open(output, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(sample_width) + wav.setframerate(16000) + wav.writeframes(samples) + return output.getvalue() + + +def typed_wav_fixture(audio_format: int, bits_per_sample: int, samples: bytes) -> bytes: + sample_width = bits_per_sample // 8 + channels = 1 + sample_rate = 16000 + block_align = channels * sample_width + byte_rate = sample_rate * block_align + fmt = struct.pack( + " bytes: + padding = b"\x00" if len(data) & 1 else b"" + return name + struct.pack(" ModelRegistry: + chat = tmp_path / "chat.bundle" + asr = tmp_path / "asr.bundle" + chat.write_bytes(b"chat") + asr.write_bytes(b"asr") + return ModelRegistry( + [ + ModelSpec("chat", chat, "chat"), + ModelSpec("asr", asr, "transcription"), + ], + trtmc_binary=FAKE_TRTMC, + load_options=WorkerLoadOptions(runtime_root=str(tmp_path)), + startup_timeout=1, + request_timeout=1, + model_replicas=model_replicas, + required_streaming_transcription=required_streaming, + ) + + +def make_single_chat_registry( + tmp_path: Path, + bundle_stem: str, + *, + request_timeout: float, + replicas: int = 1, +) -> ModelRegistry: + bundle = tmp_path / f"{bundle_stem}.bundle" + bundle.write_bytes(b"chat") + return ModelRegistry( + [ModelSpec("chat", bundle, "chat")], + trtmc_binary=FAKE_TRTMC, + load_options=WorkerLoadOptions(runtime_root=str(tmp_path)), + startup_timeout=1, + request_timeout=request_timeout, + model_replicas={"chat": replicas}, + ) + + +def make_single_asr_registry( + tmp_path: Path, + bundle_stem: str, + *, + require_streaming: bool = False, + replicas: int = 1, +) -> ModelRegistry: + bundle = tmp_path / f"{bundle_stem}.bundle" + bundle.write_bytes(b"asr") + return ModelRegistry( + [ModelSpec("asr", bundle, "transcription")], + trtmc_binary=FAKE_TRTMC, + load_options=WorkerLoadOptions(runtime_root=str(tmp_path)), + startup_timeout=1, + request_timeout=1, + model_replicas={"asr": replicas}, + required_streaming_transcription=("asr",) if require_streaming else (), + ) + + +def authorization() -> dict[str, str]: + return {"Authorization": "Bearer test-token"} + + +def chat_request(prompt: str, **options: object) -> dict[str, object]: + return { + "messages": [{"role": "user", "content": prompt}], + **options, + } + + +def chat_text(response: Response) -> str: + return str(response.json()["choices"][0]["message"]["content"]) + + +def test_worker_request_releases_lane_when_submit_fails() -> None: + class FailingSession: + closed = False + + def submit(self, *_args: object, **_kwargs: object) -> object: + raise RuntimeError("executor is unavailable") + + def close(self) -> None: + self.closed = True + + session = FailingSession() + with pytest.raises(RuntimeError, match="executor is unavailable"): + asyncio.run(_worker_request(session, "generate", {})) + assert session.closed + + +def test_worker_request_cancellation_retains_lane_until_native_completion() -> None: + class PendingSession: + def __init__(self) -> None: + self.future: Future[object] = Future() + self.closed = False + + def submit(self, *_args: object, **_kwargs: object) -> Future[object]: + return self.future + + def close(self) -> None: + self.closed = True + + async def exercise() -> None: + session = PendingSession() + request = asyncio.create_task(_worker_request(session, "generate", {})) # type: ignore[arg-type] + await asyncio.sleep(0) + + request.cancel() + with pytest.raises(asyncio.CancelledError): + await request + assert session.closed is False + + session.future.set_result({"text": "completed after disconnect"}) + for _ in range(10): + if session.closed: + break + await asyncio.sleep(0) + assert session.closed is True + + asyncio.run(exercise()) + + +def test_health_readiness_models_and_bearer_auth(tmp_path: Path) -> None: + app = create_app(make_registry(tmp_path), config=ServerConfig(api_key="test-token")) + with TestClient(app) as client: + health = client.get("/healthz") + assert health.status_code == 200 + assert health.json() == {"status": "ok", "degraded": False} + assert client.get("/readyz").status_code == 401 + + readiness = client.get("/readyz", headers=authorization()) + assert readiness.status_code == 200 + readiness_payload = readiness.json() + assert readiness_payload["ready"] is True + assert readiness_payload["degraded"] is False + assert set(readiness_payload["models"]) == {"chat", "asr"} + for model_status in readiness_payload["models"].values(): + assert set(model_status) == { + "ready", + "degraded", + "replicas", + "ready_replicas", + "idle_replicas", + "busy", + } + assert {"pid", "pids", "returncode", "error"}.isdisjoint(model_status) + + denied = client.get("/v1/models") + assert denied.status_code == 401 + assert denied.headers["www-authenticate"] == "Bearer" + + non_ascii = client.get( + "/v1/models", + headers=[(b"authorization", b"Bearer \xff")], + ) + assert non_ascii.status_code == 401 + + response = client.get("/v1/models", headers=authorization()) + assert response.status_code == 200 + models = response.json()["data"] + assert [model["id"] for model in models] == ["chat", "asr"] + assert models[0]["capabilities"] == ["text_generation"] + assert models[0]["metadata"] == { + "default_max_new_tokens": 64, + "streaming_transcription": False, + } + assert models[1]["capabilities"] == [ + "transcription", + "transcription_streaming", + ] + assert models[1]["metadata"] == { + "streaming_transcription": True, + } + + preflight = client.options( + "/v1/chat/completions", + headers={ + "Origin": "null", + "Access-Control-Request-Method": "POST", + "Access-Control-Request-Headers": "authorization,content-type", + }, + ) + assert preflight.status_code == 200 + assert preflight.headers["access-control-allow-origin"] == "null" + assert "Authorization" in preflight.headers["access-control-allow-headers"] + + local_origin = client.options( + "/v1/chat/completions", + headers={ + "Origin": "http://127.0.0.1:4173", + "Access-Control-Request-Method": "POST", + }, + ) + assert local_origin.status_code == 200 + assert local_origin.headers["access-control-allow-origin"] == ("http://127.0.0.1:4173") + + denied_origin = client.options( + "/v1/chat/completions", + headers={ + "Origin": "https://example.com", + "Access-Control-Request-Method": "POST", + }, + ) + assert denied_origin.status_code == 400 + assert "access-control-allow-origin" not in denied_origin.headers + + +def test_chat_generation_and_unsupported_semantics_are_explicit(tmp_path: Path) -> None: + app = create_app(make_registry(tmp_path)) + with TestClient(app) as client: + assert client.get("/readyz").json() == { + "status": "ready", + "ready": True, + "degraded": False, + } + stopped = client.post( + "/v1/chat/completions", + json=chat_request("hello STOP tail", stop="STOP", max_tokens=8), + ) + assert stopped.status_code == 400 + assert stopped.json()["error"]["param"] == "stop" + + thinking_disabled = client.post( + "/v1/chat/completions", + json=chat_request("summary", enable_thinking=False), + ) + assert thinking_disabled.status_code == 200 + assert chat_text(thinking_disabled) == ("enable_thinking=false:generated:summary") + + multi_turn = client.post( + "/v1/chat/completions", + json={ + "messages": [ + {"role": "system", "content": "Be brief"}, + {"role": "user", "content": "Hi"}, + ] + }, + ) + assert multi_turn.status_code == 400 + assert multi_turn.json()["error"]["param"] == "messages" + + thinking_enabled = client.post( + "/v1/chat/completions", + json={ + "messages": [{"role": "user", "content": "reason"}], + "enable_thinking": True, + }, + ) + assert thinking_enabled.status_code == 200 + assert chat_text(thinking_enabled) == ("enable_thinking=true:generated:reason") + + streamed = client.post( + "/v1/chat/completions", + json=chat_request("stream", stream=True), + ) + assert streamed.status_code == 400 + assert streamed.json()["error"]["code"] == "streaming_not_supported" + + single = prepare_chat_prompt([{"role": "user", "content": "question"}]) + assert single == "question" + with pytest.raises(ValueError, match="exactly one"): + prepare_chat_prompt( + [ + {"role": "system", "content": "policy"}, + {"role": "user", "content": "question"}, + ] + ) + + +def test_chat_rejects_empty_messages(tmp_path: Path) -> None: + app = create_app(make_registry(tmp_path)) + with TestClient(app) as client: + response = client.post("/v1/chat/completions", json={"messages": []}) + + assert response.status_code == 422 + assert response.json()["error"]["code"] == "invalid_request" + assert response.json()["error"]["param"] == "messages" + + +@pytest.mark.parametrize( + ("parameter", "value"), + [ + ("n", 2), + ("best_of", 2), + ("logprobs", True), + ("top_logprobs", 2), + ("frequency_penalty", 0.5), + ("presence_penalty", 0.5), + ("logit_bias", {"42": 1}), + ("ignore_eos", True), + ("tools", [{"type": "function"}]), + ("tool_choice", "auto"), + ("parallel_tool_calls", True), + ("response_format", {"type": "json_object"}), + ("stream_options", {"include_usage": True}), + ("functions", [{"name": "lookup"}]), + ("function_call", "auto"), + ("modalities", ["text", "audio"]), + ("audio", {"format": "wav"}), + ("prediction", {"type": "content", "content": "expected"}), + ("reasoning_effort", "high"), + ("future_semantic_option", True), + ], +) +def test_chat_rejects_meaningful_unsupported_parameters( + tmp_path: Path, + parameter: str, + value: object, +) -> None: + app = create_app(make_registry(tmp_path)) + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json=chat_request("hello", **{parameter: value}), + ) + + assert response.status_code == 400 + assert response.json()["error"] == { + "message": f"{parameter} is not supported", + "type": "invalid_request_error", + "param": parameter, + "code": "unsupported_parameter", + } + + +def test_chat_accepts_no_op_and_metadata_fields(tmp_path: Path) -> None: + app = create_app(make_registry(tmp_path)) + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json=chat_request( + "hello", + n=1, + best_of=1, + logprobs=False, + top_logprobs=0, + frequency_penalty=0, + presence_penalty=0, + logit_bias={}, + ignore_eos=False, + tools=[], + tool_choice="none", + parallel_tool_calls=False, + response_format={"type": "text"}, + stream_options={}, + user="client-metadata", + metadata={"request_class": "interactive"}, + ), + ) + + assert response.status_code == 200 + + +@pytest.mark.parametrize( + "message", + [ + {"role": "tool", "content": "tool output"}, + {"role": "assistant", "content": "prior answer"}, + {"role": "developer", "content": "policy"}, + {"role": "system", "content": "policy"}, + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}} + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "hello", "unsupported": True}], + }, + {"role": "user", "content": "hello", "name": "named-participant"}, + ], +) +def test_chat_rejects_unsupported_message_semantics( + tmp_path: Path, + message: dict[str, object], +) -> None: + app = create_app(make_registry(tmp_path)) + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json={"messages": [message]}, + ) + + assert response.status_code == 400 + assert response.json()["error"]["code"] == "unsupported_parameter" + assert response.json()["error"]["param"] == "messages" + + +@pytest.mark.parametrize("part_type", [[], {}, None, 7]) +def test_chat_rejects_non_string_content_part_types( + tmp_path: Path, + part_type: object, +) -> None: + app = create_app(make_registry(tmp_path)) + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json={ + "messages": [ + { + "role": "user", + "content": [{"type": part_type, "text": "hello"}], + } + ] + }, + ) + + assert response.status_code == 400 + assert response.json()["error"]["param"] == "messages" + + +def test_model_capability_and_unknown_model_errors_are_openai_shaped( + tmp_path: Path, +) -> None: + app = create_app(make_registry(tmp_path)) + with TestClient(app) as client: + wrong_kind = client.post("/v1/chat/completions", json=chat_request("hello", model="asr")) + assert wrong_kind.status_code == 400 + assert wrong_kind.json()["error"]["code"] == "model_capability_mismatch" + + missing = client.post("/v1/chat/completions", json=chat_request("hello", model="missing")) + assert missing.status_code == 404 + assert missing.json()["error"]["param"] == "model" + + invalid = client.post("/v1/chat/completions", json={"max_tokens": 1}) + assert invalid.status_code == 422 + assert invalid.json()["error"]["code"] == "invalid_request" + assert invalid.json()["error"]["param"] == "messages" + + +def test_oversized_prompt_is_rejected_before_worker_protocol(tmp_path: Path) -> None: + registry = make_registry(tmp_path) + app = create_app(registry, config=ServerConfig(max_prompt_bytes=5)) + with TestClient(app) as client: + oversized = client.post("/v1/chat/completions", json=chat_request("123456")) + assert oversized.status_code == 413 + assert oversized.json()["error"]["code"] == "prompt_too_large" + + healthy = client.post("/v1/chat/completions", json=chat_request("ok")) + assert healthy.status_code == 200 + assert registry.ready + + +def test_generation_token_hard_cap_applies_to_explicit_and_default_requests( + tmp_path: Path, +) -> None: + app = create_app(make_registry(tmp_path), config=ServerConfig(max_generation_tokens=4)) + with TestClient(app) as client: + rejected = client.post("/v1/chat/completions", json=chat_request("hello", max_tokens=5)) + assert rejected.status_code == 400 + assert rejected.json()["error"]["code"] == "max_tokens_exceeded" + + bounded_default = client.post("/v1/chat/completions", json=chat_request("hello")) + assert bounded_default.status_code == 200 + assert bounded_default.json()["trtmc"]["effective_max_tokens"] == 4 + + +def test_ready_capabilities_control_exposed_capability_and_startup(tmp_path: Path) -> None: + registry = make_registry(tmp_path, required_streaming=("asr",)) + app = create_app(registry) + with TestClient(app) as client: + models = client.get("/v1/models").json()["data"] + asr = next(model for model in models if model["id"] == "asr") + assert asr["metadata"]["streaming_transcription"] is True + + unsupported = make_single_asr_registry(tmp_path, "no-stream-asr", require_streaming=True) + with pytest.raises(WorkerProtocolError, match="does not advertise transcription_streaming"): + unsupported.start() + assert not unsupported.ready + + old_protocol = make_single_asr_registry(tmp_path, "protocol-v2-asr") + with pytest.raises(WorkerProtocolError, match="unsupported protocol version"): + old_protocol.start() + assert not old_protocol.ready + + +def test_required_streaming_probe_failure_rolls_back_every_replica(tmp_path: Path) -> None: + bundle = tmp_path / "advertises-stream-but-probe-fails-asr.bundle" + bundle.write_bytes(b"asr") + workers: list[WorkerProcess] = [] + + def make_worker(spec: ModelSpec) -> WorkerProcess: + worker = WorkerProcess( + name=f"{spec.name}-{len(workers) + 1}", + bundle=spec.bundle, + trtmc_binary=FAKE_TRTMC, + startup_timeout=1, + request_timeout=1, + load_options=WorkerLoadOptions(runtime_root=str(tmp_path)), + ) + workers.append(worker) + return worker + + registry = ModelRegistry( + [ModelSpec("asr", bundle, "transcription")], + trtmc_binary=FAKE_TRTMC, + model_replicas={"asr": 2}, + required_streaming_transcription=("asr",), + worker_factory=make_worker, + ) + with pytest.raises(WorkerRemoteError): + registry.start() + assert len(workers) == 2 + assert all(worker.state == "closed" for worker in workers) + assert not registry.ready + + +def test_saturated_worker_returns_429_without_waiting_for_active_request( + tmp_path: Path, +) -> None: + registry = make_single_chat_registry( + tmp_path, + "saturation-chat", + request_timeout=1, + ) + app = create_app(registry) + with TestClient(app) as client, ThreadPoolExecutor(max_workers=1) as executor: + active = executor.submit(client.post, "/v1/chat/completions", json=chat_request("first")) + deadline = time.monotonic() + 1 + while not registry.status()["models"]["chat"]["busy"] and time.monotonic() < deadline: + time.sleep(0.005) + assert registry.status()["models"]["chat"]["busy"] + assert client.get("/healthz").status_code == 200 + saturated = client.post("/v1/chat/completions", json=chat_request("second")) + assert saturated.status_code == 429 + assert not active.done() + assert saturated.json()["error"]["code"] == "server_busy" + assert saturated.headers["retry-after"] == "1" + assert active.result().status_code == 200 + + +def test_model_replicas_execute_requests_in_parallel_and_bound_overload( + tmp_path: Path, +) -> None: + registry = make_single_chat_registry( + tmp_path, + "saturation-chat", + request_timeout=1, + replicas=2, + ) + app = create_app(registry) + with TestClient(app) as client, ThreadPoolExecutor(max_workers=2) as executor: + active = [ + executor.submit(client.post, "/v1/chat/completions", json=chat_request(prompt)) + for prompt in ("first", "second") + ] + deadline = time.monotonic() + 1 + while registry.status()["models"]["chat"]["idle_replicas"] and time.monotonic() < deadline: + time.sleep(0.005) + assert registry.status()["models"]["chat"]["idle_replicas"] == 0 + + saturated = client.post("/v1/chat/completions", json=chat_request("third")) + assert saturated.status_code == 429 + assert saturated.json()["error"]["code"] == "server_busy" + assert [response.result().status_code for response in active] == [200, 200] + + +def test_health_remains_ok_while_one_replica_can_serve(tmp_path: Path) -> None: + registry = make_single_chat_registry( + tmp_path, + "crash-request-chat", + request_timeout=1, + replicas=2, + ) + app = create_app(registry) + with TestClient(app) as client: + failed = client.post("/v1/chat/completions", json=chat_request("trigger")) + assert failed.status_code == 503 + assert registry.status()["models"]["chat"]["ready_replicas"] == 1 + assert client.get("/healthz").json() == {"status": "ok", "degraded": True} + assert client.get("/readyz").status_code == 200 + assert client.get("/readyz").json()["degraded"] is True + + +def test_health_and_registry_readiness_have_distinct_failure_scopes(tmp_path: Path) -> None: + chat = tmp_path / "crash-request-chat.bundle" + asr = tmp_path / "healthy-asr.bundle" + chat.write_bytes(b"chat") + asr.write_bytes(b"asr") + registry = ModelRegistry( + [ + ModelSpec("chat", chat, "chat"), + ModelSpec("asr", asr, "transcription"), + ], + trtmc_binary=FAKE_TRTMC, + load_options=WorkerLoadOptions(runtime_root=str(tmp_path)), + request_timeout=1, + ) + app = create_app(registry) + with TestClient(app) as client: + failed = client.post("/v1/chat/completions", json=chat_request("trigger")) + assert failed.status_code == 503 + assert client.get("/healthz").status_code == 200 + assert client.get("/readyz").status_code == 503 + + +@pytest.mark.parametrize( + ( + "bundle_stem", + "expected_status", + "expected_code", + "expected_message", + "private_detail", + ), + [ + ( + "slow-request-chat", + 504, + "worker_timeout", + "The model worker timed out", + "slow-request-chat", + ), + ( + "crash-request-chat", + 503, + "worker_crashed", + "The model worker is unavailable", + "intentional fake crash", + ), + ], +) +def test_worker_failures_are_http_errors_and_clear_readiness( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + bundle_stem: str, + expected_status: int, + expected_code: str, + expected_message: str, + private_detail: str, +) -> None: + registry = make_single_chat_registry(tmp_path, bundle_stem, request_timeout=0.05) + app = create_app(registry, config=ServerConfig(api_key="test-token")) + with TestClient(app) as client: + failed = client.post( + "/v1/chat/completions", + json=chat_request("trigger"), + headers=authorization(), + ) + assert failed.status_code == expected_status + public_error = failed.json()["error"] + assert public_error["code"] == expected_code + assert public_error["message"] == expected_message + assert private_detail not in failed.text + assert "worker-secret" not in failed.text + assert any(record.message == "Model worker request failed" for record in caplog.records) + assert private_detail not in caplog.text + assert "worker-secret" not in caplog.text + + health = client.get("/healthz") + assert health.status_code == 503 + assert health.json() == {"status": "unavailable", "degraded": False} + assert private_detail not in health.text + assert "worker-secret" not in health.text + + readiness = client.get("/readyz", headers=authorization()) + assert readiness.status_code == 503 + assert private_detail not in readiness.text + assert "worker-secret" not in readiness.text + registry_status = str(registry.status()) + assert private_detail not in registry_status + assert "worker-secret" not in registry_status + for model_status in readiness.json()["models"].values(): + assert {"pid", "pids", "returncode", "error"}.isdisjoint(model_status) + + +def test_native_invalid_request_maps_to_400_without_crashing_worker( + tmp_path: Path, +) -> None: + registry = make_single_chat_registry(tmp_path, "invalid-request-chat", request_timeout=1) + app = create_app(registry) + with TestClient(app) as client: + failed = client.post("/v1/chat/completions", json=chat_request("trigger")) + assert failed.status_code == 400 + assert failed.json()["error"]["code"] == "invalid_request" + assert "invalid fake generation request" in failed.json()["error"]["message"] + assert client.get("/readyz").status_code == 200 + + +def test_multipart_ingress_limit_precedes_upload_processing(tmp_path: Path) -> None: + registry = make_registry(tmp_path) + app = create_app(registry, config=ServerConfig(max_upload_bytes=4)) + oversized = b"x" * (_HTTP_ENVELOPE_OVERHEAD_BYTES + 5) + + with TestClient(app) as client: + response = client.post( + "/v1/audio/transcriptions", + files={"file": ("oversized.wav", oversized, "audio/wav")}, + ) + assert registry.ready + + assert response.status_code == 413 + assert response.json()["error"]["code"] == "request_body_too_large" + assert response.headers["connection"] == "close" + + +def test_multipart_audio_transcription_json_text_and_empty_rejection( + tmp_path: Path, +) -> None: + app = create_app(make_registry(tmp_path)) + with TestClient(app) as client: + audio = wav_fixture() + response = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", audio, "audio/wav")}, + data={"model": "asr"}, + ) + assert response.status_code == 200 + assert response.json() == {"text": f"transcribed {len(audio)} bytes"} + + shorter_audio = wav_fixture(b"\x01\x00") + text = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", shorter_audio, "audio/wav")}, + data={"response_format": "text"}, + ) + assert text.status_code == 200 + assert text.text == f"transcribed {len(shorter_audio)} bytes" + + empty = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", b"", "audio/wav")}, + ) + assert empty.status_code == 400 + assert empty.json()["error"]["code"] == "empty_audio" + + unsupported = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.mp3", b"not an mp3 either", "audio/mpeg")}, + ) + assert unsupported.status_code == 415 + assert unsupported.json()["error"]["code"] == "unsupported_media_type" + assert unsupported.json()["error"]["param"] == "file" + + unsupported_wav = client.post( + "/v1/audio/transcriptions", + files={ + "file": ( + "sample.wav", + wav_fixture(b"\x01\x02", sample_width=1), + "audio/wav", + ) + }, + ) + assert unsupported_wav.status_code == 415 + assert "PCM16 or IEEE float32" in unsupported_wav.json()["error"]["message"] + + +def test_audio_transcription_accepts_ieee_float32_but_rejects_pcm32( + tmp_path: Path, +) -> None: + float32_wav = typed_wav_fixture(3, 32, struct.pack(" None: + async def native_result( + session: WorkerSession, + _operation: str, + _payload: object, + ) -> object: + session.close() + return { + "text": "public transcript", + "model_path": "/private/models/asr.bundle", + "token_ids": [11, 12], + "setup_ms": 1.25, + "prefill_ms": 2.5, + "decode_ms": 3.75, + "segments": [ + { + "start_seconds": 0, + "end_seconds": 1.5, + "text": "public transcript", + "token_ids": [11, 12], + "provider_debug": "private provider state", + } + ], + } + + monkeypatch.setattr("trtmc_server.app._worker_request", native_result) + app = create_app(make_registry(tmp_path)) + with TestClient(app) as client: + response = client.post( + "/v1/audio/transcriptions", + files={"file": ("sample.wav", wav_fixture(), "audio/wav")}, + data={"model": "asr", "response_format": "verbose_json"}, + ) + + assert response.status_code == 200 + assert response.json() == { + "text": "public transcript", + "model": "asr", + "segments": [ + { + "start_seconds": 0.0, + "end_seconds": 1.5, + "text": "public transcript", + } + ], + } + for private_field in ( + "model_path", + "token_ids", + "setup_ms", + "prefill_ms", + "decode_ms", + "provider_debug", + ): + assert private_field not in response.text + + +def test_realtime_transcription_events_use_cumulative_transcript(tmp_path: Path) -> None: + app = create_app(make_registry(tmp_path), config=ServerConfig(api_key="test-token")) + with TestClient(app) as client: + with client.websocket_connect( + "/v1/realtime?intent=transcription&access_token=test-token" + ) as websocket: + created = websocket.receive_json() + assert created["type"] == "session.created" + assert created["session"]["model"] == "asr" + + websocket.send_json( + { + "type": "session.update", + "session": { + "model": "asr", + "input_audio_format": "pcm16", + "trtmc": {"sample_rate_hz": 16000}, + }, + } + ) + updated = websocket.receive_json() + assert updated["type"] == "session.updated" + assert updated["session"]["trtmc"]["sample_rate_hz"] == 16000 + + websocket.send_json( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(b"\x01\x00\x02\x00").decode(), + } + ) + delta = websocket.receive_json() + assert delta["type"] == ("conversation.item.input_audio_transcription.delta") + assert delta["transcript"] == "2 samples" + + websocket.send_json( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(b"\x03\x00").decode(), + } + ) + second_delta = websocket.receive_json() + assert second_delta["transcript"] == "3 samples" + + websocket.send_json({"type": "input_audio_buffer.commit"}) + completed = websocket.receive_json() + assert completed["type"] == ("conversation.item.input_audio_transcription.completed") + assert completed["transcript"] == "3 samples" + + websocket.send_json({"type": "input_audio_buffer.clear"}) + assert websocket.receive_json()["type"] == "input_audio_buffer.cleared" + + +def test_realtime_reconfiguration_starts_a_fresh_item(tmp_path: Path) -> None: + app = create_app(make_registry(tmp_path), config=ServerConfig(api_key="test-token")) + with TestClient(app) as client: + with client.websocket_connect( + "/v1/realtime?intent=transcription&access_token=test-token" + ) as websocket: + assert websocket.receive_json()["type"] == "session.created" + + audio = base64.b64encode(b"\x01\x00").decode() + websocket.send_json({"type": "input_audio_buffer.append", "audio": audio}) + first_delta = websocket.receive_json() + assert first_delta["delta"] == "1 samples" + first_item_id = first_delta["item_id"] + + websocket.send_json( + { + "type": "session.update", + "session": {"trtmc": {"sample_rate_hz": 16000}}, + } + ) + assert websocket.receive_json()["type"] == "session.updated" + + websocket.send_json({"type": "input_audio_buffer.append", "audio": audio}) + next_delta = websocket.receive_json() + assert next_delta["delta"] == "1 samples" + assert next_delta["transcript"] == "1 samples" + assert next_delta["item_id"] != first_item_id + + +def test_realtime_reset_clears_local_item_without_active_lease() -> None: + class Registry: + default_transcription_model = "asr" + + connection = RealtimeTranscriptionConnection( + object(), # type: ignore[arg-type] + Registry(), # type: ignore[arg-type] + ) + connection.transcript = "stale transcript" + previous_item_id = connection.item_id + connection.total_audio_bytes = 2 + + asyncio.run(connection._release_stream(reset=True)) # noqa: SLF001 + + assert connection.transcript == "" + assert connection.item_id != previous_item_id + assert connection.total_audio_bytes == 2 + + +def test_realtime_forwards_validated_audio_wire_text_without_reencoding( + tmp_path: Path, +) -> None: + encoded = "AQB=" + assert base64.b64encode(base64.b64decode(encoded, validate=True)).decode() != encoded + app = create_app( + make_single_asr_registry(tmp_path, "echo-wire-asr"), + ) + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime?intent=transcription") as websocket: + assert websocket.receive_json()["type"] == "session.created" + websocket.send_json( + { + "type": "input_audio_buffer.append", + "audio": encoded, + } + ) + delta = websocket.receive_json() + assert delta["transcript"] == encoded + + +def test_realtime_rejects_bad_token_and_invalid_audio(tmp_path: Path) -> None: + app = create_app(make_registry(tmp_path), config=ServerConfig(api_key="test-token")) + with TestClient(app) as client: + with pytest.raises(WebSocketDisconnect) as denied: + with client.websocket_connect("/v1/realtime?intent=transcription"): + pass + assert denied.value.code == 4401 + + with pytest.raises(WebSocketDisconnect) as non_ascii_token: + with client.websocket_connect("/v1/realtime?intent=transcription&access_token=%C3%BF"): + pass + assert non_ascii_token.value.code == 4401 + + with pytest.raises(WebSocketDisconnect) as denied_origin: + with client.websocket_connect( + "/v1/realtime?intent=transcription&access_token=test-token", + headers={"Origin": "https://example.com"}, + ): + pass + assert denied_origin.value.code == 4403 + + with client.websocket_connect( + "/v1/realtime?intent=transcription&access_token=test-token", + headers={"Origin": "http://localhost:4173"}, + ) as websocket: + websocket.receive_json() + websocket.send_json({"type": "input_audio_buffer.append", "audio": "not base64"}) + error = websocket.receive_json() + assert error["type"] == "error" + assert error["error"]["code"] == "invalid_audio" + + +def test_realtime_handles_distinct_asyncio_timeout_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class LegacyAsyncioTimeoutError(Exception): + pass + + class TimeoutWebSocket: + def __init__(self) -> None: + self.events: list[dict[str, object]] = [] + + async def send_json(self, event: dict[str, object]) -> None: + self.events.append(event) + + async def receive_json(self) -> dict[str, object]: + raise LegacyAsyncioTimeoutError + + class Registry: + default_transcription_model = "asr" + + websocket = TimeoutWebSocket() + monkeypatch.setattr( + realtime_module, + "_TIMEOUT_ERRORS", + (TimeoutError, LegacyAsyncioTimeoutError), + ) + connection = RealtimeTranscriptionConnection( + websocket, # type: ignore[arg-type] + Registry(), # type: ignore[arg-type] + idle_timeout_seconds=1, + ) + + asyncio.run(connection.run()) + + assert websocket.events[0]["type"] == "session.created" + assert websocket.events[1]["error"]["code"] == "session_idle_timeout" # type: ignore[index] + + +def test_realtime_unexpected_error_logs_only_safe_context( + caplog: pytest.LogCaptureFixture, +) -> None: + private_detail = "access_token=client-secret /private/model client-prompt" + + class FailingWebSocket: + def __init__(self) -> None: + self.events: list[dict[str, object]] = [] + + async def send_json(self, event: dict[str, object]) -> None: + self.events.append(event) + + async def receive_json(self) -> dict[str, object]: + raise RuntimeError(private_detail) + + class Registry: + default_transcription_model = "asr" + + websocket = FailingWebSocket() + connection = RealtimeTranscriptionConnection( + websocket, # type: ignore[arg-type] + Registry(), # type: ignore[arg-type] + ) + caplog.set_level(logging.ERROR, logger="trtmc_server.realtime") + + asyncio.run(connection.run()) + + failure = websocket.events[-1]["error"] # type: ignore[index] + assert failure == { + "type": "server_error", + "code": "internal_error", + "message": "realtime transcription session failed", + } + record = next( + record + for record in caplog.records + if record.message.startswith("Unexpected realtime session failure") + ) + assert connection.connection_id in record.message + assert "RuntimeError" in record.message + assert "receive_json" in record.message + assert record.exc_info is None + rendered_events = repr(websocket.events) + for sensitive in ( + "access_token", + "authorization", + "client-secret", + "/private/model", + "client-prompt", + ): + assert sensitive not in caplog.text + assert sensitive not in rendered_events + assert record.message.count(" <- ") <= 7 + assert "/" not in record.message + assert "Traceback" not in caplog.text + + +def test_realtime_worker_diagnostics_are_not_returned_to_clients( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + app = create_app(make_single_asr_registry(tmp_path, "crash-asr")) + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime?intent=transcription") as websocket: + websocket.receive_json() + websocket.send_json( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(b"\x01\x00").decode(), + } + ) + failure = websocket.receive_json() + + assert failure["type"] == "conversation.item.input_audio_transcription.failed" + assert failure["error"]["code"] == "worker_crashed" + assert failure["error"]["message"] == "The model worker is unavailable" + assert "intentional fake crash" not in str(failure) + assert any( + record.message == "Realtime model worker request failed" for record in caplog.records + ) + assert "intentional fake crash" not in caplog.text + assert "worker-secret" not in caplog.text + + +def test_realtime_chunk_error_resets_native_session_for_next_turn(tmp_path: Path) -> None: + registry = make_single_asr_registry(tmp_path, "chunk-error-asr") + app = create_app(registry) + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime?intent=transcription") as websocket: + websocket.receive_json() + audio = base64.b64encode(b"\x01\x00").decode() + websocket.send_json({"type": "input_audio_buffer.append", "audio": audio}) + error = websocket.receive_json() + assert error["type"] == "error" + assert error["error"]["code"] == "invalid_request" + assert error["error"]["message"] == "invalid fake audio chunk" + + websocket.send_json({"type": "input_audio_buffer.append", "audio": audio}) + recovered = websocket.receive_json() + assert recovered["type"] == ("conversation.item.input_audio_transcription.delta") + assert recovered["transcript"] == "1 samples" + websocket.send_json({"type": "input_audio_buffer.commit"}) + assert websocket.receive_json()["type"] == ( + "conversation.item.input_audio_transcription.completed" + ) + + +def test_realtime_sessions_lease_distinct_replicas_and_bound_overload( + tmp_path: Path, +) -> None: + registry = make_single_asr_registry(tmp_path, "asr-replicas", replicas=2) + app = create_app(registry) + audio = base64.b64encode(b"\x01\x00").decode() + with TestClient(app) as client: + with ( + client.websocket_connect("/v1/realtime?intent=transcription") as first, + client.websocket_connect("/v1/realtime?intent=transcription") as second, + client.websocket_connect("/v1/realtime?intent=transcription") as third, + ): + for websocket in (first, second, third): + assert websocket.receive_json()["type"] == "session.created" + for websocket in (first, second): + websocket.send_json({"type": "input_audio_buffer.append", "audio": audio}) + assert websocket.receive_json()["type"].endswith(".delta") + + third.send_json({"type": "input_audio_buffer.append", "audio": audio}) + rejected = third.receive_json() + assert rejected["type"] == "error" + assert rejected["error"]["code"] == "server_busy" + + for websocket in (first, second): + websocket.send_json({"type": "input_audio_buffer.commit"}) + assert websocket.receive_json()["type"].endswith(".completed") + + +def test_realtime_disconnect_during_stream_start_resets_lane(tmp_path: Path) -> None: + registry = make_single_asr_registry(tmp_path, "asr-delayed-start") + app = create_app(registry) + audio = base64.b64encode(b"\x01\x00").decode() + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime?intent=transcription") as websocket: + websocket.receive_json() + websocket.send_json({"type": "input_audio_buffer.append", "audio": audio}) + deadline = time.monotonic() + 1 + while not registry.status()["models"]["asr"]["busy"] and time.monotonic() < deadline: + time.sleep(0.005) + assert registry.status()["models"]["asr"]["busy"] + + with client.websocket_connect("/v1/realtime?intent=transcription") as recovered: + recovered.receive_json() + recovered.send_json({"type": "input_audio_buffer.append", "audio": audio}) + assert recovered.receive_json()["type"].endswith(".delta") + recovered.send_json({"type": "input_audio_buffer.commit"}) + assert recovered.receive_json()["type"].endswith(".completed") + + +@pytest.mark.parametrize( + ("config", "expected_code"), + [ + ({"realtime_idle_timeout_seconds": 0.05}, "session_idle_timeout"), + ( + { + "realtime_idle_timeout_seconds": 1, + "realtime_max_session_seconds": 0.05, + }, + "session_duration_exceeded", + ), + ], +) +def test_realtime_time_limits_emit_structured_failure_and_cleanup( + tmp_path: Path, + config: dict[str, float], + expected_code: str, +) -> None: + # Delaying reset makes client-close cancellation overlap cleanup + # deterministically. The worker lock must still be released before the + # websocket context exits and the next operation starts. + registry = make_single_asr_registry(tmp_path, f"asr-delayed-reset-{expected_code}") + app = create_app(registry, config=ServerConfig(**config)) + audio = base64.b64encode(b"\x01\x00").decode() + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime?intent=transcription") as websocket: + websocket.receive_json() + websocket.send_json({"type": "input_audio_buffer.append", "audio": audio}) + assert websocket.receive_json()["type"].endswith(".delta") + failure = websocket.receive_json() + assert failure["type"].endswith(".failed") + assert failure["error"]["code"] == expected_code + + _spec, session = registry.acquire_session("transcription", "asr") + with session: + probe = session.request( + "probe_transcription_stream", + { + "config": { + "sample_rate_hz": 16000, + "channels": 1, + "audio_format": "pcm16le", + } + }, + ) + assert probe["supported"] is True + + +def test_realtime_audio_limit_is_cumulative_across_clear(tmp_path: Path) -> None: + registry = make_single_asr_registry(tmp_path, "asr-byte-limit") + app = create_app(registry, config=ServerConfig(max_realtime_session_bytes=2)) + audio = base64.b64encode(b"\x01\x00").decode() + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime?intent=transcription") as websocket: + websocket.receive_json() + websocket.send_json({"type": "input_audio_buffer.append", "audio": audio}) + assert websocket.receive_json()["type"].endswith(".delta") + websocket.send_json({"type": "input_audio_buffer.clear"}) + assert websocket.receive_json()["type"] == "input_audio_buffer.cleared" + websocket.send_json({"type": "input_audio_buffer.append", "audio": audio}) + failure = websocket.receive_json() + assert failure["type"].endswith(".failed") + assert failure["error"]["code"] == "audio_session_too_large" diff --git a/server/tests/test_serve_cli.py b/server/tests/test_serve_cli.py new file mode 100644 index 0000000000..f3756bccb0 --- /dev/null +++ b/server/tests/test_serve_cli.py @@ -0,0 +1,611 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import base64 +import json +import logging +import os +import queue +import socket +import subprocess +import sys +import threading +import time +import urllib.error +import urllib.request +from pathlib import Path +from urllib.parse import parse_qs + +import pytest +import websockets + +import trtmc_server +from trtmc_server import cli as serve_cli_module +from trtmc_server.cli import ( + _RedactAccessToken, + bind_socket, + build_parser, + is_loopback_host, + main as serve_main, + parse_model_assignment, + parse_replica_assignment, + resolve_runtime_root, + resolve_trtmc_binary, + validate_bind_policy, +) + + +FAKE_TRTMC = Path(__file__).with_name("fake_serve_worker.py") +REPOSITORY = Path(__file__).resolve().parents[2] +SERVER_SOURCE_ROOT = REPOSITORY / "server" / "python" +_REQUIRES_PROC = pytest.mark.skipif( + not Path("/proc").is_dir(), + reason="process-lifecycle assertions require Linux /proc", +) + + +def _server_process_environment() -> dict[str, str]: + environment = dict(os.environ) + if environment.get("TRTMC_TEST_INSTALLED_WHEEL") == "1": + return environment + source = str(SERVER_SOURCE_ROOT) + existing = environment.get("PYTHONPATH") + environment["PYTHONPATH"] = source if not existing else f"{source}{os.pathsep}{existing}" + return environment + + +def test_server_import_and_subprocess_environment_follow_selected_runtime( + monkeypatch: pytest.MonkeyPatch, +) -> None: + package_path = Path(trtmc_server.__file__).resolve() + installed_wheel = os.environ.get("TRTMC_TEST_INSTALLED_WHEEL") == "1" + assert package_path.is_relative_to(SERVER_SOURCE_ROOT) is not installed_wheel + + monkeypatch.setenv("TRTMC_TEST_INSTALLED_WHEEL", "1") + monkeypatch.setenv("PYTHONPATH", "/selected/site-packages:/source") + selected = _server_process_environment() + assert selected["PYTHONPATH"] == "/selected/site-packages:/source" + assert str(SERVER_SOURCE_ROOT) not in selected["PYTHONPATH"] + + monkeypatch.delenv("TRTMC_TEST_INSTALLED_WHEEL") + source = _server_process_environment() + assert source["PYTHONPATH"].split(os.pathsep, maxsplit=1)[0] == str(SERVER_SOURCE_ROOT) + + +def test_model_assignment_and_bind_policy(tmp_path: Path) -> None: + bundle = tmp_path / "model=revision.bundle" + bundle.write_bytes(b"fixture") + spec = parse_model_assignment(f"asr={bundle}", kind="transcription") + assert spec.name == "asr" + assert spec.bundle == bundle + assert spec.kind == "transcription" + + assert is_loopback_host("127.0.0.1") + assert is_loopback_host("::1") + assert not is_loopback_host("localhost") + assert not is_loopback_host("0.0.0.0") + with pytest.raises(ValueError, match="loopback IP literal"): + validate_bind_policy("0.0.0.0") + with pytest.raises(ValueError, match="loopback IP literal"): + validate_bind_policy("localhost") + validate_bind_policy("127.0.0.1") + + +def test_prebound_port_zero_returns_actual_port() -> None: + listener = bind_socket("127.0.0.1", 0) + try: + assert listener.getsockname()[0] == "127.0.0.1" + assert listener.getsockname()[1] > 0 + assert listener.getblocking() is False + finally: + listener.close() + + +def test_cli_reports_bind_failures_without_a_traceback( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + bundle = tmp_path / "chat.bundle" + bundle.write_bytes(b"fixture") + + def fail_to_bind(_host: str, _port: int) -> socket.socket: + raise OSError("cannot bind test listener") + + monkeypatch.setattr(serve_cli_module, "bind_socket", fail_to_bind) + with pytest.raises(SystemExit) as failure: + serve_main( + [ + "--chat-model", + f"chat={bundle}", + "--trtmc-binary", + str(FAKE_TRTMC), + "--runtime-root", + str(tmp_path), + "--api-key", + "test-token", + ] + ) + + assert failure.value.code == 2 + stderr = capsys.readouterr().err + assert "error: cannot bind test listener" in stderr + assert "Traceback" not in stderr + + +def test_model_replicas_are_explicit_and_default_to_one() -> None: + assert build_parser().parse_args([]).model_replicas == [] + assert parse_replica_assignment("chat=3") == ("chat", 3) + for value in ("chat=0", "chat=-1", "chat=many", "chat", "=2"): + with pytest.raises(ValueError, match="model-replicas|positive integer"): + parse_replica_assignment(value) + + +def test_cli_log_level_is_a_fixed_non_debug_allowlist() -> None: + parser = build_parser() + for level in ("critical", "error", "warning", "info"): + assert parser.parse_args(["--log-level", level]).log_level == level + for level in ("debug", "trace", "notset"): + with pytest.raises(SystemExit): + parser.parse_args(["--log-level", level]) + + +@pytest.mark.parametrize( + "option", + ( + "--startup-timeout", + "--request-timeout", + "--realtime-idle-timeout", + "--realtime-max-session-seconds", + ), +) +@pytest.mark.parametrize("value", ("nan", "+inf", "-inf")) +def test_cli_timeouts_require_finite_values(option: str, value: str) -> None: + with pytest.raises(SystemExit): + build_parser().parse_args([f"{option}={value}"]) + + +def test_log_filter_recursively_redacts_transport_credentials() -> None: + record = logging.LogRecord( + "uvicorn.access", + logging.INFO, + __file__, + 1, + { + "headers": [ + ("Authorization", "Bearer authorization-secret"), + (b"cookie", b"session=cookie-secret"), + ], + "url": "ws://localhost/v1/realtime?access_token=query-secret&intent=transcription", + "nested": {"ACCESS_TOKEN": "mapping-secret", "safe": "visible"}, + }, + ( + b"authorization: Bearer bytes-secret", + ["cookie: sid=sequence-secret", {"access-token": "hyphen-secret"}], + ), + None, + ) + assert _RedactAccessToken().filter(record) + rendered = repr((record.msg, record.args)) + for secret in ( + "authorization-secret", + "cookie-secret", + "query-secret", + "mapping-secret", + "bytes-secret", + "sequence-secret", + "hyphen-secret", + ): + assert secret not in rendered + assert rendered.count("") >= 7 + assert "visible" in rendered + + +def test_log_filter_drops_query_when_access_token_key_is_percent_encoded() -> None: + query = "intent=transcription&access%5Ftoken=encoded-secret" + assert parse_qs(query)["access_token"] == ["encoded-secret"] + record = logging.LogRecord( + "uvicorn.access", + logging.INFO, + __file__, + 1, + "WebSocket /v1/realtime?" + query, + (), + None, + ) + + assert _RedactAccessToken().filter(record) + assert record.msg == "WebSocket /v1/realtime?" + assert "encoded-secret" not in record.getMessage() + + +def test_cli_startup_errors_do_not_expose_bundle_or_binary_paths(tmp_path: Path) -> None: + missing_bundle = tmp_path / "private" / "missing.bundle" + with pytest.raises(ValueError, match="does not exist or is not a file") as bundle_failure: + parse_model_assignment(f"chat={missing_bundle}", kind="chat") + assert str(tmp_path) not in str(bundle_failure.value) + assert str(missing_bundle) not in str(bundle_failure.value) + + missing_binary = tmp_path / "private" / "missing-trtmc" + with pytest.raises(ValueError, match="does not exist or is not a file") as binary_failure: + resolve_trtmc_binary(str(missing_binary)) + assert str(tmp_path) not in str(binary_failure.value) + assert str(missing_binary) not in str(binary_failure.value) + + missing_runtime = tmp_path / "private" / "missing-runtime" + with pytest.raises(ValueError, match="existing directory") as runtime_failure: + resolve_runtime_root(str(missing_runtime)) + assert str(tmp_path) not in str(runtime_failure.value) + assert str(missing_runtime) not in str(runtime_failure.value) + + +def test_cli_requires_a_nonempty_authentication_token( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + bundle = tmp_path / "chat.bundle" + bundle.write_bytes(b"chat") + common = [ + "--chat-model", + f"chat={bundle}", + "--trtmc-binary", + str(FAKE_TRTMC), + "--runtime-root", + str(tmp_path), + ] + monkeypatch.delenv("TRTMC_SERVE_TOKEN", raising=False) + + with pytest.raises(SystemExit) as missing: + serve_main(common) + assert missing.value.code == 2 + assert "authentication requires --api-key or TRTMC_SERVE_TOKEN" in capsys.readouterr().err + + with pytest.raises(SystemExit) as empty: + serve_main([*common, "--api-key", " "]) + assert empty.value.code == 2 + assert "authentication requires --api-key or TRTMC_SERVE_TOKEN" in capsys.readouterr().err + + +def test_cli_returns_failure_when_uvicorn_does_not_start( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import uvicorn + + bundle = tmp_path / "chat.bundle" + bundle.write_bytes(b"chat") + + class Listener: + closed = False + + @staticmethod + def getsockname() -> tuple[str, int]: + return ("127.0.0.1", 54321) + + def close(self) -> None: + self.closed = True + + listener = Listener() + + def fail_startup(server: uvicorn.Server, *, sockets: object) -> None: + assert sockets == [listener] + assert server.started is False + + monkeypatch.setattr(serve_cli_module, "bind_socket", lambda _host, _port: listener) + monkeypatch.setattr(uvicorn.Server, "run", fail_startup) + + status = serve_main( + [ + "--chat-model", + f"chat={bundle}", + "--trtmc-binary", + str(FAKE_TRTMC), + "--runtime-root", + str(tmp_path), + "--api-key", + "test-token", + ] + ) + + assert status == 1 + assert listener.closed is True + + +def test_cli_rejects_duplicate_and_unknown_replica_assignments(tmp_path: Path) -> None: + bundle = tmp_path / "chat.bundle" + bundle.write_bytes(b"chat") + common = [ + "--chat-model", + f"chat={bundle}", + "--trtmc-binary", + str(FAKE_TRTMC), + "--runtime-root", + str(tmp_path), + "--api-key", + "test-token", + ] + with pytest.raises(SystemExit): + serve_main([*common, "--model-replicas", "missing=2"]) + with pytest.raises(SystemExit): + serve_main( + [ + *common, + "--model-replicas", + "chat=2", + "--model-replicas", + "chat=3", + ] + ) + + +@_REQUIRES_PROC +def test_cli_port_zero_emits_single_machine_readable_ready_record( + tmp_path: Path, +) -> None: + bundle = tmp_path / "asr.bundle" + bundle.write_bytes(b"fixture") + environment = _server_process_environment() + environment["TRTMC_SERVE_TOKEN"] = "environment-token" + process = subprocess.Popen( + [ + sys.executable, + "-m", + "trtmc_server", + "--transcription-model", + f"asr={bundle}", + "--trtmc-binary", + str(FAKE_TRTMC), + "--host", + "127.0.0.1", + "--port", + "0", + "--require-streaming-transcription", + "asr", + "--model-replicas", + "asr=2", + "--runtime-root", + str(tmp_path), + "--kv-cache-size", + "1024", + "--runtime-cache", + "/tmp/runtime.cache", + "--cuda-graphs", + "--access-log", + ], + cwd=REPOSITORY, + env=environment, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + assert process.stdout is not None + assert process.stderr is not None + records: queue.Queue[str] = queue.Queue() + reader = threading.Thread(target=lambda: records.put(process.stdout.readline()), daemon=True) + reader.start() + try: + line = records.get(timeout=10) + ready = json.loads(line) + assert ready["event"] == "ready" + assert ready["host"] == "127.0.0.1" + assert ready["port"] > 0 + assert "pid" not in ready + assert ready["models"] == ["asr"] + worker_pids = _wait_for_child_pids(process.pid, expected=2) + + with pytest.raises(urllib.error.HTTPError) as unauthorized: + urllib.request.urlopen(f"http://127.0.0.1:{ready['port']}/v1/models", timeout=3) + assert unauthorized.value.code == 401 + + request = urllib.request.Request( + f"http://127.0.0.1:{ready['port']}/v1/models", + headers={"Authorization": "Bearer environment-token"}, + ) + with urllib.request.urlopen(request, timeout=3) as response: + model = json.load(response)["data"][0] + assert model["id"] == "asr" + assert model["metadata"] == { + "streaming_transcription": True, + } + assert model["capabilities"] == [ + "transcription", + "transcription_streaming", + ] + logged_query_request = urllib.request.Request( + f"http://127.0.0.1:{ready['port']}/v1/models" + "?access%5Ftoken=environment-token", + headers={"Authorization": "Bearer environment-token"}, + ) + with urllib.request.urlopen(logged_query_request, timeout=3) as response: + assert response.status == 200 + with urllib.request.urlopen( + f"http://127.0.0.1:{ready['port']}/healthz", timeout=3 + ) as response: + assert json.load(response) == {"status": "ok", "degraded": False} + with pytest.raises(urllib.error.HTTPError) as anonymous_ready: + urllib.request.urlopen(f"http://127.0.0.1:{ready['port']}/readyz", timeout=3) + assert anonymous_ready.value.code == 401 + ready_request = urllib.request.Request( + f"http://127.0.0.1:{ready['port']}/readyz", + headers={"Authorization": "Bearer environment-token"}, + ) + with urllib.request.urlopen(ready_request, timeout=3) as response: + worker_status = json.load(response)["models"]["asr"] + assert set(worker_status) == { + "ready", + "degraded", + "replicas", + "ready_replicas", + "idle_replicas", + "busy", + } + assert worker_status["replicas"] == 2 + assert worker_status["ready_replicas"] == 2 + + async def verify_websocket() -> None: + uri = ( + f"ws://127.0.0.1:{ready['port']}/v1/realtime" + "?intent=transcription&access_token=environment-token" + ) + async with websockets.connect(uri, origin="http://localhost:4173") as websocket: + assert json.loads(await websocket.recv())["type"] == "session.created" + await websocket.send( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(b"\x01\x00").decode(), + } + ) + ) + delta = json.loads(await websocket.recv()) + assert delta["type"] == ("conversation.item.input_audio_transcription.delta") + assert delta["transcript"] == "1 samples" + await websocket.send(json.dumps({"type": "input_audio_buffer.commit"})) + completed = json.loads(await websocket.recv()) + assert completed["type"] == ( + "conversation.item.input_audio_transcription.completed" + ) + + asyncio.run(verify_websocket()) + + process.terminate() + process.wait(timeout=5) + assert process.stdout.read() == "" + stderr = process.stderr.read() + assert "environment-token" not in stderr + assert "/v1/models?" in stderr + assert "access_token" not in stderr + for worker_pid in worker_pids: + _assert_pid_disappears(worker_pid) + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=5) + + +@_REQUIRES_PROC +def test_parent_liveness_stdin_eof_gracefully_stops_server_and_worker( + tmp_path: Path, +) -> None: + process, _ready = _start_test_server(tmp_path, "parent-eof", parent_liveness=True) + assert process.stdin is not None + try: + worker_pid = _wait_for_child_pids(process.pid, expected=1)[0] + process.stdin.close() + assert process.wait(timeout=5) == 0 + _assert_pid_disappears(worker_pid) + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=5) + + +@_REQUIRES_PROC +def test_server_sigkill_does_not_leave_native_worker(tmp_path: Path) -> None: + process, _ready = _start_test_server(tmp_path, "sigkill", parent_liveness=False) + try: + worker_pid = _wait_for_child_pids(process.pid, expected=1)[0] + process.kill() + process.wait(timeout=5) + _assert_pid_disappears(worker_pid) + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=5) + + +def _start_test_server( + tmp_path: Path, name: str, *, parent_liveness: bool +) -> tuple[subprocess.Popen[str], dict[str, object]]: + bundle = tmp_path / f"{name}-asr.bundle" + bundle.write_bytes(b"fixture") + environment = _server_process_environment() + environment["TRTMC_SERVE_TOKEN"] = "environment-token" + command = [ + sys.executable, + "-m", + "trtmc_server", + "--transcription-model", + f"asr={bundle}", + "--trtmc-binary", + str(FAKE_TRTMC), + "--runtime-root", + str(tmp_path), + "--port", + "0", + ] + if parent_liveness: + command.append("--parent-liveness-stdin") + process = subprocess.Popen( + command, + cwd=REPOSITORY, + env=environment, + text=True, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + assert process.stdout is not None + records: queue.Queue[str] = queue.Queue() + threading.Thread(target=lambda: records.put(process.stdout.readline()), daemon=True).start() + try: + line = records.get(timeout=10) + except queue.Empty: + process.kill() + process.wait(timeout=5) + stderr = process.stderr.read() if process.stderr is not None else "" + raise AssertionError(f"server did not emit ready: {stderr}") from None + if not line: + stderr = process.stderr.read() if process.stderr is not None else "" + raise AssertionError(f"server exited before ready: {stderr}") + return process, json.loads(line) + + +def _wait_for_child_pids(parent_pid: int, *, expected: int) -> list[int]: + task_root = Path(f"/proc/{parent_pid}/task") + deadline = time.monotonic() + 3 + children: list[int] = [] + while time.monotonic() < deadline: + try: + children = sorted( + { + int(value) + for children_file in task_root.glob("*/children") + for value in children_file.read_text().split() + } + ) + except FileNotFoundError: + children = [] + if len(children) == expected: + return children + time.sleep(0.02) + raise AssertionError( + f"server {parent_pid} has {len(children)} direct children; expected {expected}" + ) + + +def _assert_pid_disappears(pid: int) -> None: + deadline = time.monotonic() + 3 + while _pid_exists(pid) and time.monotonic() < deadline: + time.sleep(0.02) + assert not _pid_exists(pid), f"native worker {pid} survived parent termination" + + +def _pid_exists(pid: int) -> bool: + try: + state = Path(f"/proc/{pid}/stat").read_text().rsplit(")", maxsplit=1)[1].split()[0] + if state == "Z": + return False + except (FileNotFoundError, IndexError): + return False + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True diff --git a/server/tests/test_serve_worker.cpp b/server/tests/test_serve_worker.cpp new file mode 100644 index 0000000000..4a5c9b4037 --- /dev/null +++ b/server/tests/test_serve_worker.cpp @@ -0,0 +1,515 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// The native protocol is tested with an injected ITask, so every case is CPU-only. + +#include "native/entrypoint.h" +#include "native/worker.h" +#include "trtmc/task.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using Json = nlohmann::json; + +int failures = 0; + +void check(bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ++failures; + } +} + +struct StreamObservations { + int accept_calls{0}; + int finish_calls{0}; + int reset_calls{0}; + std::vector last_samples; +}; + +class FakeStream final : public trtmc::ITranscriptionStream { + public: + FakeStream(trtmc::TranscriptionStreamConfig config, + std::shared_ptr observations) + : config_(std::move(config)), observations_(std::move(observations)) {} + + trtmc::TranscriptionStreamResult accept_audio(const float* samples, std::int32_t count, + bool is_final) override { + ++observations_->accept_calls; + observations_->last_samples.assign(samples, samples + count); + return {"partial transcript", {4, 5}, is_final, + observations_->accept_calls, count, config_.input_sample_rate}; + } + + trtmc::TranscriptionStreamResult finish() override { + ++observations_->finish_calls; + return {"final transcript", + {4, 5, 6}, + true, + observations_->accept_calls, + static_cast(observations_->last_samples.size()), + config_.input_sample_rate}; + } + + void reset() override { ++observations_->reset_calls; } + trtmc::TranscriptionStreamConfig config() const override { return config_; } + + private: + trtmc::TranscriptionStreamConfig config_; + std::shared_ptr observations_; +}; + +class FakeTask final : public trtmc::ITextGeneration, + public trtmc::ITranscription, + public trtmc::IStreamingTranscription { + public: + const char* task() const noexcept override { return "test_multi_capability"; } + std::int32_t default_max_new_tokens() const override { return 77; } + + trtmc::TextResult generate(const std::string& prompt, + const trtmc::TextGenerationConfig& config) override { + ++generate_calls; + last_prompt = prompt; + last_generate_config = config; + if (generate_invalid_argument) + throw std::invalid_argument( + "sensitive provider invalid argument at /tmp/private-provider.bundle"); + if (generate_runtime_error) + throw std::runtime_error("sensitive native runtime detail at /tmp/private.bundle"); + if (generate_json_exception) { + const auto ignored = Json::parse("sensitive-provider-json-detail"); + (void)ignored; + } + trtmc::TextResult result{"generated: " + prompt, {8, 9}, 1.25, 2.5}; + if (generate_invalid_utf8) + result.text.push_back(static_cast(0xFF)); + result.setup_ms = 0.5; + return result; + } + + trtmc::TextResult transcribe(const float* samples, std::int32_t count, + const trtmc::TranscriptionConfig& config) override { + ++transcribe_calls; + last_transcription_config = config; + last_transcription_samples.assign(samples, samples + count); + trtmc::TranscriptionSegment segment; + segment.end_seconds = 0.25; + segment.text = "hello"; + segment.token_ids = {10}; + return {"hello from wav", {10, 11}, 0.0, 3.0, {segment}}; + } + + std::unique_ptr + create_transcription_stream(const trtmc::TranscriptionStreamConfig& config) override { + if (!stream_supported) + throw std::runtime_error("streaming transcription unsupported"); + ++create_stream_calls; + stream_configs.push_back(config); + auto observations = std::make_shared(); + stream_history.push_back(observations); + return std::make_unique(config, std::move(observations)); + } + + int generate_calls{0}; + int transcribe_calls{0}; + int create_stream_calls{0}; + bool stream_supported{true}; + bool generate_invalid_argument{false}; + bool generate_invalid_utf8{false}; + bool generate_runtime_error{false}; + bool generate_json_exception{false}; + std::string last_prompt; + trtmc::TextGenerationConfig last_generate_config; + trtmc::TranscriptionConfig last_transcription_config; + std::vector last_transcription_samples; + std::vector stream_configs; + std::vector> stream_history; +}; + +class TextOnlyTask final : public trtmc::ITextGeneration { + public: + std::int32_t default_max_new_tokens() const override { return 5; } + trtmc::TextResult generate(const std::string& prompt, + const trtmc::TextGenerationConfig&) override { + return {prompt, {1}}; + } +}; + +std::filesystem::path make_temp_dir() { + char pattern[] = "/tmp/trtmc_serve_worker_test_XXXXXX"; + char* directory = mkdtemp(pattern); + if (directory == nullptr) + throw std::runtime_error("mkdtemp failed"); + return std::filesystem::path(directory); +} + +struct TempDir { + std::filesystem::path path{make_temp_dir()}; + ~TempDir() { + std::error_code error; + std::filesystem::remove_all(path, error); + } +}; + +std::vector parse_output_lines(const std::string& output) { + std::vector messages; + std::istringstream stream(output); + for (std::string line; std::getline(stream, line);) { + if (!line.empty()) + messages.push_back(Json::parse(line)); + } + return messages; +} + +void append_request(std::ostringstream& input, const Json& request) { + input << request.dump() << '\n'; +} + +template +void write_wav(const std::filesystem::path& path, const std::vector& samples, + std::uint16_t format, std::uint16_t channels, std::uint32_t sample_rate) { + std::ofstream output(path, std::ios::binary); + const std::uint16_t bits_per_sample = static_cast(sizeof(Sample) * 8U); + const std::uint16_t block_align = static_cast(channels * sizeof(Sample)); + const std::uint32_t byte_rate = sample_rate * block_align; + const std::uint32_t data_size = static_cast(samples.size() * sizeof(Sample)); + const std::uint32_t file_size = 36U + data_size; + const std::uint32_t fmt_size = 16; + output.write("RIFF", 4); + output.write(reinterpret_cast(&file_size), 4); + output.write("WAVEfmt ", 8); + output.write(reinterpret_cast(&fmt_size), 4); + output.write(reinterpret_cast(&format), 2); + output.write(reinterpret_cast(&channels), 2); + output.write(reinterpret_cast(&sample_rate), 4); + output.write(reinterpret_cast(&byte_rate), 4); + output.write(reinterpret_cast(&block_align), 2); + output.write(reinterpret_cast(&bits_per_sample), 2); + output.write("data", 4); + output.write(reinterpret_cast(&data_size), 4); + output.write(reinterpret_cast(samples.data()), data_size); +} + +std::vector run(FakeTask& task, const std::string& input, int* status = nullptr) { + std::istringstream requests(input); + std::ostringstream output; + const int result = trtmc::serve::run_worker_protocol(task, requests, output); + if (status != nullptr) + *status = result; + return parse_output_lines(output.str()); +} + +void test_ready_schema_uses_actual_task_capabilities() { + std::ostringstream requests; + append_request(requests, {{"id", "stop"}, {"op", "shutdown"}}); + FakeTask task; + const auto messages = run(task, requests.str()); + check(messages.size() == 2, "ready plus shutdown response"); + if (messages.size() == 2) { + const auto& ready = messages[0]; + check(ready == Json{{"event", "ready"}, + {"protocol_version", 3}, + {"capabilities", Json::array({"text_generation", "transcription", + "transcription_streaming"})}, + {"default_max_new_tokens", 77}}, + "ready schema is minimal and capability based"); + } + + TextOnlyTask text; + std::istringstream input(requests.str()); + std::ostringstream output; + check(trtmc::serve::run_worker_protocol(text, input, output) == 0, + "text-only task runs protocol"); + const auto text_messages = parse_output_lines(output.str()); + check(text_messages.size() == 2 && + text_messages[0]["capabilities"] == Json::array({"text_generation"}), + "ready does not infer unsupported task capabilities"); +} + +void test_generation_and_stream_lifecycle() { + std::ostringstream requests; + append_request(requests, {{"id", "generate"}, + {"op", "generate"}, + {"prompt", "summarize this"}, + {"config", + {{"max_new_tokens", 12}, + {"temperature", 0.25}, + {"top_p", 0.8}, + {"min_p", 0.05}, + {"top_k", 7}, + {"seed", 42}, + {"use_chat_template", true}, + {"enable_thinking", false}}}}); + append_request(requests, {{"id", "probe"}, + {"op", "probe_transcription_stream"}, + {"config", + {{"sample_rate_hz", 16000}, + {"channels", 1}, + {"audio_format", "pcm16le"}, + {"language", "en-US"}}}}); + append_request( + requests, + {{"id", "start"}, + {"op", "stream_start"}, + {"config", {{"sample_rate_hz", 16000}, {"channels", 1}, {"audio_format", "pcm16le"}}}}); + append_request(requests, {{"id", "chunk"}, {"op", "stream_chunk"}, {"audio", "AAAAQACA"}}); + append_request(requests, {{"id", "reset"}, {"op", "stream_reset"}}); + append_request( + requests, + {{"id", "restart"}, {"op", "stream_start"}, {"config", {{"sample_rate_hz", 16000}}}}); + append_request(requests, {{"id", "finish"}, {"op", "stream_finish"}}); + append_request(requests, {{"id", "stop"}, {"op", "shutdown"}}); + + FakeTask task; + int status = -1; + const auto messages = run(task, requests.str(), &status); + check(status == 0 && messages.size() == 9, "full worker lifecycle completes"); + if (messages.size() != 9) + return; + check(messages[1]["result"].value("text", "") == "generated: summarize this" && + messages[1]["result"]["token_ids"] == Json::array({8, 9}) && + messages[1]["result"].value("completion_tokens", 0) == 2, + "generation result preserves public fields"); + check(task.generate_calls == 1 && task.last_generate_config.max_new_tokens == 12 && + std::abs(task.last_generate_config.temperature - 0.25F) < 1e-6F && + task.last_generate_config.top_k == 7 && task.last_generate_config.use_chat_template && + !task.last_generate_config.enable_thinking, + "generation config maps to ITextGeneration"); + check(messages[2]["result"] == Json{{"supported", true}}, "stream capability probe succeeds"); + check(messages[3]["result"] == Json::object() && + messages[4]["result"] == Json{{"text", "partial transcript"}} && + messages[5]["result"] == Json::object() && + messages[7]["result"] == Json{{"text", "final transcript"}}, + "stream lifecycle returns canonical payloads"); + check(task.create_stream_calls == 3 && task.stream_history.size() == 3, + "probe and each stream use independent native state"); + if (task.stream_history.size() == 3) { + const auto& samples = task.stream_history[1]->last_samples; + check(samples.size() == 3 && std::abs(samples[0]) < 1e-6F && + std::abs(samples[1] - 0.5F) < 1e-6F && std::abs(samples[2] + 1.0F) < 1e-6F, + "stream PCM16 base64 is decoded exactly once"); + check(task.stream_history[0]->reset_calls == 1 && + task.stream_history[1]->reset_calls == 1 && + task.stream_history[2]->finish_calls == 1, + "stream ownership terminates by probe, reset, or finish"); + } +} + +void test_server_owned_wav_decode() { + TempDir temporary; + const auto pcm_path = temporary.path / "pcm.wav"; + write_wav( + pcm_path, {16384, 8192, 0, -8192, 0, 8192, 16384, 24576, -16384, 0, 16384, 0}, 1, 4, 22050); + const auto float_path = temporary.path / "float.wav"; + write_wav(float_path, {0.25F, -0.5F}, 3, 1, 16000); + + std::ostringstream requests; + append_request(requests, {{"id", "pcm"}, + {"op", "transcribe"}, + {"audio_path", pcm_path.string()}, + {"config", {{"language", "fr"}}}}); + append_request(requests, + {{"id", "float"}, {"op", "transcribe"}, {"audio_path", float_path.string()}}); + append_request(requests, {{"id", "stop"}, {"op", "shutdown"}}); + FakeTask task; + const auto messages = run(task, requests.str()); + check(messages.size() == 4 && messages[1].value("ok", false) && messages[2].value("ok", false), + "PCM16 and float32 WAV requests succeed"); + check(task.transcribe_calls == 2 && task.last_transcription_samples.size() == 2 && + std::abs(task.last_transcription_samples[0] - 0.25F) < 1e-6F && + std::abs(task.last_transcription_samples[1] + 0.5F) < 1e-6F, + "server-owned float32 WAV decoder preserves samples"); + + std::ostringstream pcm_only; + append_request(pcm_only, + {{"id", "pcm"}, {"op", "transcribe"}, {"audio_path", pcm_path.string()}}); + append_request(pcm_only, {{"id", "stop"}, {"op", "shutdown"}}); + FakeTask pcm_task; + (void)run(pcm_task, pcm_only.str()); + check(pcm_task.last_transcription_samples.size() == 3 && + std::abs(pcm_task.last_transcription_samples[0] - 0.125F) < 1e-6F && + std::abs(pcm_task.last_transcription_samples[1] - 0.375F) < 1e-6F && + std::abs(pcm_task.last_transcription_samples[2]) < 1e-6F && + pcm_task.last_transcription_config.input_sample_rate == 22050, + "server-owned PCM16 WAV decoder downmixes and preserves sample rate"); +} + +void test_protocol_rejects_edges_without_poisoning_task() { + TempDir temporary; + const auto invalid_path = temporary.path / "invalid.wav"; + { + std::ofstream output(invalid_path, std::ios::binary); + output << "not a wav"; + } + const auto missing_path = temporary.path / "missing.wav"; + + std::ostringstream requests; + requests << "not-json\n"; + append_request(requests, {{"id", 7}, {"op", "generate"}, {"prompt", "bad id"}}); + append_request(requests, {{"id", "bad-config"}, + {"op", "generate"}, + {"prompt", "extra"}, + {"config", {{"num_samples", 2}}}}); + append_request( + requests, {{"id", "bad-wav"}, {"op", "transcribe"}, {"audio_path", invalid_path.string()}}); + append_request( + requests, + {{"id", "missing-wav"}, {"op", "transcribe"}, {"audio_path", missing_path.string()}}); + append_request(requests, {{"id", "start"}, {"op", "stream_start"}}); + append_request(requests, {{"id", "duplicate"}, {"op", "stream_start"}}); + append_request(requests, {{"id", "bad-audio"}, {"op", "stream_chunk"}, {"audio", "%%%"}}); + append_request(requests, {{"id", "reset"}, {"op", "stream_reset"}}); + append_request(requests, {{"id", "healthy"}, {"op", "generate"}, {"prompt", "alive"}}); + append_request(requests, {{"id", "stop"}, {"op", "shutdown"}}); + + FakeTask task; + std::ostringstream diagnostics; + auto* previous = std::cerr.rdbuf(diagnostics.rdbuf()); + const auto messages = run(task, requests.str()); + std::cerr.rdbuf(previous); + check(messages.size() == 12, "each malformed request receives one response"); + if (messages.size() != 12) + return; + check(messages[1]["id"].is_null() && + messages[1]["error"].value("type", "") == "invalid_request_error" && + messages[2]["id"].is_null(), + "malformed JSON and ids stay structured"); + check(messages[3]["error"].value("message", "").find("config.num_samples") != std::string::npos, + "noncanonical generation fields fail closed"); + check(messages[4]["error"].value("code", "") == "unsupported_media_type" && + messages[4]["error"].value("param", "") == "file", + "malformed WAV is a stable media error"); + check(messages[5]["error"].value("type", "") == "runtime_error" && + messages[5]["error"].value("message", "") == "native worker operation failed" && + messages[5].dump().find(missing_path.string()) == std::string::npos, + "I/O paths remain private runtime details"); + check(!messages[7].value("ok", true) && !messages[8].value("ok", true) && + messages[9].value("ok", false) && messages[10].value("ok", false), + "state and base64 errors do not poison later canonical requests"); + check(task.generate_calls == 1, "only canonical generation reaches the task"); + check(diagnostics.str().find("cannot open uploaded audio file") != std::string::npos, + "private I/O diagnostic remains on stderr"); +} + +void test_runtime_failures_are_redacted_and_worker_survives() { + for (int failure_kind = 0; failure_kind < 3; ++failure_kind) { + std::ostringstream requests; + append_request(requests, {{"id", "failure"}, {"op", "generate"}, {"prompt", "private"}}); + append_request(requests, {{"id", "stop"}, {"op", "shutdown"}}); + FakeTask task; + task.generate_runtime_error = failure_kind == 0; + task.generate_invalid_argument = failure_kind == 1; + task.generate_json_exception = failure_kind == 2; + std::istringstream input(requests.str()); + std::ostringstream output; + std::ostringstream diagnostics; + auto* previous = std::cerr.rdbuf(diagnostics.rdbuf()); + const int status = trtmc::serve::run_worker_protocol(task, input, output); + std::cerr.rdbuf(previous); + const auto messages = parse_output_lines(output.str()); + check(status == 0 && messages.size() == 3 && + messages[1]["error"].value("type", "") == "runtime_error" && + messages[1]["error"].value("message", "") == "native worker operation failed", + "provider exception is a generic runtime error"); + const bool private_detail_logged = + failure_kind == 2 ? diagnostics.str().find("json.exception") != std::string::npos + : diagnostics.str().find("sensitive") != std::string::npos; + check(output.str().find("sensitive") == std::string::npos && private_detail_logged, + "provider diagnostics remain stderr-only"); + } +} + +void test_invalid_utf8_and_oversized_records_are_bounded() { + std::ostringstream utf8_requests; + append_request(utf8_requests, {{"id", "generate"}, {"op", "generate"}, {"prompt", "text"}}); + append_request(utf8_requests, {{"id", "stop"}, {"op", "shutdown"}}); + FakeTask utf8_task; + utf8_task.generate_invalid_utf8 = true; + std::istringstream utf8_input(utf8_requests.str()); + std::ostringstream utf8_output; + check(trtmc::serve::run_worker_protocol(utf8_task, utf8_input, utf8_output) == 0, + "invalid UTF-8 result does not stop worker"); + check(utf8_output.str().find(static_cast(0xFF)) == std::string::npos && + utf8_output.str().find("\xEF\xBF\xBD") != std::string::npos, + "invalid UTF-8 is replaced on JSONL output"); + + constexpr std::size_t limit = 16U * 1024U * 1024U; + std::ostringstream oversized; + oversized << std::string(limit + 4096U, 'x') << '\n'; + append_request(oversized, {{"id", "healthy"}, {"op", "generate"}, {"prompt", "alive"}}); + append_request(oversized, {{"id", "stop"}, {"op", "shutdown"}}); + FakeTask task; + const auto messages = run(task, oversized.str()); + check(messages.size() == 4 && messages[1]["id"].is_null() && + messages[1]["error"].value("message", "") == + "request exceeds the 16 MiB JSONL limit" && + messages[2].value("ok", false) && task.generate_calls == 1, + "oversized record is discarded before the next request"); +} + +void test_entrypoint_requires_explicit_runtime_and_keeps_stdout_clean() { + TempDir temporary; + const std::string missing = (temporary.path / "private.bundle").string(); + trtmc::server::NativeWorkerOptions options; + options.bundle_path = missing; + options.runtime_root = temporary.path.string(); + std::ostringstream protocol; + std::ostringstream diagnostics; + auto* previous_stdout = std::cout.rdbuf(protocol.rdbuf()); + auto* previous_stderr = std::cerr.rdbuf(diagnostics.rdbuf()); + const int status = trtmc::server::run_native_worker(options); + std::cout.rdbuf(previous_stdout); + std::cerr.rdbuf(previous_stderr); + check(status == EXIT_FAILURE && protocol.str().empty(), + "startup failure never enters protocol stdout"); + check(diagnostics.str().find("Error: native worker failed:") != std::string::npos && + diagnostics.str().find(missing) != std::string::npos, + "startup detail remains on private stderr"); + + char command[] = "bundle.bundle"; + char* argv[] = {command}; + std::ostringstream parse_diagnostics; + previous_stderr = std::cerr.rdbuf(parse_diagnostics.rdbuf()); + const int parse_status = trtmc::server::run_native_worker(1, argv); + std::cerr.rdbuf(previous_stderr); + check(parse_status == EXIT_FAILURE && + parse_diagnostics.str().find("requires --runtime-root") != std::string::npos, + "worker CLI requires explicit runtime placement"); +} + +} // namespace + +int main() { + try { + test_ready_schema_uses_actual_task_capabilities(); + test_generation_and_stream_lifecycle(); + test_server_owned_wav_decode(); + test_protocol_rejects_edges_without_poisoning_task(); + test_runtime_failures_are_redacted_and_worker_survives(); + test_invalid_utf8_and_oversized_records_are_bounded(); + test_entrypoint_requires_explicit_runtime_and_keeps_stdout_clean(); + } catch (const std::exception& error) { + std::cerr << "Unhandled test exception: " << error.what() << '\n'; + return 1; + } + if (failures != 0) + std::cerr << failures << " serve worker test(s) failed\n"; + return failures == 0 ? 0 : 1; +} diff --git a/server/tests/test_serve_worker.py b/server/tests/test_serve_worker.py new file mode 100644 index 0000000000..fcc539e0df --- /dev/null +++ b/server/tests/test_serve_worker.py @@ -0,0 +1,432 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest + +from trtmc_server import worker as worker_module +from trtmc_server.errors import ( + WorkerCrashedError, + WorkerProtocolError, + WorkerRemoteError, + WorkerRequestTooLargeError, + WorkerSaturatedError, + WorkerStartupError, + WorkerTimeoutError, +) +from trtmc_server.worker import WorkerGroup, WorkerLoadOptions, WorkerProcess + + +FAKE_TRTMC = Path(__file__).with_name("fake_serve_worker.py") + + +def make_worker( + tmp_path: Path, + name: str = "chat", + *, + startup_timeout: float = 1.0, + request_timeout: float = 1.0, + max_request_line_bytes: int = 16 * 1024 * 1024, + load_options: WorkerLoadOptions | None = None, +) -> WorkerProcess: + bundle = tmp_path / f"{name}.bundle" + bundle.write_bytes(b"fixture") + return WorkerProcess( + name=name, + bundle=bundle, + trtmc_binary=FAKE_TRTMC, + startup_timeout=startup_timeout, + request_timeout=request_timeout, + max_request_line_bytes=max_request_line_bytes, + load_options=load_options or WorkerLoadOptions(runtime_root=str(tmp_path)), + ) + + +def wait_for_stderr(worker: WorkerProcess, text: str, timeout: float = 1.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if any(text in line for line in worker.stderr_tail): + return True + time.sleep(0.005) + return any(text in line for line in worker.stderr_tail) + + +def test_worker_ready_request_ids_metadata_stderr_and_cleanup(tmp_path: Path) -> None: + worker = make_worker(tmp_path) + worker.start() + + assert worker.ready + assert worker.pid is not None + assert worker.ready_payload["protocol_version"] == 3 + assert worker.ready_payload["capabilities"] == ["text_generation"] + assert wait_for_stderr(worker, "fake worker ready") + + worker.close() + assert worker.state == "closed" + assert not worker.ready + + +def test_worker_serializes_concurrent_requests(tmp_path: Path) -> None: + worker = make_worker(tmp_path) + worker.start() + try: + with ThreadPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(worker.request, "generate", {"prompt": f"p{index}"}) + for index in range(12) + ] + assert {future.result()["text"] for future in futures} == { + f"generated:p{index}" for index in range(12) + } + assert worker.ready + finally: + worker.close() + + +def test_worker_session_excludes_other_requests_without_poisoning_worker( + tmp_path: Path, +) -> None: + worker = make_worker(tmp_path, "stream-asr", request_timeout=0.1) + worker.start() + session = worker.acquire_session() + try: + session.request( + "stream_start", + {"config": {}}, + ) + with pytest.raises(WorkerTimeoutError, match="remained busy"): + worker.request("generate", {"prompt": "busy"}, timeout=0.05) + assert worker.ready + session.request("stream_reset") + finally: + session.close() + worker.close() + + +@pytest.mark.parametrize( + ("mode", "message"), + [("bad-ready", "invalid JSONL"), ("legacy-ready", "string request id")], +) +def test_worker_startup_rejects_invalid_ready_shape( + tmp_path: Path, mode: str, message: str +) -> None: + worker = make_worker(tmp_path, mode, startup_timeout=0.5) + with pytest.raises(WorkerStartupError, match=message): + worker.start() + assert worker.state == "failed" + worker.close() + + +def test_worker_startup_timeout_terminates_process_without_exposing_stderr( + tmp_path: Path, +) -> None: + worker = make_worker(tmp_path, "startup-timeout", startup_timeout=1.0) + with pytest.raises(WorkerStartupError, match="did not become ready") as failure: + worker.start() + assert "startup detail" not in str(failure.value) + assert str(tmp_path) not in str(failure.value) + assert wait_for_stderr(worker, "startup detail") + assert worker.state == "failed" + assert not worker.ready + worker.close() + + +def test_worker_binary_startup_error_does_not_expose_absolute_path(tmp_path: Path) -> None: + bundle = tmp_path / "model.bundle" + bundle.write_bytes(b"fixture") + missing_binary = tmp_path / "private" / "missing-trtmc" + worker = WorkerProcess( + name="chat", + bundle=bundle, + trtmc_binary=missing_binary, + startup_timeout=0.1, + request_timeout=0.1, + load_options=WorkerLoadOptions(runtime_root=str(tmp_path)), + ) + with pytest.raises(WorkerStartupError, match="failed to launch worker") as failure: + worker.start() + assert str(tmp_path) not in str(failure.value) + assert str(missing_binary) not in str(failure.value) + worker.close() + + +def test_worker_structured_remote_error(tmp_path: Path) -> None: + worker = make_worker(tmp_path) + worker.start() + try: + with pytest.raises(WorkerRemoteError, match="unsupported op") as failure: + worker.request("does-not-exist") + assert failure.value.details["type"] == "invalid_request_error" + assert worker.ready + finally: + worker.close() + + +def test_worker_rejects_reserved_payload_fields_without_losing_sync(tmp_path: Path) -> None: + worker = make_worker(tmp_path) + worker.start() + try: + with pytest.raises(ValueError, match="cannot override 'id'"): + worker.request("generate", {"id": "caller-controlled"}) + assert worker.request("generate", {"prompt": "healthy"})["text"] == ("generated:healthy") + assert worker.ready + finally: + worker.close() + + +def test_worker_rejects_oversized_serialized_line_without_poisoning_worker( + tmp_path: Path, +) -> None: + worker = make_worker(tmp_path, max_request_line_bytes=256) + worker.start() + try: + with pytest.raises(WorkerRequestTooLargeError, match="maximum is 256 bytes"): + worker.request("generate", {"prompt": '"' * 300}) + assert worker.request("generate", {"prompt": "healthy"})["text"] == ("generated:healthy") + assert worker.ready + finally: + worker.close() + + +def test_worker_environment_is_an_explicit_runtime_allowlist( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7") + for name in ( + "ACCESS_TOKEN", + "AUTHORIZATION", + "AWS_SECRET_ACCESS_KEY", + "COOKIE", + "GITHUB_TOKEN", + "HF_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "NVIDIA_API_KEY", + "TRTMC_SERVE_TOKEN", + "TRTMC_MODEL_PLUGIN_DIR", + "TRTMC_MODEL_PLUGIN_STRICT", + "TRTMC_TRT_LIBRARY_DIR", + "UNLISTED_ENVIRONMENT", + ): + monkeypatch.setenv(name, f"secret-{name.lower()}") + for name in ( + "OMPI_COMM_WORLD_JOBID", + "OMPI_COMM_WORLD_LOCAL_RANK", + "OMPI_COMM_WORLD_RANK", + "OMPI_COMM_WORLD_SIZE", + "PMI_RANK", + "PMI_SIZE", + "PMIX_NAMESPACE", + "RANK", + "SLURM_PROCID", + "TRTMC_NCCL_RENDEZVOUS", + "TRTMC_NCCL_SKIP_DESTROY", + "WORLD_SIZE", + ): + monkeypatch.setenv(name, "distributed-test-value") + worker = make_worker(tmp_path, "inspect-environment-chat") + worker.start() + try: + assert worker.ready_payload["serve_token_present"] is False + assert worker.ready_payload["secret_environment_present"] is False + assert worker.ready_payload["distributed_environment_present"] is False + assert worker.ready_payload["allowed_cuda_visible_devices"] == "7" + finally: + worker.close() + + +def test_worker_environment_adds_windows_process_basics_only_on_windows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + windows_values = { + "COMSPEC": r"C:\Windows\System32\cmd.exe", + "PATHEXT": ".COM;.EXE;.BAT;.CMD", + "SYSTEMROOT": r"C:\Windows", + "TEMP": r"C:\Temp", + "TMP": r"C:\Temp", + "USERPROFILE": r"C:\Users\trtmc", + } + for name, value in windows_values.items(): + monkeypatch.setenv(name, value) + monkeypatch.setenv("TRTMC_SERVE_TOKEN", "must-stay-private") + + with monkeypatch.context() as windows: + windows.setattr(worker_module.os, "name", "nt") + environment = worker_module._worker_environment() + + assert {name: environment[name] for name in windows_values} == windows_values + assert "TRTMC_SERVE_TOKEN" not in environment + if worker_module.os.name != "nt": + posix_environment = worker_module._worker_environment() + assert windows_values.keys().isdisjoint(posix_environment) + + +def test_worker_forwards_native_load_options_verbatim(tmp_path: Path) -> None: + options = WorkerLoadOptions( + runtime_root="/opt/trtmc/lib", + kv_cache_size_bytes=1024, + runtime_cache="/tmp/runtime.cache", + cuda_graphs=True, + ) + worker = make_worker(tmp_path, "inspect-argv-chat", load_options=options) + worker.start() + try: + assert worker.ready_payload["worker_args"] == [ + "--runtime-root", + "/opt/trtmc/lib", + "--kv-cache-size", + "1024", + "--runtime-cache", + "/tmp/runtime.cache", + "--cuda-graphs", + ] + finally: + worker.close() + + +def test_worker_group_scales_across_replicas(tmp_path: Path) -> None: + workers = [ + make_worker(tmp_path, f"saturation-chat-{replica}", request_timeout=1) + for replica in range(2) + ] + for worker in workers: + worker.start() + group = WorkerGroup("chat", workers) + try: + with ThreadPoolExecutor(max_workers=2) as executor: + results = list( + executor.map( + lambda prompt: _group_request(group, "generate", {"prompt": prompt}), + ("first", "second"), + ) + ) + assert {result["text"] for result in results} == { + "generated:first", + "generated:second", + } + finally: + group.close() + + +def test_worker_group_rejects_when_all_replicas_are_busy(tmp_path: Path) -> None: + workers = [make_worker(tmp_path, f"chat-{replica}") for replica in range(2)] + for worker in workers: + worker.start() + group = WorkerGroup("chat", workers) + first = group.acquire_session() + second = group.acquire_session() + try: + with pytest.raises(WorkerSaturatedError, match="all 2 replicas"): + group.acquire_session() + finally: + first.close() + second.close() + group.close() + + +def test_worker_group_drops_failed_lane_and_keeps_healthy_lane(tmp_path: Path) -> None: + crashed = make_worker(tmp_path, "crash-request-chat") + healthy = make_worker(tmp_path, "healthy-chat") + crashed.start() + healthy.start() + group = WorkerGroup("chat", (crashed, healthy)) + try: + with pytest.raises(WorkerCrashedError): + _group_request(group, "generate", {"prompt": "crash"}) + + assert _group_request(group, "generate", {"prompt": "still-running"})["text"] == ( + "generated:still-running" + ) + assert group.ready + status = group.status() + assert set(status) == { + "state", + "ready", + "degraded", + "replicas", + "ready_replicas", + "idle_replicas", + "busy", + } + assert status["ready_replicas"] == 1 + assert status["degraded"] is True + finally: + group.close() + + +def _group_request(group: WorkerGroup, operation: str, payload: dict[str, str]) -> object: + with group.acquire_session() as session: + return session.request(operation, payload) + + +def test_worker_crash_keeps_stderr_out_of_exception_and_status(tmp_path: Path) -> None: + worker = make_worker(tmp_path, "crash-chat") + worker.start() + with pytest.raises(WorkerCrashedError) as failure: + worker.request("generate", {"prompt": "boom"}) + assert "intentional fake crash" not in str(failure.value) + assert "worker-secret" not in str(failure.value) + assert str(tmp_path) not in str(failure.value) + assert wait_for_stderr(worker, "intentional fake crash") + assert worker.state == "failed" + assert not worker.ready + worker.close() + + +def test_worker_request_timeout_terminates_desynchronized_worker(tmp_path: Path) -> None: + worker = make_worker(tmp_path, "slow-chat", request_timeout=0.1) + worker.start() + with pytest.raises(WorkerTimeoutError, match="timed out"): + worker.request("generate", {"prompt": "wait"}) + assert worker.state == "failed" + assert not worker.ready + worker.close() + + +def test_worker_close_interrupts_an_active_request(tmp_path: Path) -> None: + worker = make_worker(tmp_path, "slow-chat", request_timeout=5) + worker.start() + with ThreadPoolExecutor(max_workers=1) as executor: + request = executor.submit(worker.request, "generate", {"prompt": "wait"}) + deadline = time.monotonic() + 1 + while not worker.busy and time.monotonic() < deadline: + time.sleep(0.005) + assert worker.busy + + started = time.monotonic() + worker.close(grace_period=0.05) + assert time.monotonic() - started < 1.5 + with pytest.raises(WorkerCrashedError): + request.result(timeout=1) + + assert worker.state == "closed" + assert not worker.ready + + +def test_worker_unknown_response_id_is_protocol_failure(tmp_path: Path) -> None: + worker = make_worker(tmp_path, "bad-response-chat") + worker.start() + with pytest.raises(WorkerProtocolError, match="unknown id"): + worker.request("generate", {"prompt": "bad"}) + assert worker.state == "failed" + worker.close() + + +@pytest.mark.parametrize( + ("mode", "message"), + [ + ("non-bool-ok-chat", "boolean ok"), + ("missing-result-chat", "missing result"), + ], +) +def test_worker_rejects_malformed_v3_response(tmp_path: Path, mode: str, message: str) -> None: + worker = make_worker(tmp_path, mode) + worker.start() + with pytest.raises(WorkerProtocolError, match=message): + worker.request("generate", {"prompt": "bad"}) + assert worker.state == "failed" + worker.close() diff --git a/tools/ci/package.py b/tools/ci/package.py index d1b097d187..8c51bfbf37 100644 --- a/tools/ci/package.py +++ b/tools/ci/package.py @@ -30,6 +30,20 @@ def family_ids(repository: Path) -> tuple[str, ...]: return families +def server_python_files(repository: Path) -> dict[str, Path]: + """Return the exact optional server package payload by wheel path.""" + + root = repository / "server/python" + files = { + path.relative_to(root).as_posix(): path + for path in root.rglob("*.py") + if "__pycache__" not in path.parts + } + if "trtmc_server/__init__.py" not in files: + raise CiError("repository has no trtmc_server Python package") + return dict(sorted(files.items())) + + class SourceArchiveValidator: """Require the source archive to carry the physical dependency declarations.""" @@ -48,14 +62,29 @@ def validate(self, archives: list[Path]) -> None: if len(roots) != 1: raise CiError(f"{archive}: source archive must have one root directory") packaged = {Path(*path.parts[1:]).as_posix() for path in members if len(path.parts) > 1} - expected = {"requirements/base.txt"} | { + expected_dependencies = {"requirements/base.txt"} | { path.relative_to(self.context.repository).as_posix() for path in (self.context.repository / "families").glob("*/requirements.txt") } - missing = sorted(expected - packaged) - if missing: - raise CiError(f"{archive}: dependency declarations are missing: {missing}") - print(f"validated source archive={archive} family_requirements={len(expected) - 1}") + missing_dependencies = sorted(expected_dependencies - packaged) + if missing_dependencies: + raise CiError( + f"{archive}: dependency declarations are missing: {missing_dependencies}" + ) + expected_server = { + path.relative_to(self.context.repository).as_posix() + for path in (self.context.repository / "server").rglob("*") + if path.is_file() + and "__pycache__" not in path.parts + and path.suffix != ".pyc" + } + missing_server = sorted(expected_server - packaged) + if missing_server: + raise CiError(f"{archive}: server files are missing: {missing_server}") + print( + f"validated source archive={archive} " + f"family_requirements={len(expected_dependencies) - 1}" + ) def load_native_libraries(bin_dir: Path, families: tuple[str, ...]) -> None: @@ -119,6 +148,33 @@ def validate(self, wheels: list[Path]) -> None: raise CiError(f"{wheel}: Python core package is missing") if "trtmc_benchmark/__init__.py" not in names: raise CiError(f"{wheel}: Python benchmark application is missing") + expected_server = server_python_files(self.context.repository) + missing_server = sorted(set(expected_server) - set(names)) + if missing_server: + raise CiError(f"{wheel}: Python server package is missing: {missing_server}") + mismatched_server = sorted( + name + for name, source in expected_server.items() + if archive.read(name) != source.read_bytes() + ) + if mismatched_server: + raise CiError(f"{wheel}: Python server files differ from Source: {mismatched_server}") + invalid_server = [] + for name in expected_server: + try: + compile(archive.read(name), name, "exec") + except (SyntaxError, UnicodeError): + invalid_server.append(name) + if invalid_server: + raise CiError(f"{wheel}: Python server files do not compile: {invalid_server}") + legacy_server = sorted( + name + for name in names + if name == "tensorrt_model_connect/serve.py" + or name.startswith("tensorrt_model_connect/serve/") + ) + if legacy_server: + raise CiError(f"{wheel}: legacy in-package server namespace is present") source_suffixes = { ".c", ".cc", @@ -152,7 +208,7 @@ def validate(self, wheels: list[Path]) -> None: for line in metadata.splitlines() if line.startswith("Provides-Extra:") ) - if extras != ["cutedsl", "test"]: + if extras != ["cutedsl", "serve", "test"]: raise CiError(f"{wheel}: expected only application extras, found {extras}") packaged_python = tuple( sorted( @@ -294,9 +350,11 @@ def validate(self, wheel: Path) -> None: core = importlib.import_module("tensorrt_model_connect") families = importlib.import_module("families") +server = importlib.import_module("trtmc_server") print(json.dumps({ "core": str(Path(core.__file__).resolve()), "families": str(Path(families.__file__).resolve()), + "server": str(Path(server.__file__).resolve()), "family_requirements": sorted( path.parent.name for path in Path(families.__file__).resolve().parent.glob("*/requirements.txt") ), @@ -315,7 +373,7 @@ def validate(self, wheel: Path) -> None: env=environment, ) payload = json.loads(completed.stdout) - imported = (Path(payload["core"]), Path(payload["families"])) + imported = (Path(payload["core"]), Path(payload["families"]), Path(payload["server"])) if any(path.is_relative_to(self.repository.resolve()) for path in imported): raise CiError("installed wheel validation imported the source checkout") bin_dir = Path(payload["bin"]) @@ -369,6 +427,16 @@ def validate(self, wheel: Path) -> None: ) if not version.stdout.startswith("trtmc "): raise CiError("installed trtmc CLI returned an invalid version") + serve_help = subprocess.run( + [executable, "serve", "--help"], + check=True, + capture_output=True, + text=True, + cwd=Path("/tmp"), + env=environment, + ) + if "Serve TensorRT-Model-Connect bundles" not in serve_help.stdout: + raise CiError("installed trtmc serve CLI returned invalid help") with tempfile.TemporaryDirectory(prefix="trtmc-installed-wheel-") as directory: bundle = Path(directory) / "inspect.bundle" subprocess.run( diff --git a/tools/ci/quality.py b/tools/ci/quality.py index dbc9d07069..c28d86a0ce 100644 --- a/tools/ci/quality.py +++ b/tools/ci/quality.py @@ -82,6 +82,7 @@ def complexity(self) -> None: "python", "tools/check_cyclomatic_complexity.py", "core/runtime", + "server/native", "--max-ccn", "10", "--top", @@ -124,12 +125,14 @@ def architecture_contracts(self) -> None: "tools/tests/test_public_source_hygiene.py", "tools/tests/test_new_ci.py", "tools/tests/test_pr_metadata.py", + "server/tests/test_dependency_direction.py", "-q", "-p", "no:cacheprovider", ], updates={ "PYTHONPATH": ( + f"{self.context.repository / 'server/python'}:" f"{self.context.repository / 'core/builder'}:" f"{self.context.repository / 'apps/benchmark'}:" f"{self.context.repository}" @@ -165,6 +168,15 @@ def __init__(self, context: CiContext): def premerge(self) -> None: EnvironmentVerifier(self.context).verify() + pytest_options = [ + "-q", + "-x", + "-m", + "not gpu and not trt", + "-p", + "no:cacheprovider", + ] + python_timeout = self.context.env.get("PYTHON_UNIT_TIMEOUT", "20m") self.context.run( [ "python", @@ -179,12 +191,7 @@ def premerge(self) -> None: "test_voicechat_full_duplex_source.py" ), "tools/tests", - "-q", - "-x", - "-m", - "not gpu and not trt", - "-p", - "no:cacheprovider", + *pytest_options, ], updates={ "PYTHONPATH": ( @@ -194,7 +201,19 @@ def premerge(self) -> None: ), "PYTHONDONTWRITEBYTECODE": "1", }, - limit=self.context.env.get("PYTHON_UNIT_TIMEOUT", "20m"), + limit=python_timeout, + ) + self.context.run( + ["python", "-m", "pytest", "server/tests", *pytest_options], + updates={ + "PYTHONPATH": ( + "/opt/trtmc-server-test-deps:" + f"{self.context.repository / 'server/python'}:" + f"{self.context.repository}" + ), + "PYTHONDONTWRITEBYTECODE": "1", + }, + limit=python_timeout, ) build = Path( self.context.env.get( diff --git a/tools/test_impact.py b/tools/test_impact.py index 456b9180a2..17b4817283 100644 --- a/tools/test_impact.py +++ b/tools/test_impact.py @@ -26,6 +26,10 @@ MODEL_PROOF_NEUTRAL_FILES = { "apps/benchmark/performance/release.yaml", } +# The server is an application over public runtime contracts. Its complete +# Python and native contract suites run in the CPU unit stage, while unrelated +# model-family E2E does not validate server behavior. +UNIT_ONLY_PREFIXES = ("server/",) SHARED_PREFIXES = ( ".github/", "apps/", @@ -92,6 +96,7 @@ def classify(repo: Path, files: Sequence[str]) -> Impact: selected: set[str] = set() shared = False docs = False + unit_only = False unknown: list[str] = [] for path in changed: @@ -113,6 +118,9 @@ def classify(repo: Path, files: Sequence[str]) -> Impact: continue if path in MODEL_PROOF_NEUTRAL_FILES: continue + if path.startswith(UNIT_ONLY_PREFIXES): + unit_only = True + continue if len(parts) == 1 and path.endswith(".py"): shared = True continue @@ -131,6 +139,8 @@ def classify(repo: Path, files: Sequence[str]) -> Impact: return Impact("all", tuple(sorted(known)), direct_families, changed, True, docs) if selected: return Impact("families", direct_families, direct_families, changed, True, docs) + if unit_only: + return Impact("none", (), (), changed, True, docs) return Impact("docs" if docs else "none", (), (), changed, False, docs) diff --git a/tools/tests/test_architecture.py b/tools/tests/test_architecture.py index f3969e24db..be568edea1 100644 --- a/tools/tests/test_architecture.py +++ b/tools/tests/test_architecture.py @@ -7,6 +7,7 @@ import importlib.util import json import re +import tomllib from pathlib import Path @@ -384,6 +385,7 @@ def test_applications_depend_only_on_public_model_connect_surfaces() -> None: application_roots = ( REPO / "apps", REPO / "examples", + REPO / "server", ) application_files = [path for root in application_roots for path in root.rglob("*")] + [ REPO / "tools/perf_matrix.py" @@ -426,7 +428,7 @@ def test_applications_depend_only_on_public_model_connect_surfaces() -> None: source = path.read_text(encoding="utf-8", errors="ignore") if path.suffix in {".cpp", ".h", ".hpp", ".cu"}: for include in re.findall(r'#include\s+[<"]([^>"]+)', source): - if include.startswith(("apps/", "examples/")): + if include.startswith(("apps/", "examples/", "server/")): violations.append(f"{path.relative_to(REPO)}:reverse-include:{include}") for path in (REPO / "core/builder/tensorrt_model_connect").rglob("*.py"): tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) @@ -437,7 +439,12 @@ def test_applications_depend_only_on_public_model_connect_surfaces() -> None: elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module: modules.append(node.module) for module in modules: - if module == "trtmc_benchmark" or module.startswith("trtmc_benchmark."): + if ( + module == "trtmc_benchmark" + or module.startswith("trtmc_benchmark.") + or module == "trtmc_server" + or module.startswith("trtmc_server.") + ): violations.append(f"{path.relative_to(REPO)}:{node.lineno}:reverse:{module}") assert violations == [] @@ -764,11 +771,19 @@ def dependency_lines(path: Path) -> list[str]: pyproject = (REPO / "pyproject.toml").read_text(encoding="utf-8") assert 'requires-python = ">=3.12"' in pyproject assert "tomli" not in pyproject - optional = pyproject.split("[project.optional-dependencies]", 1)[1].split("\n[", 1)[0] - assert set(re.findall(r"^([a-z][a-z0-9_-]*)\s*=", optional, re.MULTILINE)) == { + optional = tomllib.loads(pyproject)["project"]["optional-dependencies"] + assert set(optional) == { "cutedsl", + "serve", "test", } + assert optional["serve"] == [ + "fastapi>=0.115,<0.142", + "pydantic>=2.11,<3", + "python-multipart>=0.0.9,<1", + "uvicorn>=0.30,<0.53", + "websockets>=13,<17", + ] requirements = sorted(FAMILIES.glob("*/requirements.txt")) assert requirements diff --git a/tools/tests/test_coderabbit_config.py b/tools/tests/test_coderabbit_config.py index dade94f6ac..b63b2cc63c 100644 --- a/tools/tests/test_coderabbit_config.py +++ b/tools/tests/test_coderabbit_config.py @@ -32,6 +32,10 @@ def test_coderabbit_covers_only_current_shared_paths() -> None: assert "model-agnostic" in instructions["core/**"] assert "public core" in instructions["apps/**"] + assert "one-way" in instructions["server/**"] + assert "fixed" in instructions["server/**"] + assert "self-healing" in instructions["server/**"] + assert "weakening" in instructions["server/tests/**"] assert "timed regions" in instructions["apps/benchmark/**"] assert not { "python/tensorrt_model_connect/families/**", diff --git a/tools/tests/test_family_impact.py b/tools/tests/test_family_impact.py index 2f6b395527..98d24cdce6 100644 --- a/tools/tests/test_family_impact.py +++ b/tools/tests/test_family_impact.py @@ -72,6 +72,37 @@ def test_shared_contract_selects_all_directly(tmp_path: Path) -> None: assert impact.families == ("alpha", "beta") +def test_server_change_runs_units_without_unrelated_family_proofs(tmp_path: Path) -> None: + impact = test_impact.classify(_repo(tmp_path), ["server/python/trtmc_server/app.py"]) + + assert impact.scope == "none" + assert impact.families == () + assert impact.direct_families == () + assert impact.run_core_tests is True + assert impact.run_docs is False + + +def test_server_and_family_change_selects_only_the_family(tmp_path: Path) -> None: + impact = test_impact.classify( + _repo(tmp_path), + ["server/native/worker.cpp", "families/alpha/model.py"], + ) + + assert impact.scope == "families" + assert impact.families == ("alpha",) + assert impact.direct_families == ("alpha",) + + +def test_server_does_not_narrow_a_shared_change(tmp_path: Path) -> None: + impact = test_impact.classify( + _repo(tmp_path), + ["server/native/worker.cpp", "core/runtime/include/trtmc/task.h"], + ) + + assert impact.scope == "all" + assert impact.families == ("alpha", "beta") + + def test_shared_change_preserves_directly_changed_family(tmp_path: Path) -> None: repo = _repo(tmp_path) impact = test_impact.classify( diff --git a/tools/tests/test_new_ci.py b/tools/tests/test_new_ci.py index cab2bc05d3..ab008a32c9 100644 --- a/tools/tests/test_new_ci.py +++ b/tools/tests/test_new_ci.py @@ -21,6 +21,7 @@ from tools.ci.docker_image import DockerImageManager from tools.ci.e2e import E2ERunner, _require_e2e_junit, _require_passing_junit from tools.ci.package import ( + InstalledWheelValidator, SourceArchiveValidator, WheelArchiveValidator, WheelPackageManager, @@ -463,6 +464,14 @@ def test_source_quality_runs_complexity_before_other_checks() -> None: assert source.index("self.complexity()") < source.index("self.lint_changed_files()") assert source.index("self.lint_changed_files()") < source.index("self.architecture_contracts()") + complexity = inspect.getsource(SourceQualityChecks.complexity) + assert '"core/runtime"' in complexity + assert '"server/native"' in complexity + + architecture = inspect.getsource(SourceQualityChecks.architecture_contracts) + assert '"server/tests/test_dependency_direction.py"' in architecture + assert "server/python" in architecture + def test_source_quality_lints_only_files_that_still_exist(tmp_path: Path) -> None: (tmp_path / "kept.py").write_text("", encoding="utf-8") @@ -643,6 +652,15 @@ def test_dev_images_install_the_pinned_requirements_without_deleted_docs() -> No assert "source-build.md" not in source +def test_server_dependencies_are_isolated_from_the_shared_runtime_python() -> None: + dockerfile = (Path(__file__).resolve().parents[2] / "Dockerfile").read_text( + encoding="utf-8" + ) + + assert "pip install --target /opt/trtmc-server-test-deps" in dockerfile + assert "PYTHONPATH=/opt/trtmc-server-test-deps" not in dockerfile + + def test_gpu_free_unit_scope_keeps_family_python_in_physical_jobs(tmp_path: Path) -> None: context = RecordingContext( tmp_path, @@ -651,9 +669,9 @@ def test_gpu_free_unit_scope_keeps_family_python_in_physical_jobs(tmp_path: Path UnitTestRunner(context).premerge() - python_command = context.calls[2][0] - assert python_command[:3] == ["python", "-m", "pytest"] - assert python_command[3:9] == [ + core_python_command = context.calls[2][0] + assert core_python_command[:3] == ["python", "-m", "pytest"] + assert core_python_command[3:9] == [ "core/builder/tests", "apps/benchmark/trtmc_benchmark/tests", "examples/audio_streaming/test_audio_streaming.py", @@ -661,9 +679,16 @@ def test_gpu_free_unit_scope_keeps_family_python_in_physical_jobs(tmp_path: Path ("examples/models/nemotron_voicechat/full_duplex/test_voicechat_full_duplex_source.py"), "tools/tests", ] - assert "families" not in python_command + assert "families" not in core_python_command + assert "server/python" not in context.calls[2][1]["updates"]["PYTHONPATH"] - ctest_command = context.calls[5][0] + server_python_command = context.calls[3][0] + assert server_python_command[:4] == ["python", "-m", "pytest", "server/tests"] + server_pythonpath = context.calls[3][1]["updates"]["PYTHONPATH"] + assert server_pythonpath.startswith("/opt/trtmc-server-test-deps:") + assert "server/python" in server_pythonpath + + ctest_command = context.calls[6][0] assert ctest_command[-2:] == ["--label-exclude", "gpu"] source = inspect.getsource(UnitTestRunner.premerge) @@ -776,6 +801,9 @@ def test_source_archive_carries_base_and_family_requirements(tmp_path: Path) -> family.mkdir(parents=True) (family / "model.py").write_text("def build(request, writer): pass\n") (family / "requirements.txt").write_text("family-dependency\n") + server_readme = tmp_path / "server/README.md" + server_readme.parent.mkdir() + server_readme.write_text("server boundary\n") archive_path = tmp_path / "package.tar.gz" def write_archive(paths: tuple[str, ...]) -> None: @@ -791,6 +819,12 @@ def write_archive(paths: tuple[str, ...]) -> None: SourceArchiveValidator(CiContext(tmp_path, {})).validate([archive_path]) write_archive(("requirements/base.txt", "families/alpha/requirements.txt")) + with pytest.raises(CiError, match="server files are missing"): + SourceArchiveValidator(CiContext(tmp_path, {})).validate([archive_path]) + + write_archive( + ("requirements/base.txt", "families/alpha/requirements.txt", "server/README.md") + ) SourceArchiveValidator(CiContext(tmp_path, {})).validate([archive_path]) @@ -800,10 +834,19 @@ def test_wheel_validation_requires_exact_new_payload(tmp_path: Path) -> None: root = tmp_path / "families" / family root.mkdir(parents=True) (root / "model.py").write_text("def build(request, writer): pass\n") + server = tmp_path / "server/python/trtmc_server" + server.mkdir(parents=True) + (server / "__init__.py").write_text('"""Server package."""\n') + (server / "worker.py").write_text("READY = True\n") wheel = tmp_path / "package.whl" with zipfile.ZipFile(wheel, "w") as archive: archive.writestr("tensorrt_model_connect/__init__.py", "") archive.writestr("trtmc_benchmark/__init__.py", "") + archive.writestr( + "trtmc_server/__init__.py", + (server / "__init__.py").read_bytes(), + ) + archive.writestr("trtmc_server/worker.py", (server / "worker.py").read_bytes()) archive.writestr("families/__init__.py", "") archive.writestr("tensorrt_model_connect/bin/trtmc", "") archive.writestr("tensorrt_model_connect/bin/trtmc_benchmark_worker", "") @@ -822,6 +865,7 @@ def test_wheel_validation_requires_exact_new_payload(tmp_path: Path) -> None: "Name: package\n" "Version: 0.1\n" "Provides-Extra: cutedsl\n" + "Provides-Extra: serve\n" "Provides-Extra: test\n", ) archive.writestr("package-0.1.data/scripts/trtmc", "") @@ -836,6 +880,30 @@ def test_wheel_validation_requires_exact_new_payload(tmp_path: Path) -> None: WheelArchiveValidator(CiContext(tmp_path, {})).validate([wheel]) + server_corrupt = tmp_path / "server-corrupt.whl" + with zipfile.ZipFile(wheel) as source, zipfile.ZipFile(server_corrupt, "w") as output: + for entry in source.infolist(): + payload = source.read(entry.filename) + if entry.filename == "trtmc_server/worker.py": + payload = b"READY = False\n" + output.writestr(entry, payload) + with pytest.raises(CiError, match="Python server files differ from Source"): + WheelArchiveValidator(CiContext(tmp_path, {})).validate([server_corrupt]) + + legacy = tmp_path / "legacy-server.whl" + with zipfile.ZipFile(wheel) as source, zipfile.ZipFile(legacy, "w") as output: + for entry in source.infolist(): + output.writestr(entry, source.read(entry.filename)) + output.writestr("tensorrt_model_connect/serve/__init__.py", "") + with pytest.raises(CiError, match="legacy in-package server namespace"): + WheelArchiveValidator(CiContext(tmp_path, {})).validate([legacy]) + + server_helper = server / "new_helper.py" + server_helper.write_text("VALUE = 1\n") + with pytest.raises(CiError, match="Python server package is missing"): + WheelArchiveValidator(CiContext(tmp_path, {})).validate([wheel]) + server_helper.unlink() + corrupt = tmp_path / "corrupt.whl" with zipfile.ZipFile(wheel) as source, zipfile.ZipFile(corrupt, "w") as output: for entry in source.infolist(): @@ -876,6 +944,14 @@ def test_wheel_validation_requires_exact_new_payload(tmp_path: Path) -> None: WheelArchiveValidator(CiContext(tmp_path, {})).validate([wheel]) +def test_installed_wheel_validation_checks_server_import_and_cli() -> None: + source = inspect.getsource(InstalledWheelValidator.validate) + + assert 'import_module("trtmc_server")' in source + assert '[executable, "serve", "--help"]' in source + assert "installed trtmc serve CLI returned invalid help" in source + + def test_native_validation_rejects_unresolved_family_symbols(tmp_path: Path) -> None: def compile_library(name: str, source: str) -> None: source_path = tmp_path / f"{name}.c" diff --git a/website/docs/api/cli-reference.md b/website/docs/api/cli-reference.md index 062f946e57..87cdc0ad60 100644 --- a/website/docs/api/cli-reference.md +++ b/website/docs/api/cli-reference.md @@ -123,5 +123,40 @@ speech commands expose only the options listed by their Task contracts in Unknown commands, unknown command-specific options, duplicate options, task interface mismatches, invalid values, and missing DSOs fail with a nonzero exit -status. Run `trtmc help` for the compiled executable's concise synopsis and +status. + +## Serve local bundles + +Install the optional control-plane dependencies, then register one or more +bundles behind the local process API: + +```bash +python -m pip install "tensorrt-model-connect[serve]" + +trtmc serve \ + --runtime-root /opt/trtmc/lib \ + --chat-model chat=/models/qwen.bundle \ + --host 127.0.0.1 \ + --port 8000 +``` + +`--chat-model NAME=PATH` and `--transcription-model NAME=PATH` are repeatable. +`--model-replicas NAME=N` sets a fixed number of independent native execution +lanes for one registered model. `--require-streaming-transcription MODEL` +fails startup unless every configured replica for that model passes the native +streaming probe. The server has no waiting queue or dynamic worker placement; +saturation returns HTTP 429, and a failed lane is not restarted. See +[Serve Local Models](../user-guides/serve-local-models.md) for the HTTP, +Realtime, readiness, security, and scaling boundaries. + +The initial chat endpoint accepts exactly one text-only `user` message. Stop +sequences, multi-turn messages, text streaming, and tool calling are rejected. + +Use `--model-replicas` above `1` only for bundles that can be loaded as +independent single-process workers. MPI/NCCL distributed bundles are not +supported by `trtmc serve`. To use multiple GPUs, run one independent +single-process server instance per GPU, pin each process with +`CUDA_VISIBLE_DEVICES`, and put any routing outside this server. + +Run `trtmc help` for the compiled executable's concise synopsis and `python -m tensorrt_model_connect build --help` for the exact build parser. diff --git a/website/docs/api/overview.md b/website/docs/api/overview.md index 6a6aa7d032..17027c0b0b 100644 --- a/website/docs/api/overview.md +++ b/website/docs/api/overview.md @@ -8,13 +8,14 @@ Reference pages are for exact lookup. Begin with the or use the [User Guides](../user-guides/overview.md) for goal-oriented procedures. -TensorRT-Model-Connect exposes three public entry layers: +TensorRT-Model-Connect exposes four public entry layers: | API | Entry point | Best for | | --- | --- | --- | | Python build API | `python -m tensorrt_model_connect build` and `tensorrt_model_connect.build()` | Resolving a supported checkpoint and building a `.bundle`. | | Native CLI | `trtmc inspect` and task commands such as `trtmc run` | Inspecting a bundle or invoking one abstract Task interface. | | C++ Task API | `#include ` and `trtmc::load_task()` | Native applications that need task-specific results. | +| Local serving process | `trtmc serve` | Local applications that need persistent bundle workers behind HTTP or WebSocket contracts. | The build and runtime entry points are intentionally separate. The Python builder resolves exactly one `families//support.py`, imports only that @@ -30,6 +31,8 @@ Hugging Face model ID or local snapshot -> task-specific output ``` -There is no Python runtime wrapper, central model registry, runtime-strategy -switch, backend search path, or fallback runtime discovery in the current -architecture. +The local server is a downstream application over the same loader and Task +contracts. Its Python control plane never owns TensorRT objects or becomes a +library dependency. There is no general Python runtime wrapper, central model +registry, runtime-strategy switch, backend search path, or fallback runtime +discovery in the current architecture. diff --git a/website/docs/architecture/ai-native-horizontal-scaling.md b/website/docs/architecture/ai-native-horizontal-scaling.md index f975ff1c79..5951875433 100644 --- a/website/docs/architecture/ai-native-horizontal-scaling.md +++ b/website/docs/architecture/ai-native-horizontal-scaling.md @@ -281,6 +281,7 @@ interfaces cannot express. | Family runtime | factory contract, Task APIs, BundleReader, Engine API, model-owned custom plugin API when needed | sibling family, loader implementation, concrete backend implementation | | Engine backend | Engine API | family, Task behavior, model policy | | Examples and benchmark | public build, load, Task, and BYOK APIs | family/backend private implementation, reverse core dependency | +| Local server | public loader and Task APIs | family/backend private implementation, reverse core dependency | ### Application dependency is one-way @@ -292,13 +293,21 @@ flowchart BT Benchmark["Benchmark"] --> BuildAPI Benchmark --> LoadAPI Benchmark --> TaskAPI + Server["Local server"] --> LoadAPI + Server --> TaskAPI ByokExample["BYOK example"] --> ByokAPI["Public BYOK API"] ByokAPI --> TVMFFI["TVM-FFI C ABI"] ByokAPI --> TRTPlugin["TensorRT plugin API"] ``` Core, families, and backend must not import, include, or link `examples/`, -`apps/benchmark/`, or the benchmark Python package. +`apps/benchmark/`, `server/`, or their Python packages. + +The local server's replicas are fixed process-local execution lanes. This is +bounded concurrency within one placement domain, not the family-level +horizontal scaling architecture described on this page. The server does not +perform cluster discovery, cross-host routing, autoscaling, or worker restart; +those remain external deployment concerns. The following dependencies are always forbidden: diff --git a/website/docs/architecture/build-system.md b/website/docs/architecture/build-system.md index 0e88a85da3..648b716f64 100644 --- a/website/docs/architecture/build-system.md +++ b/website/docs/architecture/build-system.md @@ -12,6 +12,7 @@ family DSOs. | `trtmc_backend_trt` | Standard TensorRT Engine implementation. | | `trtmc_backend_rtx` | Optional TensorRT-RTX Engine implementation. | | `trtmc` | Native application under `apps/cli/`. | +| `trtmc_server_native` | Private JSONL worker adapter under `server/native/`; linked downstream of the runtime loader. | | `trtmc_benchmark_worker` | Benchmark application under `apps/benchmark/`. | | `trtmc_model_` | One family's complete native runtime. | @@ -23,8 +24,8 @@ private dependencies, warnings, tests, output name, and install rule. Adding a normal family never changes a central model source list. The wheel packages `core/builder/tensorrt_model_connect`, the top-level -`families` package, benchmark Python code, and installed native binaries/DSOs. -Optional family dependencies remain in each +`families` package, benchmark Python code, the optional `trtmc_server` control +plane, and installed native binaries/DSOs. Optional family dependencies remain in each `families//requirements.txt`; package validation does not import every family implementation. diff --git a/website/docs/architecture/overview.md b/website/docs/architecture/overview.md index eba93fad56..cb39153758 100644 --- a/website/docs/architecture/overview.md +++ b/website/docs/architecture/overview.md @@ -45,8 +45,22 @@ second strategy dispatch. ## Applications stay above the public boundary The native CLI in `apps/cli/`, benchmark application in `apps/benchmark/`, -examples, and TVM-FFI BYOK use public build, load, Task, and Engine contracts. -Core and families never depend on those applications. +the optional local server in `server/`, examples, and TVM-FFI BYOK use public +build, load, Task, and Engine contracts. Core and families never depend on +those applications. + +The server keeps HTTP and WebSocket handling in a Python control-plane process +and loads bundles only in native child workers. Each configured replica is one +fixed serial execution lane. Replicas are created at startup; there is no +hidden request queue, dynamic placement, distributed scheduler, or worker +self-healing. A failed lane is removed from scheduling and readiness becomes +degraded while another lane remains healthy. One server process is one local +placement domain, not a cluster-level scaling system. + +Replica fan-out applies only to independently loadable single-process bundles. +MPI/NCCL distributed bundles are not supported by `trtmc serve`. Multi-GPU +deployment uses independent single-process server instances pinned with +`CUDA_VISIBLE_DEVICES`; placement and routing remain external concerns. See [AI-Native Horizontal Scaling Architecture](ai-native-horizontal-scaling.md) for the complete rules and [Source Layout](../reference/source-layout.md) for diff --git a/website/docs/architecture/units-and-ownership.md b/website/docs/architecture/units-and-ownership.md index b71e92fcd8..ac1e229893 100644 --- a/website/docs/architecture/units-and-ownership.md +++ b/website/docs/architecture/units-and-ownership.md @@ -30,12 +30,14 @@ duplicated across families until a stable model-agnostic contract exists. | `core/runtime/tensorrt/` | Standard TensorRT and optional TensorRT-RTX Engine implementations. | | `apps/cli/` | Native command parsing and private image/WAV/file adapters. | | `apps/benchmark/` | Benchmark catalog, workers, reference runners, and reporting policy. | +| `server/` | Optional local HTTP/WebSocket control plane, native worker process, and server-owned tests. | | `examples/` | Optional applications over public APIs. | ## Allowed dependency direction ```text applications -> public build/load/Task/BYOK contracts +local server -> public loader + Task contracts family build -> BuildRequest + BundleWriter + TensorRT build API family runtime -> BundleReader + Task + Engine contracts runtime loader -> bundle + factory + backend contracts @@ -47,3 +49,8 @@ family dependencies, backend-to-family dependencies, and family/core dependencies on applications. A family-local custom TensorRT plugin is allowed only when the family graph genuinely requires it; TVM-FFI BYOK remains the model-agnostic custom-kernel boundary. + +The server boundary is deliberately one-way. `server/` may link or call public +library contracts; `core/` and `families/` must not import, include, link, or +load server implementation. Applications integrate through the `trtmc serve` +process and its HTTP/WebSocket contracts, not by importing `trtmc_server`. diff --git a/website/docs/architecture/validation-design.md b/website/docs/architecture/validation-design.md index fe454b07ea..c4ac290a6c 100644 --- a/website/docs/architecture/validation-design.md +++ b/website/docs/architecture/validation-design.md @@ -14,6 +14,7 @@ and official-reference adapters below `families//tests/`. | Family native tests | `families//tests/cpp/` | The owning DSO's C++ contracts. | | Family E2E | `families//tests/test_e2e.py` | Checkpoint-to-bundle-to-Task behavior and model-owned oracle. | | Application tests | `apps/*/tests/`, `examples/**/test_*.py` | Public API consumers without reverse dependencies. | +| Local server | `server/tests/` | HTTP/Realtime contracts, bounded worker lifecycle, native JSONL protocol, and one-way dependency checks. | Each family manifest declares its exact checkpoint inputs, task, precision, topology, premerge selection, and case-specific contract. Optional threshold @@ -31,13 +32,19 @@ family's owner-selected cases; nightly runs the complete declared inventory. checkpoint evidence without importing unrelated families. - Package validation finds every family dependency file and DSO while avoiding unselected family imports. +- Server-only changes run the complete CPU server suite and package checks + without selecting unrelated family E2E. A mixed server/family or server/core + change retains the broader owner or shared-contract scope. +- The installed-wheel gate imports `trtmc_server` and runs `trtmc serve --help` + outside the source checkout; a source-only import is not packaging evidence. ## Local checks ```bash python3 -m tools.model_ci validate python3 tools/test_impact.py --validate -python3 -m pytest core/builder/tests tools/tests +PYTHONPATH=server/python:core/builder:. python3 -m pytest \ + core/builder/tests server/tests tools/tests python3 -m pytest families/qwen/tests ``` diff --git a/website/docs/reference/source-layout.md b/website/docs/reference/source-layout.md index 7e4c2925d5..6e4be156f6 100644 --- a/website/docs/reference/source-layout.md +++ b/website/docs/reference/source-layout.md @@ -13,10 +13,14 @@ core/runtime/loader/ libtrtmc_runtime.so exact DSO loader core/runtime/tensorrt/ libtrtmc_backend_trt.so implementation apps/cli/ native CLI and private image/audio file I/O apps/benchmark/ benchmark application, workers, and performance policy +server/python/trtmc_server/ optional local HTTP/WebSocket control plane +server/native/ native JSONL worker over public runtime contracts +server/tests/ server-owned Python and native contract tests tools/model_ci.py family inventory and impact website/ documentation generated from family ownership ``` `core/builder/` contains only Python. `core/runtime/` contains only C++ headers -and sources. No production source lives under the retired `python/`, `src/`, -`include/`, or `tests/` roots. +and sources. The server is a separate downstream unit rather than part of +either core tree. No production source lives under the retired `python/`, +`src/`, `include/`, or `tests/` roots. diff --git a/website/docs/reference/testing.md b/website/docs/reference/testing.md index b091fadfba..45b1417c9b 100644 --- a/website/docs/reference/testing.md +++ b/website/docs/reference/testing.md @@ -14,8 +14,8 @@ contracts and mechanics. Run the structural validators before selecting expensive model tests: ```bash -PYTHONPATH=core/builder:. python3 -m tools.model_ci validate -PYTHONPATH=core/builder:. python3 tools/test_impact.py --validate +PYTHONPATH=server/python:core/builder:. python3 -m tools.model_ci validate +PYTHONPATH=server/python:core/builder:. python3 tools/test_impact.py --validate git diff --check ``` @@ -27,9 +27,20 @@ Run host-side Python tests with the builder package and repository root on the import path: ```bash -PYTHONPATH=core/builder:. python3 -m pytest core/builder/tests tools/tests -q +PYTHONPATH=server/python:core/builder:apps/benchmark:. python3 -m pytest \ + core/builder/tests \ + apps/benchmark/trtmc_benchmark/tests \ + server/tests \ + tools/tests \ + -q -m "not gpu and not trt" ``` +`server/tests/` covers both the Python process/API contracts and the static +one-way dependency boundary. Its native JSONL worker test is compiled into the +same CTest tree and remains CPU-runnable. Package validation separately proves +that the sdist carries `server/`, the wheel carries `trtmc_server`, and the +installed native CLI can dispatch `trtmc serve --help` outside the checkout. + After configuring a native build, run its compiled tests with CTest: ```bash diff --git a/website/docs/user-guides/overview.md b/website/docs/user-guides/overview.md index 3b44534ce9..727e8dccfc 100644 --- a/website/docs/user-guides/overview.md +++ b/website/docs/user-guides/overview.md @@ -12,6 +12,7 @@ command as a support claim. | Create an artifact | [Build a Bundle](build-a-bundle.md) | A format-1 `.bundle` built by exactly one family. | | Diagnose an artifact | [Inspect a Bundle](inspect-a-bundle.md) | Family, task, backend, and section inventory. | | Execute a task | [Run Inference](run-inference.md) | Correct Task command and typed JSON/media output. | +| Keep local models loaded | [Serve Local Models](serve-local-models.md) | Local HTTP and Realtime APIs backed by fixed native worker lanes. | | Place a setting correctly | [Configure Runtime Behavior](configure-runtime.md) | Build, load, or request input at its typed boundary. | | Establish evidence | [Validate & Benchmark](validate-benchmark.md) | Reproducible correctness or performance evidence. | diff --git a/website/docs/user-guides/serve-local-models.md b/website/docs/user-guides/serve-local-models.md new file mode 100644 index 0000000000..39c55f6b99 --- /dev/null +++ b/website/docs/user-guides/serve-local-models.md @@ -0,0 +1,150 @@ +--- +title: Serve Local Models +description: Keep text and speech bundles loaded behind local HTTP and Realtime APIs. +--- + +`trtmc serve` starts an optional Python HTTP/WebSocket control plane and a +fixed group of native worker processes. Each worker loads one bundle through +the public runtime loader and owns one serial execution lane for its lifetime. + +## Install the optional control plane + +```bash +python -m pip install "tensorrt-model-connect[serve]" +``` + +The native runtime remains C++. The Python package owns transport, validation, +authentication, and worker lifecycle only; TensorRT objects stay in the native +workers. + +## Start a local server + +```bash +export TRTMC_SERVE_TOKEN="replace-with-a-random-local-token" + +trtmc serve \ + --runtime-root /opt/trtmc/lib \ + --chat-model chat=/models/qwen.bundle \ + --transcription-model asr=/models/whisper.bundle \ + --default-chat-model chat \ + --default-transcription-model asr \ + --host 127.0.0.1 \ + --port 8000 +``` + +A non-empty token is required. Prefer `TRTMC_SERVE_TOKEN` to `--api-key` so the +credential does not appear in the process argument list. + +The initial server accepts only loopback IP literals such as `127.0.0.1` and +`::1`. Hostnames such as `localhost` and every non-loopback bind are rejected. +Access logs are disabled by default because browser WebSocket clients may use +an `access_token` query parameter. Other diagnostics continue on stderr with +transport credentials redacted. + +Check process health, whole-registry readiness, and registered models: + +```bash +curl http://127.0.0.1:8000/healthz +curl http://127.0.0.1:8000/readyz \ + -H "Authorization: Bearer $TRTMC_SERVE_TOKEN" +curl http://127.0.0.1:8000/v1/models \ + -H "Authorization: Bearer $TRTMC_SERVE_TOKEN" +``` + +`/healthz` is unauthenticated and detail-free. It remains healthy while at +least one native worker is usable. Authenticated `/readyz` reports registry and +replica state and fails when any configured model has no ready lane. + +## Generate text + +```bash +curl http://127.0.0.1:8000/v1/chat/completions \ + -H "Authorization: Bearer $TRTMC_SERVE_TOKEN" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "chat", + "messages": [{"role": "user", "content": "Summarize this transcript."}], + "max_completion_tokens": 160, + "stream": false + }' +``` + +The initial endpoint accepts exactly one text-only `user` message. Multi-turn +messages, stop sequences, text streaming, tool calling, log probabilities, +native token callbacks, and cooperative cancellation are not implemented. +Unsupported execution options fail explicitly rather than being silently +ignored. + +## Transcribe a WAV file + +```bash +curl http://127.0.0.1:8000/v1/audio/transcriptions \ + -H "Authorization: Bearer $TRTMC_SERVE_TOKEN" \ + -F model=asr \ + -F response_format=verbose_json \ + -F file=@meeting.wav +``` + +The endpoint accepts PCM16 and IEEE float32 WAV data supported by the native +audio adapter. Oversized request bodies are rejected before multipart parsing +or temporary-file spooling. +`verbose_json` includes timestamp segments when the selected Task returns +them. + +## Use Realtime transcription + +Connect to: + +```text +ws://127.0.0.1:8000/v1/realtime?intent=transcription&access_token=TOKEN +``` + +The client sends `session.update`, `input_audio_buffer.append`, +`input_audio_buffer.commit`, and `input_audio_buffer.clear`. The server returns +session, transcription, and structured error events. True partial results +require native streaming support. The `--require-streaming-transcription MODEL` +option requires every configured replica for that model to pass the native +streaming startup probe. Other speech bundles remain usable through the offline +endpoint without that requirement. + +## Configure bounded concurrency + +The default is one worker lane per model. Add replicas only after confirming +that the model copies fit in available GPU memory: + +```bash +trtmc serve \ + --runtime-root /opt/trtmc/lib \ + --chat-model chat=/models/qwen.bundle \ + --model-replicas chat=2 +``` + +Each replica is a separate native process and may duplicate model and KV-cache +memory. There is no server-side waiting queue: a request atomically leases one +idle lane, and the server returns HTTP 429 when every lane is busy. + +Set a replica count above `1` only for a bundle that can be loaded independently +in each single-process worker. MPI/NCCL distributed bundles are not supported +by `trtmc serve`. + +To use multiple GPUs, run one independent single-process server instance per +GPU and pin each process before startup, for example with +`CUDA_VISIBLE_DEVICES=0` and `CUDA_VISIBLE_DEVICES=1`. Each instance needs a +distinct port; any routing across instances remains external to this server. + +## Scaling and failure boundary + +This release intentionally implements one simple placement domain: + +- replica counts and model assignments are fixed at process startup; +- one replica handles one request or Realtime session at a time; +- failed replicas are removed from scheduling and are not restarted; +- a model is degraded while some, but not all, configured replicas are ready; +- Realtime sessions hold one lane until commit, clear, failure cleanup, or + disconnect; +- there is no continuous batching, dynamic model loading, cluster membership, + cross-host routing, autoscaling, or rolling replacement. + +Operational restart and multi-process placement belong to an external +supervisor. The loopback-only bind means this initial server is for local +applications; it is not a network-facing multi-node serving tier. diff --git a/website/sidebars.js b/website/sidebars.js index f7f437b673..7b616afb4c 100644 --- a/website/sidebars.js +++ b/website/sidebars.js @@ -62,6 +62,7 @@ module.exports = { 'user-guides/build-a-bundle', 'user-guides/inspect-a-bundle', 'user-guides/run-inference', + 'user-guides/serve-local-models', { type: 'category', label: 'Task Guides',