[CI] Multi-GPU pytest: dynamic work-stealing + shutdown-hang stack capture

[hujc7]

[CI] Multi-GPU pytest: dynamic work-stealing + shutdown-hang stack capture

Summary

• Adds dynamic work-stealing across multi-GPU pytest shards: a single job spawns N concurrent docker containers (one per non-default GPU), each pulls test files from a shared flock-locked work queue and runs them on its assigned GPU. Beats fixed round-robin sharding for wall time when test durations vary.
• Auto-discovers opt-in test files (any test_*.py whose test_devices() scope advertises a non-default GPU) so adding a test to multi-GPU coverage needs no workflow edit.
• Adds a file-level MULTI_GPU_SKIP_REASON opt-out mechanism for files with known concurrency-only failures; both files keep running in single-GPU CI.
• Adds py-spy + gdb stack capture in tools/conftest.py on shutdown_hang / startup_hang / timeout so the upstream Kit shutdown hang (#3475) is observable in CI logs. Walks the process group; safe no-op when py-spy/gdb are missing.
• Honors ISAACLAB_SIM_DEVICE env var in AppLauncher so the workflow can boot Kit on cuda:1+ without editing each test's AppLauncher(device=...) call site.

1. Scope

This PR is multi-GPU CI infrastructure only. The actual fixes that unblock the skipped tests live in companion PRs:

Fix PR	What it does	When [MGPU] skip marker drops
#5883 — [MGPU] Tests: make newton-only tests kitless (drop AppLauncher boot)	Drops AppLauncher boot from newton-only tests; SimulationContext's has_kit() gate carries them. No Kit lifecycle exposure.	newton test_articulation
#5886 — [MGPU] App: bounded shutdown — SIGHUP handler + force-exit on hang	Adds SIGHUP handler + opt-in ISAACLAB_FORCE_EXIT_TIMEOUT (force-exit timer).	physx test_articulation (when CI sets the env var and the hang is bounded)
#5881 — [MGPU] Sim: honor device kwarg over sim_cfg.device in build_simulation_context	Fixes device-drift in build_simulation_context so newton kernels actually run on the requested cuda:N.	n/a (correctness fix; unblocks downstream tests but no marker drops)

2. CI mechanics

2.1 Dynamic work-stealing

.github/workflows/test-multi-gpu-pytest.yaml runs N parallel docker shards on the same multi-GPU runner. A shared flock-locked queue file in a runtime dir holds test paths; each shard flocks the file, pops the next path, releases, and runs that file on its GPU. No work duplication; balanced wall time even when file durations vary 10x.

2.2 Discover step + opt-in / opt-out

# In .github/workflows/test-multi-gpu-pytest.yaml
mapfile -t candidates < <(grep -rlE 'test_devices\(\)|test_devices\("[^"]*X"\)' source/ --include='test_*.py' | sort -u)
discovered=()
skipped=()
for f in "${candidates[@]}"; do
  if grep -q '^MULTI_GPU_SKIP_REASON' "$f"; then
    skipped+=("$f")
  else
    discovered+=("$f")
  fi
done

Any test that calls test_devices() argless or with a non-default-GPU mask (e.g. "..X") is in scope. To exclude a file from multi-GPU lane while keeping it in single-GPU CI, declare a module-level MULTI_GPU_SKIP_REASON = "...". The discover step prints each excluded file as a ::notice with the reason.

2.3 Hang stack capture

tools/conftest.py:_capture_hang_stacks(pid, pgid, kill_reason) runs before the kill path's SIGKILL. Invokes py-spy dump --pid and gdb -batch -ex "thread apply all bt" -p <pid> against every pid in the process group (cap 8). Each tool is optional; missing tools report inline rather than crashing the diagnostic.

The output gets attached to the JUnit error report so a CI failure surfaces both Python and C++ frames at the moment of the hang — critical because Kit core is closed-source and the hang otherwise terminates with no actionable diagnostic.

Workflow adds --cap-add=SYS_PTRACE to the per-shard docker run (required for both tools to attach) and py-spy to the in-container pip install line (gdb is already in the image).

2.4 AppLauncher honors ISAACLAB_SIM_DEVICE

When the test's AppLauncher() call site doesn't pin a device, the env var becomes the device. Used by the workflow so each shard's container can -e ISAACLAB_SIM_DEVICE=cuda:N without editing test source.

3. Files currently skipped from multi-GPU lane

Both test_articulation.py variants carry the MULTI_GPU_SKIP_REASON marker pointing at #3475:

• source/isaaclab_newton/test/assets/test_articulation.py — marker is removed once #5883 lands (test goes kitless, no Kit lifecycle exposure).
• source/isaaclab_physx/test/assets/test_articulation.py — needs Kit (physx physics is a Kit extension); marker drops once the upstream hang is fixed OR #5886 is set with ISAACLAB_FORCE_EXIT_TIMEOUT=10 on the runner and proven to clear the failure.

4. Validation

Three consecutive CI green at HEAD (before the revert of the cross-PR cherry-picks; the reverts add net-zero file changes vs the green state plus the conftest diagnostics, which only fire on the kill path).

5. Non-scope

• The 3 fix PRs above are NOT included. The PR description for each makes their independent scope explicit.
• The upstream Kit hang itself is not fixed here; it's just instrumented and worked-around at the harness layer.

[isaaclab-review-bot]

---

Update (commit e7d1af7): New commit implements a 2-shard cap as a concurrency-threshold workaround for the Kit lifecycle bug.

Changes in This Commit

1. .github/workflows/test-multi-gpu-pytest.yaml

Shard count is now capped at 2, regardless of available GPUs:

# TEMP EXPERIMENT: cap shards at 2 to test concurrency-threshold workaround for
# the Kit lifecycle bug (SIGHUP/shutdown-hang fires at >=3 concurrent Kit
# instances; 2 is below the threshold per local 2-MIG cross-GPU repro).
shards=$(python3 -c "import json; cap=2; print(json.dumps(list(range(1, min($N, cap+1)))))")
count=$(python3 -c "cap=2; print(min($N-1, cap))")

2. Test Files Re-enabled

Both test_articulation.py files (isaaclab_newton and isaaclab_physx) now have their MULTI_GPU_SKIP_REASON markers commented out to test whether the 2-shard cap avoids the Kit lifecycle bug.

Assessment

✅ Hypothesis-driven debugging — The approach systematically tests the theory that ≥3 concurrent Kit instances trigger the lifecycle bug, while ≤2 are stable.

✅ Conservative cap — 2 shards still provides multi-GPU coverage (cuda:1, cuda:2) while staying below the suspected threshold.

⚠️ GPU Utilization — On 4+ GPU systems, only 2 of the non-default GPUs will be utilized. This is acceptable as a diagnostic step but should be reverted once the root cause is understood.

Comments on Implementation

1. Inline Python for cap logic: The cap calculation uses inline Python:

   count=$(python3 -c "cap=2; print(min($N-1, cap))")

   This works but is slightly fragile. Consider extracting cap=2 to a workflow-level variable for easier adjustment:

   env:
     SHARD_CAP: 2  # TEMP: workaround for Kit lifecycle bug

2. Experiment lifecycle: This is now the third commit in the diagnostic cycle (enable → test → observe). The clear comments document the experiment purpose, which is good practice.

Outstanding Items (Cumulative)

Item	Status	Location
run_docker_tests: 'false' hardcoding	Pre-existing	build.yaml
45-minute timeout (target 30)	Pre-existing	test-multi-gpu-pytest.yaml
2-shard cap	New (diagnostic)	test-multi-gpu-pytest.yaml
Commented MULTI_GPU_SKIP_REASON markers	New (experiment)	Both test_articulation.py

Verdict

LGTM for experiment ✅ — Clean diagnostic change. The 2-shard cap is a reasonable approach to isolate the concurrency threshold. Results from this CI run will inform whether the cap resolves the issue or if the problem lies elsewhere.

Pre-merge checklist (once experiments conclude):

• Revert or formalize the shard cap based on findings
• Restore or remove MULTI_GPU_SKIP_REASON markers accordingly
• Revert run_docker_tests: 'false' in build.yaml
• Adjust 45-minute timeout if validated

---

Update (commit 4bc8b33): The 2-shard cap experiment has concluded. This commit:

1. Removes the shard cap — all available GPUs are now used (shards=$(python3 -c "...range(1, $N)"), count=$((N - 1))). The inline Python is simplified and the bash arithmetic is cleaner.
2. Re-enables MULTI_GPU_SKIP_REASON in both test_articulation.py files — with improved documentation explaining the opt-out mechanism. These files are excluded from concurrent multi-GPU CI while the upstream Kit lifecycle bug persists.

Previous Concerns Status

• ⚠️ Inline Python for cap logic → ✅ Moot — the cap is removed entirely. The remaining inline Python is minimal and the count uses native bash arithmetic.
• ⚠️ GPU Utilization → ✅ Fixed — all non-default GPUs are now utilized.
• 📋 2-shard cap → ✅ Reverted (experiment concluded).
• 📋 Commented MULTI_GPU_SKIP_REASON markers → ✅ Markers restored and properly documented.

Assessment

Clean resolution of the diagnostic cycle. The approach is sound: use all GPUs but exclude known-bad test files via a well-documented constant. The improved comments on MULTI_GPU_SKIP_REASON clearly explain the mechanism and removal criteria. No new issues introduced.

---

Update (commit 3d3f136): New commit adds a fix for build_simulation_context silently ignoring the device kwarg when sim_cfg is also provided.

Changes

1. simulation_context.py — device parameter default changed from "cuda:0" to None. When device is explicitly passed alongside sim_cfg, it now overrides sim_cfg.device. This fixes warp kernel device-mismatch errors on non-default GPUs.
2. Changelog entry added documenting the fix.
3. Test updates — both headless and non-headless test files split the test case to separately verify: (a) sim_cfg values are preserved when no device kwarg is given, and (b) explicit device kwarg wins when both are passed.

Assessment

✅ Correct fix — The None default + explicit override pattern is clean and backward-compatible. Callers not passing device get SimulationCfg's default; callers passing it explicitly get what they asked for.

✅ Well-tested — Both branches (sim_cfg-only and device-override) are covered.

✅ Good documentation — The inline comment explains the rationale and the failure mode it prevents.

⚠️ Scope note — This commit fixes a real bug but is unrelated to the CI multi-GPU sharding work in the PR title. Consider whether this should be a separate PR for cleaner git history, though it is understandable if the bug was discovered during multi-GPU testing.

No blocking issues. LGTM.

---

Update (commit 3bd01c4): Removes MULTI_GPU_SKIP_REASON from both test_articulation.py files (isaaclab_newton and isaaclab_physx), fully re-enabling them in concurrent multi-GPU CI.

This is the logical conclusion of the diagnostic cycle: the build_simulation_context device-override fix from the previous commit likely resolved the warp kernel failures that were triggering Kit lifecycle issues under concurrent execution. With the root cause addressed, the opt-out markers are no longer needed.

Changes

Both files simply remove the 6-line MULTI_GPU_SKIP_REASON constant block:

# File-level opt-out from concurrent multi-GPU CI. The multi-GPU workflow's
# discover step skips any test file declaring this module-level constant.
# ...
MULTI_GPU_SKIP_REASON = "Kit lifecycle bug: ..."

Assessment

✅ Clean removal — No code changes beyond deleting the opt-out marker.

✅ Logical progression — The diagnostic cycle is complete: 2-shard cap (experiment) → device fix (root cause) → re-enable tests (validation).

⚠️ Risk note — If the Kit lifecycle issues resurface in CI, these markers can be quickly re-added. The workflow's MULTI_GPU_SKIP_REASON discovery mechanism remains in place.

Outstanding Pre-merge Items

Item	Status	Notes
run_docker_tests: 'false'	⚠️ Still hardcoded	build.yaml — revert before final merge
45-minute timeout	ℹ️ Diagnostic	Can tighten to 30 once validated
Cherry-picked #5881	ℹ️ Noted in PR description	Will drop on rebase after #5881 merges

No new issues introduced. LGTM. ✅

Update (commit 7d04e41): Adds SIGHUP signal handling and fixes _abort_signal_handle_callback to properly exit after cleanup.

Changes

1. app_launcher.py — Registers SIGHUP handler so child processes shut down cleanly when the parent shell (supervising sibling shards) exits. Rewrites the callback to call sys.exit(128 + signum) after app.close(), preventing Python from resuming with half-torn-down Kit state. Uses contextlib.suppress(Exception) for robustness in signal context.
2. Changelog entry — Documents both fixes clearly.

Assessment

✅ Correct fix — SIGHUP cascading to child shards was the missing piece causing "Stage X already attached" failures. The explicit sys.exit() after cleanup is the right pattern for replaced signal dispositions.

✅ Clean implementation — Both contextlib and sys were already imported. Parameter rename from signal to signum avoids shadowing the module. Comments are thorough.

✅ No new issues introduced.

Outstanding Pre-merge Items

Item	Status	Notes
run_docker_tests: 'false'	⚠️ Still hardcoded	build.yaml — revert before final merge
45-minute timeout	ℹ️ Diagnostic	Can tighten to 30 once validated

---

Update (commit 2b4530d): Two new diagnostic commits (7251b74 + 2b4530d). The latest restores 3 shards and replaces --gpus all with --gpus "device=$cuda" to give each container exclusive access to a single physical GPU. This eliminates the cross-GPU visibility that was triggering the SIGHUP cascade / "Stage already attached" pattern.

Assessment: Clean diagnostic change. The per-container GPU isolation via --gpus device=N is the correct approach — it mirrors MIG-level hardware isolation and prevents Kit processes from interfering with each other's GPU contexts. The explanatory comment is clear and well-reasoned.

No new issues introduced. Previous outstanding items unchanged:

Item	Status	Notes
run_docker_tests: 'false'	⚠️ Still hardcoded	build.yaml — revert before final merge
45-minute timeout	ℹ️ Diagnostic	Can tighten to 30 once validated

---

Update (commit 6137b1b): Significant pivot — instead of per-container GPU isolation (--gpus device=N), the new approach eliminates Kit entirely from newton tests (kitless mode) and re-adds MULTI_GPU_SKIP_REASON for the physx test that still requires Kit.

Changes

1. Workflow — Reverts --gpus "device=$cuda" → --gpus all. The per-container isolation is no longer needed since newton tests no longer boot Kit.

2. schemas.py — New _create_fixed_joint_to_world() helper replaces the from omni.physx.scripts import utils import with pure pxr.UsdPhysics calls. Well-implemented: handles instance proxy / prototype prim climbing, unique joint naming, and correctly sets body0 (world) / body1 (articulation) with proper transforms via UsdGeom.XformCache.

3. Newton test files — test_articulation.py and test_rigid_object_collection.py remove AppLauncher boot and test_devices() import, replacing with hardcoded ["cuda:0", "cpu"] or ["cuda:0"]. This makes them fully kitless, avoiding the Kit lifecycle bug entirely.

4. PhysX test_articulation.py — Re-adds MULTI_GPU_SKIP_REASON to opt out of concurrent multi-GPU CI (this file still requires Kit).

5. Changelog entries — Properly added for the kitless fix.

Assessment

✅ Better root-cause solution — Rather than working around Kit concurrency bugs with container isolation, this removes the Kit dependency where it isn't needed. Kitless newton tests run faster (~30s saved per file) and are immune to Kit lifecycle issues.

✅ _create_fixed_joint_to_world() is well-written — Correctly mirrors the omni.physx.scripts.utils.createJoint single-selection Fixed branch. The prototype/instance climb, unique naming, and transform extraction are all correct.

✅ Clean test migration — Removing AppLauncher is safe for newton-only tests since SimulationContext gates Kit-specific paths on has_kit().

Outstanding Pre-merge Items

Item	Status	Notes
run_docker_tests: 'false'	⚠️ Still hardcoded	build.yaml — revert before final merge
45-minute timeout	ℹ️ Diagnostic	Can tighten to 30 once validated
MULTI_GPU_SKIP_REASON in physx test	ℹ️ Intentional	Stays until upstream Kit fix (IsaacLab #3475)

No new issues introduced. LGTM. ✅

---

Update (commit 6744a2d): Major direction change — reverts the kitless newton approach from the previous commit and instead adds robust hang diagnostics (py-spy + gdb stack capture) to the CI workflow.

Changes (vs 6137b1b)

1. Workflow (.github/workflows/test-multi-gpu-pytest.yaml)

• Adds --cap-add=SYS_PTRACE to the Docker run command, enabling py-spy/gdb to attach to hung Kit processes inside the container.
• Adds py-spy to the in-container pip install list.

2. tools/conftest.py — New _capture_hang_stacks() Function

Captures Python (py-spy) and C++ (gdb) stack traces from all processes in the test's process group (capped at 8 pids) before SIGKILL erases them. Called on shutdown_hang, startup_hang, and timeout detection. Output is appended to the JUnit diagnostic report. Gracefully degrades when py-spy/gdb are unavailable.

Implementation is solid:

• Enumerates process group via ps -o pid= -g <pgid>
• Per-pid captures with sensible timeouts (10s py-spy, 20s gdb)
• Truncates gdb output at 8KB to avoid flooding CI logs
• Safe no-op when tools are missing

3. Reverts kitless newton changes

• app_launcher.py — Reverts to simple self._app.close() callback (removes SIGHUP handler, removes contextlib.suppress + sys.exit(128+signum) pattern)
• schemas.py — Removes the _create_fixed_joint_to_world() helper, reverts to from omni.physx.scripts import utils import
• simulation_context.py — device param reverts to "cuda:0" default, removes the explicit-device-override logic
• Newton test files — Restore AppLauncher boot, test_devices() usage, and MULTI_GPU_SKIP_REASON marker
• Changelog entries for the reverted features are deleted

4. MULTI_GPU_SKIP_REASON in newton test_articulation.py

Re-added with documentation that it should be removed once PR #5883 (kitless newton conversion) lands separately.

Assessment

✅ Good separation of concerns — The kitless newton approach (which is a larger refactor) is being split into its own PR (#5883), while this PR focuses on what it says in the title: multi-GPU CI infrastructure (dynamic work-stealing + hang diagnostics).

✅ Hang diagnostics are well-implemented — _capture_hang_stacks() is robust: handles missing tools, process enumeration failures, and output size limits gracefully. The --cap-add=SYS_PTRACE addition is necessary and minimal.

✅ Clean revert — The AppLauncher signal handling reverts to the simpler pre-existing pattern. The build_simulation_context device override is removed (presumably moving to #5883 as well).

⚠️ Concern: AppLauncher signal handling regression — The previous commit introduced valuable hardening (SIGHUP handler, sys.exit(128+signum) to prevent post-handler resume, contextlib.suppress for robustness in signal context). Reverting this means the Kit half-torn-down issue persists. Presumably this will be addressed in #5883 or a dedicated fix, but worth tracking.

⚠️ build_simulation_context device-override revert — The device kwarg no longer overrides sim_cfg.device when both are passed. The original bug (warp kernel device-mismatch on non-default GPUs) may resurface in multi-GPU test scenarios. The test changes confirm this is intentional ("values don't get overridden"), but callers passing both kwargs will silently get sim_cfg's device.

Outstanding Pre-merge Items

Item	Status	Notes
run_docker_tests: 'false'	⚠️ Still hardcoded	build.yaml — revert before final merge
45-minute timeout	ℹ️ Diagnostic	Can tighten to 30 once validated
MULTI_GPU_SKIP_REASON in newton test	ℹ️ Temporary	Drops when #5883 lands
AppLauncher SIGHUP / exit hardening	⚠️ Reverted	Track for #5883 or separate PR
build_simulation_context device override	⚠️ Reverted	Track for #5883 or separate PR

Verdict

LGTM for the focused scope ✅ — The PR now cleanly addresses its stated goal: multi-GPU CI improvements with hang diagnostics. The kitless newton refactor is correctly being separated into #5883. The diagnostic stack capture is a strong addition that will make Kit shutdown hangs observable. The reverts are intentional scope reduction, not regressions (assuming the features land via #5883).

Update (commit 6f09e08): Adds ISAACLAB_PIN_KIT_GPU environment variable support — a clean, targeted fix for the Kit multi-GPU renderer interference that was causing the shutdown hangs.

Changes

1. app_launcher.py — New ISAACLAB_PIN_KIT_GPU env var handling

When ISAACLAB_PIN_KIT_GPU is set to a truthy value (anything not in {"", "0", "false", "no", "off"}), appends three Kit command-line flags:

• --/renderer/multiGpu/enabled=False
• --/renderer/multiGpu/autoEnable=False
• --/renderer/multiGpu/maxGpuCount=1

This pins each Kit process to its assigned GPU only, preventing the shared cubric / PhysX-fabric GPU-interop context that was causing [Error] [omni.physx.plugin] Stage X already attached and SimulationApp.close hangs.

Implementation is clean: Placed logically after the _resolve_device_settings method's existing GPU assignment. The truthy-value check handles common falsy strings. Logger info message aids debugging.

2. Workflow (.github/workflows/test-multi-gpu-pytest.yaml)

Adds -e ISAACLAB_PIN_KIT_GPU=1 to the Docker run command. Simple, non-invasive.

3. test_articulation.py (isaaclab_physx)

Removes MULTI_GPU_SKIP_REASON — re-enabling the test in concurrent multi-GPU CI. With Kit pinned to a single GPU per process, the concurrency bug should no longer trigger.

4. Changelog entry (jichuanh-mgpu-pin-kit-resources.rst)

Well-written: documents the env var, explains the Kit renderer behavior it overrides, references the issue (#3475) and NVBug, and clarifies it's off by default.

Assessment

✅ Excellent root-cause fix — Instead of working around the problem (shard caps, container isolation, removing Kit entirely), this directly addresses the Kit renderer's multi-GPU enumeration that was causing cross-process interference. Kelly Guo's documented WAR is properly attributed.

✅ Opt-in design — Off by default, only activated in CI via env var. Single-GPU and user-facing rendering paths are completely unchanged. No risk to existing workflows.

✅ Minimal footprint — 4 lines of functional code in app_launcher.py + 1 env var in the workflow. Clean and auditable.

✅ Enables test re-enablement — The physx test_articulation.py can now run in concurrent multi-GPU CI since the underlying cause is addressed.

Minor Observations

1. Comment references internal person — The inline comment in app_launcher.py mentions "Kelly Guo's documented WAR (#omni-kit thread, 2024-2025)". Consider whether this internal reference is appropriate for an open-source codebase, or if just the NVBug reference suffices.

2. Newton test_articulation.py unchanged — The newton equivalent still has MULTI_GPU_SKIP_REASON (presumably waiting for #5883 kitless conversion). This is fine and consistent with the PR's scope.

Outstanding Pre-merge Items (Updated)

Item	Status	Notes
run_docker_tests: 'false'	⚠️ Still hardcoded	build.yaml — revert before final merge
45-minute timeout	ℹ️ Diagnostic	Can tighten to 30 once validated
Newton MULTI_GPU_SKIP_REASON	ℹ️ Expected	Drops when #5883 lands
AppLauncher SIGHUP / exit hardening	ℹ️ Tracked	For #5883 or separate PR
build_simulation_context device override	ℹ️ Tracked	For #5883 or separate PR

Verdict

LGTM ✅ — This is the cleanest solution in the PR's evolution. Rather than elaborate workarounds, it applies the known Kit WAR (disable multi-GPU renderer enumeration) via an opt-in env var. The implementation is minimal, well-documented, and correctly scoped. Ready for merge pending CI validation of the re-enabled physx test.