Fixed bug when using catch + retry

Netflix · Mar 19, 2022 · 5a93d9f · 5a93d9f
1 parent 62bc8df
commit 5a93d9f
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/metaflow/plugins/airflow/airflow_decorator.py b/metaflow/plugins/airflow/airflow_decorator.py
@@ -40,12 +40,11 @@ def task_pre_step(
         ]
         # Register book-keeping metadata for debugging.
         metadata.register_metadata(run_id, step_name, task_id, entries)
-        if retry_count == 0:
-            push_xcom_values(
-                {
-                    TASK_ID_XCOM_KEY: os.environ["METAFLOW_AIRFLOW_TASK_ID"],
-                }
-            )
+        push_xcom_values(
+            {
+                TASK_ID_XCOM_KEY: os.environ["METAFLOW_AIRFLOW_TASK_ID"],
+            }
+        )
 
     def task_finished(
         self, step_name, flow, graph, is_task_ok, retry_count, max_user_code_retries