From de8ffe0c129cb166215046a9fbc8bfec34cf08d3 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Fri, 14 Jan 2022 16:36:33 +0800 Subject: [PATCH] refactor(imperative): unify interpreter option setting GitOrigin-RevId: 53510445cc38866a4bae1059aec56e4c91e9ca4d --- imperative/python/megengine/__init__.py | 1 + imperative/python/megengine/core/_config.py | 14 +++---- imperative/python/megengine/dtr/dtr.py | 2 +- imperative/python/src/tensor.cpp | 12 +----- .../integration/test_converge_with_drop.py | 6 +-- .../python/test/unit/core/test_interpreter.py | 40 +++++++++---------- 6 files changed, 32 insertions(+), 43 deletions(-) diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py index a7c8c9ad..a980024a 100644 --- a/imperative/python/megengine/__init__.py +++ b/imperative/python/megengine/__init__.py @@ -79,6 +79,7 @@ from .core._imperative_rt.core2 import close as _close from .core._imperative_rt.core2 import full_sync as sync from .core._imperative_rt.core2 import sync as _sync from .core._imperative_rt.utils import _set_fork_exec_path_for_timed_func +from .config import * from .device import * from .logger import enable_debug_log, get_logger, set_log_file, set_log_level from .serialization import load, save diff --git a/imperative/python/megengine/core/_config.py b/imperative/python/megengine/core/_config.py index 70d663ab..b4451320 100644 --- a/imperative/python/megengine/core/_config.py +++ b/imperative/python/megengine/core/_config.py @@ -9,12 +9,12 @@ import os from contextlib import contextmanager +from ._imperative_rt.core2 import get_option, set_option + __compute_mode = "default" __conv_format = "default" _benchmark_kernel = False _deterministic_kernel = False -_async_level = os.getenv("MEGENGINE_INTERP_ASYNC_LEVEL", 2) - __all__ = [ "benchmark_kernel", @@ -77,13 +77,13 @@ def async_level(mod) -> int: import megengine as mge mge.config.async_level = 2 """ - return _async_level + return get_option("async_level") @async_level.setter def async_level(mod, level: int): - global _async_level - _async_level = level + assert level >= 0 and level <= 2, "async_level should be 0, 1 or 2" + set_option("async_level", level) @property @@ -148,7 +148,7 @@ def _reset_execution_config( orig_flags = ( _benchmark_kernel, _deterministic_kernel, - _async_level, + get_option("async_level"), __compute_mode, __conv_format, ) @@ -157,7 +157,7 @@ def _reset_execution_config( if deterministic_kernel is not None: _deterministic_kernel = deterministic_kernel if async_level is not None: - _async_level = async_level + set_option("async_level", async_level) if compute_mode is not None: __compute_mode = compute_mode if conv_format is not None: diff --git a/imperative/python/megengine/dtr/dtr.py b/imperative/python/megengine/dtr/dtr.py index 382bbec3..aa2b445c 100644 --- a/imperative/python/megengine/dtr/dtr.py +++ b/imperative/python/megengine/dtr/dtr.py @@ -9,8 +9,8 @@ import re from typing import Union +from ..core import set_option as _set_option from ..core._imperative_rt.core2 import clear_candidates as _clear_candidates -from ..core._imperative_rt.core2 import set_option as _set_option _eviction_threshold = 0 _evictee_minimum_size = 1024 ** 2 diff --git a/imperative/python/src/tensor.cpp b/imperative/python/src/tensor.cpp index af0f8cae..4d7b6767 100644 --- a/imperative/python/src/tensor.cpp +++ b/imperative/python/src/tensor.cpp @@ -599,7 +599,7 @@ void init_tensor(py::module m) { auto val2 = py_async_error.py::object::operator()( "An async error is reported. See above for the actual cause." " Hint: This is where it is reported, not where it happened." - " You may call `megengine.core.set_option('async_level', 0)` " + " You may call `megengine.config.async_level = 0 " "to get better error reporting."); PyException_SetCause( val2.ptr(), val); // PyException_SetCause steals reference @@ -698,20 +698,12 @@ void init_tensor(py::module m) { py_task_q.wait_all_task_finish(); }; + m.def("clear_candidates", [channel]() { channel->clear_candidates(); }); m.def("set_option", [channel](std::string name, size_t value) { channel->set_option(name, value); }); - m.def("clear_candidates", [channel]() { channel->clear_candidates(); }); m.def("get_option", [channel](std::string name) { return channel->get_option(name); }); - m.def("_set_drop_flag", - [channel](bool flag) { channel->set_option("enable_drop", flag); }); - m.def("config_async_level", [channel](int level) { - mgb_assert(level >= 0 and level <= 2, "async_level should be 0, 1 or 2"); - channel->set_option("async_level", level); - }); - m.def("get_async_level", - [channel]() { return channel->get_option("async_level"); }); m.def("set_buffer_length", [channel](int length) { mgb_assert(length >= 0 and length < 100, "buffer_length should be in [0, 100)"); channel->set_option("buffer_length", length); diff --git a/imperative/python/test/integration/test_converge_with_drop.py b/imperative/python/test/integration/test_converge_with_drop.py index 5e31e365..a0d09b4f 100644 --- a/imperative/python/test/integration/test_converge_with_drop.py +++ b/imperative/python/test/integration/test_converge_with_drop.py @@ -14,7 +14,7 @@ import megengine as mge import megengine.autodiff as ad import megengine.functional as F from megengine import Tensor -from megengine.core._imperative_rt.core2 import _set_drop_flag, get_option, set_option +from megengine.core import get_option, set_option from megengine.module import Linear, Module from megengine.optimizer import SGD @@ -75,7 +75,7 @@ class XORNet(Module): def test_training_converge_with_drop(): - _set_drop_flag(True) + set_option("enable_drop", 1) old_buffer_length = get_option("buffer_length") set_option("buffer_length", 0) net = XORNet() @@ -118,5 +118,5 @@ def test_training_converge_with_drop(): precision ) - _set_drop_flag(False) + set_option("enable_drop", 0) set_option("buffer_length", old_buffer_length) diff --git a/imperative/python/test/unit/core/test_interpreter.py b/imperative/python/test/unit/core/test_interpreter.py index 2513c36b..00052096 100644 --- a/imperative/python/test/unit/core/test_interpreter.py +++ b/imperative/python/test/unit/core/test_interpreter.py @@ -6,23 +6,19 @@ import pytest import megengine as mge import megengine.functional as F -from megengine.core._imperative_rt.core2 import ( - AsyncError, - _set_drop_flag, - config_async_level, - get_async_level, -) +from megengine.core import set_option +from megengine.core._imperative_rt.core2 import AsyncError def test_basic(): - config_async_level(2) - assert get_async_level() == 2 - with pytest.raises(RuntimeError): - config_async_level(3) + mge.config.async_level = 2 + assert mge.config.async_level == 2 + with pytest.raises(AssertionError): + mge.config.async_level = 3 def test_level1_infer_value(): - config_async_level(1) + mge.config.async_level = 1 a = mge.tensor([[1, 2], [2, 3], [3, 4]], dtype="float32") b = mge.tensor([1, 1], dtype="float32") identity = mge.tensor(np.array([[1, 0], [0, 1]]), dtype="float32") @@ -30,11 +26,11 @@ def test_level1_infer_value(): c = F.matmul(b, identity) with pytest.raises(RuntimeError): d = F.reshape(a, c) - config_async_level(2) + mge.config.async_level = 2 def test_level1_infer_shape_with_unknown(): - config_async_level(2) + mge.config.async_level = 2 a = mge.tensor([[1, 2, 2, 3]], dtype="float32") b = mge.tensor([1, 1], dtype="float32") multi2 = mge.tensor(np.array([[2, 0], [0, 2]]), dtype="float32") @@ -42,13 +38,13 @@ def test_level1_infer_shape_with_unknown(): # make DepType::SHAPE unknown d = F.reshape(a, c) e = mge.tensor([[1, 2]], dtype="float32") - config_async_level(1) + mge.config.async_level = 1 # test src no shape, throw in level1 with pytest.raises(RuntimeError): f = F.reshape(d, b) with pytest.raises(RuntimeError): g = F.matmul(d, e) - config_async_level(2) + mge.config.async_level = 2 def test_host_compute_elemwise(): @@ -61,7 +57,7 @@ def test_host_compute_elemwise(): def test_drop_basic(): - _set_drop_flag(True) + set_option("enable_drop", True) # test xpu compute x = mge.tensor(np.ones((3, 3)), dtype=np.float32) y = mge.tensor(np.ones((3, 3)), dtype=np.float32) @@ -74,7 +70,7 @@ def test_drop_basic(): z = x + y z._drop() z.numpy() - _set_drop_flag(False) + set_option("enable_drop", False) def test_finalize(): @@ -107,21 +103,21 @@ def test_async_error_check(): # NOTE: DO NOT REMOVE THIS TEST # This is also a compatibility test for -# mge.core.set_option('async_level', 0). +# mge.config.async_level = 0. # If you change the canonical API to set async level, # update the error message of AsyncError as well. def test_async_error(): - orig_lvl = mge.core.get_option("async_level") + orig_lvl = mge.config.async_level try: - mge.core.set_option("async_level", 1) + mge.config.async_level = 1 x = F.utils._simulate_error() try: x.numpy() except AsyncError as e: assert isinstance(e.__cause__, RuntimeError) - mge.core.set_option("async_level", 0) + mge.config.async_level = 0 with pytest.raises(RuntimeError): F.utils._simulate_error() finally: - mge.core.set_option("async_level", orig_lvl) + mge.config.async_level = orig_lvl