From 15c86b9db509edd14600b9a017865bc0348aba78 Mon Sep 17 00:00:00 2001
From: Valay Dave <valaygaurang@gmail.com>
Date: Mon, 25 Nov 2024 18:23:57 -0800
Subject: [PATCH] [runtime] run-hbs for local runtime's `run` command (#2150)

---
 metaflow/cli.py | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/metaflow/cli.py b/metaflow/cli.py
index a318b84a3ec..498ea1b74b2 100644
--- 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):