Fix TraceEvent.__init__() got an unexpected keyword argument pool_id

[HugeEngine]

Summary:

Problem

Newer PyTorch memory snapshots (introduced in D97221533) include a pool_id field in trace events. Mosaic's TraceEvent.from_raw() classmethod splats the raw dict directly into TraceEvent(**raw_modified), causing a TypeError for any unrecognized field:

TypeError: TraceEvent.__init__() got an unexpected keyword argument 'pool_id'

This breaks all memory snapshot analysis tools (peak_memory_analysis, categorical_profiling, annotation_analysis, memory_diff) when used with snapshots from newer PyTorch versions.

Root Cause

TraceEvent.from_raw() in fbcode/mosaic/libmosaic/utils/data_utils.py individually pops known extra fields (e.g., frames, user_metadata) before splatting the remaining dict into the dataclass constructor. When PyTorch added pool_id to snapshot trace events (D97221533), this field leaked through as an unexpected kwarg.

Fix

Replaced the individual .pop() approach (previously used for user_metadata in D88310416) with generic field filtering — raw_modified is now filtered to only include keys that match actual TraceEvent dataclass fields using dataclasses.fields(). This is future-proof: any new fields PyTorch adds to snapshot trace events will be silently ignored without requiring another code change.

Changes

1. Updated import (line 13): Added fields as dataclass_fields to the dataclasses import
2. Generic field filtering (lines 158-166): Replaced individual .pop() calls with {k: v for k, v in raw_modified.items() if k in valid_fields} filtering

Reviewed By: basilwong

Differential Revision: D99884123

[meta-codesync]

@HugeEngine has exported this pull request. If you are a Meta employee, you can view the originating Diff in D99884123.

[meta-codesync]

This pull request has been merged in 67b27bb.