[None][infra] CBTS coverage data enhancement

jenkins/L0_MergeRequest.groovy

@@ -339,7 +339,10 @@ def setupPipelineEnvironment(pipeline, testFilter, globalVars)
     testFilter[(AUTO_TRIGGER_TAG_LIST)] = getAutoTriggerTagList(pipeline, testFilter, globalVars)
     testFilter[(CBTS_RESULT)] = getCbtsResult(pipeline, testFilter, globalVars)
     // Decide CBTS coverage eligibility here so L0_Test only consumes the propagated flag.
-    testFilter[(CBTS_COVERAGE)] = ENABLE_CBTS_COVERAGE && (env.JOB_NAME ==~ /.*PostMerge.*/)
+    // TEST ONLY (revert before merge): the official-post-merge gate is dropped so coverage is
+    // collected on a manual run to validate collection. Restore the JOB_NAME gate before merging:
+    //   testFilter[(CBTS_COVERAGE)] = ENABLE_CBTS_COVERAGE && (env.JOB_NAME ==~ /.*PostMerge.*/)
+    testFilter[(CBTS_COVERAGE)] = ENABLE_CBTS_COVERAGE
     pipeline.echo("CBTS coverage eligible: ${testFilter[(CBTS_COVERAGE)]}")
     getContainerURIs().each { k, v ->
         globalVars[k] = v

jenkins/scripts/cbts/coverage_utils/README.md

@@ -13,10 +13,10 @@ not lines executed (far cheaper than line tracing).
 
 | File | Role |
 |---|---|
-| `cbts_pystart.py` | The tracker: a `sys.monitoring` (tool id 4) `PY_START` tool that records, per test context, the set of product `(file, qualname)` entered. Writes one `.cbtscov.<stage>.<suffix>.sqlite` per process; these leave the node only inside compressed tarballs, so the publish-artifacts guardword/secret scanner (which byte-matches product paths in raw files but does not recurse into archives) never sees the paths stored inside. |
-| `sitecustomize.py` | Starts the tracker in each Python process under `CBTS_COVERAGE_CONFIG` (except dependency build/install tools — `pip`, `setup.py`, `cmake`, … — which opt out themselves and their spawned subtree). Reads `source` + `data_file` from the rcfile. Long-lived non-pytest processes (e.g. `trtllm-serve`) poll a marker file to switch context; `mpi4py.futures` pool workers use the inherited `CBTS_TEST_ID` context plus the atexit save. |
-| `cbts_plugin.py` | Pytest plugin (`-p cbts_plugin`): per test, writes the marker file, sets `CBTS_TEST_ID`, and switches the tracker context via `sitecustomize.switch_test_context`; also patches `mpi_session._start_mpi_pool` so workers inherit the coverage env. |
-| `pystart_report.py` | Merges all `.cbtscov.*.sqlite` (union per test; legacy `.json`/`.json.gz` also accepted) and emits any of: `--out-sqlite` (indexed `touch(test, file, qualname)` DB — the selector artifact), `--out-dir` (per-file split HTML report: index + one page per file), `--out-json` (full `test -> [file::qualname]` map). With `--source-root` also computes the file/function coverage rate. Prints a one-line touch-count summary. |
+| `cbts_pystart.py` | The tracker: a `sys.monitoring` (tool id 4) `PY_START` tool that records, per test context, the set of product `(file, qualname)` entered. Writes one `.cbtscov.<stage>.<suffix>.sqlite` per process, each carrying a `proc_meta` (stage) row and, in the coordinator, a `test_meta` (per-test outcome + spawned-worker count); these leave the node only inside compressed tarballs, so the publish-artifacts guardword/secret scanner (which byte-matches product paths in raw files but does not recurse into archives) never sees the paths stored inside. |
+| `sitecustomize.py` | Starts the tracker in each Python process under `CBTS_COVERAGE_CONFIG` (except dependency build/install tools — `pip`, `setup.py`, `cmake`, … — which opt out themselves and their spawned subtree). Reads `source` + `data_file` from the rcfile. Every instrumented process saves periodically so its coverage survives a non-clean exit (`mpi4py.futures` pool workers included — their pool is torn down at test end); long-lived non-pytest processes (e.g. `trtllm-serve`) additionally poll a marker file to switch context. |
+| `cbts_plugin.py` | Pytest plugin (`-p cbts_plugin`): per test, writes the marker file, sets `CBTS_TEST_ID`, switches the tracker context, and records the test outcome via `sitecustomize`; patches `mpi_session._start_mpi_pool` so workers inherit the coverage env, and patches `MPIPoolExecutor.__init__` to count the subprocess workers each test spawns. |
+| `pystart_report.py` | Merges all `.cbtscov.*.sqlite` (union per test; legacy `.json`/`.json.gz` also accepted) and emits any of: `--out-sqlite` (indexed `touch(test, file, qualname, stage)` + `test_meta` DB — the selector artifact), `--out-dir` (per-file split HTML report: index + one page per file), `--out-json` (full `test -> [file::qualname]` map). With `--source-root` also computes the file/function coverage rate. Prints touch-count and completeness summaries. |
 | `coveragerc.template` | Template for the runtime rcfile; only `[run] source` + `data_file` are used. |
 | `make_coveragerc.sh` | Substitutes `@...@` placeholders in the template; writes `$JOB_WORKSPACE/.coveragerc`. |
 
@@ -59,10 +59,22 @@ c.execute("SELECT DISTINCT test FROM touch WHERE file = ?",
 # function-level (phase 2): tests that entered a specific function/method
 c.execute("SELECT DISTINCT test FROM touch WHERE file = ? AND qualname = ?",
           ("tensorrt_llm/_torch/pyexecutor/py_executor.py", "PyExecutor.forward")).fetchall()
-# coverage rate
+# per-stage: coverage is attributed to the single-GPU stage it came from
+c.execute("SELECT DISTINCT test FROM touch WHERE file = ? AND stage = ?",
+          ("tensorrt_llm/_torch/pyexecutor/py_executor.py", "RTXPro6000D-PyTorch-1")).fetchall()
+# coverage rate + schema version
 dict(c.execute("SELECT key, value FROM meta"))
 ```
 
+A `(test, stage)` is only safe to skip when its coverage is complete — the test passed and
+every process it spawned saved. Force-run anything else:
+
+```python
+# rows NOT safe to skip: non-passed, or a spawned worker/server never saved its coverage
+c.execute("SELECT test, stage FROM test_meta WHERE test != '' AND "
+          "(outcome IS NULL OR outcome != 'passed' OR saved_procs < expected_workers + 1)").fetchall()
+```
+
 ## Smoke test
 
 ```bash

jenkins/scripts/cbts/coverage_utils/cbts_plugin.py

@@ -23,6 +23,8 @@
 
 _PATCHED_MARKER = "_cbts_patched_start_mpi_pool"
 
+_POOL_PATCHED_MARKER = "_cbts_patched_pool_init"
+
 
 def install_mpi_pool_patch(*, raise_on_refactor=True):
     """Widen ``MpiPoolSession._start_mpi_pool``'s env whitelist; idempotent."""
@@ -68,26 +70,53 @@ def _patched_start_mpi_pool(self):
     return True
 
 
-def _switch_test_context(nodeid):
-    """Switch the per-test tracking context via the sitecustomize bootstrap, if active."""
+def install_expected_workers_patch():
+    """Count subprocess pool workers per test (any ``MPIPoolExecutor``) for the completeness signal; idempotent."""
+    try:
+        from mpi4py.futures import MPIPoolExecutor
+    except ImportError:
+        return False
+
+    init = MPIPoolExecutor.__init__
+    if getattr(init, _POOL_PATCHED_MARKER, False):
+        return False
+
+    def _patched_init(self, *args, **kwargs):
+        # Attribute the workers to the test running now; disagg's raw pool is counted here too.
+        try:
+            max_workers = kwargs.get("max_workers", args[0] if args else None)
+            n = int(max_workers) if max_workers else 1
+            _sitecustomize_call("note_expected_workers", os.environ.get("CBTS_TEST_ID", ""), n)
+        except Exception:
+            pass
+        return init(self, *args, **kwargs)
+
+    setattr(_patched_init, _POOL_PATCHED_MARKER, True)
+    MPIPoolExecutor.__init__ = _patched_init
+    return True
+
+
+def _sitecustomize_call(func_name, *args):
+    """Forward to a sitecustomize bootstrap hook (context switch / outcome / worker count), if active."""
     try:
         import sitecustomize
 
-        switch = getattr(sitecustomize, "switch_test_context", None)
+        fn = getattr(sitecustomize, func_name, None)
     except ImportError:
-        switch = None
-    if switch is not None:
-        switch(nodeid)
+        fn = None
+    if fn is not None:
+        fn(*args)
 
 
 # Bind pytest only when already loaded, so importing this module for install_mpi_pool_patch stays cheap.
 if "pytest" in sys.modules:
     import pytest
 
     def pytest_configure(config):  # noqa: D401 - pytest hook
-        """Apply ``mpi_session`` monkeypatch with a compatibility guard."""
+        """Apply the ``mpi_session`` env monkeypatch and the pool-worker accounting patch."""
         del config
         install_mpi_pool_patch(raise_on_refactor=True)
+        install_expected_workers_patch()
 
     @pytest.hookimpl(hookwrapper=True)
     def pytest_runtest_protocol(item, nextitem):  # noqa: D401 - pytest hook
@@ -105,6 +134,16 @@ def pytest_runtest_protocol(item, nextitem):  # noqa: D401 - pytest hook
         # Propagate nodeid via env so subprocesses pick it up in sitecustomize.py.
         os.environ["CBTS_TEST_ID"] = nodeid
 
-        _switch_test_context(nodeid)
+        _sitecustomize_call("switch_test_context", nodeid)
 
         yield
+
+    @pytest.hookimpl(hookwrapper=True)
+    def pytest_runtest_makereport(item, call):  # noqa: D401 - pytest hook
+        """Record each test's outcome so the merge can flag coverage that isn't safe to trust."""
+        del call
+        outcome = yield
+        report = outcome.get_result()
+        # The call phase is the test body; a non-passing setup is the test's effective outcome.
+        if report.when == "call" or (report.when == "setup" and report.outcome != "passed"):
+            _sitecustomize_call("record_test_outcome", item.nodeid, report.outcome)

jenkins/scripts/cbts/coverage_utils/cbts_pystart.py

@@ -33,6 +33,8 @@ def __init__(self, source_roots, data_dir, stage="stage", tool_id=_DEFAULT_TOOL_
         self.tool_id = tool_id
         self._ctx = os.environ.get("CBTS_TEST_ID", "") or ""
         self._data = {}  # context -> set((filename, qualname))
+        self._outcomes = {}  # context -> pytest outcome (filled only in the outer pytest process)
+        self._expected = {}  # context -> pool workers the coordinator spawned for it
         self._file_ok = {}  # co_filename -> bool (cached source-membership)
         self._active = False
         self._new_suffix()
@@ -106,10 +108,21 @@ def switch_test_context(self, nodeid):
         if self._active:
             _MON.restart_events()
 
+    def record_outcome(self, nodeid, outcome):
+        """Record a test's pytest outcome (passed/failed/skipped) for the completeness signal."""
+        self._outcomes[nodeid or ""] = outcome
+
+    def note_expected_workers(self, nodeid, n):
+        """Add to the count of subprocess pool workers the coordinator spawned for a test."""
+        key = nodeid or ""
+        self._expected[key] = self._expected.get(key, 0) + int(n)
+
     def save(self):
         # Write a per-process SQLite the downstream merge reads directly; uploaded compressed only.
         snap = self._data.copy()  # atomic shallow copy; each set snapshotted below
-        if not snap:
+        outcomes = dict(self._outcomes)
+        expected = dict(self._expected)
+        if not snap and not outcomes and not expected:
             return None
         os.makedirs(self.data_dir, exist_ok=True)
         path = os.path.join(self.data_dir, f".cbtscov.{self.stage}.{self._suffix}.sqlite")
@@ -121,6 +134,18 @@ def save(self):
             con.execute("CREATE TABLE touch (test TEXT, file TEXT, qualname TEXT)")
             rows = ((ctx, f, q) for ctx, fs in snap.items() for (f, q) in fs.copy())
             con.executemany("INSERT INTO touch VALUES (?, ?, ?)", rows)
+            # Stage rides in the file content so the merge attributes rows without parsing the filename.
+            con.execute("CREATE TABLE proc_meta (stage TEXT)")
+            con.execute("INSERT INTO proc_meta VALUES (?)", (self.stage,))
+            # Per-test completeness signal; only the coordinator process fills outcome / expected.
+            con.execute(
+                "CREATE TABLE test_meta "
+                "(test TEXT PRIMARY KEY, outcome TEXT, expected_workers INTEGER)"
+            )
+            con.executemany(
+                "INSERT OR REPLACE INTO test_meta VALUES (?, ?, ?)",
+                [(k, outcomes.get(k), expected.get(k, 0)) for k in set(outcomes) | set(expected)],
+            )
             con.commit()
         finally:
             con.close()

jenkins/scripts/cbts/coverage_utils/pystart_report.py

@@ -59,6 +59,35 @@ def _safe_rows(fp):
         return
 
 
+def _read_stage(fp):
+    """Return the stage recorded inside a per-process sqlite, or '' for legacy files."""
+    if not fp.endswith(".sqlite"):
+        return ""
+    try:
+        con = sqlite3.connect(f"file:{fp}?mode=ro", uri=True)
+        try:
+            row = con.execute("SELECT stage FROM proc_meta LIMIT 1").fetchone()
+            return row[0] if row and row[0] else ""
+        finally:
+            con.close()
+    except sqlite3.Error:
+        return ""
+
+
+def _read_test_meta(fp):
+    """Yield (test, outcome, expected_workers) from a per-process sqlite; nothing for legacy files."""
+    if not fp.endswith(".sqlite"):
+        return
+    try:
+        con = sqlite3.connect(f"file:{fp}?mode=ro", uri=True)
+        try:
+            yield from con.execute("SELECT test, outcome, expected_workers FROM test_meta")
+        finally:
+            con.close()
+    except sqlite3.Error:
+        return
+
+
 def merge_to_sqlite(pattern, out_path):
     """Stream every per-process file into a deduped touch DB. Returns (connection, n_data_files)."""
     files = sorted(glob.glob(pattern))
@@ -68,22 +97,63 @@ def merge_to_sqlite(pattern, out_path):
     con.execute("PRAGMA journal_mode=OFF")
     con.execute("PRAGMA synchronous=OFF")
     con.execute(
-        "CREATE TABLE touch (test TEXT, file TEXT, qualname TEXT, UNIQUE(test, file, qualname))"
+        "CREATE TABLE touch (test TEXT, file TEXT, qualname TEXT, stage TEXT, "
+        "UNIQUE(test, file, qualname, stage))"
     )
     con.execute("CREATE TABLE meta (key TEXT PRIMARY KEY, value TEXT)")
+
+    # Per-(test, stage) completeness signal: saved_procs counts the per-process files that
+    # contributed rows; outcome / expected_workers come from the coordinator's test_meta.
+    test_procs = defaultdict(set)
+    test_outcome = {}
+    test_expected = defaultdict(int)
+
     for fp in files:
+        stage = _read_stage(fp)
         batch = []
+        seen_tests = set()
         # Input errors are recovered in _safe_rows; destination-write errors abort the merge.
-        for row in _safe_rows(fp):
-            batch.append(row)
+        for test, file, qual in _safe_rows(fp):
+            seen_tests.add(test)
+            batch.append((test, file, qual, stage))
             if len(batch) >= 20000:
-                con.executemany("INSERT OR IGNORE INTO touch VALUES (?, ?, ?)", batch)
+                con.executemany("INSERT OR IGNORE INTO touch VALUES (?, ?, ?, ?)", batch)
                 batch = []
         if batch:
-            con.executemany("INSERT OR IGNORE INTO touch VALUES (?, ?, ?)", batch)
+            con.executemany("INSERT OR IGNORE INTO touch VALUES (?, ?, ?, ?)", batch)
+        for test in seen_tests:
+            if test:
+                test_procs[(test, stage)].add(fp)
+        for test, outcome, expected in _read_test_meta(fp):
+            if not test:
+                continue
+            if outcome is not None:
+                test_outcome[(test, stage)] = outcome
+            if expected:
+                test_expected[(test, stage)] += int(expected)
+
+    con.execute(
+        "CREATE TABLE test_meta (test TEXT, stage TEXT, outcome TEXT, "
+        "expected_workers INTEGER, saved_procs INTEGER, PRIMARY KEY(test, stage))"
+    )
+    con.executemany(
+        "INSERT OR REPLACE INTO test_meta VALUES (?, ?, ?, ?, ?)",
+        [
+            (
+                t,
+                s,
+                test_outcome.get((t, s)),
+                test_expected.get((t, s), 0),
+                len(test_procs.get((t, s), ())),
+            )
+            for (t, s) in set(test_procs) | set(test_outcome) | set(test_expected)
+        ],
+    )
+
     con.execute("CREATE INDEX ix_file ON touch(file)")
     con.execute("CREATE INDEX ix_func ON touch(file, qualname)")
     con.execute("CREATE INDEX ix_test ON touch(test)")
+    con.execute("CREATE INDEX ix_stage ON touch(stage)")
     con.commit()
     return con, len(files)
 
@@ -144,7 +214,8 @@ def write_html_tree(out_dir, con, rate_line, n_tests, n_data_files):
         # one file's (qualname -> tests) held at a time
         funcs = defaultdict(list)
         for qual, test in con.execute(
-            "SELECT qualname, test FROM touch WHERE file = ? AND test != '' ORDER BY qualname, test",
+            "SELECT DISTINCT qualname, test FROM touch WHERE file = ? AND test != '' "
+            "ORDER BY qualname, test",
             (file,),
         ):
             funcs[qual].append(test)
@@ -208,7 +279,8 @@ def write_json(con, out_path):
         cur_test = None
         touched = []
         rows = con.execute(
-            "SELECT test, file, qualname FROM touch WHERE test != '' ORDER BY test, file, qualname"
+            "SELECT DISTINCT test, file, qualname FROM touch WHERE test != '' "
+            "ORDER BY test, file, qualname"
         )
         for test, file, qual in rows:
             if test != cur_test:
@@ -273,7 +345,24 @@ def main():
         f"{n_tests} tests, {n_files} product files, {n_funcs} functions touched"
     )
 
-    meta = {"tests": str(n_tests), "files": str(n_files), "functions": str(n_funcs)}
+    # Completeness: a (test, stage) is safe to skip only if it passed and every expected
+    # process saved (saved_procs >= expected_workers + 1, the +1 being the coordinator).
+    n_meta = con.execute("SELECT COUNT(*) FROM test_meta WHERE test != ''").fetchone()[0]
+    n_unsafe = con.execute(
+        "SELECT COUNT(*) FROM test_meta WHERE test != '' AND "
+        "(outcome IS NULL OR outcome != 'passed' OR saved_procs < expected_workers + 1)"
+    ).fetchone()[0]
+    print(
+        f"CBTS completeness: {n_meta} (test, stage) rows, "
+        f"{n_unsafe} not safe to skip (non-passed or saved_procs < expected_workers + 1)"
+    )
+
+    meta = {
+        "schema_version": "2",
+        "tests": str(n_tests),
+        "files": str(n_files),
+        "functions": str(n_funcs),
+    }
     rate_line = ""
     if a.source_root:
         coverable_files, all_funcs = enumerate_defs(a.source_root)