From d24f198cfa23862fb42a72c4a258c04144b0efa6 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Thu, 13 Jan 2022 14:40:46 +0800
Subject: [PATCH] fix(imperative): clear custom op build cache

GitOrigin-RevId: cb145fe1ada24cae7e1553bcb1ce8a371895b7c4
---
 .../python/megengine/utils/custom_op_tools.py      |  9 +++---
 imperative/python/test/unit/core/test_custom_op.py | 34 +++++++++++++---------
 2 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/imperative/python/megengine/utils/custom_op_tools.py b/imperative/python/megengine/utils/custom_op_tools.py
index d9150fef..6b09fde0 100644
--- a/imperative/python/megengine/utils/custom_op_tools.py
+++ b/imperative/python/megengine/utils/custom_op_tools.py
@@ -16,7 +16,7 @@ import sys
 import time
 from typing import List, Optional, Union
 
-from ..core.ops.custom import load
+from ..core.ops.custom import get_custom_op_abi_tag, load
 from ..logger import get_logger
 
 
@@ -730,9 +730,10 @@ def build(
     """
 
     # phase 1: prepare config
-    if abi_tag != None:
-        global MGE_ABI_VER
-        MGE_ABI_VER = abi_tag
+    if abi_tag == None:
+        abi_tag = get_custom_op_abi_tag()
+    global MGE_ABI_VER
+    MGE_ABI_VER = abi_tag
 
     def strlist(args, name):
         assert isinstance(args, str) or isinstance(
diff --git a/imperative/python/test/unit/core/test_custom_op.py b/imperative/python/test/unit/core/test_custom_op.py
index e2a9e4b2..c5ace4b5 100644
--- a/imperative/python/test/unit/core/test_custom_op.py
+++ b/imperative/python/test/unit/core/test_custom_op.py
@@ -37,7 +37,12 @@ def compare(ref, real):
 def build_and_clean(test_func):
     def wrapper():
         cur_dir_path = os.path.dirname(os.path.abspath(__file__))
-        build_path = os.path.join(cur_dir_path, "custom_opsrc", "build")
+        build_root_dir = custom_op_tools._get_default_build_root()
+        build_path = os.path.join(build_root_dir, "custom_opsrc", "build")
+
+        if os.path.exists(build_path):
+            shutil.rmtree(build_path)
+
         mgb_root_path = os.path.dirname(
             os.path.dirname(
                 os.path.dirname(os.path.dirname(os.path.dirname(cur_dir_path)))
@@ -67,20 +72,21 @@ def build_and_clean(test_func):
         else:
             custom_opsrc = [os.path.join(cur_dir_path, "custom_opsrc", "elem_add.cpp")]
 
-        lib_path = custom_op_tools.build_and_load(
-            "test_op",
-            custom_opsrc,
-            extra_include_paths=extra_include_paths,
-            extra_ldflags=extra_ld_flags,
-            build_dir=build_path,
-            verbose=False,
-            abi_tag=custom.get_custom_op_abi_tag(),
-        )
-        test_func()
+        try:
+            lib_path = custom_op_tools.build_and_load(
+                "test_op",
+                custom_opsrc,
+                extra_include_paths=extra_include_paths,
+                extra_ldflags=extra_ld_flags,
+                build_dir=build_path,
+                verbose=False,
+            )
+            test_func()
+            custom.unload(lib_path)
 
-        custom.unload(lib_path)
-        if os.path.exists(build_path):
-            shutil.rmtree(build_path)
+        finally:
+            if os.path.exists(build_path):
+                shutil.rmtree(build_path)
 
     return wrapper