refactor(orchestrator): fold packing.py into batch.py

src/prime_rl/orchestrator/batch.py

@@ -5,8 +5,11 @@
 from typing import Any
 
 import numpy as np
+from transformers import AutoConfig
 
+from prime_rl.configs.orchestrator import OrchestratorConfig
 from prime_rl.transports.batch.types import MicroBatch, MMImageRef, MMRefs, RoutedExperts, SamplingMask, TrainingSample
+from prime_rl.utils.logger import get_logger
 
 # Backfill value per component weight stream when a packed sample doesn't
 # carry it: absent rl means weight 1.0 on the loss mask, absent ce/ref_kl
@@ -858,3 +861,32 @@ def prepare_batch(
             batches_per_gpu[worker_idx].extend(worker_batches)
 
     return batches_per_gpu
+
+
+class BatchPacker:
+    """Bin-packs a step's training samples into one micro-batch list per trainer DP rank."""
+
+    def __init__(self, config: OrchestratorConfig):
+        self.seq_len = config.seq_len
+        self.num_train_workers = config.num_train_workers
+        self.pad_to_multiple_of = config.pad_to_multiple_of
+        try:
+            model_config = AutoConfig.from_pretrained(
+                config.model.name, trust_remote_code=config.tokenizer.trust_remote_code
+            )
+        except Exception as e:
+            get_logger().warning(
+                f"Could not load model config for {config.model.name} ({e}) - "
+                "packing balances by token count instead of estimated FLOPs"
+            )
+            model_config = None
+        self.bin_cost = build_bin_cost(model_config)
+
+    def pack(self, samples: list[TrainingSample]) -> list[list[MicroBatch]]:
+        return prepare_batch(
+            rollouts=samples,
+            seq_len=self.seq_len,
+            num_train_workers=self.num_train_workers,
+            bin_cost=self.bin_cost,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+        )

src/prime_rl/orchestrator/orchestrator.py

@@ -38,6 +38,7 @@
 from prime_rl.configs.orchestrator import OrchestratorConfig
 from prime_rl.orchestrator.algo.routing import is_trainable
 from prime_rl.orchestrator.annotations import stamp_arrival, stamp_batch
+from prime_rl.orchestrator.batch import BatchPacker
 from prime_rl.orchestrator.ckpt import setup_ckpt_manager
 from prime_rl.orchestrator.clients import AdminPlane, InferenceClient, setup_admin_plane
 from prime_rl.orchestrator.concurrency import ConcurrencyController
@@ -47,7 +48,6 @@
 from prime_rl.orchestrator.eval_source import EvalSource
 from prime_rl.orchestrator.inference_metrics import InferenceMetricsCollector
 from prime_rl.orchestrator.metrics import TrainEpisodes, dispatch_failure_metrics
-from prime_rl.orchestrator.packing import BatchPacker
 from prime_rl.orchestrator.periodic_logger import PeriodicLogger
 from prime_rl.orchestrator.train_sink import TrainSink
 from prime_rl.orchestrator.train_source import TrainSource