[skyrl-train][inference] Inference Server Refactor (1/N)

[kouroshHakha]

Why

We're building HTTP-based inference serving for RL workloads. This enables:

• Decoupled training ↔ inference via standard HTTP (no Ray object refs)
• Flexible backends (vLLM now, SGLang later via protocol)
• Dynamic weight sync between trainer and inference servers

What

skyrl_train/inference_servers/
├── router.py             # HTTP proxy: load balancing + control plane fan-out
├── server_group.py       # Manages N vLLM actors on a placement group
├── vllm_server_actor.py  # Ray actor wrapping vLLM OpenAI server
├── protocols.py          # ServerActorProtocol interface
├── vllm_worker.py        # Worker extension for NCCL weight sync
└── common.py             # Utilities (get_open_port, get_node_ip)

Key Design Decisions

Decision	Rationale
ServerActorProtocol	Swap vLLM for SGLang without changing ServerGroup
Router as HTTP proxy	Compatible with any HTTP client; no vLLM-specific deps
Control plane fan-out	/pause, /resume, /init_weight_transfer hit all servers
Session-aware hashing	Sticky routing for multi-turn conversations

Review Guide

1. Start here: protocols.py (30 lines) - the interface contract
2. Core logic: router.py - proxy routing + fan-out behavior
3. Glue code: server_group.py + vllm_server_actor.py - Ray actor lifecycle
4. Tests: test_weight_sync.py - trainer→inference weight sync flow

Testing

# GPU CI (requires 4 GPUs)
uv run pytest tests/gpu/gpu_ci/test_inference_server_group.py -v
uv run pytest tests/gpu/gpu_ci/test_weight_sync.py -v

Next

RemoteInferenceClient - typed client for trainer to call inference servers

[gemini-code-assist]

Code Review

This pull request introduces a significant and well-designed refactoring for inference serving. The new architecture, with its clear separation of concerns into a router, server group, actor pool, and a server actor protocol, is robust and extensible. The implementation of the vLLM server actor and the comprehensive test suite, including complex scenarios like weight synchronization, are particularly impressive. I have a few suggestions to improve the router's lifecycle management and fix a test case.

[kouroshHakha]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant refactor of the inference server architecture, moving towards a more decoupled and flexible design. Key components include a new InferenceRouter for load balancing and control plane fan-out, ServerGroup for managing server actors with placement groups, and a ServerActorProtocol to enable interchangeable backends like vLLM and SGLang. The changes also include a robust weight synchronization mechanism between trainers and inference servers. New utility functions for environment variables and port management are added, along with comprehensive unit and GPU integration tests. The overall structure is well-thought-out, enhancing maintainability and scalability for RL workloads.

[CharlieFRuan]

Thank you so much for the PR and pushing the refactor of the rollout stack. Learned a lot when reviewing, and cannot wait to see it land!

[CharlieFRuan]

The current skyrl-train/skyrl_train/inference_engines/inference_engine_client_http_endpoint.py allows the users to, in their agent harness, instead using a say OpenAI endpoint, to use the HTTP endpoint that SkyRL spins up.

With this refactoring, the user will use http://0.0.0.0/8080/v1/chat/completions for the router router = InferenceRouter(server_urls, host="0.0.0.0", port=8080).

My question is, for the sticky routing, how will the user add the header? Would it be simpler if the session ID is in the body?

Currently for terminal bench, it will have a session_id field in the request:

• SkyRL/skyrl-train/examples/terminal_bench/generator/terminal_bench_generator.py

  Line 105 in b9a6307

  "session_id": session_id,

  • TBench side will forward this kwarg in the request body
• Then on SkyRL, upon receiving this request, pops the session_id here:

  SkyRL/skyrl-train/skyrl_train/inference_engines/inference_engine_client.py

  Lines 367 to 377 in b9a6307

  	async def chat_completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]:
  	session_id = request_payload["json"].pop("session_id", None)
  	if session_id is None:
  	# if session_id is not provided, we'll use a random engine
  	engine_idx = random.randint(0, len(self.engines) - 1)
  	else:
  	assert isinstance(session_id, (str, int)), "Session ID must be an integer or string for `/chat/completions`"
  	engine_idx = hash_with_sha256(str(session_id)) % len(self.engines)
  	# Always use the retry loop which also issues the first request inside
  	return await self._chat_completion_with_retry(engine_idx, request_payload)

[kouroshHakha]

so we need to make it compatible with industry standards. if you look at sgl router, vllm-router, etc. the dominant way to pass in the session id is through passing a request header that has X-Session-ID. I don't know what the benefit of this is over keeping it in the body, but I think this is more of a standard for decoupling middle-ware / routing ops from the request body.

[kouroshHakha]

session_id is not part of chat or completion api standard btw. So I'd say we need to change the agent harness / generator to the header instead.

[CharlieFRuan]

if you look at sgl router, vllm-router, etc. the dominant way to pass in the session id is through passing a request header that has X-Session-ID

Thanks for the info, didn't know before! Let's keep it as X-Session-ID for now! We just need to remember properly telling users about how to do sticky routing somewhere. Made an issue here: #939

[CharlieFRuan]

if you look at sgl router, vllm-router, etc. the dominant way to pass in the session id is through passing a request header that has X-Session-ID

Thanks for the info, didn't know before! Let's keep it as X-Session-ID for now! We just need to remember properly telling users about how to do sticky routing somewhere. Made an issue here: #939

[CharlieFRuan]

Thanks for addressing the comments! Last 3 nits!

[CharlieFRuan]

What is the semantic of _proxy_to_all() for /health and /is_paused? If one of the engines returns false, does the router return false altogether? Could we add a unit test for/health when one of the engine is not healthy?

I feel like the semantic of /health and /is_paused (GET methods) are different from the other control plane APIs (POST) for _proxy_to_all().

[kouroshHakha]

good q.

/is_paused --> routers returns a mapping from server_urls to is_paused state of each engine (same as any other control plane ops)

/health --> I think we should actually make it not a control plane endpoint. Basically if at least one engine exists that is healthy it should return true, By making /health go through the routing, if it returns true you are guaranteed that at least one engine is healhty. The router should smart enough to not route to the engines that are unhealthy (that's part of the responsibility of the router that we can add to our simple one if we want later)

I removed /health from the control plane endpoints because of above.

[CharlieFRuan]

sounds good, thanks!

routers returns a mapping from server_urls to is_paused state of each engine (same as any other control plane ops)

Got it. Would be good to document this behavior somewhere! We can do that in another PR

[CharlieFRuan]

On /health, I made an issue here: #958

[vercel]

@kouroshHakha is attempting to deploy a commit to the Tyler's projects Team on Vercel.

A member of the Team first needs to authorize it.

[CharlieFRuan]

LGTM! Thank you!