feat(nemo-gym): route rollouts through vLLM Router

.github/workflows/lockfile-check.yml

@@ -41,6 +41,7 @@ jobs:
       - name: Set up uv
         uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5
         with:
+          version: "0.11.28"
           enable-cache: true
 
       # Building megatron-bridge's metadata provisions its build requirements

examples/configs/evals/eval.yaml

@@ -55,6 +55,25 @@ env:
   math:
     num_workers: 8
 
+logger:
+  log_dir: "logs/eval"
+  wandb_enabled: false
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  swanlab_enabled: false
+  monitor_gpus: false # set true to log cluster GPU/host-memory metrics at eval_step
+  wandb:
+    project: "nemo-rl-eval"
+    name: "eval"
+  tensorboard: {}
+  swanlab:
+    project: "nemo-rl-eval"
+    name: "eval"
+  mlflow: {}
+  gpu_monitoring:
+    collection_interval: 10
+    flush_interval: 10
+
 cluster:
   gpus_per_node: 1
   num_nodes: 1

examples/nemo_gym/eval_workplace_assistant_nemotron_nano_v2_9b.yaml

@@ -0,0 +1,70 @@
+defaults: ../configs/evals/eval.yaml
+
+eval:
+  metric: "mean_reward"
+  num_tests_per_prompt: 1
+  save_path: "results/eval-workplace-assistant-nemotron-nano-v2-9b"
+
+generation:
+  max_new_tokens: 8192
+  model_name: "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+  top_k: null
+  num_prompts_per_step: 32
+  vllm_cfg:
+    async_engine: true
+    expose_http_server: true
+    enable_vllm_metrics_logger: true
+    vllm_metrics_logger_interval: 0.5
+    tensor_parallel_size: 1
+    gpu_memory_utilization: 0.8
+    max_model_len: 8192
+    skip_tokenizer_init: false
+    http_server_serving_chat_kwargs:
+      enable_auto_tools: true
+      tool_parser: nemotron_json
+  vllm_kwargs:
+    compilation_config:
+      backend: eager
+    mamba_ssm_cache_dtype: "float32"
+
+tokenizer:
+  name: ${generation.model_name}
+
+data:
+  max_input_seq_length: null
+  dataset_name: "NemoGymDataset"
+  # The checked-in example split makes this recipe runnable after submodule init.
+  # Override with resources_servers/workplace_assistant/data/validation.jsonl
+  # after downloading the full validation artifact through NeMo Gym.
+  data_path: "3rdparty/Gym-workspace/Gym/resources_servers/workplace_assistant/data/example.jsonl"
+  processor: "nemo_gym_data_processor"
+  env_name: "nemo_gym"
+  repeat: 1
+
+env:
+  should_use_nemo_gym: true
+  nemo_gym:
+    port_range_low: 5000
+    port_range_high: 5999
+    rollout_max_attempts_to_avoid_lp_nan: 1
+    config_paths:
+      - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+      - resources_servers/workplace_assistant/configs/workplace_assistant.yaml
+    policy_model:
+      responses_api_models:
+        vllm_model:
+          uses_reasoning_parser: false
+          extra_body:
+            chat_template_kwargs:
+              enable_thinking: false
+
+logger:
+  log_dir: "logs/eval-workplace-assistant-nemotron-nano-v2-9b"
+  wandb_enabled: true
+  wandb:
+    project: "nemo-rl-eval"
+    name: "workplace-assistant-nemotron-nano-v2-9b"
+
+cluster:
+  gpus_per_node: 8
+  num_nodes: 1

examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml

@@ -132,6 +132,7 @@ policy:
     #gives ~20% training perf speedup with sequence packing
     apply_rope_fusion: True
     defer_fp32_logits: false
+    moe_per_layer_logging: false
     moe_permute_fusion: true
     moe_enable_deepep: false
     moe_token_dispatcher_type: "alltoall"
@@ -309,6 +310,27 @@ env:
     # RL (3000-4999) or vLLM (7000-8999). See ray.sub for the full port layout.
     port_range_low: 5000
     port_range_high: 5999
+    # Optional owned Router. Set enabled=true for cache_aware or consistent_hash
+    # Phase 2 arms; false is the direct-engine baseline.
+    vllm_router:
+      enabled: false
+      policy: consistent_hash
+      # vLLM Router 0.1.15 does not increment its declared cache counters.
+      # Use debug_log_compat only with policy=cache_aware when those counters
+      # are required; it adds DEBUG-log parsing overhead.
+      cache_metrics_mode: native
+      cache_threshold: 0.3
+    # Optional strict RL-Insight/Prometheus evidence collection. A formal run
+    # should set enabled=true, required=true, RL_INSIGHT_SERVER_URL, and a stable
+    # NEMO_RL_RUN_ID. See docs/design-docs/nemo-gym-integration.md.
+    prometheus:
+      enabled: false
+      required: false
+      scrape_interval_s: 10
+      initial_scrape_wait_s: 12
+      final_scrape_wait_s: 12
+      # Formal required runs must use a separate RL-Insight instance per run.
+      target_lifecycle: shared
     is_trajectory_collection: false  # Set this to true to enable trajectory collection (no training). You may also want to increase `policy.generation.vllm_cfg.gpu_memory_utilization`
     config_paths:
     - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml  # Required! And it must be *for_training

examples/nemo_gym/phase2_experiment_metadata.example.json

@@ -0,0 +1,63 @@
+{
+  "schema_version": 1,
+  "engine": {
+    "fresh": true,
+    "launch_id": "replace-with-unique-slurm-job-and-launch-id"
+  },
+  "workload_replay": {
+    "faithful": true,
+    "workload_sha256": "0000000000000000000000000000000000000000000000000000000000000000",
+    "seed": "42",
+    "num_prompts": 64,
+    "num_generations_per_prompt": 4
+  },
+  "warmup": {
+    "completed": true,
+    "workload_sha256": "0000000000000000000000000000000000000000000000000000000000000000",
+    "requests": 32
+  },
+  "software": {
+    "nemo_rl_commit": "0000000000000000000000000000000000000000",
+    "container_digest": "sha256:0000000000000000000000000000000000000000000000000000000000000000"
+  },
+  "model": {
+    "name": "Qwen/Qwen2.5-1.5B-Instruct",
+    "revision": "replace-with-model-revision",
+    "tokenizer": "Qwen/Qwen2.5-1.5B-Instruct",
+    "tokenizer_revision": "replace-with-tokenizer-revision",
+    "chat_template_sha256": "0000000000000000000000000000000000000000000000000000000000000000"
+  },
+  "topology": {
+    "tensor_parallel_size": 1,
+    "data_parallel_size": 8,
+    "num_nodes": 1,
+    "gpus_per_node": 8
+  },
+  "generation": {
+    "sampling_parameters": {
+      "temperature": 0.7,
+      "top_p": 0.95,
+      "top_k": null
+    },
+    "concurrency": 64,
+    "max_context_tokens": 8192,
+    "max_output_tokens": 256
+  },
+  "backend": {
+    "prefix_caching_enabled": true,
+    "scheduler_parameters": {
+      "scheduling_policy": "fcfs"
+    },
+    "batching_parameters": {
+      "max_num_seqs": 256,
+      "max_num_batched_tokens": 8192
+    }
+  },
+  "router": {
+    "enabled": true,
+    "policy": "cache_aware",
+    "session_affinity_header": "X-Session-ID",
+    "cache_metrics_mode": "debug_log_compat",
+    "cache_threshold": 0.3
+  }
+}

examples/nemo_gym/rl_insight_phase2/config.yaml

@@ -0,0 +1,20 @@
+server:
+  port: 18080
+  install_dir: ${oc.env:PHASE2_RL_INSIGHT_ROOT}/services
+  runtime_dir: ${oc.env:PHASE2_RL_INSIGHT_ROOT}/runtime
+  data_dir: ${oc.env:PHASE2_RL_INSIGHT_ROOT}/data
+  state_file: ${oc.env:PHASE2_RL_INSIGHT_ROOT}/run/services.json
+
+prometheus:
+  binary_path: ${oc.env:PHASE2_PROMETHEUS_BIN}
+  prometheus_port: 19090
+  retention_time: 1d
+  config_file: ${service_root}/prometheus.yml
+
+# Phase 2 needs the target API and Prometheus only. Disabling the other
+# services reduces moving parts in the per-run dedicated instance.
+tempo:
+  enable: false
+
+grafana:
+  enable: false

examples/nemo_gym/rl_insight_phase2/prometheus.yml

@@ -0,0 +1,5 @@
+global:
+  scrape_interval: 1s
+  evaluation_interval: 1s
+
+scrape_configs: []