[None][test] CBTS sitecustomize: fix nested-pytest hang + argv detection

Two fixes in jenkins/scripts/cbts/coverage_utils/sitecustomize.py.

1. Use sys.orig_argv instead of sys.argv for pytest detection.
   When invoked as `python -m pytest`, sys.argv at sitecustomize-run
   time is ['-m', ...] -- the pytest path is injected later by the
   runpy machinery. The previous check `any("pytest" in a for a in
   sys.argv[:2])` therefore returned False for every `python -m pytest`
   process, including pytest main itself, causing daemon threads to
   spawn where they shouldn't.

2. Detect the nested-pytest layer used by
   tests/integration/defs/test_unittests.py::test_unittests_v2 (an
   outer pytest test that subprocess-spawns `python -m pytest
   <unittest_file>`), and in that layer skip the periodic-save /
   marker-poll / mpi-watcher daemon threads while still synchronously
   applying install_mpi_pool_patch. The 100ms-loop daemons in the
   middle layer were contending with the zmq IPC threads talking to
   MPI workers, pushing LLM() fixture setup past the 120s
   pytest-timeout on A30-PyTorch-1 (128 timeouts in
   unittest/_torch/sampler/test_logits_logprobs.py). MPI workers
   still get the widened env via the synchronous patch, so worker
   coverage (the bulk of py_executor.py data) is preserved.

Validated locally on H100 with a minimal repro: nested pytest spawns
LLM(TinyLlama), inner pytest has zero cbts-* daemon threads, mpi patch
is applied, MPI worker starts full coverage, LLM init + generate
complete normally.

Signed-off-by: Ivy Zhang <25222398+crazydemo@users.noreply.github.com>

Author: crazydemo

jenkins/scripts/cbts/coverage_utils/sitecustomize.py

@@ -21,16 +21,15 @@
   * Starts a coverage.Coverage instance with the given config in every
     Python process under our PYTHONPATH (pytest main + MPI workers +
     trtllm-serve subprocesses).
-  * In subprocesses (pytest is NOT on argv[:2]), runs a daemon thread
-    that calls cov.save() every 5s so processes killed via SIGKILL
-    (e.g. trtllm-serve teardown in disaggregated tests) still leave
-    most of their data on disk. Pytest main exits cleanly via atexit,
-    so it doesn't need periodic save.
-  * In the same subprocesses, runs a second daemon thread that polls a
-    marker file and calls cov.switch_context(<test_nodeid>) whenever
-    the marker changes; the marker file is written by cbts_plugin in
-    the pytest main process. This is how worker code gets attributed
-    to the test that triggered it.
+  * In non-pytest subprocesses (MPI workers, trtllm-bench, trtllm-serve)
+    runs a daemon thread that calls cov.save() every 5s so processes
+    killed via SIGKILL still leave most of their data on disk. Pytest
+    layers (main or nested) skip this — they exit cleanly via atexit.
+  * In the same non-pytest subprocesses, runs a second daemon thread
+    that polls a marker file and calls cov.switch_context(<test_nodeid>)
+    whenever the marker changes; the marker file is written by
+    cbts_plugin in the pytest main process. This is how worker code
+    gets attributed to the test that triggered it.
   * Registers a single atexit handler that signals the periodic thread
     to stop, then stops coverage and saves once. This avoids the
     periodic thread racing with atexit's own cov.save().
@@ -41,6 +40,27 @@
 import os
 import sys
 
+
+def _parent_is_pytest():
+    """Return True if our parent process is also running pytest.
+
+    Detects the nested-pytest pattern used by
+    ``tests/integration/defs/test_unittests.py::test_unittests_v2``: the
+    outer pytest spawns ``python -m pytest <unittest_file>`` as a
+    subprocess for each unittest module. In that nested case we keep
+    coverage running but skip the daemon threads in this layer (they
+    contend with the zmq IPC threads that talk to MPI workers spawned
+    by LLM() fixtures, and the inner pytest exits cleanly via atexit so
+    periodic save is unnecessary).
+    """
+    try:
+        with open(f"/proc/{os.getppid()}/cmdline", "rb") as f:
+            parent_cmdline = f.read().split(b"\x00")
+    except OSError:
+        return False
+    return any(b"pytest" in part for part in parent_cmdline)
+
+
 if os.getenv("CBTS_COVERAGE_CONFIG"):
     import atexit
 
@@ -70,7 +90,22 @@
     )
     cov.start()
 
-    _is_pytest_main = any("pytest" in a for a in sys.argv[:2])
+    # Use sys.orig_argv (Python 3.10+) instead of sys.argv: at the moment
+    # sitecustomize runs, Python has not yet rewritten argv[0] from "-m"
+    # to the pytest __main__.py path for `python -m pytest` invocations,
+    # so sys.argv looks like ['-m', '--co', ...] -- pytest is not in
+    # there yet. sys.orig_argv preserves the full launching cmdline,
+    # e.g. ['python', '-m', 'pytest', '--co', ...].
+    _orig_argv = getattr(sys, "orig_argv", sys.argv)
+    _is_pytest_main = any("pytest" in a for a in _orig_argv[:4])
+    _is_nested_pytest = _parent_is_pytest() and _is_pytest_main
+    # Skip daemon threads in pytest main AND in nested-pytest layers
+    # (see _parent_is_pytest docstring). Both exit cleanly via atexit
+    # so periodic save is unnecessary, and in the nested case the
+    # 100ms-loop daemon threads contend for the GIL with the zmq IPC
+    # threads talking to MPI workers, which on slower hosts (A30) can
+    # push LLM() fixture setup past the per-test pytest-timeout.
+    _skip_daemons = _is_pytest_main or _is_nested_pytest
 
     # Periodic save protects against SIGKILL'd subprocesses (e.g.
     # trtllm-serve cleanup in disaggregated tests). Pytest main exits
@@ -112,7 +147,7 @@ def _final_save():
 
     atexit.register(_final_save)
 
-    if not _is_pytest_main:
+    if not _skip_daemons:
         threading.Thread(
             target=_periodic_save,
             daemon=True,
@@ -139,12 +174,29 @@ def _final_save():
     # actually run py_executor.py -- which is why those workers ended
     # up with no coverage even though sitecustomize ran here. Without
     # this, the 21 SUBPROCESS_LOST tests stay SUBPROCESS_LOST.
-    if not _is_pytest_main:
-        import time as _time
+    _initial_nodeid = os.environ.get("CBTS_TEST_ID", "").strip()
+    if _initial_nodeid and not _is_pytest_main:
+        cov.switch_context(_initial_nodeid)
+
+    if _is_nested_pytest:
+        # Inner pytest: apply the mpi_session env-whitelist patch
+        # synchronously so MPI workers spawned by LLM() inside the
+        # inner test fixtures still get the widened env. The watcher
+        # thread used in non-pytest subprocesses would contend with
+        # zmq IPC threads here, so we just do it inline.
+        try:
+            sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+            from cbts_plugin import install_mpi_pool_patch
 
-        _initial_nodeid = os.environ.get("CBTS_TEST_ID", "").strip()
-        if _initial_nodeid:
-            cov.switch_context(_initial_nodeid)
+            install_mpi_pool_patch(raise_on_refactor=False)
+        except Exception as _exc:
+            print(
+                f"[cbts] nested-pytest mpi patch skipped in pid {os.getpid()}: {_exc!r}",
+                file=sys.stderr,
+            )
+
+    if not _skip_daemons:
+        import time as _time
 
         def _watch_mpi_session():
             try: