Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions transformer_engine/common/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -245,11 +245,13 @@ def _nvidia_cudart_include_dir() -> str:
return ""

# Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia"
# above doesn't through. However, they don't set "__file__" attribute.
if nvidia.__file__ is None:
return ""
# above doesn't throw. However, they don't set "__file__" attribute.
if nvidia.__file__ is not None:
nvidia_root = Path(nvidia.__file__).parent
else:
nvidia_root = Path(nvidia.__path__[0]) # namespace package
Comment on lines +250 to +252
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unconditional __path__[0] indexing
If nvidia.__path__ exists but is empty (or contains a non-path-like entry), nvidia.__path__[0] will raise IndexError/TypeError, which would regress the “clear error message” intent. Consider guarding for emptiness and using the same explicit RuntimeError path.

Suggested change
nvidia_root = Path(nvidia.__file__).parent
else:
nvidia_root = Path(nvidia.__path__[0]) # namespace package
nvidia_dir = list(nvidia.__path__)[0]

→ check nvidia.__path__ is truthy/non-empty before indexing and raise the intended error when it isn’t.


include_dir = Path(nvidia.__file__).parent / "cuda_runtime"
include_dir = nvidia_root / "cuda_runtime"
return str(include_dir) if include_dir.exists() else ""


Expand Down
Loading