From 9178b3459a300212f3bba1100ddbe1f3d8c34963 Mon Sep 17 00:00:00 2001 From: Dun Liang Date: Fri, 10 Apr 2020 16:10:01 +0800 Subject: [PATCH] add lock_scope and unlock_scope --- python/jittor/__init__.py | 18 ++++---- python/jittor/compile_extern.py | 5 --- python/jittor/compiler.py | 21 ++++----- python/jittor/dataset/mnist.py | 3 -- python/jittor/dataset/utils.py | 3 +- python/jittor/lock.py | 59 +++++++++++++++++++----- python/jittor/test/test_lock.py | 6 +-- python/jittor_utils/__init__.py | 4 +- src/lock.cc | 80 +++++++++++++-------------------- src/lock.h | 22 +++++---- src/op_compiler.cc | 5 +-- 11 files changed, 123 insertions(+), 103 deletions(-) diff --git a/python/jittor/__init__.py b/python/jittor/__init__.py index 015eb7e1..6598fd5e 100644 --- a/python/jittor/__init__.py +++ b/python/jittor/__init__.py @@ -7,14 +7,16 @@ # This file is subject to the terms and conditions defined in # file 'LICENSE.txt', which is part of this source code package. # *************************************************************** -from . import compiler -from .compiler import LOG, has_cuda -from .compiler import compile_custom_ops, compile_custom_op -import jittor_core as core -from jittor_core import * -from jittor_core.ops import * -from . import compile_extern -from .compile_extern import mkl_ops +from . import lock +with lock.lock_scope(): + from . import compiler + from .compiler import LOG, has_cuda + from .compiler import compile_custom_ops, compile_custom_op + import jittor_core as core + from jittor_core import * + from jittor_core.ops import * + from . import compile_extern + from .compile_extern import mkl_ops import contextlib import numpy as np diff --git a/python/jittor/compile_extern.py b/python/jittor/compile_extern.py index f02d323e..fd31178f 100644 --- a/python/jittor/compile_extern.py +++ b/python/jittor/compile_extern.py @@ -7,7 +7,6 @@ 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: @@ -374,8 +373,6 @@ 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() @@ -383,5 +380,3 @@ setup_cutt() setup_mkl() setup_cuda_extern() - -jittor_lock.unlock() diff --git a/python/jittor/compiler.py b/python/jittor/compiler.py index 2107c91a..6216c6da 100644 --- a/python/jittor/compiler.py +++ b/python/jittor/compiler.py @@ -17,7 +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 +from . import lock def find_jittor_path(): return os.path.dirname(__file__) @@ -555,6 +555,7 @@ def compile_custom_op(header, source, op_name, warp=True): m = compile_custom_ops([hname, ccname]) return getattr(m, op_name) +@lock.lock_scope() def compile_custom_ops( filenames, extra_flags="", @@ -644,10 +645,10 @@ def compile_custom_ops( lib_path = os.path.join(cache_path, "custom_ops") if lib_path not in os.sys.path: os.sys.path.append(lib_path) - jittor_lock.unlock() - with jit_utils.import_scope(dlopen_flags): - exec(f"import {gen_name}") - jittor_lock.lock() + # unlock scope when initialize + with lock.unlock_scope(): + with jit_utils.import_scope(dlopen_flags): + exec(f"import {gen_name}") mod = locals()[gen_name] if return_module: return mod @@ -801,20 +802,16 @@ 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" @@ -856,13 +853,11 @@ 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 @@ -887,7 +882,6 @@ if has_cuda: return nvcc_flags nvcc_flags = convert_nvcc_flags(nvcc_flags) -jittor_lock.lock() # build core gen_jit_flags() gen_jit_tests() @@ -977,4 +971,5 @@ 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 + +core.set_lock_path(lock.lock_path) diff --git a/python/jittor/dataset/mnist.py b/python/jittor/dataset/mnist.py index af37adf2..ccc1c832 100644 --- a/python/jittor/dataset/mnist.py +++ b/python/jittor/dataset/mnist.py @@ -15,7 +15,6 @@ 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): @@ -65,6 +64,4 @@ 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/dataset/utils.py b/python/jittor/dataset/utils.py index 5c86a128..ce40057f 100644 --- a/python/jittor/dataset/utils.py +++ b/python/jittor/dataset/utils.py @@ -15,6 +15,7 @@ from tqdm import tqdm import numpy as np from collections.abc import Sequence, Mapping from PIL import Image +from .. import lock def ensure_dir(dir_path): if not os.path.isdir(dir_path): @@ -36,7 +37,7 @@ def _progress(): return bar_update - +@lock.lock_scope() def download_url_to_local(url, filename, root_folder, md5): ensure_dir(root_folder) file_path = os.path.join(root_folder, filename) diff --git a/python/jittor/lock.py b/python/jittor/lock.py index b64996f0..3510460e 100644 --- a/python/jittor/lock.py +++ b/python/jittor/lock.py @@ -1,24 +1,63 @@ import fcntl import os -from jittor_utils import cache_path +from jittor_utils import cache_path, LOG class Lock: def __init__(self, filename): - self.handle = open(filename, 'w') - print(f'Create lock for {filename}, PID {os.getpid()}') + self.handle = open(filename, 'w') + LOG.v(f'OPEN LOCK path: {filename} PID: {os.getpid()}') + self.is_locked = False def lock(self): - ret = fcntl.flock(self.handle, fcntl.LOCK_EX) - print(f'Add lock success {ret}, PID {os.getpid()}') + fcntl.flock(self.handle, fcntl.LOCK_EX) + self.is_locked = True + LOG.vv(f'LOCK PID: {os.getpid()}') def unlock(self): - ret = fcntl.flock(self.handle, fcntl.LOCK_UN) - print(f'Release lock success {ret}, PID {os.getpid()}') + fcntl.flock(self.handle, fcntl.LOCK_UN) + self.is_locked = False + LOG.vv(f'UNLOCK PID: {os.getpid()}') def __del__(self): self.handle.close() -lock_path = os.path.join(cache_path, "../jittor.lock") + +class _base_scope: + '''base_scope for support @xxx syntax''' + def __enter__(self): pass + def __exit__(self, *exc): pass + def __call__(self, func): + def inner(*args, **kw): + with self: + ret = func(*args, **kw) + return ret + return inner + +class lock_scope(_base_scope): + def __enter__(self): + self.is_locked = jittor_lock.is_locked + if not self.is_locked: + jittor_lock.lock() + + def __exit__(self, *exc): + if not self.is_locked: + jittor_lock.unlock() + +class unlock_scope(_base_scope): + def __enter__(self): + self.is_locked = jittor_lock.is_locked + if self.is_locked: + jittor_lock.unlock() + + def __exit__(self, *exc): + if self.is_locked: + jittor_lock.lock() + +lock_path = os.path.abspath(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 + LOG.i("Create lock file:", lock_path) + try: + os.mknod(lock_path) + except: + pass +jittor_lock = Lock(lock_path) diff --git a/python/jittor/test/test_lock.py b/python/jittor/test/test_lock.py index 9d8a2e0e..65c93131 100644 --- a/python/jittor/test/test_lock.py +++ b/python/jittor/test/test_lock.py @@ -11,11 +11,11 @@ import os, sys import jittor as jt from pathlib import Path -@unittest.skipIf(jt.compile_extern.mpi_ops is None, "no mpi found") +@unittest.skipIf(not jt.compile_extern.has_mpi, "no mpi found") class TestLock(unittest.TestCase): def test(self): mpi = jt.compile_extern.mpi - mpirun_path = jt.compiler.env_or_try_find('mpirun_path', 'mpirun') + mpirun_path = jt.compile_extern.mpicc_path.replace("mpicc", "mpirun") if os.environ.get('lock_full_test', '0') == '1': cache_path = os.path.join(str(Path.home()), ".cache", "jittor", "lock") cmd = f"rm -rf {cache_path} && cache_name=lock {mpirun_path} -np 2 {sys.executable} -m jittor.test.test_example" @@ -23,7 +23,7 @@ class TestLock(unittest.TestCase): cache_path = os.path.join(str(Path.home()), ".cache", "jittor") cmd = f"{mpirun_path} -np 2 {sys.executable} -m jittor.test.test_example" print("run cmd", cmd) - jt.compiler.run_cmd(cmd) + assert os.system(cmd) == 0 if __name__ == "__main__": diff --git a/python/jittor_utils/__init__.py b/python/jittor_utils/__init__.py index a9ccb8a1..710a25c3 100644 --- a/python/jittor_utils/__init__.py +++ b/python/jittor_utils/__init__.py @@ -158,9 +158,10 @@ def find_cache_path(): r = sp.run(["git","branch"], cwd=os.path.dirname(__file__), stdout=sp.PIPE, stderr=sp.PIPE) assert r.returncode == 0 - bs = r.stdout.decode() + bs = r.stdout.decode().splitlines() for b in bs: if b.startswith("* "): break + cache_name = b[2:] for c in " (){}": cache_name = cache_name.replace(c, "_") except: @@ -168,6 +169,7 @@ def find_cache_path(): for name in cache_name.split("/"): dirs.insert(-1, name) os.environ["cache_name"] = cache_name + LOG.v("cache_name", cache_name) for d in dirs: path = os.path.join(path, d) if not os.path.isdir(path): diff --git a/src/lock.cc b/src/lock.cc index 6024c95f..e0e2f6cb 100644 --- a/src/lock.cc +++ b/src/lock.cc @@ -1,66 +1,50 @@ // *************************************************************** // Copyright (c) 2020 Jittor. Authors: -// Dun Liang . -// Wenyang Zhou <576825820@qq.com>. +// Wenyang Zhou <576825820@qq.com> +// Dun Liang // 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 +#include +#include +#include + #include "lock.h" -#include "jit_compiler.h" -#include "utils/cache_compile.h" namespace jittor { -DECLARE_FLAG(string, cache_path); +static int lock_fd = -1; -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; +void set_lock_path(string path) { + lock_fd = open(path.c_str(), O_RDWR); + ASSERT(lock_fd >= 0); + LOGv << "OPEN LOCK path:" << path << "Pid:" << getpid(); } -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; +void lock() { + ASSERT(lock_fd >= 0); + struct flock lock = { + .l_type = F_WRLCK, + .l_whence = SEEK_SET, + .l_start = 0, + .l_len = 0 + }; + ASSERT(fcntl(lock_fd, F_SETLKW, &lock) == 0); + LOGvv << "LOCK Pid:" << getpid(); } -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; +void unlock() { + ASSERT(lock_fd >= 0); + struct flock lock = { + .l_type = F_UNLCK, + .l_whence = SEEK_SET, + .l_start = 0, + .l_len = 0 + }; + ASSERT(fcntl(lock_fd, F_SETLKW, &lock) == 0); + LOGvv << "UNLOCK Pid:" << getpid(); } } // jittor \ No newline at end of file diff --git a/src/lock.h b/src/lock.h index 4912d8dc..5595ea83 100644 --- a/src/lock.h +++ b/src/lock.h @@ -1,20 +1,26 @@ // *************************************************************** // Copyright (c) 2020 Jittor. Authors: -// Dun Liang . -// Wenyang Zhou <576825820@qq.com>. +// Wenyang Zhou <576825820@qq.com> +// Dun Liang // 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 +#include "common.h" namespace jittor { -int lock(); -int unlock(); +// @pyjt(set_lock_path) +void set_lock_path(string path); + +void lock(); + +void unlock(); + +struct lock_guard { + inline lock_guard() { lock(); } + inline ~lock_guard() { unlock(); } +}; } // jittor diff --git a/src/op_compiler.cc b/src/op_compiler.cc index 3c28a00a..2f556c54 100644 --- a/src/op_compiler.cc +++ b/src/op_compiler.cc @@ -673,7 +673,7 @@ string OpCompiler::get_jit_src(Op* op) { else after_include_src += src; } - ASSERT(file_exist(src_path)); + ASSERT(file_exist(src_path)) << src_path; LOGvvv << "Read from" << src_path; string src = read_all(src_path); ASSERT(src.size()) << "Source read failed:" << src_path; @@ -946,7 +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(); + jittor::lock_guard lg; OpCompiler oc(op); string* src = &oc.src; string src_after_passes; @@ -957,7 +957,6 @@ jit_op_entry_t OpCompiler::do_compile(Op* op) { src = &src_after_passes; } auto ret = oc.compile(op->get_jit_key(), *src); - unlock(); return ret; }