write to tmp file & replace, avoid writing a module per process in PointwiseDynamic#532
Merged
iclementine merged 5 commits intoflagos-ai:masterfrom May 16, 2025
Merged
Conversation
…cFunction, add test for multiprocessing & multithreading
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR refactors code caching to use file locks instead of per-process module files and adds tests for concurrent execution of pointwise dynamic functions.
- Introduce
filelockto guard writes and remove PID suffix from generated filenames inpointwise_dynamicand various ops. - Update module loading to use stable filenames and drop PID from spec names.
- Add multithreading and multiprocessing tests for
pointwise_dynamicto ensure correct behavior under concurrency.
Reviewed Changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_pointwise_dynamic.py | Add multithread/multiprocess tests for pointwise dynamic ops |
| src/flag_gems/utils/pointwise_dynamic.py | Use filelock around code cache writes; remove PID suffix |
| src/flag_gems/ops/tile.py | Add file locking and drop PID in cache filenames |
| src/flag_gems/ops/scatter.py | Add file locking and drop PID in cache filenames |
| src/flag_gems/ops/repeat.py | Add file locking and drop PID in cache filenames |
| src/flag_gems/ops/pad.py | Add file locking and drop PID in cache filenames |
| src/flag_gems/ops/index_put.py | Add file locking and drop PID in cache filenames |
| src/flag_gems/ops/gather.py | Add file locking and drop PID in cache filenames |
Comments suppressed due to low confidence (1)
src/flag_gems/utils/pointwise_dynamic.py:1263
- Since filenames no longer include a process-unique suffix, the cache may serve stale code when regeneration is needed. Consider incorporating a hash of
code.getvalue()into the filename or invalidating the cache when the code changes.
file_name = ( f"pointwise_dynamic_{self._scalar_fn_cache_key}_{kernel_name}_... .py")
Comment on lines
+1273
to
+1276
| with filelock.FileLock(lock_path): | ||
| if not os.path.exists(file_path): | ||
| with open(file_path, "wt", encoding="utf-8") as f: | ||
| f.write(code.getvalue()) |
There was a problem hiding this comment.
[nitpick] The file-locking and write logic is duplicated across multiple operators. Extract this pattern into a shared utility to reduce duplication and simplify future updates.
Suggested change
| with filelock.FileLock(lock_path): | |
| if not os.path.exists(file_path): | |
| with open(file_path, "wt", encoding="utf-8") as f: | |
| f.write(code.getvalue()) | |
| write_with_lock(lock_path, file_path, code.getvalue()) |
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
iclementine
changed the title
nicelynice
pushed a commit
to nicelynice/FlagGems
that referenced
this pull request
Feb 24, 2026
flagos-ai#532) * write to tmp file & os.replace, so as to avoid writing a module per process in PointwiseDynamicFunction, add test for multiprocessing & multithreading * update for other operators * Update tests/test_pointwise_dynamic.py * use os.replace to write the same contents to the same path concurrently --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR Category
Other
Type of Change
Other
Description
Previously, PointwiseDynamic woule generate module for pointwise functions with the process id as suffix in filename to avoid conflicts in file-writing. It has 2 drawbacks: filenames with pid prevent file reuse between different runs, it generate duplicated files with different names, which is redundant.
Issue
Progress
Performance