Skip to content

Log training metrics as MetricEvents through the event logger - #1403

Merged
fzyzcjy merged 21 commits into
mainfrom
tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger
Jul 10, 2026
Merged

Log training metrics as MetricEvents through the event logger#1403
fzyzcjy merged 21 commits into
mainfrom
tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger

Conversation

@fzyzcjy

@fzyzcjy fzyzcjy commented Jun 22, 2026

Copy link
Copy Markdown
Collaborator

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.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a new MetricEvent model and integrates it into the tracking utility to log metrics when the event logger is initialized. The review feedback highlights a potential issue where passing raw metrics (which may contain PyTorch tensors or NumPy scalars) directly to the event logger could cause validation or serialization failures. A code suggestion is provided to preprocess the metrics dictionary and convert these values to standard Python types before logging.

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.

Comment thread miles/utils/tracking_utils/__init__.py Outdated
Comment on lines +20 to +21
if is_event_logger_initialized():
get_event_logger().log(MetricEvent, {"metrics": dict(metrics)}, print_log=False)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Passing metrics directly (even wrapped in dict()) to the event logger can cause runtime crashes or validation failures. Since MetricEvent inherits from FrozenStrictBaseModel (which likely enforces strict type validation) and eventually serializes to JSON, any non-standard Python types in metrics—such as PyTorch tensors (e.g., scalar loss tensors) or NumPy arrays/scalars—will trigger Pydantic validation or JSON serialization errors.

To prevent training runs from crashing, preprocess the metrics dictionary to convert any tensor or NumPy scalar values to standard Python types using .item() before passing them to the event logger.

Suggested change
if is_event_logger_initialized():
get_event_logger().log(MetricEvent, {"metrics": dict(metrics)}, print_log=False)
if is_event_logger_initialized():
serializable_metrics = {
k: (v.item() if hasattr(v, "item") else v)
for k, v in metrics.items()
}
get_event_logger().log(MetricEvent, {"metrics": serializable_metrics}, print_log=False)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please fix. This looks like an authentic risk.

@fzyzcjy
fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-event-log-snapshot-and-restore-checkpointing branch from 2abf57a to 12a82f0 Compare June 23, 2026 07:47
@fzyzcjy
fzyzcjy requested a review from yushengsu-thu as a code owner June 23, 2026 07:47
@fzyzcjy
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
fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-event-log-snapshot-and-restore-checkpointing branch from 12a82f0 to d60490a Compare June 23, 2026 09:25
@fzyzcjy
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
fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-event-log-snapshot-and-restore-checkpointing branch from d60490a to 81809bf Compare June 23, 2026 13:29
@fzyzcjy
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
from miles.utils.event_logger.logger import get_event_logger, is_event_logger_initialized
from miles.utils.event_logger.models import MetricEvent

from .base import TrackingManager

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: let's also fix this relative import.

Comment thread miles/utils/tracking_utils/__init__.py Outdated
Comment on lines +20 to +21
if is_event_logger_initialized():
get_event_logger().log(MetricEvent, {"metrics": dict(metrics)}, print_log=False)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please fix. This looks like an authentic risk.

fzyzcjy added a commit that referenced this pull request Jul 8, 2026
Review comment on #1403: a scalar torch.Tensor in metrics fails
MetricEvent validation/JSON serialization; unwrap via .item().
@fzyzcjy
fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-event-log-snapshot-and-restore-checkpointing branch from 81809bf to dd582d1 Compare July 8, 2026 03:52
@fzyzcjy
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
fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-event-log-snapshot-and-restore-checkpointing branch from dd582d1 to c4b826c Compare July 8, 2026 05:54
@fzyzcjy
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 added 9 commits July 10, 2026 10:04
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.
fzyzcjy added 11 commits July 10, 2026 10:04
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.
@fzyzcjy
fzyzcjy force-pushed the tom/pr_chain/trainer_ft/dev_revert_reversed/add-event-log-snapshot-and-restore-checkpointing branch from c4b826c to 691e4a7 Compare July 10, 2026 02:08
@fzyzcjy
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
Base automatically changed from tom/pr_chain/trainer_ft/dev_revert_reversed/add-event-log-snapshot-and-restore-checkpointing to main July 10, 2026 03:14
…ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger
@fzyzcjy
fzyzcjy merged commit cf77275 into main Jul 10, 2026
6 checks passed
@fzyzcjy
fzyzcjy deleted the tom/pr_chain/trainer_ft/dev_revert_reversed/log-training-metrics-as-metricevents-through-the-event-logger branch July 10, 2026 03:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants