[lora][tinker] Add pause and resume for multi-tenant lora (#1657)

Addresses #1647

Today's `pause_generation()` is a **global** vLLM keep-mode pause: when
one LoRA tenant syncs new weights, every other tenant's in-flight
generation freezes for the duration of the swap. This blocks practical
multi-tenant LoRA RL training.

This PR adds a per-LoRA `lora_name` arg to pause/resume so a weight sync
for adapter A only aborts A's requests; other adapters keep generating.
The aborted requests come back with `finish_reason="abort"` and partial
tokens, and a new `sample_with_retry()` client method accumulates those
partial tokens, awaits resume, and resubmits with `prompt + accumulated`
and remaining `max_tokens` until completion.

**This is a transient fix.** Hopefully in the future if we can upstream
a lora specific pause, we delete `sample_with_retry`, the per-LoRA gate,
and the abort endpoint; the `lora_name` kwarg stays and routes to the
new vLLM API.

## What changed

### New: `/skyrl/v1/abort_lora_requests` server endpoint


[vllm_server_actor.py](skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py)
gets a small custom endpoint that iterates
`engine.output_processor.request_states`, filters by `lora_name`, and
calls `engine.abort(ids, internal=True)`. `internal=True` is
load-bearing — the states dict is keyed by internal IDs.

### `pause_generation`/`resume_generation` gain `lora_name:
Optional[str] = None`

- `lora_name=None` → unchanged: vLLM `/pause?mode=keep` global pause.
- `lora_name="X"` → clears a per-LoRA `asyncio.Event` (gates retries
client-side), sleeps a 5 s grace, then fans out to
`/skyrl/v1/abort_lora_requests`.

The new inference path (`_SKYRL_USE_NEW_INFERENCE=1`, the default) goes
through `RemoteInferenceClient` → `vllm_server_actor`, which is where
the new `/skyrl/v1/abort_lora_requests` endpoint lives — so that's the
only path that actually supports targeted pause. Legacy-path classes
(`InferenceEngineClient`, `RayWrappedInferenceEngine`,
`RemoteInferenceEngine`, `AsyncVLLMInferenceEngine`) accept the
`lora_name` kwarg as a guardrail but raise
`NotImplementedError("targeted pause is HTTP-only")` when it's non-None.
The legacy path keeps its existing global-pause behavior unchanged.

### `RemoteInferenceClient`: `sample()` refactor + new
`sample_with_retry()`

This is the only data-plane surface that gained retry logic.

**`sample()` is split**:

- `sample()` — renders the prompt and dispatches once. **Public behavior
unchanged.**
- `_sample_with_rendered_tokens()` — the post-render half, parameterized
on rendered `token_ids`. Pure refactor (regression-checked by the
existing `TestSample` tests).

**`sample_with_retry()` is new** — the only retry-bearing method.
Renders once, then runs a `while stop_reason == "abort"` loop:

1. `await _lora_pause_events[model].wait()` (no-op if no event exists).
2. Build a body with `token_ids = original_prompt + accum_tokens` and
`max_tokens = original_max_tokens - len(accum_tokens)`.
3. Dispatch `_sample_with_rendered_tokens`.
4. Extend `accum_tokens` + `accum_logprobs` with the returned segment.
5. Repeat until non-abort.

Returns the same `SampleResponse` shape as `sample()`. Asserts
`num_samples == 1` (current Tinker callers all do this; multi-sample
retry is straightforward but deferred).

`generate()`, `chat_completion()`, `completion()` are **untouched** — no
retry, no per-LoRA gating. The multi-tenant Tinker path is the only
production caller that flips to retry mode.

### `worker_dispatch.save_weights_for_sampler` threads `lora_name`


[worker_dispatch.py](skyrl/backends/skyrl_train/workers/worker_dispatch.py)
non-colocate branch forwards `model_id` (set by
`SkyRLTrainBackend.save_sampler_checkpoint` to the LoRA name for
multi-tenant, `None` for FFT) to the pause/resume calls.

## API impact

| API | Change | Migration |
|---|---|---|
| `InferenceEngineInterface.pause_generation` / `resume_generation` |
New optional `lora_name: Optional[str] = None` kwarg. | None — default
preserves existing behavior on every subclass. |
| `RemoteInferenceClient.sample()` | **Refactored** into render +
dispatch (pure refactor; public behavior unchanged). | None. |
| `RemoteInferenceClient.sample_with_retry()` | **New.** Same
signature/return as `sample()`. | Optional; auto-used by the multi-LoRA
path in `SkyRLTrainBackend`. |
| `generate()`, `chat_completion()`, `completion()` | **Unchanged.** No
retry, no per-LoRA gate. | None. |
| `worker_dispatch.save_weights_for_sampler` | Now forwards
`lora_name=model_id` to pause/resume on the non-colocate path. | None —
`model_id=None` for FFT preserves the global keep-mode pause. |

## Tests

### Unit (no GPU) —
`tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py`

7 new cases in `TestTargetedLoraPause`:

- `pause_generation(lora_name=X)` clears the event and fans out abort.
- `pause_generation()` (no lora_name) still drives keep-mode.
- `sample_with_retry` accumulates partial tokens across an abort,
`max_tokens` decrements correctly, logprobs concatenate, final response
shape OK.
- `sample_with_retry` no-abort path is a single shot (refactor
regression).
- `sample_with_retry` blocks on the per-LoRA event until
`resume_generation` is called.
- `sample_with_retry` rejects `num_samples > 1`.
- LoRAs that were never paused never block (no spurious event creation).

### GPU integration —
`tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_pause_lora.py`

4 end-to-end cases, run against a real vLLM server with two LoRA
adapters (Meow + Woof) loaded:

1. **`test_pause_lora_does_not_affect_other_lora`** — while lora-meow is
paused, 4 concurrent `sample_with_retry(model="lora-woof")` calls
complete promptly and contain "woof" content (proves the gate doesn't
spill across adapters and weights aren't mixed up).
2. **`test_sample_with_retry_recovers_from_abort`** — 4+4 concurrent
in-flight samples on meow + woof; pausing lora-meow mid-flight aborts
all 4 meow requests, retry resubmits after resume, all 8 complete with
non-abort stop reasons and correct content. **Hard-asserts** 0/8 tasks
completed pre-pause and 0/4 escaped during the pause window — without
these assertions the test would silently no-op if the LoRA emitted EOS
too fast.
3. **`test_pause_swap_weights_resume_mid_sample`** — single sample call
spans a real weight swap: starts with Meow weights, mid-flight calls
`pause → load_lora_adapter("lora-target", woof_path) → resume`, and the
merged output literally shows `meow meow meow ... woof woof woof`.
Proves the abort/retry boundary preserves accumulated state AND that the
retried request observes the newly-loaded weights.
4. **`test_global_pause_still_works`** — `pause_generation()` with no
`lora_name` still drives keep-mode pause (FFT regression).

---------

Signed-off-by: ahao-anyscale <ahao@anyscale.com>

Author: hao-aaron

skyrl/backends/skyrl_train/inference_engines/base.py

@@ -165,13 +165,22 @@ async def reset_prefix_cache(self):
         raise NotImplementedError
 
     @abstractmethod
-    async def pause_generation(self) -> None:
-        """Pause generation, freezing in-flight requests so they can be resumed later."""
+    async def pause_generation(self, lora_name: Optional[str] = None) -> None:
+        """Pause generation, freezing in-flight requests so they can be resumed later.
+
+        When ``lora_name`` is None (default), pauses all generation globally
+        (vLLM keep-mode pause). When ``lora_name`` is provided, only requests
+        targeting that specific LoRA adapter are paused (HTTP path only).
+        """
         raise NotImplementedError
 
     @abstractmethod
-    async def resume_generation(self) -> None:
-        """Resume generation after a pause, continuing any frozen in-flight requests."""
+    async def resume_generation(self, lora_name: Optional[str] = None) -> None:
+        """Resume generation after a pause, continuing any frozen in-flight requests.
+
+        ``lora_name`` must match the value used in the corresponding
+        ``pause_generation`` call.
+        """
         raise NotImplementedError
 
     @abstractmethod

skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py

@@ -366,21 +366,29 @@ def dp_size(self) -> int:
     # ----------------------------
     # Generation pause and resume
     # ----------------------------
-    async def pause_generation(self) -> None:
+    async def pause_generation(self, lora_name: Optional[str] = None) -> None:
         """
         Pauses generation for all engines using vLLM's native keep mode.
 
         In-flight requests are frozen (not aborted) and will resume from where they left off
         when `resume_generation()` is called. New requests are blocked until resume.
+
+        ``lora_name`` is accepted for interface parity with the HTTP path but
+        targeted (per-LoRA) pause is HTTP-only; passing a non-None value
+        raises ``NotImplementedError``.
         """
+        if lora_name is not None:
+            raise NotImplementedError("targeted pause is HTTP-only")
         await self._run_on_all_engines("pause_generation")
 
-    async def resume_generation(self) -> None:
+    async def resume_generation(self, lora_name: Optional[str] = None) -> None:
         """
         Resumes generation for all engines after a keep-mode pause.
 
         Frozen in-flight requests continue from where they left off, and new requests are unblocked.
         """
+        if lora_name is not None:
+            raise NotImplementedError("targeted pause is HTTP-only")
         await self._run_on_all_engines("resume_generation")
 
     # ----------------------------

skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Dict, List
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import ray
 from packaging import version
@@ -76,10 +76,14 @@ async def chat_completion(self, request_payload: Dict[str, Any]) -> Dict[str, An
     async def completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]:
         return await self.inference_engine_actor.completion.remote(request_payload)
 
-    async def pause_generation(self) -> None:
+    async def pause_generation(self, lora_name: Optional[str] = None) -> None:
+        if lora_name is not None:
+            raise NotImplementedError("targeted pause is HTTP-only")
         return await self.inference_engine_actor.pause_generation.remote()
 
-    async def resume_generation(self) -> None:
+    async def resume_generation(self, lora_name: Optional[str] = None) -> None:
+        if lora_name is not None:
+            raise NotImplementedError("targeted pause is HTTP-only")
         return await self.inference_engine_actor.resume_generation.remote()
 
 

skyrl/backends/skyrl_train/inference_engines/remote_inference_engine.py

@@ -282,8 +282,10 @@ async def reset_prefix_cache(self):
                 "body": text,
             }
 
-    async def pause_generation(self) -> None:
+    async def pause_generation(self, lora_name: Optional[str] = None) -> None:
         """Pause generation using vLLM's native keep mode, freezing in-flight requests."""
+        if lora_name is not None:
+            raise NotImplementedError("targeted pause is HTTP-only")
         async with aiohttp.ClientSession() as session:
             async with session.post(
                 f"{self.url}/pause",
@@ -293,8 +295,10 @@ async def pause_generation(self) -> None:
                 if resp.status != 200:
                     raise RuntimeError(f"Failed to pause generation: {result.get('error', result)}")
 
-    async def resume_generation(self) -> None:
+    async def resume_generation(self, lora_name: Optional[str] = None) -> None:
         """Resume generation after a keep-mode pause."""
+        if lora_name is not None:
+            raise NotImplementedError("targeted pause is HTTP-only")
         async with aiohttp.ClientSession() as session:
             async with session.post(f"{self.url}/resume") as resp:
                 result = await resp.json()

skyrl/backends/skyrl_train/inference_engines/vllm/vllm_engine.py

@@ -216,10 +216,10 @@ def reset_prefix_cache(self):
         """Reset the prefix cache. Subclasses override for async version."""
         return self.llm.llm_engine.reset_prefix_cache()
 
-    async def pause_generation(self, clear_cache: bool = False) -> None:
+    async def pause_generation(self, lora_name: Optional[str] = None, clear_cache: bool = False) -> None:
         raise NotImplementedError("pause_generation is only supported for AsyncVLLMInferenceEngine.")
 
-    async def resume_generation(self) -> None:
+    async def resume_generation(self, lora_name: Optional[str] = None) -> None:
         raise NotImplementedError("resume_generation is only supported for AsyncVLLMInferenceEngine.")
 
 
@@ -648,14 +648,18 @@ async def completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]:
         """
         return await self._handle_openai_request(request_payload, endpoint="/completions")
 
-    async def pause_generation(self, clear_cache: bool = False) -> None:
+    async def pause_generation(self, lora_name: Optional[str] = None, clear_cache: bool = False) -> None:
         """Pause generation using vLLM's native keep mode, freezing in-flight requests."""
+        if lora_name is not None:
+            raise NotImplementedError("targeted pause is HTTP-only")
         engine = self._get_engine()
         await engine.pause_generation(mode="keep", clear_cache=clear_cache)
         logger.info("pause_generation(mode='keep') finished")
 
-    async def resume_generation(self) -> None:
+    async def resume_generation(self, lora_name: Optional[str] = None) -> None:
         """Resume generation after a keep-mode pause."""
+        if lora_name is not None:
+            raise NotImplementedError("targeted pause is HTTP-only")
         engine = self._get_engine()
         await engine.resume_generation()
         logger.info("resume_generation() finished")