diff --git a/nemo_rl/models/generation/dynamo/dynamo_generation.py b/nemo_rl/models/generation/dynamo/dynamo_generation.py index 91bf3d25a32..c65380260c0 100644 --- a/nemo_rl/models/generation/dynamo/dynamo_generation.py +++ b/nemo_rl/models/generation/dynamo/dynamo_generation.py @@ -184,7 +184,7 @@ def __init__( ) self._token_wrapper_server: Optional[DynamoTokenWrapperServer] = None self._dynamo_frontend_base_url = "" - self.dp_openai_server_base_urls: list[Optional[str]] = [] + self.dp_openai_server_base_urls: list[str] = [] self._refit_channel: DynamoRefitChannel | None = None self._metrics_sampler: DynamoMetricsSampler | None = None try: @@ -217,7 +217,9 @@ def __init__( flush=True, ) else: - self.dp_openai_server_base_urls = [None] + # No token wrapper means no OpenAI-compatible server to advertise; + # stay empty rather than holding a slot open with a placeholder. + self.dp_openai_server_base_urls = [] print(f" [Dynamo] Forwarding rollouts to {url}", flush=True) if vllm_cfg.enable_vllm_metrics_logger: diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 53294bce379..eb201df4be5 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -238,7 +238,7 @@ def __init__( self._policy_config = config self.cfg: MCoreGenerationConfig = config["generation"] # Populated after the first prepare_for_generation (which starts the HTTP server). - self.dp_openai_server_base_urls: list[Optional[str]] = [] + self.dp_openai_server_base_urls: list[str] = [] # Installed by setup via create_weight_synchronizer. self.weight_synchronizer: Optional["WeightSynchronizer"] = None diff --git a/nemo_rl/models/generation/trtllm/trtllm_generation.py b/nemo_rl/models/generation/trtllm/trtllm_generation.py index 422e2099856..2bf396d35f2 100644 --- a/nemo_rl/models/generation/trtllm/trtllm_generation.py +++ b/nemo_rl/models/generation/trtllm/trtllm_generation.py @@ -261,16 +261,23 @@ def _get_tied_worker_bundle_indices( ) return tied_groups - def _report_dp_openai_server_base_urls(self) -> list[Optional[str]]: + def _report_dp_openai_server_base_urls(self) -> list[str]: """Collect HTTP server base URLs from each DP-rank-0 worker.""" if not self.cfg["trtllm_cfg"].get("expose_http_server"): - urls = [cast(Optional[str], None)] * self.dp_size - return urls + # Empty rather than one None per DP rank: with no server running there is + # no per-rank slot to hold open, and the length said otherwise. + return [] futures = self.worker_group.run_all_workers_single_data( "report_dp_openai_server_base_url", run_rank_0_only_axes=["tensor_parallel"], ) - return ray.get(futures) + results = ray.get(futures) + # All-or-nothing, like the guarded branch above: expose_http_server is one + # config value for every worker. Collapsing the whole list rather than + # filtering keeps one entry per DP rank for anything that reads it by index. + if not any(results): + return [] + return results def _report_device_id(self) -> list[list[str]]: futures = self.worker_group.run_all_workers_single_data( diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 382ba1dded5..b3d16cc1bd8 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -524,10 +524,10 @@ def _report_device_id(self) -> list[list[str]]: results = ray.get(futures) return results - def _report_dp_openai_server_base_urls(self) -> list[Optional[str]]: + def _report_dp_openai_server_base_urls(self) -> list[str]: """Report the data parallel OpenAI server base URLs of vLLM workers, only populated if it is async vLLM engine and the HTTP server is active.""" if not self.cfg["vllm_cfg"]["async_engine"]: - return [None] # Not applicable since this is sync + return [] # Not applicable since this is sync # Use run_all_workers_single_data for methods that don't need data futures = self.worker_group.run_all_workers_single_data( @@ -536,22 +536,34 @@ def _report_dp_openai_server_base_urls(self) -> list[Optional[str]]: ) # Wait for all futures to complete results = ray.get(futures) + # A worker leaves base_url unset unless expose_http_server started a server, so + # an async engine without one answers with all Nones -- report that as no URLs, + # for the same reason the sync branch returns empty. Collapse the whole list + # rather than filtering entry by entry: expose_http_server is one config value + # for every worker, so the answer is all-or-nothing, and GenerationFleetHealth + # indexes this list by shard and length-checks it against shard_count. + if not any(results): + return [] return results - def _collect_reserved_urls(self) -> list[Optional[str]]: + def _collect_reserved_urls(self) -> list[str]: """Collect reserved URLs from DP leaders before model loading. Only called when defer_model_load=True. Workers have bound ports during __init__ and can report their reserved URLs immediately. """ if not self.cfg["vllm_cfg"]["async_engine"]: - return [None] + return [] futures = self.worker_group.run_all_workers_single_data( "get_reserved_url", run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], ) results = ray.get(futures) + # Only expose_http_server reserves a socket, so without one every worker + # reports None here too. Collapsed as a whole, per the note above. + if not any(results): + return [] return results def load_and_start(self) -> None: diff --git a/tests/unit/models/generation/test_dynamo_generation.py b/tests/unit/models/generation/test_dynamo_generation.py index 7a53e212f89..c64b3c97ab2 100644 --- a/tests/unit/models/generation/test_dynamo_generation.py +++ b/tests/unit/models/generation/test_dynamo_generation.py @@ -164,7 +164,8 @@ def test_runtime_start_world_size_sender_geometry_and_shutdown(monkeypatch) -> N assert calls[:2] == ["init", "start"] assert generation.frontend_url == "http://10.0.0.1:3000/v1" - assert generation.dp_openai_server_base_urls == [None] + # No token wrapper started, so there is no OpenAI-compatible server to name. + assert generation.dp_openai_server_base_urls == [] assert generation.get_inference_world_size() == 4 sender = generation.get_collective_sender_spec() assert sender.nccl_peer == "vllm" diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 4429f26446e..a1fc6c3cce0 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -233,6 +233,66 @@ def test_vllm_generation_rejects_partial_refit_pause_and_resume( generation.resume_generation_after_refit() +def _url_reporting_generation(*, async_engine: bool) -> VllmGeneration: + generation = VllmGeneration.__new__(VllmGeneration) + generation.cfg = {"vllm_cfg": {"async_engine": async_engine}} + generation.worker_group = MagicMock() + return generation + + +def test_sync_engine_reports_no_openai_server_urls() -> None: + """A sync engine runs no HTTP server, so there is no URL to report.""" + generation = _url_reporting_generation(async_engine=False) + + assert generation._report_dp_openai_server_base_urls() == [] + assert generation._collect_reserved_urls() == [] + + +def test_async_engine_without_http_server_reports_no_urls( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Workers answer None unless expose_http_server started a server. + + The list used to keep those Nones, which made a plain truth test read as + "we have URLs" when every entry stood for the absence of one. + """ + generation = _url_reporting_generation(async_engine=True) + monkeypatch.setattr(ray, "get", MagicMock(return_value=[None, None])) + + assert generation._report_dp_openai_server_base_urls() == [] + assert generation._collect_reserved_urls() == [] + + +def test_async_engine_with_http_server_reports_served_urls( + monkeypatch: pytest.MonkeyPatch, +) -> None: + generation = _url_reporting_generation(async_engine=True) + served = ["http://10.0.0.1:8000/v1", "http://10.0.0.2:8000/v1"] + monkeypatch.setattr(ray, "get", MagicMock(return_value=list(served))) + + assert generation._report_dp_openai_server_base_urls() == served + assert generation._collect_reserved_urls() == served + + +def test_partially_reported_urls_keep_one_entry_per_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Report emptiness as a whole list, never by dropping individual entries. + + ``expose_http_server`` is one config value for every worker, so a mixed answer + is not reachable today. Asserted anyway because ``GenerationFleetHealth`` + indexes these URLs by shard and rejects a list whose length does not match + ``shard_count``: shortening one would turn a degraded shard into a startup + failure with a length error that names neither cause. + """ + generation = _url_reporting_generation(async_engine=True) + partial = ["http://10.0.0.1:8000/v1", None] + monkeypatch.setattr(ray, "get", MagicMock(return_value=list(partial))) + + assert generation._report_dp_openai_server_base_urls() == partial + assert generation._collect_reserved_urls() == partial + + def test_sampling_params_preserve_bad_words(): worker = object.__new__(VllmGenerationWorkerImpl) worker.cfg = { diff --git a/tests/unit/models/generation/trtllm/test_trtllm_generation.py b/tests/unit/models/generation/trtllm/test_trtllm_generation.py index 1123ebae58a..a91ef3ce0b9 100644 --- a/tests/unit/models/generation/trtllm/test_trtllm_generation.py +++ b/tests/unit/models/generation/trtllm/test_trtllm_generation.py @@ -170,6 +170,27 @@ def test_cross_node_tp_replicas_use_unified_placement_group(): ] +def test_no_http_server_reports_an_empty_url_list(): + """Without a server there is no per-rank slot, so the length said nothing. + + This used to be one None per DP rank, which is truthy and carries a length + that reads like real URLs were collected. + """ + generation = _bare_generation(dp_size=2) + generation.dp_size = 2 + + assert generation._report_dp_openai_server_base_urls() == [] + + +def test_exposed_http_server_reports_served_addresses(monkeypatch): + generation = _bare_generation(dp_size=2, expose_http_server=True) + generation.dp_size = 2 + served = ["http://10.0.0.1:8000/v1", "http://10.0.0.2:8000/v1"] + monkeypatch.setattr(trtllm_generation.ray, "get", lambda futures: list(served)) + + assert generation._report_dp_openai_server_base_urls() == served + + @pytest.mark.asyncio async def test_generate_async_dispatches_round_robin_and_returns_leader_index(): generation = _bare_generation(dp_size=2)