refactor: extract duplicated checkpoint interval logic into reusable helper (THUDM#1027)

Author: lancerts

slime/utils/misc.py

@@ -65,3 +65,23 @@ def get_free_port(start_port=10000, consecutive=1):
     while not all(is_port_available(port + i) for i in range(consecutive)):
         port += 1
     return port
+
+
+def should_run_periodic_action(
+    rollout_id: int,
+    interval: int | None,
+    num_rollout_per_epoch: int | None = None,
+) -> bool:
+    """
+    Return True when a periodic action (eval/save/checkpoint) should run.
+
+    Args:
+        rollout_id: The current rollout index (0-based).
+        interval: Desired cadence; disables checks when None.
+        num_rollout_per_epoch: Optional epoch boundary to treat as a trigger.
+    """
+    if interval is None:
+        return False
+
+    step = rollout_id + 1
+    return (step % interval == 0) or (num_rollout_per_epoch is not None and step % num_rollout_per_epoch == 0)

train.py

@@ -9,6 +9,7 @@
 from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
 from slime.utils.arguments import parse_args
 from slime.utils.logging_utils import configure_logger
+from slime.utils.misc import should_run_periodic_action
 from slime.utils.tracking_utils import init_tracking
 
 
@@ -61,7 +62,6 @@ def onload_rollout():
     # train loop.
     # note that for async training, one can change the position of the sync operation(ray.get).
     for rollout_id in range(args.start_rollout_id, args.num_rollout):
-        # TODO extract the duplicated eval logic
         if args.eval_interval is not None and rollout_id == 0:
             ray.get(rollout_manager.eval.remote(rollout_id))
 
@@ -78,10 +78,7 @@ def onload_rollout():
         else:
             ray.get(actor_model.async_train(rollout_id, rollout_data_ref))
 
-        if args.save_interval is not None and (
-            (rollout_id + 1) % args.save_interval == 0
-            or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0)
-        ):
+        if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch):
             if (not args.use_critic) or (rollout_id >= args.num_critic_only_steps):
                 actor_model.save_model(rollout_id)
             if args.use_critic:
@@ -98,10 +95,7 @@ def onload_rollout():
                 ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH]))
             ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE]))
 
-        if args.eval_interval is not None and (
-            (rollout_id + 1) % args.eval_interval == 0
-            or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0)
-        ):
+        if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch):
             ray.get(rollout_manager.eval.remote(rollout_id))
 
     ray.get(rollout_manager.dispose.remote())

train_async.py

@@ -3,6 +3,7 @@
 from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
 from slime.utils.arguments import parse_args
 from slime.utils.logging_utils import configure_logger
+from slime.utils.misc import should_run_periodic_action
 from slime.utils.tracking_utils import init_tracking
 
 
@@ -46,10 +47,7 @@ def train(args):
         else:
             ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref))
 
-        if args.save_interval is not None and (
-            (rollout_id + 1) % args.save_interval == 0
-            or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0)
-        ):
+        if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch):
             actor_model.save_model(rollout_id)
             if args.use_critic:
                 critic_model.save_model(rollout_id)
@@ -62,10 +60,7 @@ def train(args):
             rollout_data_next_future = None
             actor_model.update_weights()
 
-        if args.eval_interval is not None and (
-            (rollout_id + 1) % args.eval_interval == 0
-            or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0)
-        ):
+        if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch):
             ray.get(rollout_manager.eval.remote(rollout_id))
 
     ray.get(rollout_manager.dispose.remote())