Fix NotImplementedError for mo.batch when single call not implemented (#2635)

Author: wjsi

.github/workflows/core-ci.yml

@@ -50,7 +50,6 @@ jobs:
 
         source ./ci/rewrite-cov-config.sh
 
-        pip install git+https://github.com/mars-project/pytest-asyncio.git
         pip install numpy scipy cython oss2
         pip install -e ".[dev,extra]"
         pip install virtualenv flaky

.github/workflows/os-compat-ci.yml

@@ -37,7 +37,6 @@ jobs:
 
         source ./ci/rewrite-cov-config.sh
 
-        pip install git+https://github.com/mars-project/pytest-asyncio.git
         pip install numpy scipy cython oss2
         pip install -e ".[dev,extra]"
         pip install virtualenv flaky

.github/workflows/platform-ci.yml

@@ -51,7 +51,6 @@ jobs:
 
           source ./ci/rewrite-cov-config.sh
 
-          pip install git+https://github.com/mars-project/pytest-asyncio.git
           pip install numpy scipy cython
 
           pip install -e ".[dev,extra]"

mars/oscar/batch.py

@@ -82,6 +82,7 @@ class _ExtensibleCallable:
     func: Callable
     batch_func: Optional[Callable]
     is_async: bool
+    has_single_func: bool
 
     def __call__(self, *args, **kwargs):
         if self.is_async:
@@ -91,20 +92,26 @@ def __call__(self, *args, **kwargs):
 
     async def _async_call(self, *args, **kwargs):
         try:
-            return await self.func(*args, **kwargs)
+            if self.has_single_func:
+                return await self.func(*args, **kwargs)
         except NotImplementedError:
-            if self.batch_func:
-                ret = await self.batch_func([args], [kwargs])
-                return None if ret is None else ret[0]
-            raise
+            self.has_single_func = False
+
+        if self.batch_func is not None:
+            ret = await self.batch_func([args], [kwargs])
+            return None if ret is None else ret[0]
+        raise NotImplementedError
 
     def _sync_call(self, *args, **kwargs):
         try:
-            return self.func(*args, **kwargs)
+            if self.has_single_func:
+                return self.func(*args, **kwargs)
         except NotImplementedError:
-            if self.batch_func:
-                return self.batch_func([args], [kwargs])[0]
-            raise
+            self.has_single_func = False
+
+        if self.batch_func is not None:
+            return self.batch_func([args], [kwargs])[0]
+        raise NotImplementedError
 
 
 class _ExtensibleWrapper(_ExtensibleCallable):
@@ -119,6 +126,7 @@ def __init__(
         self.batch_func = batch_func
         self.bind_func = bind_func
         self.is_async = is_async
+        self.has_single_func = True
 
     @staticmethod
     def delay(*args, **kwargs):
@@ -138,7 +146,7 @@ async def _async_batch(self, *delays):
         # will be more efficient
         if len(delays) == 1:
             d = delays[0]
-            return [await self.func(*d.args, **d.kwargs)]
+            return [await self._async_call(*d.args, **d.kwargs)]
         elif self.batch_func:
             args_list, kwargs_list = self._gen_args_kwargs_list(delays)
             return await self.batch_func(args_list, kwargs_list)
@@ -184,6 +192,7 @@ def __init__(self, func: Callable):
         self.batch_func = None
         self.bind_func = build_args_binder(func, remove_self=True)
         self.is_async = asyncio.iscoroutinefunction(self.func)
+        self.has_single_func = True
 
     def batch(self, func: Callable):
         self.batch_func = func

mars/oscar/tests/test_batch.py

@@ -126,6 +126,11 @@ def __init__(self):
             self.arg_list = []
             self.kwarg_list = []
 
+        @extensible
+        @_wrap_async(use_async)
+        def not_implemented_method(self, *args, **kw):
+            raise NotImplementedError
+
         @extensible
         @_wrap_async(use_async)
         def method(self, *args, **kwargs):
@@ -141,6 +146,11 @@ def method(self, args_list, kwargs_list):
     if use_async:
         assert asyncio.iscoroutinefunction(TestClass.method)
 
+    test_inst = TestClass()
+    ret = test_inst.method.batch(test_inst.method.delay(12))
+    ret = await ret if use_async else ret
+    assert ret == [1]
+
     test_inst = TestClass()
     ret = test_inst.method.batch(test_inst.method.delay(12), test_inst.method.delay(10))
     ret = await ret if use_async else ret
@@ -149,6 +159,10 @@ def method(self, args_list, kwargs_list):
     assert test_inst.kwarg_list == [{}, {}]
 
     test_inst = TestClass()
+    for _ in range(2):
+        with pytest.raises(NotImplementedError):
+            ret = test_inst.not_implemented_method()
+            await ret if use_async else ret
     ret = test_inst.method(12, kwarg=34)
     ret = await ret if use_async else ret
     assert ret == 1