refactor(inference): extract reusable RemoteInferenceGenerator

[dyurk-lila]

Note on the diff: This PR is part of a routed-expert-replay / sampler-support series and builds on the PRs below. GitHub can't show the intermediate branches here, so the diff is cumulative on top of main — the changes new to this PR sit on top of:

• perf(r3): accelerate routed-expert transport with packed arrays #1909 — perf(r3) packed-array transport
• perf(r3): cast expert indices to int32 for pp local layers only #1910 — perf(r3) expand routes to local layers only

Reviewing in PR order (lowest number first) shows each incremental change cleanly.

Summary

Pure refactor (no functionality change) that extracts the single-request HTTP
generation path out of RemoteInferenceClient so it can be reused on its own:

• Extract the raw-token generation path into a standalone
  RemoteInferenceGenerator and a RemoteGenerateResult dataclass.
• Have RemoteInferenceClient own an internal RemoteInferenceGenerator and
  delegate session management, _post, and _generate_single to it.
• Expose optional routed-expert results without requiring callers to construct
  the full inference/control-plane client.
• Preserve existing endpoint routing, retry/backoff, cache_salt handling,
  serialization, and lifecycle behavior.

Testing

• uv run --isolated --extra dev --extra fsdp pytest tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py
  (58 passed)
• ruff and black clean on the changed files.

[gemini-code-assist]

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

[SumanthRH]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the MoE router replay mechanism and optimizes inference server communication by introducing a dedicated RemoteInferenceGenerator and a compact base64-encoded payload format for routed expert indices. It also adds a router_padding_mask to training batches and refactors token metadata alignment to correctly handle padding and context-parallel sharding. The reviewer feedback highlights critical improvements: ensuring the _generator field is properly cleared and reset during serialization/deserialization in RemoteInferenceClient, falling back to a default value instead of raising a ValueError for non-finite logprobs in vllm_server_actor.py, and handling list-type unwrapped models when virtual pipeline parallelism is enabled in replay_utils.py to prevent runtime crashes.

[erictang000]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the remote inference client by extracting the HTTP request and token generation logic into a new RemoteInferenceGenerator class and introducing a RemoteGenerateResult dataclass. The RemoteInferenceClient now delegates its generation tasks to this new generator. Feedback is provided regarding the loss of a protective try...except block around session.close() during the refactoring of the session closure logic, which could lead to unhandled exceptions during cleanup.

[gemini-code-assist]

The protective try...except block around session.close() was lost during the refactoring into RemoteInferenceGenerator.aclose(). Since aclose() is called during cleanup and teardown paths (such as in __aexit__), any exception raised during session closure could propagate and disrupt the cleanup process or mask other errors. We should restore the try...except block to safely handle and log any exceptions during session closure.

Suggested change

	async def aclose(self) -> None:
	if self._session is not None and not self._session.closed:
	await self._session.close()
	self._session = None
	async def aclose(self) -> None:
	if self._session is not None and not self._session.closed:
	try:
	await self._session.close()
	except Exception as exc:
	logger.warning(f"Encountered exception {exc} while closing client session")
	self._session = None