diff --git a/python/jittor/compile_extern.py b/python/jittor/compile_extern.py index fd31178f..f02d323e 100644 --- a/python/jittor/compile_extern.py +++ b/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() diff --git a/python/jittor/compiler.py b/python/jittor/compiler.py index 059f597d..53682eaf 100644 --- a/python/jittor/compiler.py +++ b/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() \ No newline at end of file diff --git a/python/jittor/dataset/mnist.py b/python/jittor/dataset/mnist.py index ccc1c832..af37adf2 100644 --- a/python/jittor/dataset/mnist.py +++ b/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() diff --git a/python/jittor/lock.py b/python/jittor/lock.py new file mode 100644 index 00000000..53c30cc2 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/python/jittor_utils/__init__.py b/python/jittor_utils/__init__.py index ab909880..a9ccb8a1 100644 --- a/python/jittor_utils/__init__.py +++ b/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) diff --git a/src/lock.cc b/src/lock.cc new file mode 100644 index 00000000..152668e8 --- /dev/null +++ b/src/lock.cc @@ -0,0 +1,66 @@ +// *************************************************************** +// Copyright (c) 2020 Jittor. Authors: +// Dun Liang . +// 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 \ No newline at end of file diff --git a/src/lock.h b/src/lock.h new file mode 100644 index 00000000..4912d8dc --- /dev/null +++ b/src/lock.h @@ -0,0 +1,20 @@ +// *************************************************************** +// Copyright (c) 2020 Jittor. Authors: +// Dun Liang . +// 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 +#include +#include +#include + +namespace jittor { +int lock(); + +int unlock(); + +} // jittor diff --git a/src/op_compiler.cc b/src/op_compiler.cc index 435f93b0..3c28a00a 100644 --- a/src/op_compiler.cc +++ b/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; } - } \ No newline at end of file