Support var-len params in func passed to executor. (#1457)

Fixes #1456

Author: chuckwondo

lithops/tests/test_map.py

@@ -125,3 +125,31 @@ def test_lithops_return_futures_map_multiple(self):
         fexec.wait()
         result = fexec.get_result()
         assert result == [1, 2, 3, 1, 2, 3]
+
+    def test_lithops_return_futures_map_over_decorator(self):
+        def doubled(f):
+            def wrapper(*args, **kwargs):
+                return 2 * f(*args, **kwargs)
+
+            return wrapper
+
+        @doubled
+        def total(a, b, c=0, *args, d, e=5, **kwargs):
+            return sum([a, b, c, d, e, *args, *kwargs.values()])
+
+        with lithops.FunctionExecutor(config=pytest.lithops_config) as fexec:
+            fexec.map(
+                total,
+                [
+                    {"a": 1, "b": 2, "d": 3},
+                    {"args": (1, 2), "d": 3},
+                    {"args": (1, 2), "d": 3, "e": 4},
+                    {"args": (1, 2), "d": 3, "f": 4},
+                    {"a": 1, "b": 2, "d": 3, "e": 6, "f": 4},
+                    {"args": (1, 2), "d": 3, "e": 6, "f": 4},
+                    {"args": (1, 2), "kwargs": {"d": 3, "e": 6, "f": 4}},
+                ],
+            )
+            result = fexec.get_result()
+
+        assert result == [22, 22, 20, 30, 32, 32, 32]

lithops/utils.py

@@ -533,26 +533,40 @@ def verify_args(func, iterdata, extra_args):
     non_verify_args = ['ibm_cos', 'storage', 'id', 'rabbitmq']
     func_sig = inspect.signature(func)
 
-    new_parameters = list()
-    for param in func_sig.parameters:
-        if param not in non_verify_args:
-            new_parameters.append(func_sig.parameters[param])
-
+    new_parameters = [
+        param
+        for name, param in func_sig.parameters.items()
+        if name not in non_verify_args
+    ]
     new_func_sig = func_sig.replace(parameters=new_parameters)
 
-    new_data = list()
+    # Detect presence of **kwargs (with any name)
+    has_var_keyword = any(
+        p.kind == inspect.Parameter.VAR_KEYWORD
+        for p in new_func_sig.parameters.values()
+    )
+
+    new_data = []
+
     for elem in data:
         if isinstance(elem, dict):
-            if set(list(new_func_sig.parameters.keys())) <= set(elem):
+            # If the function accepts **kwargs (any name), we cannot reliably
+            # enforce exact param name matching here, and we *want* to allow
+            # passing through arbitrary dicts (e.g., original function args)
+            # even when a decorator wrapper has **kwargs, etc.
+            if has_var_keyword:
+                new_data.append(elem)
+            elif set(expected_keys := list(new_func_sig.parameters)) <= set(elem):
+                # No **kwargs: enforce that the dict contains at least all
+                # required user parameters (excluding reserved ones).
                 new_data.append(elem)
             else:
-                raise ValueError("Check the args names in the data. "
-                                 "You provided these args: {}, and "
-                                 "the args must be: {}"
-                                 .format(list(elem.keys()),
-                                         list(new_func_sig.parameters.keys())))
+                raise ValueError(
+                    "Check the args names in the data. You provided these args: ",
+                    f"{list(elem)}, and the args must be: {expected_keys}",
+                )
         elif isinstance(elem, tuple):
-            new_elem = dict(new_func_sig.bind(*list(elem)).arguments)
+            new_elem = dict(new_func_sig.bind(*elem).arguments)
             new_data.append(new_elem)
         else:
             # single value (list, string, integer, dict, etc)

lithops/worker/jobrunner.py

@@ -234,7 +234,8 @@ def run(self):
             logger.info(f"Going to execute '{str(fn_name)}()'")
             print('---------------------- FUNCTION LOG ----------------------')
             function_start_tstamp = time.time()
-            result = func(**data)
+            args, kwargs = _prepare_args(func, data)
+            result = func(*args, **kwargs)
             function_end_tstamp = time.time()
             print('----------------------------------------------------------')
             logger.info("Success function execution")
@@ -310,3 +311,28 @@ def run(self):
                 self.stats.write("worker_result_upload_time", round(output_upload_end_tstamp - output_upload_start_tstamp, 8))
             self.jobrunner_conn.send("Finished")
             logger.info("Process finished")
+
+
+def _prepare_args(func, data):
+    # Convert the "data" envelope into normal *args/**kwargs,
+    # respecting the actual var-length parameter names of `func`.
+    func_sig = inspect.signature(func)
+    var_pos_name = None
+    var_kw_name = None
+
+    for name, param in func_sig.parameters.items():
+        if param.kind == inspect.Parameter.VAR_POSITIONAL:
+            var_pos_name = name
+        elif param.kind == inspect.Parameter.VAR_KEYWORD:
+            var_kw_name = name
+
+    payload = dict(data)
+
+    # Extract var-positional argument value if present
+    args = payload.pop(var_pos_name) or () if var_pos_name in payload else ()
+    # Extract var-keyword argument value if present
+    kwargs = payload.pop(var_kw_name) or {} if var_kw_name in payload else {}
+    # Any remaining keys become normal keyword arguments
+    kwargs.update(payload)
+
+    return args, kwargs