[ziyixiong-nv]

@coderabbitai summary

Description

When TLLM_AUTOTUNER_CACHE_PATH points at a shared filesystem, every rank takes one exclusive lock to load the cache when an autotune() session starts, and takes it again to save at the end. Each save reads the whole JSON file, merges, and writes it back with an fsync. These calls run between collectives, so a stalled rank stalls all of them. During filesystem trouble, single saves took minutes, the ranks queued behind them, and the autotuner part of startup took close to an hour.

This change:

• Loads take no lock. Writers only ever replace the file by rename, so a reader always sees one complete version. A read that races a replacement on NFS (ESTALE) is retried. An empty, unreadable or unparsable file is skipped (with a warning for the latter two) instead of failing startup.
• Saves no longer fsync. The cache needs an atomic replacement, not durability. A file torn by a crash is ignored by the reader above.
• The lock wait is bounded. A save waits for the writer lock at most TLLM_AUTOTUNER_CACHE_LOCK_TIMEOUT seconds (default 120, enough for a healthy cold-start queue of about 100 rank saves). After that it is skipped with one warning, startup continues, and the next run tunes what was not saved.
• Saves with nothing new do no I/O. If the file already holds every entry the rank would write (e.g. a run that reused the cache), the save takes no lock and writes nothing.
• One rank writes the shared section. The ranks of a tuning group hold the same BROADCAST/MERGE/PARALLEL entries, so only the group's first rank writes them. Every rank still writes its own rank_N section.
• Loads merge into memory. Instead of replacing the in-memory cache, a load merges the file into it; on a key both hold, the file wins. A skipped save therefore cannot drop this process's entries at the next autotune() session.

The file format is unchanged: older code reads files written by this version, and this version reads older files. Older readers still take the lock, which writers now hold only for the read-modify-write itself.

Test Coverage

Two new tests in tests/unittest/_torch/misc/test_autotuner.py:

• test_cache_saves_and_loads_from_concurrent_processes
  • Three processes (two MPI workers and the test process) each save 50 new entries to one cache file and reload it after every save.
  • Every reload sees that process's latest entry.
  • The final file holds all 150 per-rank entries, plus the 50 shared entries written by one rank; shared entries that only the other ranks hold never reach the file.
• test_cache_save_is_skipped_while_another_process_holds_the_lock, while a second process holds the writer lock:
  • Loads still work, and a save with nothing new returns without waiting.
  • A save with a new entry gives up after the 1 s timeout with one warning and leaves the file unchanged; the next load keeps the unsaved entry.
  • Once the lock is free, the entry is saved.

Both tests fail when the writer lock is removed: entries are lost and the save does not wait. The first test also fails when every rank writes the shared section.

Existing tests:

• tests/unittest/_torch/misc/test_autotuner.py passes on one H100 (94 passed, 5 skipped).
• test_autotuner_distributed_strategy also passes for all four strategies with its two-GPU skip lifted and both ranks on one GPU (98 passed, 1 skipped). It covers the shared section written by one rank.
• Files written by this change load with the previous code, and files written by the previous code load with this change.
• tests/unittest/_torch/pyexecutor/engine/runners/decoder/test_warmup.py: same results as on main (79 passed; the 5 KV-cache-manager tests fail on both because the test image's compiled bindings are older than main).

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.