[sdf] Fix race condition by filelock

iory · Jan 28, 2025 · aad0d72 · aad0d72
1 parent e92b970
commit aad0d72
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,6 @@
 cached-property;python_version>="3.0"
 cached-property<=1.5.2;python_version<"3.0"
+filelock
 future
 gdown
 lxml

diff --git a/skrobot/sdf/signed_distance_function.py b/skrobot/sdf/signed_distance_function.py
@@ -1,5 +1,6 @@
 from __future__ import division
 
+import filelock
 from logging import getLogger
 from math import floor
 import os
@@ -520,13 +521,18 @@ def from_objfile(obj_filepath, dim_grid=100, padding_grid=5, **kwargs):
             checksum_md5(obj_filepath), dim_grid, padding_grid)
 
         sdf_cache_path = os.path.join(sdf_cache_dir, hashed_filename + '.sdf')
-        if not os.path.exists(sdf_cache_path):
-            logger.info(
-                'pre-computing sdf and making a cache at {0}.'
-                .format(sdf_cache_path))
-            pysdfgen.obj2sdf(str(obj_filepath), dim_grid, padding_grid,
-                             output_filepath=sdf_cache_path)
-            logger.info('finish pre-computation')
+        lock = filelock.FileLock(sdf_cache_path + '.lock')
+        with lock:
+            if not os.path.exists(sdf_cache_path):
+                logger.info(
+                     'trying to acquire lock for {0}...'
+                     .format(sdf_cache_path))
+                logger.info(
+                    'pre-computing sdf and making a cache at {0}.'
+                    .format(sdf_cache_path))
+                pysdfgen.obj2sdf(str(obj_filepath), dim_grid, padding_grid,
+                                 output_filepath=sdf_cache_path)
+                logger.info('finish pre-computation')
         return GridSDF.from_file(sdf_cache_path, **kwargs)