NVIDIA-NeMo / NVIDIA-NeMo/Gym

Use FastAPI-compatible ORJSON encoding across Gym servers

Open
#3,003 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

core-infra
Dominant language
Python
Stars
1.2k
Forks
349
Avg merge
1d 21h
Merged PRs (30d)
318

Description

Parent workstream: #2999

Slowdown path

flowchart LR
    A[Typed handler result] --> B[FastAPI response contract]
    B --> C[Generic Python traversal]
    C --> D[Standard JSON encoding]
    D --> E[Large response bytes]

Current behavior

Agent and resources servers create ordinary FastAPI applications and register typed handlers without a shared fast JSON response path. See SimpleResponsesAPIAgent.setup_webserver and SimpleResourcesServer.setup_webserver.

The model-server-only _orjson_dispatch_response fast path turns a completed model directly into a Response. It does not provide an app-wide implementation that retains FastAPI's response-model validation, field filtering, aliases, custom serializers, injected status and headers, cookies, background tasks, ready Response objects, or OpenAPI behavior. See _orjson_dispatch_response.

Impact and provenance

Measured

The evidence report measured the parity-preserving app-wide prototype at 7.5 to 9.5 rollouts/s for a training-sized end-to-end body, a 27% increase. Its isolated agent/resources hop improved from 22 to 45 requests/s. These results came from the report's loopback full-rollout and hop rigs; they are not production-cluster measurements.

Expected

A shared encoder should reduce CPU time on non-streaming JSON responses across agents and resources servers. The gain will depend on payload size and concurrency, so no production throughput improvement is assumed.

Implementation constraints

  • Preserve FastAPI's complete response contract and OpenAPI schemas; change the final byte encoder rather than bypassing response processing.
  • Preserve ready Response and streaming/SSE pass-through behavior.
  • Keep telemetry wrappers, session cookies, MCP-exposed routes, status codes, headers, and background tasks unchanged.
  • Use the shared server setup so eligible agent and resources routes inherit the behavior without per-server forks.

Acceptance criteria

  • Eligible non-streaming agent and resources responses use the fast encoder after normal FastAPI response processing.
  • Contract tests cover response-model filtering, aliases, custom serializers, include/exclude options, status codes, headers, cookies, background tasks, ready Response objects, OpenAPI schemas, telemetry wrappers, and MCP routes.
  • Streaming and non-JSON responses remain pass-through.
  • Request-validation and error response shapes remain compatible.

Benchmark plan

Run alternating baseline/candidate loopback benchmarks for small control payloads and training-shaped responses at concurrency 1 and 32. Measure requests/s, p50/p99 latency, bytes, and CPU time for agent and resources hops, then repeat the full driver → agent → model → verifier rollout rig. Report hardware, dependency versions, payload construction, sample counts, and variance.

Related work/PRs

  • Parent: #2999
  • Model-server fast path: #2867

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with SimpleResponsesAPIAgent.setup_webserver in nemo_gym/base_responses_api_agent.py, SimpleResourcesServer.setup_webserver in nemo_gym/base_resources_server.py, and _orjson_dispatch_response in nemo_gym/base_responses_api_model.py. Compare these entry points with the shared server setup, then use the contract tests and benchmark plan to verify FastAPI behavior, pass-through responses, telemetry and MCP routes, and the reported performance measurements.

Written by the indexing model from the issue text.

Assessment

Tech stack
fastapi, python
Domain
api, backend, performance, testing
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.