add lock

python/jittor/compile_extern.py

@@ -7,6 +7,7 @@ import os, sys, shutil
 from .compiler import *
 from jittor_utils import run_cmd, get_version
 from jittor.dataset.utils import download_url_to_local
+from jittor.lock import jittor_lock
 
 def search_file(dirs, name):
     for d in dirs:
@@ -373,6 +374,8 @@ def setup_mpi():
     LOG.vv("Get mpi: "+str(mpi.__dict__.keys()))
     LOG.vv("Get mpi_ops: "+str(mpi_ops.__dict__.keys()))
 
+jittor_lock.lock()
+
 setup_mpi()
 setup_nccl()
 
@@ -380,3 +383,5 @@ setup_cutt()
 setup_mkl()
 
 setup_cuda_extern()
+
+jittor_lock.unlock()

python/jittor/compiler.py

@@ -17,6 +17,7 @@ from ctypes.util import find_library
 import jittor_utils as jit_utils
 from jittor_utils import LOG, run_cmd, cache_path, find_exe, cc_path, cc_type, cache_path
 from . import pyjt_compiler
+from jittor.lock import jittor_lock
 
 def find_jittor_path():
     return os.path.dirname(__file__)
@@ -798,15 +799,20 @@ import_flags = os.RTLD_NOW | os.RTLD_GLOBAL | os.RTLD_DEEPBIND
 #     import_flags = os.RTLD_NOW | os.RTLD_GLOBAL
 dlopen_flags = os.RTLD_NOW | os.RTLD_GLOBAL | os.RTLD_DEEPBIND
 
+jittor_lock.lock()
 with jit_utils.import_scope(import_flags):
     jit_utils.try_import_jit_utils_core()
+jittor_lock.unlock()
+
 jittor_path = find_jittor_path()
 check_debug_flags()
 
 sys.path.append(cache_path)
 
+jittor_lock.lock()
 with jit_utils.import_scope(import_flags):
     jit_utils.try_import_jit_utils_core()
+jittor_lock.unlock()
 
 python_path = sys.executable
 py3_config_path = sys.executable+"-config"
@@ -848,11 +854,13 @@ make_cache_dir(os.path.join(cache_path, "jit"))
 make_cache_dir(os.path.join(cache_path, "obj_files"))
 make_cache_dir(os.path.join(cache_path, "gen"))
 
+jittor_lock.lock()
 # build cache_compile
 cc_flags += pybind_include
 cc_flags += f" -I{jittor_path}/src "
 check_cache_compile()
 LOG.v(f"Get cache_compile: {jit_utils.cc}")
+jittor_lock.unlock()
 
 # check cuda
 has_cuda = 0
@@ -877,6 +885,7 @@ if has_cuda:
         return nvcc_flags
     nvcc_flags = convert_nvcc_flags(nvcc_flags)
 
+jittor_lock.lock()
 # build core
 gen_jit_flags()
 gen_jit_tests()
@@ -966,3 +975,4 @@ flags.jittor_path = jittor_path
 flags.gdb_path = gdb_path
 flags.addr2line_path = addr2line_path
 flags.has_pybt = has_pybt
+jittor_lock.unlock()

python/jittor/dataset/mnist.py

@@ -15,6 +15,7 @@ from jittor.dataset.dataset import Dataset, dataset_root
 from jittor.dataset.utils import ensure_dir, download_url_to_local
 import jittor as jt 
 import jittor.transform as trans
+from jittor.lock import jittor_lock
 
 class MNIST(Dataset):
     def __init__(self, data_root=dataset_root+"/mnist_data/", train=True ,download=True, transform=None):
@@ -64,4 +65,6 @@ class MNIST(Dataset):
 
         for url, md5 in resources:
             filename = url.rpartition('/')[2]
+            jittor_lock.lock()
             download_url_to_local(url, filename, self.data_root, md5)
+            jittor_lock.unlock()

python/jittor/lock.py

@@ -0,0 +1,24 @@
+import fcntl
+import os
+from jittor_utils import cache_path
+
+class Lock:   
+    def __init__(self, filename):  
+        self.handle = open(filename, 'w') 
+        print(f'创建锁 {filename}') 
+      
+    def lock(self):  
+        ret = fcntl.flock(self.handle, fcntl.LOCK_EX)
+        print(f'加锁成功 {ret}')
+        
+    def unlock(self):  
+        ret = fcntl.flock(self.handle, fcntl.LOCK_UN)  
+        print(f'释放锁成功 {ret}')
+          
+    def __del__(self):  
+        self.handle.close()
+
+lock_path = os.path.join(cache_path, "../jittor.lock")
+if not os.path.exists(lock_path):
+    os.mknod(lock_path)
+jittor_lock = Lock(lock_path)

python/jittor_utils/__init__.py

@@ -171,7 +171,10 @@ def find_cache_path():
     for d in dirs:
         path = os.path.join(path, d)
         if not os.path.isdir(path):
-            os.mkdir(path)
+            try:
+                os.mkdir(path)
+            except:
+                pass
         assert os.path.isdir(path)
     if path not in sys.path:
         sys.path.append(path)

src/lock.cc

@@ -0,0 +1,66 @@
+// ***************************************************************
+// Copyright (c) 2020 Jittor. Authors: 
+// Dun Liang <randonlang@gmail.com>. 
+// Wenyang Zhou <576825820@qq.com>. 
+// All Rights Reserved.
+// This file is subject to the terms and conditions defined in
+// file 'LICENSE.txt', which is part of this source code package.
+// ***************************************************************
+#include "lock.h"
+#include "jit_compiler.h"
+#include "utils/cache_compile.h"
+
+namespace jittor {
+
+DECLARE_FLAG(string, cache_path);
+
+void lock_init(struct flock *lock, short type, short whence, off_t start, off_t len)
+{
+    if (lock == NULL)
+        return;
+ 
+    lock->l_type = type;
+    lock->l_whence = whence;
+    lock->l_start = start;
+    lock->l_len = len;
+}
+ 
+int lock()
+{
+    auto lock_path = jittor::jit_compiler::join(cache_path, "../jittor.lock");
+    const char* lockfilepath = lock_path.c_str();
+    int fd = open(lockfilepath, O_RDWR);
+    if (fd < 0)
+    {
+        return -1;
+    }
+    struct flock lock;
+    lock_init(&lock, F_WRLCK, SEEK_SET, 0, 0);
+    if (fcntl(fd, F_SETLKW, &lock) != 0)
+    {
+        return -1;
+    }
+    printf("Pid: %ld process lock to write the file.\n", (long)getpid());
+    return 0;
+}
+ 
+int unlock()
+{
+    auto lock_path = jittor::jit_compiler::join(cache_path, "../jittor.lock");
+    const char* lockfilepath = lock_path.c_str();
+    int fd = open(lockfilepath, O_RDWR);
+    if (fd < 0)
+    {
+        return -1;
+    }
+    struct flock lock;
+    lock_init(&lock, F_UNLCK, SEEK_SET, 0, 0);
+    if (fcntl(fd, F_SETLKW, &lock) != 0)
+    {
+        return -1;
+    }
+    printf("Pid: %ld process release the file.\n", (long)getpid());
+    return 0;
+}
+
+} // jittor

src/lock.h

@@ -0,0 +1,20 @@
+// ***************************************************************
+// Copyright (c) 2020 Jittor. Authors: 
+// Dun Liang <randonlang@gmail.com>. 
+// Wenyang Zhou <576825820@qq.com>. 
+// All Rights Reserved.
+// This file is subject to the terms and conditions defined in
+// file 'LICENSE.txt', which is part of this source code package.
+// ***************************************************************
+#pragma once
+#include <stdio.h>
+#include <unistd.h>
+#include <fcntl.h>
+#include <errno.h>
+
+namespace jittor {
+int lock();
+
+int unlock();
+
+} // jittor

src/op_compiler.cc

@@ -14,6 +14,7 @@
 #include "misc/str_utils.h"
 #include "ops/op_register.h"
 #include "ops/array_op.h"
+#include "lock.h"
 
 namespace jittor {
 
@@ -945,6 +946,7 @@ jit_op_entry_t OpCompiler::compile(const string& jit_key, const string& src) {
 }
 
 jit_op_entry_t OpCompiler::do_compile(Op* op) {
+    lock();
     OpCompiler oc(op);
     string* src = &oc.src;
     string src_after_passes;
@@ -954,8 +956,9 @@ jit_op_entry_t OpCompiler::do_compile(Op* op) {
         src_after_passes = tm.tune();
         src = &src_after_passes;
     }
-    return oc.compile(op->get_jit_key(), *src);
+    auto ret = oc.compile(op->get_jit_key(), *src);
+    unlock();
+    return ret;
 }
 
-
 }