Fix AutoMate SoftDTW import fallback

[ooctipus]

Summary

• Make the AutoMate SoftDTW module importable when the optional numba.cuda backend fails to import.
• Fall back to the existing CPU SoftDTW implementation when SoftDTW(use_cuda=True) is requested but the CUDA backend is unavailable.
• Add a focused regression test that simulates the reported numba.cuda.types.NPDatetime import failure.

Reproduction / Verification

• Checked the provided failure path: assembly_env.py imports SoftDTW from soft_dtw_cuda.py, which previously imported and decorated numba.cuda kernels at module import time.
• In an existing beta2 virtualenv with broken Numba imports, verified the patched soft_dtw_cuda.py imports successfully and a small SoftDTW(use_cuda=True) call falls back to CPU.

Tests

• git diff --cached --check
• Parsed modified Python files with ast.parse
• /home/zhengyuz/Projects/IsaacLab.wt/develop/.venv/bin/python -m pytest source/isaaclab_tasks/test/contrib/test_automate_soft_dtw_cuda.py -q
• /home/zhengyuz/Projects/IsaacLab.wt/beta2-automate-collision-stack/.venv/bin/python -m pytest source/isaaclab_tasks/test/contrib/test_automate_soft_dtw_cuda.py -q
• Direct patched-module import/fallback check in beta2-automate-collision-stack/.venv

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

PR #6039 — Fix AutoMate SoftDTW import fallback

Summary

This PR adds a graceful fallback when numba.cuda fails to import, allowing the AutoMate SoftDTW module to still load and fall back to the CPU implementation. The approach is sound and addresses a real import failure scenario.

✅ Strengths

1. Layered import handling — Separates numba (jit/prange) and numba.cuda imports, catching both ImportError and AttributeError. This correctly handles the reported NPDatetime attribute error.
2. Graceful degradation — When CUDA is unavailable, SoftDTW(use_cuda=True) emits a one-time RuntimeWarning and silently falls back to CPU instead of crashing.
3. _identity_jit design — Correctly handles both bare @jit and @jit(...) decorator patterns.
4. Regression test — Uses monkeypatch to simulate the failure without requiring a broken Numba install. Properly asserts the warning and fallback behavior.
5. Changelog fragment — Included per project conventions.

⚠️ Findings (2 minor)

1. CI: ruff format failure — The pre-commit check failed because 1 file needs reformatting. Please run ruff format on the changed files and push the fix.

2. Global mutable state for warning suppression — _CUDA_FALLBACK_WARNED uses global to ensure the warning fires only once. This works but consider using warnings.warn with the same message text (Python\u2019s warning filter deduplicates by default via default action), which would eliminate the need for the global flag entirely. This is a style suggestion, not a blocker.

🏁 Verdict

Approve (with minor formatting fix needed)

The logic is correct and well-structured. The only blocking issue is the pre-commit formatting failure — once that\u2019s resolved, this is ready to merge.

---

Automated review by IsaacLab Review Bot • Commit: cdda369

[greptile-apps]

Greptile Summary

This PR makes the AutoMate soft_dtw_cuda.py module importable when numba or numba.cuda fails to load (e.g. the reported NPDatetime attribute error), and silently degrades to the CPU _SoftDTW implementation rather than crashing at import time.

• Graceful import fallback: numba.jit/prange and numba.cuda are each wrapped in try/except (AttributeError, ImportError) blocks; stubs (_identity_jit, _UnavailableCuda) keep the module-level @cuda.jit decorators functional without Numba.
• Runtime warning + CPU fallback: _get_func_dtw checks _CUDA_IMPORT_ERROR at call time and emits a one-shot RuntimeWarning before switching use_cuda to False, so users are informed once without spam.
• Regression test: A focused pytest loads the module fresh via importlib, monkeypatches _CUDA_IMPORT_ERROR, and asserts the warning and CPU-path selection are both correct.

Confidence Score: 4/5

Safe to merge; the fallback path is well-guarded and the CUDA code paths are unchanged when Numba is fully available.

The only finding is that stacklevel=2 on the warning points to the internal forward() call rather than the user's call site — a minor usability gap that does not affect correctness or the fallback behavior itself.

No files require special attention; soft_dtw_cuda.py is the only non-trivial change and its fallback logic is consistent throughout.

Important Files Changed

Filename	Overview
source/isaaclab_tasks/isaaclab_tasks/contrib/automate/soft_dtw_cuda.py	Wraps numba/numba.cuda imports in try/except with identity-JIT stubs and an _UnavailableCuda sentinel; adds a one-shot RuntimeWarning and CPU fallback in _get_func_dtw when CUDA import failed. Logic is sound; stacklevel on the warning points at the internal forward() frame rather than user code.
source/isaaclab_tasks/test/contrib/test_automate_soft_dtw_cuda.py	Regression test that loads the module fresh via importlib, patches _CUDA_IMPORT_ERROR to simulate the reported failure, and asserts both the warning and the CPU fallback are applied correctly.
source/isaaclab_tasks/changelog.d/fix-automate-softdtw-numba-cuda-import.rst	Changelog fragment describing the fix; no code concerns.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Module import] --> B{import numba jit/prange}
    B -- success --> C[use real jit / prange]
    B -- AttributeError / ImportError --> D[jit = _identity_jit\nprange = range]
    C --> E{import numba.cuda}
    D --> E
    E -- success --> F[cuda = real numba.cuda\n_CUDA_IMPORT_ERROR = None]
    E -- AttributeError / ImportError --> G[cuda = _UnavailableCuda\n_CUDA_IMPORT_ERROR = exc]
    F --> H["@cuda.jit decorators applied\n(real CUDA kernels compiled)"]
    G --> I["@cuda.jit decorators applied\n(identity no-op, plain Python fns)"]
    H --> J[SoftDTW._get_func_dtw called]
    I --> J
    J --> K{use_cuda and\n_CUDA_IMPORT_ERROR?}
    K -- Yes --> L[emit one-shot RuntimeWarning\nset use_cuda = False]
    L --> M[return _SoftDTW.apply\nCPU path]
    K -- No --> N{use_cuda and\nseq_len > 1024?}
    N -- Yes --> O[print msg\nset use_cuda = False\nreturn _SoftDTW.apply]
    N -- No, use_cuda=True --> P[return _SoftDTWCUDA.apply\nGPU path]
    N -- No, use_cuda=False --> M

Loading

_{Reviews (1): Last reviewed commit: "Fix AutoMate SoftDTW import fallback" | Re-trigger Greptile}

[greptile-apps]

stacklevel=2 makes the warning point at the internal forward() call inside this module rather than the user's code that invoked the criterion. The call stack is: user code → forward() → _get_func_dtw() → warnings.warn(). Using stacklevel=3 would surface the warning at the correct site in user code.

Suggested change

	warnings.warn(
	"SoftDTW CUDA backend is unavailable because numba.cuda failed to import "
	f"({type(_CUDA_IMPORT_ERROR).__name__}: {_CUDA_IMPORT_ERROR}). "
	"Falling back to CPU SoftDTW.",
	RuntimeWarning,
	stacklevel=2,
	)
	warnings.warn(
	"SoftDTW CUDA backend is unavailable because numba.cuda failed to import "
	f"({type(_CUDA_IMPORT_ERROR).__name__}: {_CUDA_IMPORT_ERROR}). "
	"Falling back to CPU SoftDTW.",
	RuntimeWarning,
	stacklevel=3,
	)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[ooctipus]

Closing this implementation. The runtime fallback was too broad; the root cause is the uv dependency override allowing an incompatible pre-release Numba/NumPy stack. Replacing with a narrower dependency-constraint fix.