[5/7][multi-lora]: rollout request routing - per sample lora_path/rid, per-adapter rewards, prefill logprob grouping

miles/rollout/generate_utils/prefill_logprobs.py

@@ -7,10 +7,20 @@
 from miles.rollout.generate_utils.generate_endpoint_utils import compute_routing_headers, policy_uses_routing_key
 from miles.utils.http_utils import post
 from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled
+from miles.utils.multi_lora import slot_lora_name
 from miles.utils.processing_utils import encode_image_for_rollout_engine
 from miles.utils.types import Sample
 
 
+def _lora_path_for_sample(args: Any, sample: Sample) -> str | None:
+    """Adapter name to score under: the sample slot's __miles_slot_{N} name, the fixed single-LoRA name, or None."""
+    if sample.adapter is not None:
+        return slot_lora_name(sample.adapter.slot)
+    if is_lora_enabled(args):
+        return LORA_ADAPTER_NAME
+    return None
+
+
 def _build_prefill_scoring_payload(
     args: Any,
     sample: Sample,
@@ -38,8 +48,8 @@ def _build_prefill_scoring_payload(
         "logprob_start_len": prompt_len - 1,
     }
 
-    if is_lora_enabled(args):
-        payload["lora_path"] = LORA_ADAPTER_NAME
+    if (lora_path := _lora_path_for_sample(args, sample)) is not None:
+        payload["lora_path"] = lora_path
 
     if sample.multimodal_inputs and sample.multimodal_inputs.get("images"):
         image_data = sample.multimodal_inputs["images"]
@@ -64,14 +74,19 @@ def _build_batch_prefill_scoring_payload(
     if any(payload["logprob_start_len"] != logprob_start_len for payload in payloads):
         raise ValueError("Batched SGLang prefill scoring requires a shared logprob_start_len")
 
+    lora_paths = {payload.get("lora_path") for payload in payloads}
+    if len(lora_paths) > 1:
+        # A batch payload carries a single lora_path; callers must group samples by adapter first.
+        raise ValueError("Batched SGLang prefill scoring requires a shared lora_path")
+
     batch_payload: dict[str, Any] = {
         "input_ids": [payload["input_ids"] for payload in payloads],
         "sampling_params": payloads[0]["sampling_params"],
         "return_logprob": True,
         "logprob_start_len": logprob_start_len,
     }
-    if "lora_path" in payloads[0]:
-        batch_payload["lora_path"] = payloads[0]["lora_path"]
+    if (lora_path := next(iter(lora_paths))) is not None:
+        batch_payload["lora_path"] = lora_path
     return batch_payload
 
 
@@ -138,12 +153,13 @@ async def recompute_samples_rollout_logprobs_via_prefill(
     flush_url = url.rsplit("/", 1)[0] + "/flush_cache"
 
     if _can_batch_prefill_score(args, samples_to_score):
-        samples_by_logprob_start_len: dict[int, list[Sample]] = defaultdict(list)
+        # A batch must share one logprob_start_len and one lora_path, so group by both.
+        samples_by_batch_key: dict[tuple[int, str | None], list[Sample]] = defaultdict(list)
         for sample in samples_to_score:
             prompt_len = len(sample.tokens) - sample.response_length
-            samples_by_logprob_start_len[prompt_len - 1].append(sample)
+            samples_by_batch_key[(prompt_len - 1, _lora_path_for_sample(args, sample))].append(sample)
 
-        for batch_samples in samples_by_logprob_start_len.values():
+        for batch_samples in samples_by_batch_key.values():
             # SGLang can serve scoring requests from radix/KV cache. Flush before
             # each scoring group so every group uses the same clean-prefill path.
             await post(flush_url, {})

miles/rollout/rm_hub/__init__.py

@@ -5,6 +5,7 @@
 import aiohttp
 
 from miles.utils.misc import load_function
+from miles.utils.multi_lora import is_multi_lora_enabled
 from miles.utils.types import Sample
 
 from .deepscaler import get_deepscaler_rule_based_reward, get_gemma_math_reward
@@ -28,15 +29,27 @@ async def remote_rm(args, sample: Sample):
             return await resp.json()
 
 
+def _resolve_reward_config(args, sample: Sample) -> tuple[str | None, str]:
+    # Spec fields win when set; unset/empty fields fall back to sample metadata and process-wide args.
+    spec = sample.reward_spec
+    metadata = sample.metadata if isinstance(sample.metadata, dict) else {}
+    custom_rm_path = (spec.custom_rm_path if spec is not None else None) or getattr(args, "custom_rm_path", None)
+    rm_type = (
+        (spec.rm_type if spec is not None else None) or metadata.get("rm_type") or getattr(args, "rm_type", None) or ""
+    ).strip()
+    return custom_rm_path, rm_type
+
+
 async def async_rm(args, sample: Sample, **kwargs):
-    if args.custom_rm_path is not None:
-        rm_function = load_function(args.custom_rm_path)
+    custom_rm_path, rm_type = _resolve_reward_config(args, sample)
+
+    if custom_rm_path is not None:
+        rm_function = load_function(custom_rm_path)
         return await rm_function(args, sample, **kwargs)
 
-    metadata = sample.metadata if isinstance(sample.metadata, dict) else {}
-    rm_type = (metadata.get("rm_type") or args.rm_type or "").strip()
     response = sample.response
     label = sample.label
+    metadata = sample.metadata if isinstance(sample.metadata, dict) else {}
     if rm_type.startswith("boxed_"):
         response = extract_boxed_answer(response) or ""
         rm_type = rm_type[len("boxed_") :]
@@ -88,8 +101,7 @@ async def batched_async_rm(
             sample.reward = reward
         return None
 
-    if args.custom_rm_path is not None:
-        # Ensure the custom reward function is implemented in batch mode
+    if args.custom_rm_path is not None and not is_multi_lora_enabled(args):
         rm_function = load_function(args.custom_rm_path)
         return await rm_function(args, samples, **kwargs)
     tasks = [async_rm(args, sample, **kwargs) for sample in samples]

miles/rollout/sglang_rollout.py

@@ -23,6 +23,7 @@
 from miles.utils.http_utils import get, post
 from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled
 from miles.utils.misc import SingletonMeta, call_agent_abort_hook, load_function
+from miles.utils.multi_lora import make_rid, slot_lora_name
 from miles.utils.processing_utils import (
     call_processor,
     encode_image_for_rollout_engine,
@@ -176,7 +177,22 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
     if getattr(args, "use_opd", False) and opd_top_k > 0 and opd_top_k_strategy != "only-teacher":
         payload["top_logprobs_num"] = opd_top_k
 
-    if is_lora_enabled(args):
+    if sample.adapter is not None:
+        from miles.ray.multi_lora.controller import AdaptersCache
+
+        if (adapter := await AdaptersCache().get(sample.adapter.name)) is None:
+            # Adapter deregistered: don't POST, or an orphan the abort round can't see
+            # would keep decoding under the slot's next tenant and pollute its group.
+            logger.warning(
+                f"Dropping generation for adapter '{sample.adapter.name}' (slot {sample.adapter.slot}): "
+                "adapter is no longer sampleable"
+            )
+            sample.status = Sample.Status.ABORTED
+            return sample
+        payload["lora_path"] = slot_lora_name(sample.adapter.slot)
+        payload["rid"] = make_rid(sample.adapter.name)
+        payload["extra_key"] = f"{sample.adapter.name}:v{adapter.version}"
+    elif is_lora_enabled(args):
         payload["lora_path"] = LORA_ADAPTER_NAME
 
     if args.use_rollout_routing_replay:

tests/fast/rollout/generate_utils/test_prefill_logprobs.py

@@ -3,7 +3,7 @@
 import pytest
 
 from miles.rollout.generate_utils import prefill_logprobs
-from miles.utils.types import Sample
+from miles.utils.types import AdapterRef, Sample
 
 
 @pytest.mark.asyncio
@@ -113,6 +113,105 @@ async def fake_post(url, payload, action="post", headers=None):
     assert calls[1][1]["logprob_start_len"] == 1
 
 
+@pytest.mark.asyncio
+async def test_recompute_uses_per_sample_adapter_lora_path(monkeypatch):
+    """Multi-LoRA: the scoring request must go to the sample's own slot adapter,
+    not the single-adapter name (which is never registered on those engines)."""
+    sample = Sample(
+        tokens=[10, 11, 20],
+        response_length=1,
+        status=Sample.Status.COMPLETED,
+        adapter=AdapterRef(name="run-a", slot=3),
+    )
+    # Multi-LoRA forces lora_rank > 0, so is_lora_enabled(args) is always true.
+    args = SimpleNamespace(recompute_logprobs_via_prefill=True, lora_rank=8)
+    seen = {}
+
+    async def fake_post(url, payload, headers=None):
+        seen["payload"] = payload
+        return {"meta_info": {"input_token_logprobs": [(None, 11), (-0.5, 20)]}}
+
+    monkeypatch.setattr(prefill_logprobs, "post", fake_post)
+
+    await prefill_logprobs.recompute_rollout_logprobs_via_prefill(
+        args,
+        sample,
+        url="http://localhost/generate",
+        sampling_params={},
+    )
+
+    assert seen["payload"]["lora_path"] == "__miles_slot_3"
+    assert sample.rollout_log_probs == [-0.5]
+
+
+@pytest.mark.asyncio
+async def test_recompute_samples_batches_group_by_adapter(monkeypatch):
+    """Samples from different adapters must not share a batch payload: each
+    batch carries a single lora_path."""
+    samples = [
+        Sample(
+            tokens=[10, 11, 20],
+            response_length=1,
+            status=Sample.Status.COMPLETED,
+            adapter=AdapterRef(name="run-a", slot=0),
+        ),
+        Sample(
+            tokens=[10, 11, 21],
+            response_length=1,
+            status=Sample.Status.COMPLETED,
+            adapter=AdapterRef(name="run-b", slot=1),
+        ),
+        Sample(
+            tokens=[10, 11, 22],
+            response_length=1,
+            status=Sample.Status.COMPLETED,
+            adapter=AdapterRef(name="run-a", slot=0),
+        ),
+    ]
+    args = SimpleNamespace(
+        recompute_logprobs_via_prefill=True,
+        lora_rank=8,
+        sglang_router_policy="round_robin",
+    )
+    generate_payloads = []
+
+    async def fake_post(url, payload, action="post", headers=None):
+        if url.endswith("/flush_cache"):
+            return {}
+        generate_payloads.append(payload)
+        return [
+            {"meta_info": {"input_token_logprobs": [(None, 11), (-float(tokens[-1]), tokens[-1])]}}
+            for tokens in payload["input_ids"]
+        ]
+
+    monkeypatch.setattr(prefill_logprobs, "post", fake_post)
+
+    await prefill_logprobs.recompute_samples_rollout_logprobs_via_prefill(
+        args,
+        samples,
+        url="http://localhost/generate",
+        sampling_params={"max_new_tokens": 32},
+    )
+
+    assert [sample.rollout_log_probs for sample in samples] == [[-20.0], [-21.0], [-22.0]]
+    by_lora_path = {payload["lora_path"]: payload["input_ids"] for payload in generate_payloads}
+    assert by_lora_path == {
+        "__miles_slot_0": [[10, 11, 20], [10, 11, 22]],
+        "__miles_slot_1": [[10, 11, 21]],
+    }
+
+
+def test_batch_payload_rejects_mixed_lora_paths():
+    samples = [
+        Sample(tokens=[10, 11, 20], response_length=1, adapter=AdapterRef(name="run-a", slot=0)),
+        Sample(tokens=[10, 11, 21], response_length=1, adapter=AdapterRef(name="run-b", slot=1)),
+    ]
+    args = SimpleNamespace(lora_rank=8)
+
+    with pytest.raises(ValueError, match="shared lora_path"):
+        prefill_logprobs._build_batch_prefill_scoring_payload(args, samples, {})
+
+
 @pytest.mark.asyncio
 async def test_recompute_samples_batches_by_logprob_start_len(monkeypatch):
     samples = [

tests/fast/rollout/rm_hub/test_reward_config_resolution.py

@@ -0,0 +1,39 @@
+"""_resolve_reward_config precedence: reward_spec fields win when set, and
+unset fields fall through to sample metadata and process-wide args. A
+reward_spec with empty fields (an adapter config without its own rm_type)
+must not disable the fallbacks."""
+
+from types import SimpleNamespace
+
+from miles.rollout.rm_hub import _resolve_reward_config
+from miles.utils.types import RewardSpec, Sample
+
+
+def _args(**kwargs) -> SimpleNamespace:
+    return SimpleNamespace(rm_type=None, custom_rm_path=None, **{k: v for k, v in kwargs.items()})
+
+
+def test_spec_fields_win_over_args():
+    sample = Sample(prompt="p", reward_spec=RewardSpec(rm_type="math", custom_rm_path="pkg.fn"))
+    args = SimpleNamespace(rm_type="deepscaler", custom_rm_path="other.fn")
+    assert _resolve_reward_config(args, sample) == ("pkg.fn", "math")
+
+
+def test_empty_spec_falls_back_to_args():
+    sample = Sample(prompt="p", reward_spec=RewardSpec(rm_type=None, custom_rm_path=None))
+    args = SimpleNamespace(rm_type="math", custom_rm_path=None)
+    assert _resolve_reward_config(args, sample) == (None, "math")
+
+
+def test_empty_spec_falls_back_to_sample_metadata_before_args():
+    sample = Sample(prompt="p", reward_spec=RewardSpec(rm_type=None, custom_rm_path=None))
+    sample.metadata = {"rm_type": "gpqa"}
+    args = SimpleNamespace(rm_type="math", custom_rm_path=None)
+    assert _resolve_reward_config(args, sample) == (None, "gpqa")
+
+
+def test_no_spec_keeps_metadata_then_args_chain():
+    sample = Sample(prompt="p")
+    sample.metadata = {}
+    args = SimpleNamespace(rm_type="math", custom_rm_path=None)
+    assert _resolve_reward_config(args, sample) == (None, "math")