From 350dd89627bc75c8362ea800b999bf04b576066a Mon Sep 17 00:00:00 2001 From: kailash Date: Fri, 17 Jul 2026 20:56:34 +0000 Subject: [PATCH] complete forward_backward futures as soon as each model batch finishes --- skyrl/tinker/engine.py | 42 +++++++++++++++++++++++------------ tests/tinker/test_engine.py | 44 +++++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 14 deletions(-) diff --git a/skyrl/tinker/engine.py b/skyrl/tinker/engine.py index bab1998d1a..49da179ec7 100644 --- a/skyrl/tinker/engine.py +++ b/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,6 +729,7 @@ 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. @@ -735,21 +737,31 @@ def process_batch_requests( 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) diff --git a/tests/tinker/test_engine.py b/tests/tinker/test_engine.py index 584e046876..118b2b2381 100644 --- a/tests/tinker/test_engine.py +++ b/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())