[Shr3yash]

What does this PR do ?

Return the vLLM prefix-cache reset result so invalidate_kv_cache() can report failure.

reset_prefix_cache and reset_prefix_cache_async dropped the bool from vLLM's reset_prefix_cache, so every worker result was None. invalidate_kv_cache then did all([]), which is True, and callers that raise or warn on failure could not fire. Both workers now return that bool, and the aggregator returns False when no worker reports a result.

Open PR #3929 removes one caller of this function but leaves invalidate_kv_cache and the worker returns unchanged. Open PR #3530 touches the same worker lines to pass reset_connector, so there may be a trivial conflict.

Issues

Fixes #4632

Usage

uv run pytest tests/unit/models/generation/test_vllm_generation.py::test_invalidate_kv_cache_requires_a_reported_success tests/unit/models/generation/test_vllm_generation.py::test_reset_prefix_cache_returns_engine_bool tests/unit/models/generation/test_vllm_generation.py::test_reset_prefix_cache_async_returns_engine_bool

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• CPU unit tests above pass, along with the existing tests in test_vllm_generation.py. Functional tests were not run.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[terrykong]

Thanks for keeping this small and adding tests. The fix is right and needed: today invalidate_kv_cache can only say "success", so every caller's failure branch is dead code. It is still needed after #3929 lands: #3929's fallback raise, the legacy loop's warning, and the sparse synchronizer all read this result, and so will the message #4018 plans to add.

One blocker before merging: with this change, the sparse weight synchronizer crashes runs that keep KV cache across in-flight refits (see the inline comment on vllm_generation.py#L1678). Please fix it in this PR, since this PR is what turns the crash on. Please also list in the PR description the callers whose result changes: VllmRemoteSparseWeightSynchronizer.sync_weights and non-colocated finish_generation (nobody reads its return), and note that #3929 moves the single-controller call into a fallback that raises on False rather than removing it.

Merge order: this merges cleanly with both #3929 and #4018 (checked with git merge-tree). Landing before or together with #4018 is better: #4018's drained-refit message only means something once vLLM's False reaches the caller.

Router: not affected. Async vLLM pauses on its own, and the router path never calls this.

Follow-up (not for this PR): TRT-LLM's reset_prefix_cache_async throws away the result of collective_rpc and always returns True. Same kind of bug as #4632; worth its own issue.

Generated by Claude Code

[terrykong]

nemo_rl/models/generation/vllm/vllm_generation.py#L1678

1 action item — please fix in this PR, since this PR turns the crash on.
TL;DR — VllmRemoteSparseWeightSynchronizer raises when invalidate_kv_cache() returns False, and vLLM always returns False while paused requests hold KV blocks. So async GRPO with sparse refit that keeps KV cache across in-flight refits now dies at its first refit.

Config (nothing rejects it; the design doc says async engines work): async GRPO, async vLLM engine, in_flight_weight_updates: true, recompute_kv_cache_after_weight_updates: false, policy.generation.refit_transport: vllm_zmq_sparse (or vllm_s3_sparse).

1. The collector pauses with clear_cache=False (trajectory_collector.py, vllm_worker_async.py), so paused requests keep their KV blocks (vLLM core.py).
2. sync_weights calls invalidate_kv_cache() before the transfer and raises on False (L133-L137).
3. vLLM refuses the reset: "Failed to reset prefix cache because some blocks (%d) are not freed yet" (block_pool.py).
4. This line now passes that False through: RuntimeError: vLLM KV cache invalidation failed before zmq weight update.

AI-1

In VllmRemoteSparseWeightSynchronizer.sync_weights, delete the invalidate_kv_cache() check and its raise. Add a one-line comment that cache policy belongs to finish_generation and recompute_kv_cache_after_weight_updates. Replace test_fails_before_transfer_when_kv_cache_invalidation_fails with a test that sync_weights does not call invalidate_kv_cache. (Not a suggestion block: both files are outside this diff.)

Context — no action. Please don't ban the config at setup instead: the doc says it works, and this reset only tidies the cache.

Why delete the check, and who hits this

• The sparse apply path never touches the KV/prefix cache: vllm_sparse_refit.py and vllm_sparse_delta.py have no cache use.
• No other vLLM transport resets inside its synchronizer. Sync GRPO resets in finish_generation (sparse requires non-colocated); async GRPO follows recompute_kv_cache_after_weight_updates. This one transport sets its own cache policy, and it contradicts false.
• No shipped recipe has this config (the only sparse recipe is sync GRPO) and no test covers it. But recompute_kv_cache_after_weight_updates: false is common in production configs, so adding sparse refit to one of them hits this. The single controller never builds this synchronizer.

[terrykong]

tests/unit/models/generation/test_vllm_generation.py#L668-L681

1 action item.
TL;DR — both worker tests only feed the engine False, so a worker that hard-codes return False still passes them.

A mutation run confirmed it: hard-coding return False in the worker keeps both tests green. Running each test with True and False catches it (also checked by running the mutant).

AI-1

Suggested change

	def test_reset_prefix_cache_returns_engine_bool(
	monkeypatch: pytest.MonkeyPatch,
	) -> None:
	worker = VllmGenerationWorkerImpl.__new__(VllmGenerationWorkerImpl)
	worker.cfg = {"vllm_cfg": {"async_engine": False}}
	worker.llm = MagicMock()
	worker.llm.llm_engine.reset_prefix_cache.return_value = False
	monkeypatch.setattr(
	"nemo_rl.models.generation.vllm.vllm_worker.torch.cuda.empty_cache",
	lambda: None,
	)
	assert worker.reset_prefix_cache() is False
	worker.llm.llm_engine.reset_prefix_cache.assert_called_once_with()
	@pytest.mark.parametrize("engine_result", [True, False])
	def test_reset_prefix_cache_returns_engine_bool(
	monkeypatch: pytest.MonkeyPatch, engine_result: bool
	) -> None:
	worker = VllmGenerationWorkerImpl.__new__(VllmGenerationWorkerImpl)
	worker.cfg = {"vllm_cfg": {"async_engine": False}}
	worker.llm = MagicMock()
	worker.llm.llm_engine.reset_prefix_cache.return_value = engine_result
	monkeypatch.setattr(
	"nemo_rl.models.generation.vllm.vllm_worker.torch.cuda.empty_cache",
	lambda: None,
	)
	assert worker.reset_prefix_cache() is engine_result
	worker.llm.llm_engine.reset_prefix_cache.assert_called_once_with()

Make the same change to the async test: parametrize engine_result, use AsyncMock(return_value=engine_result), and assert is engine_result.

[terrykong]

nemo_rl/models/generation/vllm/vllm_worker.py#L1544

1 action item.
TL;DR — reset_prefix_cache now returns vLLM's bool, but the docstring doesn't say that False means nothing was cleared.

vLLM returns False and skips the reset while any request still holds KV blocks (block_pool.py#L777-L784). A caller reading only the docstring can't tell that.

AI-1

Suggested change

	"""Reset the prefix cache of vLLM engine."""
	"""Reset the prefix cache of vLLM engine.
	Returns:
	vLLM's result: False when vLLM skipped the reset because a request
	still holds KV blocks, so nothing was cleared.
	"""

Make the same change to the docstring of reset_prefix_cache_async.

[Shr3yash]

Dropped the invalidate_kv_cache() check in VllmRemoteSparseWeightSynchronizer.sync_weights. Cache policy stays with finish_generation and recompute_kv_cache_after_weight_updates, and the config is still allowed at setup. The replacement test checks that sync_weights does not call invalidate_kv_cache. Both worker tests now run for True and False, and the reset docstrings say False means vLLM left the cache in place because a request still holds KV blocks.