Keep prewarmed inference alive after model unload

[j316chuck]

Summary

• Start base inference during Tinker service startup when --prewarm-inference is enabled.
• On temporary model unload, stop only that model's training workers while retaining the service-owned vLLM/router runtime.
• Preserve existing full teardown behavior for services that were not prewarmed.

Testing

python -m py_compile skyrl/backends/skyrl_train_backend.py skyrl/backends/skyrl_train/workers/worker.py skyrl/backends/skyrl_train/workers/worker_dispatch.py
git diff --check HEAD^ HEAD

Both passed. The focused pytest requires the SkyRL CUDA environment, which is not installed in this worktree.

[gemini-code-assist]

Code Review

This pull request introduces a prewarm_inference configuration option to start base-model inference engines before the first model is created, allowing short-lived adapters to be loaded and unloaded without tearing down the shared inference runtime. While the implementation is well-tested, the review highlights three critical issues: first, resetting self._colocate_pg to None during model deletion when keeping inference warm loses the active placement group reference; second, self._colocate_pg is unconditionally overwritten during model creation, which breaks colocation; and third, calling prewarm_inference on the backend without a defensive check could cause an AttributeError on backends that do not support it.

[gemini-code-assist]

In delete_model, setting self._colocate_pg = None when self._keep_inference_warm is True causes the reference to the active placement group to be lost.

When a new model is subsequently created via create_model, it will unconditionally call self._create_colocate_pg() (at line 502), creating a new placement group. This breaks colocation because the prewarmed inference engines (running on the old placement group) and the new training workers (running on the new placement group) will be on different placement groups. This can also lead to resource allocation failures since the old placement group still holds the GPUs.

To fix this, we should preserve self._colocate_pg here, and update create_model to only initialize self._colocate_pg if it is currently None.

Suggested change

	self._cfg = None
	self._dispatch = None
	self._renderer = None
	self._colocate_pg = None
	self._base_lora_signature = None
	self._cfg = None
	self._dispatch = None
	self._renderer = None
	self._base_lora_signature = None

[gemini-code-assist]

When prewarm_inference is enabled, self._colocate_pg is already initialized during prewarming. However, in create_model (at line 502), self._colocate_pg is unconditionally overwritten:

self._colocate_pg = self._create_colocate_pg() if self._cfg.trainer.placement.colocate_all else None

This causes a new placement group to be created even on the very first model creation after prewarming, which breaks colocation immediately because the prewarmed inference engines and the newly created policy workers will run on different placement groups.

Please update line 502 in create_model to only create the placement group if it is not already set:

if self._colocate_pg is None and self._cfg.trainer.placement.colocate_all:
    self._colocate_pg = self._create_colocate_pg()

[gemini-code-assist]

If config.prewarm_inference is enabled but the selected backend (such as the jax backend) does not implement prewarm_inference, this will raise an AttributeError and crash the engine on startup. Adding a defensive check using hasattr prevents this crash and logs a warning instead.

Suggested change

	if config.prewarm_inference:
	self.backend.prewarm_inference()
	if config.prewarm_inference:
	if hasattr(self.backend, "prewarm_inference"):
	self.backend.prewarm_inference()
	else:
	logger.warning(f"Backend {type(self.backend).__name__} does not support prewarm_inference.")