Skip to content

Commit 0de35f9

Browse files
authored
[profiler] Suppress Callable field from tyro CLI parsing (#3038)
## Summary Fixes a crash introduced by #2926 where `trace_post_processor: Function.Config | None` in `Profiler.Config` causes `tyro.cli()` to fail because tyro cannot parse `Callable[..., Any]` types from CLI arguments. **Error:** ``` Invalid input to tyro.cli() Unsupported type annotation for field profiler.trace-post-processor.fn with type collections.abc.Callable[..., typing.Any] ``` **Fix:** Annotate `trace_post_processor` with `tyro.conf.Suppress`, matching the existing pattern used for: - `model_spec` in `trainer.py:62` - `sample_processor` in `text_datasets.py:502` The field defaults to `None` and is only set programmatically by `CUDAGraphWrapper`, never via CLI. ## Test plan - [ ] Verify `python -m torchtitan.train ...` no longer crashes at config parsing - [ ] Verify graph_trainer still works (trace_post_processor set programmatically)
1 parent 1a2fef0 commit 0de35f9

1 file changed

Lines changed: 6 additions & 2 deletions

File tree

torchtitan/tools/profiler.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,10 @@
88
import pickle
99
import time
1010
from dataclasses import dataclass
11+
from typing import Annotated
1112

1213
import torch
13-
14+
import tyro
1415
from torchtitan.config import Configurable
1516
from torchtitan.config.function import Function
1617
from torchtitan.tools.logging import logger
@@ -169,10 +170,13 @@ class Config(Configurable.Config):
169170
save_memory_snapshot_folder: str = "memory_snapshot"
170171
"""Memory snapshot files location."""
171172

172-
trace_post_processor: Function.Config | None = None
173+
trace_post_processor: Annotated[
174+
Function.Config | None, tyro.conf.Suppress
175+
] = None
173176
"""Optional hook invoked with the trace path after each export.
174177
175178
Wraps ``fn(trace_path: str) -> None``.
179+
Set programmatically (not via CLI) — tyro cannot parse Callable types.
176180
"""
177181

178182
def __init__(

0 commit comments

Comments
 (0)