Add the witness id allocator by fzyzcjy · Pull Request #1404 · radixark/miles

fzyzcjy · 2026-06-22T10:36:25Z

Add the witness id allocator and WitnessInfo carrier used to assign and track
witness ids for fault-tolerance verification.

miles/utils/witness/allocator.py and tests.

gemini-code-assist

Code Review

This pull request introduces a witness ID allocator mechanism, adding CLI arguments (--enable-witness and --witness-buffer-size), a WitnessIdAllocator class to manage and recycle witness IDs, a utility to recover the allocator state from persisted events, and a comprehensive test suite. The review feedback suggests adding defensive assertions to ensure buffer_size is strictly positive to avoid potential division-by-zero errors, and ensuring num_ids is non-negative during allocation to prevent the internal counter from decrementing.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-22T10:37:42Z

+    def __init__(self, *, buffer_size: int) -> None:
+        self._buffer_size = buffer_size
+        self._counter: int = 0


If buffer_size is zero or negative, it will cause a ZeroDivisionError during modulo operations in _compute_stale_ids. We should defensively assert that buffer_size > 0 during initialization.

Suggested change

def __init__(self, *, buffer_size: int) -> None:

self._buffer_size = buffer_size

self._counter: int = 0

def __init__(self, *, buffer_size: int) -> None:

assert buffer_size > 0, f"buffer_size ({buffer_size}) must be positive."

self._buffer_size = buffer_size

self._counter: int = 0

This is good to have. I would raise a ValueError instead of using assert.

gemini-code-assist · 2026-06-22T10:37:42Z

+    def allocate(self, num_ids: int) -> WitnessInfo:
+        assert num_ids <= self._buffer_size, (
+            f"num_ids ({num_ids}) exceeds buffer_size ({self._buffer_size}). " f"Increase --witness-buffer-size."
+        )


If num_ids is negative, range(num_ids) is empty but self._counter will be decremented, which violates the monotonic non-negative counter assumption. We should defensively assert that num_ids >= 0.

Suggested change

def allocate(self, num_ids: int) -> WitnessInfo:

assert num_ids <= self._buffer_size, (

f"num_ids ({num_ids}) exceeds buffer_size ({self._buffer_size}). " f"Increase --witness-buffer-size."

)

def allocate(self, num_ids: int) -> WitnessInfo:

assert num_ids >= 0, f"num_ids ({num_ids}) must be non-negative."

assert num_ids <= self._buffer_size, (

f"num_ids ({num_ids}) exceeds buffer_size ({self._buffer_size}). " f"Increase --witness-buffer-size."

)

Same good to have.

Shi-Dong

LGTM. Good to add the value checks suggested by Gemini.

Review comments on #1404: reject non-positive buffer_size (would hit a ZeroDivisionError later) and negative num_ids (would silently decrement the monotonic counter). Raise ValueError per reviewer preference.

Add a `deterministic_random` reward that hashes the sample tokens + response to produce a stable pseudo-random 0/1 reward, used for reproducible fault-tolerance / CI tests. - rm_hub/__init__.py (+ test).

Add `inplace_modify_args`, a context manager that temporarily overrides args attributes and restores them on exit (asserting they weren't clobbered), used to scope per-attempt argument overrides in the fault-tolerant trainer. - argparse_utils.py (+ test).

Small shared-utility additions used by the fault-tolerant trainer: hash non-contiguous tensors safely (reshape before viewing as bytes), an `enable_experimental_ft_trainer` env flag, forward NCCL_DEBUG/NCCL_DEBUG_FILE to worker environments, and a `filter_keys` helper. - ci_utils.py / environ.py / external_utils/command_utils.py / misc.py.

Thread the original backend through ReloadableProcessGroup so that, when a process group is rebuilt (e.g. after a reconfigure/heal), it is recreated with the same backend instead of hard-coding NCCL. - reloadable_process_group.py: carry `backend` in the reload group info.

Add small foundation utilities used across the fault-tolerance trainer: a strict pydantic base model, a retry helper, a tensor checksum helper, a per-cell megatron world-size computation, the TrainStepOutcome enum, and the IndepDPInfo dataclass describing a cell's independent-DP identity. - pydantic_utils.py / retry_utils.py / checksum_utils.py / megatron_args_utils.py / types.py / indep_dp.py and tests.

Add a `log_structured` helper that emits logfmt-style key/value log lines, used by the fault-tolerance components for greppable structured logs. - structured_log.py (+ test).

Add a small `Clock` interface (`RealClock` plus a controllable fake clock) so time-dependent fault-tolerance code (health checks, heartbeats) can be driven deterministically in tests. - miles/utils/clock.py and tests.

Add a fault-injector test utility used to deterministically exercise fault-tolerance code paths. - miles/utils/test_utils/fault_injector.py.

Add the shared data models for the fault-tolerance control server (e.g. the `TriState` health value), used by the health checker and later by the HTTP control server. - miles/utils/control_server/models.py.

Add the periodic health checker (debounced TriState status driven by a Clock) and heartbeat utilities used to monitor train-cell liveness. - miles/utils/health_checker.py, miles/utils/heartbeat_utils.py and tests.

Add the nvidia-resiliency-ext dependency, the "ft" CI test label, the FT test fixtures in the rollout conftest, and route startup logging through configure_logger_raw. The fault-tolerance CLI arguments themselves now live with the features that consume them (distributed across the per-feature commits).

Unconditionally disconnect-then-reconnect the model-update process group when (re)connecting rollout engines, guarding the destroy against a missing group, so a reconfigured/healed engine set can rebuild the NCCL group from scratch. - broadcast.py: drop the "only disconnect if group exists" short-circuit; guard `destroy_process_group` against None.

Expose an `inject_fault` Ray method on TrainRayActor (in its own concurrency group) that triggers a configured failure mode via the fault injector, so fault-tolerance tests can crash/hang specific actors on demand. - train_actor.py: `inject_fault` RPC.

Extract the DP split into a witness-aware split_train_data_by_dp_raw helper (with unit tests; the key list also carries seq_witness_ids) and use it to split the training data on the actor side when delay_split_train_data_by_dp is set, deferring the DP split from the rollout side to actor-side processing. split_train_data_by_dp stays a thin wrapper that ray.puts each partition. - miles/ray/rollout/train_data_conversion.py (+ tests), miles/utils/data.py, actor_group.py, rollout_manager.py.

Add an opt-in deterministic NCCL process-group backend (`--debug-deterministic-collective`) that folds order-sensitive SUM/AVG reductions into a fixed order so training collectives are bit-reproducible, registering it as the training world's distributed backend and requiring synchronous grad sync. - det_process_group.py (+ GPU test, dist test helper). - train_actor.py: register the backend and select it when enabled. - initialize.py: assert synchronous grad reduce under the deterministic backend.

Add a per-process identity helper that uniquely keys each training process, used to attribute structured fault-tolerance events to their originating process. - miles/utils/process_identity.py and tests.

Add the structured event models (Event / EventBase hierarchy) for the fault-tolerance event log, each tagged with the originating ProcessIdentity. - miles/utils/event_logger/models.py and tests.

Add the structured event logger that records typed events keyed by per-process identity, wire it through the logging helper and CLI argument, and start it from the train entrypoints. - miles/utils/event_logger/logger.py, logging_utils.py, arguments.py and entrypoint wiring, with tests.

Add snapshot/restore for the structured event log so the event history survives cell restarts during fault-tolerant training. - miles/utils/event_logger/checkpoint.py and tests.

Add the `MetricEvent` model (a discriminated-union member) and emit every tracking metric into the structured event log: `tracking_utils.log` now forwards `{metrics}` to `get_event_logger().log(MetricEvent, ...)` when the event logger is initialized.

Add the witness id allocator and `WitnessInfo` carrier used to assign and track witness ids for fault-tolerance verification. - miles/utils/witness/allocator.py and tests.

…ft/dev_revert_reversed/add-the-witness-id-allocator

fzyzcjy requested review from Shi-Dong, Zhichenzzz, guapisolo, jybsuper, maocheng23 and yueming-yuan as code owners June 22, 2026 10:36

gemini-code-assist Bot reviewed Jun 22, 2026

View reviewed changes

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger branch from d1ab021 to f547c3e Compare June 23, 2026 07:47

fzyzcjy requested a review from yushengsu-thu as a code owner June 23, 2026 07:47

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-the-witness-id-allocator branch from 64b5e9b to 477bf8f Compare June 23, 2026 07:47

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger branch from f547c3e to b776abd Compare June 23, 2026 09:26

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-the-witness-id-allocator branch from 477bf8f to 56ce1fd Compare June 23, 2026 09:26

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger branch from b776abd to 51b40a7 Compare June 23, 2026 13:29

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-the-witness-id-allocator branch from 56ce1fd to 6c8f9cc Compare June 23, 2026 13:29

Shi-Dong approved these changes Jun 25, 2026

View reviewed changes

yueming-yuan approved these changes Jun 26, 2026

View reviewed changes

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger branch from 51b40a7 to 79deb5a Compare July 8, 2026 03:52

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-the-witness-id-allocator branch from 6c8f9cc to 9e840d3 Compare July 8, 2026 03:52

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger branch from 79deb5a to 7f8afe5 Compare July 8, 2026 05:55

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-the-witness-id-allocator branch from 9e840d3 to 5e59784 Compare July 8, 2026 05:55

fzyzcjy added 9 commits July 10, 2026 10:04

Add a deterministic_random reward type

96eecb4

Add a `deterministic_random` reward that hashes the sample tokens + response to produce a stable pseudo-random 0/1 reward, used for reproducible fault-tolerance / CI tests. - rm_hub/__init__.py (+ test).

Add structured logfmt logging helper

8dc8ddd

Add a `log_structured` helper that emits logfmt-style key/value log lines, used by the fault-tolerance components for greppable structured logs. - structured_log.py (+ test).

Add a Clock abstraction with a fake clock for tests

96d220f

Add a small `Clock` interface (`RealClock` plus a controllable fake clock) so time-dependent fault-tolerance code (health checks, heartbeats) can be driven deterministically in tests. - miles/utils/clock.py and tests.

Add a fault injector test utility

a005a2f

Add a fault-injector test utility used to deterministically exercise fault-tolerance code paths. - miles/utils/test_utils/fault_injector.py.

Add control-server data models

8365971

Add the shared data models for the fault-tolerance control server (e.g. the `TriState` health value), used by the health checker and later by the HTTP control server. - miles/utils/control_server/models.py.

fzyzcjy added 12 commits July 10, 2026 10:04

Add a cell health checker and heartbeat utilities

e308821

Add the periodic health checker (debounced TriState status driven by a Clock) and heartbeat utilities used to monitor train-cell liveness. - miles/utils/health_checker.py, miles/utils/heartbeat_utils.py and tests.

Add a per-process identity helper

a48d03f

Add a per-process identity helper that uniquely keys each training process, used to attribute structured fault-tolerance events to their originating process. - miles/utils/process_identity.py and tests.

Add structured event models keyed by per-process identity

3e85a86

Add the structured event models (Event / EventBase hierarchy) for the fault-tolerance event log, each tagged with the originating ProcessIdentity. - miles/utils/event_logger/models.py and tests.

Add event-log snapshot and restore checkpointing

691e4a7

Add snapshot/restore for the structured event log so the event history survives cell restarts during fault-tolerant training. - miles/utils/event_logger/checkpoint.py and tests.

Add the witness id allocator

34fb714

Add the witness id allocator and `WitnessInfo` carrier used to assign and track witness ids for fault-tolerance verification. - miles/utils/witness/allocator.py and tests.

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger branch from 7f8afe5 to 5a91150 Compare July 10, 2026 02:08

fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-the-witness-id-allocator branch from 5e59784 to 34fb714 Compare July 10, 2026 02:09

Base automatically changed from tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger to main July 10, 2026 03:14

Merge remote-tracking branch 'origin/main' into tom/pr_chain/trainer_…

b2ff7c0

…ft/dev_revert_reversed/add-the-witness-id-allocator

fzyzcjy merged commit edf438a into main Jul 10, 2026
5 checks passed

fzyzcjy deleted the tom/pr_chain/trainer_ft/dev_revert_reversed/add-the-witness-id-allocator branch July 10, 2026 03:15

Navigation Menu

Uh oh!

Add the witness id allocator - #1404

Add the witness id allocator#1404
fzyzcjy merged 22 commits into
mainfrom
tom/pr_chain/trainer_ft/dev_revert_reversed/add-the-witness-id-allocator

fzyzcjy commented Jun 22, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jun 22, 2026

Uh oh!

Shi-Dong Jun 25, 2026

Uh oh!

gemini-code-assist Bot Jun 22, 2026

Uh oh!

Shi-Dong Jun 25, 2026

Uh oh!

Shi-Dong left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

fzyzcjy commented Jun 22, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 22, 2026

Choose a reason for hiding this comment

Uh oh!

Shi-Dong Jun 25, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 22, 2026

Choose a reason for hiding this comment

Uh oh!

Shi-Dong Jun 25, 2026

Choose a reason for hiding this comment

Uh oh!

Shi-Dong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants