Skip to content

[fix] Fix top-level CUDAGraph export - #8

Open
matteso1 wants to merge 1 commit into
foundry-org:mainfrom
matteso1:fix-top-level-cudagraph-export
Open

[fix] Fix top-level CUDAGraph export#8
matteso1 wants to merge 1 commit into
foundry-org:mainfrom
matteso1:fix-top-level-cudagraph-export

Conversation

@matteso1

Copy link
Copy Markdown

The bug

The top-level package imports the Python CUDAGraph wrapper and then applies from .ops import *. The native extension exports the same name. Its wildcard import therefore replaces the wrapper. As a result, foundry.CUDAGraph resolves to the raw binding instead of the documented Python wrapper.

The fix

Import the native re-exports first, then explicitly re-export foundry.graph.CUDAGraph. Add a regression to the existing import test that checks the top-level name resolves to the Python wrapper.

Test plan

  • pytest -q tests/test_imports.py::test_cudagraph_reexport_uses_python_wrapper
  • ruff check python/foundry/__init__.py tests/test_imports.py
  • ruff format --check python/foundry/__init__.py tests/test_imports.py

The focused test passed in the built Linux environment. On one RTX 4090 Laptop GPU, the low-level SAVE oracle passed in 4.2557 seconds, n=1. The LOAD oracle passed in 2.3335 seconds, n=1.

I used AI assistance. I reviewed the change and can defend the implementation and verification.

Signed-off-by: Nils Matteson <nilsmatteson@icloud.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant