[Bug] configure_cpu_affinity() is not idempotent: a second call un-pins every thread

Summary

configure_cpu_affinity() is not idempotent. Calling it twice in the same process un-pins every thread it pinned on the first call, silently reverting the NUMA optimization.

Cause

The function decides whether affinity was set by someone else using only the current mask:

all_cpus = list(range(psutil.cpu_count()))
constrained_affinity = (cpu_affinity != all_cpus)

It cannot tell an externally imposed constraint from the one it applied itself. On a second call the process is "constrained" — by its own previous NUMA pin — so with TLLM_NUMA_AWARE_WORKER_AFFINITY unset it takes the un-constrain branch, logs "Removing CPU affinity constraints", and rebinds every thread back to the full CPU set. The NUMA branch is then skipped, because constrained_affinity is true.

Reproduction

Measured on a GB300 (Grace, 144 CPUs) with both functions extracted verbatim from tensorrt_llm/llmapi/utils.py, psutil/get_numa_aware_cpu_affinity shimmed onto real syscalls, six worker threads held alive across the whole run, TLLM_NUMA_AWARE_WORKER_AFFINITY unset:

start (unconstrained)   main=ALL   non-main: 0 on NUMA subset, 6 on ALL (of 6)
after 1st configure     main=NUMA  non-main: 6 on NUMA subset, 0 on ALL (of 6)
after 2nd configure     main=ALL   non-main: 0 on NUMA subset, 6 on ALL (of 6)

Emitted logs, in order:

INFO  Worker process ... CPU affinity set to [...] for optimal NUMA-aware scheduling (7/7 threads).
WARN  Worker process ... is affined to run on the following CPUs: [...]
WARN  Worker process ... has constrained CPU affinity but `TLLM_NUMA_AWARE_WORKER_AFFINITY` is not set. Removing CPU affinity constraints.

The second warning pair is the give-away: it reports the process as externally constrained when the constraint is self-inflicted.

How it is reached

Any path that runs configure_cpu_affinity() twice in one process — most directly two sequential LLM() objects on the in-process worker path (gather_generation_logits=True or TLLM_WORKER_USE_SINGLE_PROCESS=1), a shape several in-tree tests already use.

Severity

The end state matches the pre-#18701 baseline, so this is not a regression — but it silently reduces the feature to a no-op for the whole process, and the misleading "constrained CPU affinity" warning sends anyone debugging it in the wrong direction. Threads created between the two calls are a genuine small regression: they inherited the NUMA mask and are then widened, whereas previously they would have stayed pinned.

Suggested direction

Track whether this process already applied a NUMA-aware mask (or compare the current mask against the NUMA-aware set) instead of treating any non-full mask as external. Whatever the fix, constrained_affinity should stop conflating "the user pinned us" with "we pinned ourselves".

Context

Found while reviewing #18701; not introduced by it.

Start in tensorrt_llm/llmapi/utils.py at configure_cpu_affinity() and trace how it determines constrained_affinity and applies thread masks. Reproduce two sequential calls with TLLM_NUMA_AWARE_WORKER_AFFINITY unset, then add or update coverage for that path. Done means the second call preserves the NUMA-aware affinity, including for existing and newly created worker threads, without logging a false external-constraint warning.