Limit forward_backward queue drain size

examples/repro_forward_backward_queue_drain.py

@@ -0,0 +1,88 @@
+"""Reproduce the forward_backward queue shape that can create oversized batches.
+
+This script is intentionally synthetic: it does not allocate tensors or require a
+GPU. It models a client submitting many pending FORWARD_BACKWARD requests with
+uneven sequence lengths, then compares the single-drain behavior with the
+request-count limited behavior.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+REQUEST_COUNT = 47
+EXAMPLE_COUNT = 63
+MAX_SEQUENCE_LENGTH = 147_472
+TOTAL_INPUT_TOKENS = 5_153_654
+OBSERVED_CUDA_ALLOCATION_GIB = 40.02
+FORWARD_BACKWARD_MAX_REQUEST_COUNT = 4
+
+
+@dataclass(frozen=True)
+class ForwardBackwardRequest:
+    request_id: int
+    sequence_lengths: list[int]
+
+
+def build_repro_requests() -> list[ForwardBackwardRequest]:
+    other_token_budget = TOTAL_INPUT_TOKENS - MAX_SEQUENCE_LENGTH
+    other_lengths = [other_token_budget // (EXAMPLE_COUNT - 1)] * (EXAMPLE_COUNT - 1)
+    for index in range(other_token_budget % (EXAMPLE_COUNT - 1)):
+        other_lengths[index] += 1
+
+    example_lengths = [MAX_SEQUENCE_LENGTH, *other_lengths]
+    requests = []
+    cursor = 0
+    for request_id in range(1, REQUEST_COUNT + 1):
+        examples_in_request = 2 if request_id <= EXAMPLE_COUNT - REQUEST_COUNT else 1
+        requests.append(
+            ForwardBackwardRequest(
+                request_id=request_id,
+                sequence_lengths=example_lengths[cursor : cursor + examples_in_request],
+            )
+        )
+        cursor += examples_in_request
+    return requests
+
+
+def chunk_by_request_count(
+    requests: list[ForwardBackwardRequest], max_request_count: int
+) -> list[list[ForwardBackwardRequest]]:
+    return [requests[start : start + max_request_count] for start in range(0, len(requests), max_request_count)]
+
+
+def summarize_batch(batch: list[ForwardBackwardRequest]) -> tuple[int, int, int, int]:
+    lengths = [length for request in batch for length in request.sequence_lengths]
+    request_count = len(batch)
+    example_count = len(lengths)
+    max_sequence_length = max(lengths)
+    input_tokens = sum(lengths)
+    return request_count, example_count, max_sequence_length, input_tokens
+
+
+def print_batch(label: str, batch: list[ForwardBackwardRequest]) -> None:
+    request_count, example_count, max_sequence_length, input_tokens = summarize_batch(batch)
+    print(
+        f"{label}: requests={request_count}, examples={example_count}, "
+        f"max_sequence_length={max_sequence_length:,}, input_tokens={input_tokens:,}"
+    )
+
+
+def main() -> None:
+    requests = build_repro_requests()
+
+    print("Client submits pending FORWARD_BACKWARD requests with uneven sequence lengths.")
+    print_batch("single queue drain before limiting", requests)
+    print(f"observed failure symptom: CUDA OOM while trying to allocate {OBSERVED_CUDA_ALLOCATION_GIB:.2f} GiB")
+    print()
+
+    chunks = chunk_by_request_count(requests, FORWARD_BACKWARD_MAX_REQUEST_COUNT)
+    print(
+        "with forward_backward_max_request_count=" f"{FORWARD_BACKWARD_MAX_REQUEST_COUNT}: {len(chunks)} backend calls"
+    )
+    for index, chunk in enumerate(chunks, start=1):
+        print_batch(f"chunk {index:02d}", chunk)
+
+
+if __name__ == "__main__":
+    main()

skyrl/tinker/config.py

@@ -58,6 +58,15 @@ class EngineConfig(BaseModel):
         ),
         json_schema_extra={"argparse_type": lambda v: None if v == "None" else int(v)},
     )
+    forward_backward_max_request_count: int | None = Field(
+        default=None,
+        description=(
+            "Optional cap on how many pending forward_backward requests the engine "
+            "coalesces into one backend call. Default `None` preserves the current "
+            "unlimited batching behavior."
+        ),
+        json_schema_extra={"argparse_type": lambda v: None if v == "None" else int(v)},
+    )
     session_cleanup_interval_sec: int = Field(
         default=60,
         description="How often to check for stale sessions (seconds). Set to -1 to disable cleanup.",

skyrl/tinker/engine.py

@@ -751,6 +751,34 @@ def process_batch_requests(
                 results = {request_id: types.ErrorResponse(error=str(e), status="failed") for request_id in requests}
         self._complete_futures(results)
 
+    def _chunk_forward_backward_requests(
+        self,
+        requests: dict[str, tuple[str, types.ForwardBackwardInput]],
+    ) -> list[dict[str, tuple[str, types.ForwardBackwardInput]]]:
+        max_request_count = self.config.forward_backward_max_request_count
+        if not requests:
+            return []
+        if max_request_count is None or max_request_count <= 0 or len(requests) <= max_request_count:
+            return [requests]
+
+        items = list(requests.items())
+        chunks = [dict(items[start : start + max_request_count]) for start in range(0, len(items), max_request_count)]
+        logger.info(
+            "forward_backward request batch split %s requests into %s chunks "
+            "(forward_backward_max_request_count=%s)",
+            len(requests),
+            len(chunks),
+            max_request_count,
+        )
+        return chunks
+
+    def process_forward_backward_requests(
+        self,
+        requests: dict[str, tuple[str, types.ForwardBackwardInput]],
+    ) -> None:
+        for request_chunk in self._chunk_forward_backward_requests(requests):
+            self.process_batch_requests(request_chunk, self.process_forward_backward, "forward_backward")
+
     def process_pending_requests(self):
         """Main loop to process pending requests."""
         while True:
@@ -767,7 +795,7 @@ def process_pending_requests(self):
                 other_requests = self.find_single_requests(session)
 
             # Process batches outside of session context
-            self.process_batch_requests(forward_backward_requests, self.process_forward_backward, "forward_backward")
+            self.process_forward_backward_requests(forward_backward_requests)
             self.process_batch_requests(forward_requests, self.process_forward, "forward")
             self.process_batch_requests(sample_requests, self.process_sample, "sample")
 

tests/tinker/test_engine.py

@@ -180,6 +180,66 @@ def sample_input(**kwargs):
     assert batch.all_session_ids == ["sampling_abcd:3", "sampling_abcd:3", None, None]
 
 
+def _fake_forward_backward_input() -> types.ForwardBackwardInput:
+    datum = types.Datum(
+        model_input=types.ModelInput(chunks=[types.EncodedTextChunk(tokens=[1, 2, 3])]),
+        loss_fn_inputs=types.LossFnInputs(
+            target_tokens=types.TensorData(data=[2, 3, 4]),
+            weights=types.TensorData(data=[1.0, 1.0, 1.0]),
+            advantages=types.TensorData(data=[]),
+            logprobs=types.TensorData(data=[]),
+        ),
+    )
+    return types.ForwardBackwardInput(data=[datum], loss_fn="cross_entropy")
+
+
+def _fake_forward_backward_requests(count: int) -> dict[str, tuple[str, types.ForwardBackwardInput]]:
+    return {str(index): ("model_a", _fake_forward_backward_input()) for index in range(1, count + 1)}
+
+
+def test_chunk_forward_backward_requests_preserves_default_batching():
+    engine = object.__new__(TinkerEngine)
+    engine.config = EngineConfig(base_model=BASE_MODEL)
+    requests = _fake_forward_backward_requests(5)
+
+    chunks = engine._chunk_forward_backward_requests(requests)
+
+    assert chunks == [requests]
+
+
+def test_chunk_forward_backward_requests_splits_by_configured_request_count():
+    engine = object.__new__(TinkerEngine)
+    engine.config = EngineConfig(base_model=BASE_MODEL, forward_backward_max_request_count=4)
+
+    chunks = engine._chunk_forward_backward_requests(_fake_forward_backward_requests(10))
+
+    assert [list(chunk) for chunk in chunks] == [
+        ["1", "2", "3", "4"],
+        ["5", "6", "7", "8"],
+        ["9", "10"],
+    ]
+
+
+def test_process_forward_backward_requests_processes_each_chunk_in_order():
+    engine = object.__new__(TinkerEngine)
+    engine.config = EngineConfig(base_model=BASE_MODEL, forward_backward_max_request_count=2)
+    calls = []
+
+    def process_batch_requests(requests, processor, name):
+        calls.append((list(requests), processor, name))
+
+    engine.process_batch_requests = process_batch_requests
+
+    engine.process_forward_backward_requests(_fake_forward_backward_requests(5))
+
+    assert [(keys, name) for keys, _, name in calls] == [
+        (["1", "2"], "forward_backward"),
+        (["3", "4"], "forward_backward"),
+        (["5"], "forward_backward"),
+    ]
+    assert all(processor == engine.process_forward_backward for _, processor, _ in calls)
+
+
 @pytest.fixture()
 def scheduling_engine():
     """Create a TinkerEngine with only the DB initialized (no backend) for scheduling tests."""