Fix missing extra_params when constructing operands (#2999)

mars-project · May 6, 2022 · 4a9392d · 4a9392d
1 parent 263c841
commit 4a9392d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/mars/core/operand/base.py b/mars/core/operand/base.py
@@ -195,7 +195,8 @@ def __init__(self: OperandType, *args, **kwargs):
 
     @classmethod
     def _parse_kwargs(cls, kwargs: Dict[str, Any]):
-        kwargs["extra_params"] = extras = AttributeDict()
+        extra_params = kwargs.pop("extra_params", {})
+        kwargs["extra_params"] = extras = AttributeDict(extra_params)
         kwargs["scheduling_hint"] = scheduling_hint = kwargs.get(
             "scheduling_hint", SchedulingHint()
         )

diff --git a/mars/core/operand/tests/test_core.py b/mars/core/operand/tests/test_core.py
@@ -73,6 +73,8 @@ class MyOperand5(MyOperand4):
 
 
 def test_execute():
+    op = MyOperand(extra_params={"my_extra_params": 1})
+    assert op.extra_params["my_extra_params"] == 1
     MyOperand.register_executor(lambda *_: 2)
     assert execute(dict(), MyOperand(_key="1")) == 2
     assert execute(dict(), MyOperand2(_key="1")) == 2