[Driver] Make compilation more compatible with multi-processing #350

python/hidet/drivers/build_task.py

@@ -24,7 +24,7 @@
 from hidet.ir.module import IRModule
 from hidet.ir.task import Task
 from hidet.drivers.build_module import build_ir_module, build_ir_module_batch
-from hidet.drivers.utils import lazy_initialize_cuda
+from hidet.drivers.utils import lazy_initialize_cuda, CompileLock
 from hidet.runtime.compiled_module import compiled_module_exists
 from hidet.runtime.compiled_task import CompiledTask, TensorSignature, load_compiled_task, compiled_task_cache
 from hidet.runtime.device import Device
@@ -250,42 +250,43 @@ def build_task(task: Task, target='cuda', load=True) -> Optional[CompiledTask]:
         lib_path = os.path.join(task_dir, 'lib.so')
         version_path = os.path.join(task_dir, 'version.txt')
 
-        version_matched = False
-        if os.path.exists(version_path):
-            with open(version_path, 'r') as f:
-                version = f.read()
-                if version.strip() == hidet.__version__:
-                    version_matched = True
-
-        # use previously generated library when available
-        if use_cache and version_matched and compiled_module_exists(task_dir):
-            logger.debug(f"Load cached task binary {green(task.name)} from path: \n{cyan(lib_path)}")
-            if load:
-                compiled_task = load_compiled_task(task_dir)
-                compiled_task_cache.add(target, space_level, task_string, compiled_task)
-        else:
-            logger.info(f"Compiling {target} task {green(task.signature())}...")
+        with CompileLock(os.path.join(task_dir, "compile.lock"), enabled=use_cache):
+            version_matched = False
+            if os.path.exists(version_path):
+                with open(version_path, 'r') as f:
+                    version = f.read()
+                    if version.strip() == hidet.__version__:
+                        version_matched = True
+
+            # use previously generated library when available
+            if use_cache and version_matched and compiled_module_exists(task_dir):
+                logger.debug(f"Load cached task binary {green(task.name)} from path: \n{cyan(lib_path)}")
+                if load:
+                    compiled_task = load_compiled_task(task_dir)
+                    compiled_task_cache.add(target, space_level, task_string, compiled_task)
+            else:
+                logger.info(f"Compiling {target} task {green(task.signature())}...")
 
-            # build from scratch
-            os.makedirs(task_dir, exist_ok=True)
+                # build from scratch
+                os.makedirs(task_dir, exist_ok=True)
 
-            # write task
-            with open(os.path.join(task_dir, 'task.txt'), 'w') as f:
-                f.write(task_string)
+                # write task
+                with open(os.path.join(task_dir, 'task.txt'), 'w') as f:
+                    f.write(task_string)
 
-            # write version
-            with open(version_path, 'w') as f:
-                f.write(hidet.__version__)
+                # write version
+                with open(version_path, 'w') as f:
+                    f.write(hidet.__version__)
 
-            # implement task to IRModule, each task may produce multiple IRModules (candidates)
-            # they have the same functionality but different performance
-            candidates = task.implement(target=target, working_dir=task_dir)
+                # implement task to IRModule, each task may produce multiple IRModules (candidates)
+                # they have the same functionality but different performance
+                candidates = task.implement(target=target, working_dir=task_dir)
 
-            # generate meta data
-            generate_meta_data(task, task_dir, target, len(candidates))
+                # generate meta data
+                generate_meta_data(task, task_dir, target, len(candidates))
 
-            # construct the ir module for the task
-            build_task_module(task, candidates, task_dir, target)
+                # construct the ir module for the task
+                build_task_module(task, candidates, task_dir, target)
 
             if load:
                 compiled_task = load_compiled_task(task_dir)

python/hidet/drivers/utils.py

@@ -10,6 +10,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+import os
+from filelock import FileLock
 
 import hidet.cuda
 
@@ -36,3 +38,27 @@ def lazy_initialize_cuda():
         for i in range(hidet.cuda.device_count()):
             hidet.cuda.properties(i)
             hidet.cuda.compute_capability(i)
+
+
+class CompileLock:
+    # There are cases where multiple instances of the same task or IRModule will be built at the same time, especially
+    # in the cases of distributed inference. We use this lock to make sure that a single task is not built multiple
+    # times.
+
+    def __init__(self, lock_path: str, enabled=True):
+        self.enabled = enabled
+        self.lock_path = lock_path
+        self.lock = None
+
+    def __enter__(self):
+        if not self.enabled:
+            return None
+
+        os.makedirs(os.path.dirname(self.lock_path), exist_ok=True)
+        self.lock = FileLock(self.lock_path)
+        self.lock.acquire()
+        return self.lock
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        if self.lock:
+            self.lock.release()

python/hidet/graph/frontend/torch/interpreter.py

@@ -359,6 +359,10 @@ def load_arg(a, env):
                 hidet_kwargs = load_arg(node.kwargs, hidet_env)
                 try:
                     hidet_env[node.name] = exec_func(*hidet_args, **hidet_kwargs)
+                    from .register_functions import setitem
+
+                    if exec_func.functions[0] is setitem:
+                        hidet_env[str(node.args[0])] = hidet_env[node.name]
                 except Exception as e:
                     self._raise_exception(e, node.target, exec_func, hidet_args, hidet_kwargs)
             elif node.op == "call_method":
@@ -448,6 +452,10 @@ def load_arg(a, env):
 
                 try:
                     hidet_env[node.name] = hidet_func(*hidet_args, **hidet_kwargs)
+                    from .register_functions import setitem
+
+                    if hidet_func.functions[0] is setitem:
+                        hidet_env[str(node.args[0])] = hidet_env[node.name]
                 except Exception as e:
                     self._raise_exception(e, node.target, hidet_func, hidet_args, hidet_kwargs)