[Dask] Fix dask load module & handler (mlrun#1766)

Penumbra69 · Feb 21, 2022 · 4b99f57 · 4b99f57
1 parent 3da7e76
commit 4b99f57
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/mlrun/runtimes/daskjob.py b/mlrun/runtimes/daskjob.py
@@ -431,19 +431,19 @@ def _run(self, runobj: RunObject, execution):
         extra_env = self._generate_runtime_env(runobj)
         environ.update(extra_env)
 
-        if not inspect.isfunction(handler):
-            if not self.spec.command:
-                raise ValueError(
-                    "specified handler (string) without command "
-                    "(py file path), specify command or use handler pointer"
-                )
-            mod, handler = load_module(self.spec.command, handler)
         context = MLClientCtx.from_dict(
             runobj.to_dict(),
             rundb=self.spec.rundb,
             autocommit=False,
             host=socket.gethostname(),
         )
+        if not inspect.isfunction(handler):
+            if not self.spec.command:
+                raise ValueError(
+                    "specified handler (string) without command "
+                    "(py file path), specify command or use handler pointer"
+                )
+            handler = load_module(self.spec.command, handler, context=context)
         client = self.client
         setattr(context, "dask_client", client)
         sout, serr = exec_from_params(handler, runobj, context)