From 60459bd79b8076ae234ea1746d7b1f5dbf65f0a8 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Tue, 5 Jan 2021 16:22:44 +0800 Subject: [PATCH] fix(mge/pytest): remove sync in pytest teardown & add full_sync GitOrigin-RevId: ee278a8ac10f841c15cf74e086c82af055e2dd7b --- imperative/python/megengine/__init__.py | 2 +- imperative/python/src/tensor.cpp | 7 +++++++ imperative/python/test/conftest.py | 4 ---- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py index f691035d5..219fd5ef3 100644 --- a/imperative/python/megengine/__init__.py +++ b/imperative/python/megengine/__init__.py @@ -71,7 +71,7 @@ if sys.platform == "win32": kernel32.SetErrorMode(old_error_mode) -from .core._imperative_rt.core2 import release_trace_apply_func, sync +from .core._imperative_rt.core2 import release_trace_apply_func, sync, full_sync from .core._imperative_rt.utils import _set_fork_exec_path_for_timed_func from .device import * from .logger import enable_debug_log, get_logger, set_log_file, set_log_level diff --git a/imperative/python/src/tensor.cpp b/imperative/python/src/tensor.cpp index 1d96f1c10..a93031777 100644 --- a/imperative/python/src/tensor.cpp +++ b/imperative/python/src/tensor.cpp @@ -736,6 +736,13 @@ void init_tensor(py::module m) { py_task_q.wait_all_task_finish(); }, py::call_guard()); + m.def("full_sync", + []() { + interpreter_for_py->sync(); + CompNode::sync_all(); + py_task_q.wait_all_task_finish(); + }, + py::call_guard()); m.def("release_trace_apply_func", &release_trace_apply_func); diff --git a/imperative/python/test/conftest.py b/imperative/python/test/conftest.py index ae47cb831..67e0ad3f2 100644 --- a/imperative/python/test/conftest.py +++ b/imperative/python/test/conftest.py @@ -6,7 +6,3 @@ import pytest from megengine.core._imperative_rt.core2 import sync sys.path.append(os.path.join(os.path.dirname(__file__), "helpers")) - - -def pytest_runtest_teardown(): - sync() -- GitLab