[runtime] run-hbs for local runtime's run command (Netflix#2150)

madhur-ob · Nov 26, 2024 · 15c86b9 · 15c86b9
1 parent a3d087e
commit 15c86b9
Showing 1 changed file with 15 additions and 14 deletions.
diff --git a/metaflow/cli.py b/metaflow/cli.py
@@ -802,20 +802,21 @@ def run(
             "msg": "Starting run",
         },
     )
-    runtime.print_workflow_info()
-    runtime.persist_constants()
-
-    if runner_attribute_file:
-        with open(runner_attribute_file, "w", encoding="utf-8") as f:
-            json.dump(
-                {
-                    "run_id": runtime.run_id,
-                    "flow_name": obj.flow.name,
-                    "metadata": obj.metadata.metadata_str(),
-                },
-                f,
-            )
-    runtime.execute()
+    with runtime.run_heartbeat():
+        runtime.print_workflow_info()
+        runtime.persist_constants()
+
+        if runner_attribute_file:
+            with open(runner_attribute_file, "w", encoding="utf-8") as f:
+                json.dump(
+                    {
+                        "run_id": runtime.run_id,
+                        "flow_name": obj.flow.name,
+                        "metadata": obj.metadata.metadata_str(),
+                    },
+                    f,
+                )
+        runtime.execute()
 
 
 def write_file(file_path, content):