diff --git a/scripts/collectivespec/summarize_p0_phase_trace.py b/scripts/collectivespec/summarize_p0_phase_trace.py index ba9349d..41d9c6a 100644 --- a/scripts/collectivespec/summarize_p0_phase_trace.py +++ b/scripts/collectivespec/summarize_p0_phase_trace.py @@ -328,11 +328,29 @@ def target_execution_agreement(worker_events: list[dict[str, Any]]) -> dict[str, } metadata_valid_count = 0 + anchor_local_ranks: dict[int, int] = {} + for event in target_events: + dp_rank = integer(event.get("data_parallel_rank")) + local_rank = integer(event.get("local_rank")) + if dp_rank is None or local_rank is None: + continue + previous = anchor_local_ranks.get(dp_rank) + if previous is None or local_rank < previous: + anchor_local_ranks[dp_rank] = local_rank + anchor_events = [ + event + for event in target_events + if integer(event.get("data_parallel_rank")) in anchor_local_ranks + and integer(event.get("local_rank")) + == anchor_local_ranks[integer(event.get("data_parallel_rank"))] + ] metadata_padding: list[int] = [] raw_unequal_count = 0 for event in target_events: - valid, padding, raw_unequal = _target_metadata_check(event) + valid, _, _ = _target_metadata_check(event) metadata_valid_count += valid + for event in anchor_events: + _, padding, raw_unequal = _target_metadata_check(event) if padding is not None: metadata_padding.append(padding) raw_unequal_count += bool(raw_unequal) @@ -352,8 +370,9 @@ def target_execution_agreement(worker_events: list[dict[str, Any]]) -> dict[str, "valid_record_count": metadata_valid_count, "all_records_valid": len(target_events) > 0 and metadata_valid_count == len(target_events), - "raw_dp_counts_unequal_record_count": raw_unequal_count, - "post_coordinate_minus_local_raw_rows": { + "anchor_record_count": len(anchor_events), + "raw_dp_counts_unequal_anchor_record_count": raw_unequal_count, + "per_dp_replica_anchor_post_coordinate_minus_local_raw_rows": { "n": len(metadata_padding), "min": min(metadata_padding) if metadata_padding else None, "max": max(metadata_padding) if metadata_padding else None,