feat: update metadata on ec2 spot termination for batch (Netflix#2271)

* feat: update metadata on ec2 spot termination for batch

* fix: Set METAFLOW_FLOW_FILENAME before deco variables

metaflow/plugins/aws/batch/batch_cli.py

@@ -211,7 +211,7 @@ def step(
     log_driver=None,
     log_options=None,
     num_parallel=None,
-    **kwargs
+    **kwargs,
 ):
     def echo(msg, stream="stderr", batch_id=None, **kwargs):
         msg = util.to_unicode(msg)
@@ -273,11 +273,11 @@ def echo(msg, stream="stderr", batch_id=None, **kwargs):
         "metaflow_version"
     ]
 
+    env = {"METAFLOW_FLOW_FILENAME": os.path.basename(sys.argv[0])}
+
     env_deco = [deco for deco in node.decorators if deco.name == "environment"]
     if env_deco:
-        env = env_deco[0].attributes["vars"]
-    else:
-        env = {}
+        env.update(env_deco[0].attributes["vars"])
 
     # Add the environment variables related to the input-paths argument
     if split_vars:

metaflow/plugins/aws/batch/batch_decorator.py

@@ -298,6 +298,13 @@ def task_pre_step(
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
+            # Start spot termination monitor sidecar.
+            current._update_env(
+                {"spot_termination_notice": "/tmp/spot_termination_notice"}
+            )
+            self._spot_monitor_sidecar = Sidecar("spot_termination_monitor")
+            self._spot_monitor_sidecar.start()
+
         num_parallel = int(os.environ.get("AWS_BATCH_JOB_NUM_NODES", 0))
         if num_parallel >= 1 and ubf_context == UBF_CONTROL:
             # UBF handling for multinode case
@@ -350,6 +357,7 @@ def task_finished(
 
         try:
             self._save_logs_sidecar.terminate()
+            self._spot_monitor_sidecar.terminate()
         except:
             # Best effort kill
             pass