write to tmp file & replace, avoid writing a module per process in PointwiseDynamic

[iclementine]

PR Category

Other

Type of Change

Other

Description

1. use filelock to avoid writing a module per process in PointwiseDynamiccFunction,
2. add tests for multiprocessing & multithreading

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

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

[Copilot]

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 filelock to guard writes and remove PID suffix from generated filenames in pointwise_dynamic and various ops.
• Update module loading to use stable filenames and drop PID from spec names.
• Add multithreading and multiprocessing tests for pointwise_dynamic to 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")

[Copilot]

[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())

[Bowen12992]

LGTM

[Bowen12992]

LGTM