Repository: mars-project/mars

Issue:
[BUG] `NotImplementedError` raised when calling with batch
**Describe the bug**
Batch method raises `NotImplementedError` when calling with only one delay argument.

```
Traceback (most recent call last):
  File "/Users/wenjun/miniconda3/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3444, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-4-14c7f9a6baba>", line 2, in <module>
    asyncio.run(test_inst.method.batch(test_inst.method.delay(1, 2)))
  File "/Users/wenjun/miniconda3/lib/python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/Users/wenjun/miniconda3/lib/python3.9/asyncio/base_events.py", line 642, in run_until_complete
    return future.result()
  File "/Users/wenjun/Code/mars/mars/oscar/batch.py", line 141, in _async_batch
    return [await self.func(*d.args, **d.kwargs)]
  File "<ipython-input-2-e7df29e525ae>", line 10, in method
    raise NotImplementedError
NotImplementedError
```

**To Reproduce**
```python
import asyncio
import mars.oscar as mo

class TestClass:
    def __init__(self):
        self.arg_list = []
        self.kwarg_list = []

    @mo.extensible
    async def method(self, *args, **kwargs):
        raise NotImplementedError

    @method.batch
    async def method(self, args_list, kwargs_list):
        self.arg_list.extend(args_list)
        self.kwarg_list.extend(kwargs_list)
        return [len(self.kwarg_list)] * len(args_list)

test_inst = TestClass()
asyncio.run(test_inst.method.batch(test_inst.method.delay(1, 2)))
```



Discussion:

Resolving patch:
diff --git a/mars/oscar/batch.py b/mars/oscar/batch.py
--- a/mars/oscar/batch.py
+++ b/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


Test patch:
diff --git a/mars/oscar/tests/test_batch.py b/mars/oscar/tests/test_batch.py
--- a/mars/oscar/tests/test_batch.py
+++ b/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