[tinker] Fix per-tenant blocking on forward/forward_backward calls

skyrl/tinker/engine.py

@@ -2,6 +2,7 @@
 
 import argparse
 import time
+from collections import defaultdict
 from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
@@ -728,28 +729,39 @@ def process_batch_requests(
         requests: dict[str, tuple[str, BaseModel]],
         processor: Callable[[dict[str, tuple[str, BaseModel]]], dict[str, BaseModel]],
         name: str,
+        per_model: bool = False,
     ):
         """Process a batch of requests with error handling and future completion.
 
         Args:
             requests: Dict mapping request_id to (model_id, request_data) tuples
             processor: Function that processes requests and returns results dict
             name: Name for logging
+            per_model: Process one model's requests at a time, completing each
+                model's futures as soon as its sub-batch finishes (GPU execution
+                is serialized per model anyway; this only changes completion
+                granularity, not batching within a model).
         """
         if not requests:
             return
-        with log_timing(f"process_batch_requests({name}, n={len(requests)})"):
-            try:
-                error_results, valid_requests = self._filter_valid_requests(requests)
-                if valid_requests:
-                    results = processor(valid_requests)
-                    results.update(error_results)
-                else:
-                    results = error_results
-            except Exception as e:
-                logger.exception(f"Error processing batch: {e}")
-                results = {request_id: types.ErrorResponse(error=str(e), status="failed") for request_id in requests}
-        self._complete_futures(results)
+        error_results, valid_requests = self._filter_valid_requests(requests)
+        if error_results:
+            self._complete_futures(error_results)
+        if per_model:
+            grouped: dict[str, dict] = defaultdict(dict)
+            for request_id, item in valid_requests.items():
+                grouped[item[0]][request_id] = item
+            groups = list(grouped.values())
+        else:
+            groups = [valid_requests] if valid_requests else []
+        for group in groups:
+            with log_timing(f"process_batch_requests({name}, n={len(group)})"):
+                try:
+                    results = processor(group)
+                except Exception as e:
+                    logger.exception(f"Error processing batch: {e}")
+                    results = {request_id: types.ErrorResponse(error=str(e), status="failed") for request_id in group}
+            self._complete_futures(results)
 
     def process_pending_requests(self):
         """Main loop to process pending requests."""
@@ -767,8 +779,10 @@ 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_batch_requests(forward_requests, self.process_forward, "forward")
+            self.process_batch_requests(
+                forward_backward_requests, self.process_forward_backward, "forward_backward", per_model=True
+            )
+            self.process_batch_requests(forward_requests, self.process_forward, "forward", per_model=True)
             self.process_batch_requests(sample_requests, self.process_sample, "sample")
 
             # Process other request types individually (in the future we can also batch independent optim_steps)

tests/tinker/test_engine.py

@@ -296,3 +296,47 @@ def test_find_single_requests_barrier_is_per_model(scheduling_engine):
         assert list(singles.keys()) == ["2", "5"]
         assert singles["2"][0] == "model_a"
         assert singles["5"][0] == "model_b"
+
+
+def test_process_batch_requests_per_model_completes_incrementally(scheduling_engine):
+    """per_model=True must complete each model's futures before processing the next model's group."""
+    from types import SimpleNamespace
+
+    from sqlmodel import select
+
+    engine = scheduling_engine
+    engine.backend = SimpleNamespace(has_model=lambda model_id: True)
+
+    with Session(engine.db_engine) as session:
+        for model_id in ["model_a", "model_a", "model_b"]:
+            session.add(
+                FutureDB(
+                    request_type=types.RequestType.FORWARD_BACKWARD,
+                    model_id=model_id,
+                    request_data={},
+                    status=RequestStatus.PENDING,
+                )
+            )
+        session.commit()
+        futures = session.exec(select(FutureDB).order_by(FutureDB.request_id)).all()
+        adam = types.AdamParams(learning_rate=1e-4, beta1=0.9, beta2=0.95, eps=1e-8, weight_decay=0.0)
+        requests = {str(f.request_id): (f.model_id, types.OptimStepInput(adam_params=adam)) for f in futures}
+
+    processed_models = []
+
+    def processor(group):
+        (model_id,) = {mid for mid, _ in group.values()}
+        if model_id == "model_b":
+            with Session(engine.db_engine) as session:
+                statuses = [
+                    f.status for f in session.exec(select(FutureDB).where(FutureDB.model_id == "model_a")).all()
+                ]
+                assert statuses == [RequestStatus.COMPLETED] * 2
+        processed_models.append(model_id)
+        return {request_id: types.OptimStepOutput(metrics={}) for request_id in group}
+
+    engine.process_batch_requests(requests, processor, "forward_backward", per_model=True)
+
+    assert processed_models == ["model_a", "model_b"]
+    with Session(engine.db_engine) as session:
+        assert all(f.status == RequestStatus.COMPLETED for f in session.exec(select(FutureDB)).all())