fix(mge/pytest): remove sync in pytest teardown & add full_sync

GitOrigin-RevId: ee278a8ac1

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

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<py::gil_scoped_release>());
+    m.def("full_sync",
+          []() {
+              interpreter_for_py->sync();
+              CompNode::sync_all();
+              py_task_q.wait_all_task_finish();
+          },
+          py::call_guard<py::gil_scoped_release>());
 
     m.def("release_trace_apply_func", &release_trace_apply_func);
 

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()