Skip to content

misc: jit: Replace parallel_load_modules() with build_jit_specs() #1070

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 22, 2025

Conversation

abcdabcd987
Copy link
Member

Part of AOT Refactor (#1064).

Now that build_jit_specs() can build all ops in one ninja run, there's no need for parallel_load_modules() in tests.

Also fixing a bad pattern:

try:
    flashinfer.jit.parallel_load_modules(...)
except Exception as e:
    pytest.exit(str(e))
finally:
    yield

This will not stop pytest when error is raised, because there's a finally clause. Also, the default exit code of pytest.exit() is 0, so it doesn't indicate an error.

An easy fix would be just to get rid of the whole try-except-finally structure.

@abcdabcd987 abcdabcd987 requested a review from yzh119 May 19, 2025 21:48
@abcdabcd987 abcdabcd987 mentioned this pull request May 19, 2025
10 tasks
@abcdabcd987 abcdabcd987 force-pushed the lequn/split-gen-jitspec branch from 4556311 to e863698 Compare May 20, 2025 18:59
@abcdabcd987 abcdabcd987 force-pushed the lequn/remove-parallel-load-modules branch from d4d5497 to 5ad71da Compare May 20, 2025 19:05
@abcdabcd987 abcdabcd987 force-pushed the lequn/split-gen-jitspec branch from e863698 to 2d38830 Compare May 21, 2025 20:32
@abcdabcd987 abcdabcd987 force-pushed the lequn/remove-parallel-load-modules branch from 5ad71da to 20892f7 Compare May 21, 2025 20:33
@abcdabcd987 abcdabcd987 force-pushed the lequn/split-gen-jitspec branch from 2d38830 to cfd4539 Compare May 21, 2025 22:38
Base automatically changed from lequn/split-gen-jitspec to main May 22, 2025 00:33
@abcdabcd987 abcdabcd987 force-pushed the lequn/remove-parallel-load-modules branch from 20892f7 to ed480f6 Compare May 22, 2025 00:37
Copy link
Collaborator

@yzh119 yzh119 left a comment

Choose a reason for hiding this comment

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

As we implemented build_jit_specs to generate ninja for a set of specs (and then ninja could compile the files in parallel), we should no longer rely on multi-threading to build different extensions.

@yzh119 yzh119 merged commit 98995bf into main May 22, 2025
1 of 2 checks passed
@yzh119 yzh119 deleted the lequn/remove-parallel-load-modules branch May 22, 2025 00:42
Edenzzzz pushed a commit to Edenzzzz/flashinfer that referenced this pull request Jun 6, 2025
…ashinfer-ai#1070)

Part of AOT Refactor (flashinfer-ai#1064).

Now that `build_jit_specs()` can build all ops in one ninja run, there's
no need for `parallel_load_modules()` in tests.

Also fixing a bad pattern:

```python
try:
    flashinfer.jit.parallel_load_modules(...)
except Exception as e:
    pytest.exit(str(e))
finally:
    yield
```

This will not stop pytest when error is raised, because there's a
`finally` clause. Also, the default exit code of `pytest.exit()` is 0,
so it doesn't indicate an error.

An easy fix would be just to get rid of the whole `try-except-finally`
structure.
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.

2 participants