add retries

Author: savingoyal

metaflow/plugins/aws/step_functions/dynamo_db_client.py

@@ -1,5 +1,5 @@
 import os
-
+import time
 import requests
 
 from metaflow.metaflow_config import SFN_DYNAMO_DB_TABLE
@@ -27,12 +27,31 @@ def save_foreach_cardinality(self, foreach_split_task_id, foreach_cardinality, t
     def save_parent_task_id_for_foreach_join(
         self, foreach_split_task_id, foreach_join_parent_task_id
     ):
-        return self._client.update_item(
-            TableName=self.name,
-            Key={"pathspec": {"S": foreach_split_task_id}},
-            UpdateExpression="ADD parent_task_ids_for_foreach_join :val",
-            ExpressionAttributeValues={":val": {"SS": [foreach_join_parent_task_id]}},
-        )
+        ex = None
+        for attempt in range(10):
+            try:
+                return self._client.update_item(
+                    TableName=self.name,
+                    Key={"pathspec": {"S": foreach_split_task_id}},
+                    UpdateExpression="ADD parent_task_ids_for_foreach_join :val",
+                    ExpressionAttributeValues={
+                        ":val": {"SS": [foreach_join_parent_task_id]}
+                    },
+                )
+            except self._client.exceptions.ClientError as error:
+                ex = error
+                if (
+                    error.response["Error"]["Code"]
+                    == "ProvisionedThroughputExceededException"
+                ):
+                    # hopefully, enough time for AWS to scale up! otherwise
+                    # ensure sufficient on-demand throughput for dynamo db
+                    # is provisioned ahead of time
+                    sleep_time = min((2**attempt) * 10, 60)
+                    time.sleep(sleep_time)
+                else:
+                    raise
+        raise ex
 
     def get_parent_task_ids_for_foreach_join(self, foreach_split_task_id):
         response = self._client.get_item(

metaflow/plugins/aws/step_functions/step_functions.py

@@ -545,10 +545,6 @@ def _batch(self, node):
                     "Parallel steps are not supported yet with AWS step functions."
                 )
 
-            # Inherit the run id from the parent and pass it along to children.
-            attrs["metaflow.run_id.$"] = "$.Parameters.run_id"
-            attrs["run_id.$"] = "$.Parameters.run_id"
-
             # Handle foreach join.
             if (
                 node.type == "join"
@@ -572,6 +568,9 @@ def _batch(self, node):
                 env["METAFLOW_SPLIT_PARENT_TASK_ID"] = (
                     "$.Parameters.split_parent_task_id_%s" % node.split_parents[-1]
                 )
+                # Inherit the run id from the parent and pass it along to children.
+                attrs["metaflow.run_id.$"] = "$.Parameters.run_id"
+                attrs["run_id.$"] = "$.Parameters.run_id"
             else:
                 # Set appropriate environment variables for runtime replacement.
                 if len(node.in_funcs) == 1:
@@ -580,6 +579,9 @@ def _batch(self, node):
                         % node.in_funcs[0]
                     )
                     env["METAFLOW_PARENT_TASK_ID"] = "$.JobId"
+                    # Inherit the run id from the parent and pass it along to children.
+                    attrs["metaflow.run_id.$"] = "$.Parameters.run_id"
+                    attrs["run_id.$"] = "$.Parameters.run_id"
                 else:
                     # Generate the input paths in a quasi-compressed format.
                     # See util.decompress_list for why this is written the way
@@ -589,6 +591,9 @@ def _batch(self, node):
                         "${METAFLOW_PARENT_%s_TASK_ID}" % (idx, idx)
                         for idx, _ in enumerate(node.in_funcs)
                     )
+                    # Inherit the run id from the parent and pass it along to children.
+                    attrs["metaflow.run_id.$"] = "$.[0].Parameters.run_id"
+                    attrs["run_id.$"] = "$.[0].Parameters.run_id"
                     for idx, _ in enumerate(node.in_funcs):
                         env["METAFLOW_PARENT_%s_TASK_ID" % idx] = "$.[%s].JobId" % idx
                         env["METAFLOW_PARENT_%s_STEP" % idx] = (