[heartbeats] Better error handling (#2193)

* [robust heartbeats] - Ensure better error handling so that any server related flakiness can be handled without crashing the heartbeat thread. * Update heartbeat.py --------- Co-authored-by: Savin <[email protected]>
Netflix · Jan 4, 2025 · e92990f · e92990f
1 parent 0f31e25
commit e92990f
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/metaflow/metadata_provider/heartbeat.py b/metaflow/metadata_provider/heartbeat.py
@@ -53,13 +53,27 @@ def _ping(self):
                 retry_counter = 0
             except HeartBeatException as e:
                 retry_counter = retry_counter + 1
-                time.sleep(4**retry_counter)
+                time.sleep(1.5**retry_counter)
 
     def _heartbeat(self):
         if self.hb_url is not None:
-            response = requests.post(
-                url=self.hb_url, data="{}", headers=self.headers.copy()
-            )
+            try:
+                response = requests.post(
+                    url=self.hb_url, data="{}", headers=self.headers.copy()
+                )
+            except requests.exceptions.ConnectionError as e:
+                raise HeartBeatException(
+                    "HeartBeat request (%s) failed" " (ConnectionError)" % (self.hb_url)
+                )
+            except requests.exceptions.Timeout as e:
+                raise HeartBeatException(
+                    "HeartBeat request (%s) failed" " (Timeout)" % (self.hb_url)
+                )
+            except requests.exceptions.RequestException as e:
+                raise HeartBeatException(
+                    "HeartBeat request (%s) failed"
+                    " (RequestException) %s" % (self.hb_url, str(e))
+                )
             # Unfortunately, response.json() returns a string that we need
             # to cast to json; however when the request encounters an error
             # the return type is a json blob :/