[Bug] Worker CPU affinity is applied process-wide in shared-process deployments and never restored

Summary

configure_cpu_affinity() applies a process-wide CPU affinity policy, but in several supported configurations the worker does not own the process it runs in. In those deployments the mask is applied to threads that do not belong to the worker, and it is never restored when the worker shuts down.

This is a pre-existing condition, not a regression: before #18701 the code called psutil.Process().cpu_affinity(cpus), which on Linux is sched_setaffinity(pid, ...) and rebinds the main thread of whatever process it runs in — in the shared cases below, that is the caller's own main thread. Because Linux threads inherit the creating thread's mask, everything the main thread spawned afterwards was pinned too. #18701 widened the blast radius from "the main thread and its later descendants" to "every TID present at the configuration point", which is what makes it worth tracking explicitly.

Where the worker shares a process

All four reach configure_cpu_affinity() in a process that also hosts non-worker code:

1. Single-process TP1 worker. executor.py takes the use_worker=True branch when gather_generation_logits=True or TLLM_WORKER_USE_SINGLE_PROCESS=1, constructing GenerationExecutorWorker inline → worker.py setup_engine() → base_worker.py:153. Co-resident: the user's interpreter and its threads. Under trtllm-serve this is the HTTP server process itself — uvicorn/FastAPI, its uvloop event loop and the OpenAIServer thread pools all live there.
2. Externally supplied MpiCommSession, rank 0. mpi_session.py submits only n_workers - 1 tasks to the MPI pool; rank 0's worker_main runs on the session's own ThreadPoolExecutor inside the submitting process. Co-resident: the user's rank-0 script and its event loop, the proxy's ZMQ dispatch/result threads, the second pool thread.
3. External-launch VisualGen, rank 0. _torch/visual_gen/executor.py starts the _serve_forever coordinator thread before the worker thread, then runs configure_cpu_affinity() on the worker thread at executor.py:687. Co-resident: the coordinator's own asyncio loop and ZMQ sockets, and under trtllm-serve the HTTP server on the main thread.
4. Leader process of a multi-node launch. With TLLM_SPAWN_PROXY_PROCESS=1 the leader runs RemoteMpiCommSessionServer.serve() and rank 0's worker on the session thread pool. Most threads there are genuinely worker threads, but the ZMQ session server is not.

For contrast, these are dedicated and unaffected: Ray actors (ray/gpu_worker.py:370), ordinary MPI pool workers, and mgmn_worker_node ranks > 0.

Problems

• Ownership. There is no signal distinguishing "this process is mine" from "I am a guest in the caller's process", so the policy cannot be scoped correctly today.
• No restoration. Nothing restores the previous mask at shutdown — grep -rn "cpu_affinity\|sched_setaffinity" tensorrt_llm/ finds only the three configure_cpu_affinity() call sites. A user who constructs an LLM, uses it and shuts it down is left with a permanently narrowed process.
• Default-on widening. The TLLM_NUMA_AWARE_WORKER_AFFINITY-unset branch removes an externally set constraint. Under srun --cpu-bind, this now widens the whole hosting process to all CPUs rather than just its main thread.

Suggested direction

_torch/visual_gen/executor.py already plumbs an in_client_process flag through run_diffusion_worker, so VisualGen is the cheapest first landing point — but the fix should cover all four modes rather than special-casing one.

Some care is needed: gating the sweep on "dedicated worker process" would silently disable NUMA pinning in the multi-node leader, where rank 0's full engine lives. Rank 0's MPI/UCX/NCCL progress threads would stay unpinned while every other rank is pinned, reproducing exactly the rank-0 straggler imbalance the feature exists to remove. "Restrict to worker-owned threads" is also not directly implementable: the threads the feature targets are unregistered C-level threads with no ownership metadata.

Context

Raised in review of #18701, which documents the current behaviour in docs/source/deployment-guide/configuring-cpu-affinity.md but deliberately does not change it.

Start at configure_cpu_affinity() in base_worker.py:153 and trace its call sites in executor.py, worker.py, mpi_session.py, and _torch/visual_gen/executor.py. Compare the four shared-process configurations with the dedicated worker cases, including the in_client_process flow. Done means non-worker process affinity is not incorrectly changed, prior masks are restored at shutdown, and NUMA pinning remains effective for the multi-node leader.