Merge pull request #868 from Netflix/fix-catch-on-parallel

fix @catch decorator for parallel
Netflix · Dec 17, 2021 · 100e1c8 · 100e1c8
2 parents 7745b55 + 8b9c5a7
commit 100e1c8
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/metaflow/plugins/catch_decorator.py b/metaflow/plugins/catch_decorator.py
@@ -87,9 +87,10 @@ def task_exception(
 
         # If this task is a UBF control task, it will return itself as the singleton
         # list of tasks.
-        flow._control_mapper_tasks = [
-            "/".join((current.run_id, current.step_name, current.task_id))
-        ]
+        if hasattr(flow, "_parallel_ubf_iter"):
+            flow._control_mapper_tasks = [
+                "/".join((current.run_id, current.step_name, current.task_id))
+            ]
         # store the exception
         picklable = MetaflowExceptionWrapper(exception)
         flow._catch_exception = picklable