add filelock

depyf/explain/enable_debugging.py

@@ -130,10 +130,12 @@ def prepare_debug(func, dump_src_dir, clean_wild_fx_code=True, pause=True):
                             f"Please set breakpoints in {dump_src_dir}. Then press enter to continue.")
                 else:
                     from depyf.explain import dump_src
+                    from depyf.explain.utils import write_code_to_file_template
+                    from torch._dynamo.eval_frame import innermost_fn
+                    func = innermost_fn(func)
                     full_src = dump_src(func)
-                    filename = os.path.join(dump_src_dir, f"full_code.py")
-                    with open(filename, "w") as f:
-                        f.write(full_src)
+                    filename_template = os.path.join(dump_src_dir, f"full_code_for_{func.__code__.co_name}_%s.py")
+                    filename = write_code_to_file_template(full_src, filename_template)
 
                     if pause:
                         input(

depyf/explain/utils.py

@@ -249,22 +249,38 @@ def remove_indentation(code: str):
     indent = len(lines[0]) - len(lines[0].lstrip())
     return "".join([line[indent:] + "\n" for line in lines])
 
+from contextlib import contextmanager
 
-def write_code_to_file_template(src, path_template):
+@contextmanager
+def lock_on_file(lock_path):
+    from filelock import FileLock
     import os
-    count = 0
-    while True:
-        new_filepath = path_template % str(count)
-        if not os.path.exists(new_filepath):
-            with open(new_filepath, "w") as f:
-                f.write(src)
-            break
-        # might be a hash collision
-        existing_code = open(new_filepath).read()
-        if existing_code == src:
-            break
-        count += 1
-    return new_filepath
+    lock = FileLock(lock_path)
+    try:
+        with lock:
+            yield
+    finally:
+        os.remove(lock_path)
+
+
+def write_code_to_file_template(src, path_template):
+    lock_path = path_template % "lock"
+
+    with lock_on_file(lock_path):
+        import os
+        count = 0
+        while True:
+            new_filepath = path_template % str(count)
+            if not os.path.exists(new_filepath):
+                with open(new_filepath, "w") as f:
+                    f.write(src)
+                break
+            # might be a hash collision
+            existing_code = open(new_filepath).read()
+            if existing_code == src:
+                break
+            count += 1
+        return new_filepath
 
 
 def get_code_owner(fn):