feat(session): multi-process session-server (N workers + sticky-hash router)

miles/backends/experimental/fsdp_utils/checkpoint.py

@@ -12,10 +12,7 @@
 from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
 from torch.distributed.checkpoint.stateful import Stateful
 
-from miles.backends.training_utils.log_utils import (
-    init_train_step_counter,
-    save_train_step_counter,
-)
+from miles.backends.training_utils.log_utils import init_train_step_counter, save_train_step_counter
 
 logger = logging.getLogger(__name__)
 

miles/backends/megatron_utils/actor.py

@@ -49,8 +49,6 @@
 logger = logging.getLogger(__name__)
 
 import math
-from typing import Any
-
 
 
 def validate_rollout_for_grpo_training_step(
@@ -69,7 +67,6 @@ def validate_rollout_for_grpo_training_step(
     Logs useful diagnostics before raising so NCCL-desync root cause is visible
     in the first failing rank's log.
     """
-    import math
     import socket
     import traceback
 
@@ -218,10 +215,7 @@ def _summarize_vector_list(key, limit=3):
                 shapes.append(type(x).__name__)
                 dtypes.append(type(x).__name__)
                 devices.append("python")
-        return (
-            f"{key}: len={len(xs)} first_shapes={shapes} "
-            f"first_dtypes={dtypes} first_devices={devices}"
-        )
+        return f"{key}: len={len(xs)} first_shapes={shapes} " f"first_dtypes={dtypes} first_devices={devices}"
 
     def _basic_batch_summary():
         keys = sorted(list(rollout_data.keys()))
@@ -420,7 +414,9 @@ def _basic_batch_summary():
             _add_error(f"loss_masks[{i}] has no active tokens, sum={mask_sum}, response_len={resp}")
         if mask_sum > resp:
             # Warning-only: float/weighted masks can legitimately have sum > resp.
-            _add_warning(f"loss_masks[{i}] sum={mask_sum} exceeds response_len={resp} (expected for float/weighted masks)")
+            _add_warning(
+                f"loss_masks[{i}] sum={mask_sum} exceeds response_len={resp} (expected for float/weighted masks)"
+            )
 
         if torch.is_tensor(mask):
             # Binary check: warning, not fatal, because masks may be float.

miles/backends/training_utils/data.py

@@ -127,9 +127,7 @@ def get_batch(
     # same list (aggregate_train_losses keys positionally on the first microbatch).
     if has_domains:
         if not hasattr(data_iterator, "_all_domains_cache"):
-            data_iterator._all_domains_cache = sorted(
-                {d for d in data_iterator.rollout_data["domains"] if d}
-            )
+            data_iterator._all_domains_cache = sorted({d for d in data_iterator.rollout_data["domains"] if d})
         batch["all_domains"] = data_iterator._all_domains_cache
 
     tokens = batch["tokens"]

miles/backends/training_utils/log_utils.py

@@ -23,33 +23,33 @@
 # Maps bare metric names to their W&B top-level section(s).
 # Keys appearing in multiple sections (e.g. pg_loss) are emitted under each.
 _TRAIN_METRIC_GROUPS: dict[str, list[str]] = {
-    "ppo_kl":                         ["policy_shift"],
-    "ois":                            ["policy_shift"],
-    "pg_clipfrac":                    ["policy_shift"],
-    "pg_loss":                        ["policy_shift", "optimization"],
-    "log_probs":                      ["policy_shift"],   # current policy (training forward pass)
-    "old_log_probs":                  ["policy_shift"],   # old policy (rollout or FSDP rollout)
-    "ref_kl":                         ["policy_shift"],
+    "ppo_kl": ["policy_shift"],
+    "ois": ["policy_shift"],
+    "pg_clipfrac": ["policy_shift"],
+    "pg_loss": ["policy_shift", "optimization"],
+    "log_probs": ["policy_shift"],  # current policy (training forward pass)
+    "old_log_probs": ["policy_shift"],  # old policy (rollout or FSDP rollout)
+    "ref_kl": ["policy_shift"],
     "train_rollout_logprob_abs_diff": ["train_inference_mismatch"],
-    "train_rollout_logprob_diff":     ["train_inference_mismatch"],
-    "tis":                            ["train_inference_mismatch"],
-    "tis_abs":                        ["train_inference_mismatch"],
-    "tis_clipfrac":                   ["train_inference_mismatch"],
-    "loss":                           ["optimization"],
-    "entropy_loss":                   ["optimization"],
-    "kl_loss":                        ["optimization"],
-    "grad_norm":                      ["optimization"],
+    "train_rollout_logprob_diff": ["train_inference_mismatch"],
+    "tis": ["train_inference_mismatch"],
+    "tis_abs": ["train_inference_mismatch"],
+    "tis_clipfrac": ["train_inference_mismatch"],
+    "loss": ["optimization"],
+    "entropy_loss": ["optimization"],
+    "kl_loss": ["optimization"],
+    "grad_norm": ["optimization"],
 }
 
 # Maps rollout batch field names to their W&B top-level section.
 _ROLLOUT_DATA_METRIC_GROUPS: dict[str, str] = {
-    "log_probs":         "train_inference_mismatch",  # FSDP log probs at rollout time
+    "log_probs": "train_inference_mismatch",  # FSDP log probs at rollout time
     "rollout_log_probs": "train_inference_mismatch",  # inference engine log probs
-    "ref_log_probs":     "policy_shift",              # reference model log probs
-    "rewards":           "reward",
-    "raw_reward":        "reward",
-    "advantages":        "reward",
-    "returns":           "reward",
+    "ref_log_probs": "policy_shift",  # reference model log probs
+    "rewards": "reward",
+    "raw_reward": "reward",
+    "advantages": "reward",
+    "returns": "reward",
 }
 
 # Cumulative train-step counter across all rollouts. The previous formula
@@ -570,7 +570,7 @@ def log_train_step(
     for full_key, val in log_dict_out.items():
         if not full_key.startswith(prefix):
             continue
-        bare_key = full_key[len(prefix):]
+        bare_key = full_key[len(prefix) :]
         # Per-domain keys arrive as "<metric>/<domain>" — route to "<group>/<domain>/<metric>".
         metric_name, sep, domain = bare_key.rpartition("/")
         lookup = metric_name if (sep and metric_name in _TRAIN_METRIC_GROUPS) else bare_key

miles/backends/training_utils/loss.py

@@ -775,8 +775,12 @@ def policy_loss_function(
                 for dd, lm in zip(batch["domains"], batch["loss_masks"], strict=False)
             ]
             reducer = get_sum_of_sample_mean(
-                total_lengths, response_lengths, masked,
-                args.calculate_per_token_loss, args.qkv_format, max_seq_lens,
+                total_lengths,
+                response_lengths,
+                masked,
+                args.calculate_per_token_loss,
+                args.qkv_format,
+                max_seq_lens,
                 loss_agg_mode=getattr(args, "loss_agg_mode", None),
             )
             for name, t in per_token.items():