fix: only publish trace from tp 0 (#16411)

This commit is contained in:
Douglas Yang
2026-01-04 12:17:08 -08:00
committed by GitHub
parent 52c604342c
commit 87699d48eb
2 changed files with 24 additions and 3 deletions

View File

@@ -132,12 +132,25 @@ def save_results_as_pydantic_models(
profile_link_decode = None
if res.profile_link:
# Collect all trace files, preferring TP-0 to match upload behavior
# (only TP-0 traces are published to avoid duplicates)
extend_files = []
decode_files = []
for file in os.listdir(res.profile_link):
if file.endswith(".trace.json.gz") or file.endswith(".trace.json"):
if "extend" in file.lower() or "prefill" in file.lower():
profile_link_extend = os.path.join(res.profile_link, file)
extend_files.append(file)
elif "decode" in file.lower():
profile_link_decode = os.path.join(res.profile_link, file)
decode_files.append(file)
# Sort to prefer TP-0 files (TP-0 < TP-1 < TP-2... alphabetically)
extend_files.sort()
decode_files.sort()
if extend_files:
profile_link_extend = os.path.join(res.profile_link, extend_files[0])
if decode_files:
profile_link_decode = os.path.join(res.profile_link, decode_files[0])
benchmark_result = BenchmarkResult(
model_path=model_path,