fix: update heartbeat method using threads

Author: justinmerrell

requirements.txt

@@ -1,7 +1,6 @@
 aiohttp >= 3.8.4
 backoff == 2.2.1
 boto3 >= 1.26.142
-importlib-metadata == 6.6.0
 pillow >= 9.3.0
 py-cpuinfo == 9.0.0
 python-dotenv >= 0.21.0

runpod/serverless/modules/heartbeat.py

@@ -1,6 +1,7 @@
 """Enables heartbeats."""
 
 import os
+import time
 import threading
 
 import requests
@@ -12,36 +13,47 @@
 _session.headers.update({"Authorization": f"{os.environ.get('RUNPOD_AI_API_KEY')}"})
 
 
-def _send_ping(ping_params=None):
-    if PING_URL not in [None, 'PING_URL_NOT_SET']:
-        try:
-            result = _session.get(
-                PING_URL,
-                params=ping_params,
-                timeout=int(PING_INTERVAL / 1000)
-            )
-
-            log.info(f"Heartbeat Sent  URL: {PING_URL}  Status: {result.status_code}")
-            log.info(f"Heartbeat Sent  Interval: {PING_INTERVAL}ms  Params: {ping_params}")
-
-        except Exception as err:  # pylint: disable=broad-except
-            log.error(f"Heartbeat Failed  URL: {PING_URL}  Params: {ping_params}")
-            log.error(f"Heartbeat Fail  Error: {err}")
-
-
-def start_ping():
-    """
-    Pings the heartbeat endpoint at the specified interval.
-    """
-    job_id = get_current_job_id()
-
-    ping_params = {
-        'job_id': job_id,
-    } if job_id is not None else None
-
-    _send_ping(ping_params)
-
-    log.debug(f"Scheduling next heartbeat in {PING_INTERVAL}ms")
-    heartbeat_thread = threading.Timer(int(PING_INTERVAL / 1000), start_ping)
-    heartbeat_thread.daemon = True
-    heartbeat_thread.start()
+class HeartbeatSender:
+    ''' Sends heartbeats to the Runpod server. '''
+
+    def __init__(self):
+        self._thread = threading.Thread(target=self._run, daemon=True)
+
+    def start_ping(self):
+        '''
+        Starts the heartbeat thread.
+        '''
+        self._thread.start()
+
+    def _run(self):
+        '''
+        Sends heartbeats to the Runpod server.
+        '''
+        while True:
+            self._send_ping()
+            time.sleep(int(PING_INTERVAL / 1000))
+
+    def _send_ping(self):
+        '''
+        Sends a heartbeat to the Runpod server.
+        '''
+        job_id = get_current_job_id()
+
+        ping_params = {
+            'job_id': job_id,
+        } if job_id is not None else None
+
+        if PING_URL not in [None, 'PING_URL_NOT_SET']:
+            try:
+                result = _session.get(
+                    PING_URL,
+                    params=ping_params,
+                    timeout=int(PING_INTERVAL / 1000)
+                )
+
+                log.info(f"Heartbeat Sent  URL: {PING_URL}  Status: {result.status_code}")
+                log.info(f"Heartbeat Sent  Interval: {PING_INTERVAL}ms  Params: {ping_params}")
+
+            except Exception as err:  # pylint: disable=broad-except
+                log.error(f"Heartbeat Failed  URL: {PING_URL}  Params: {ping_params}")
+                log.error(f"Heartbeat Fail  Error: {err}")

runpod/serverless/modules/rp_fastapi.py

@@ -9,7 +9,10 @@
 
 from .job import run_job
 from .worker_state import set_job_id
-from .heartbeat import start_ping
+from .heartbeat import HeartbeatSender
+
+
+heartbeat = HeartbeatSender()
 
 
 class Job(BaseModel):
@@ -29,7 +32,7 @@ def __init__(self, handler=None):
         3. Sets the handler for processing jobs.
         '''
         # Start the heartbeat thread.
-        start_ping()
+        heartbeat.start_ping()
 
         # Set the handler for processing jobs.
         self.config = {"handler": handler}

runpod/serverless/utils/rp_download.py

@@ -13,15 +13,12 @@
 from typing import List, Union
 from urllib.parse import urlparse
 from email import message_from_string
-from importlib.metadata import version
 from concurrent.futures import ThreadPoolExecutor
 
 import backoff
 import requests
 
-HEADERS = {
-    "User-Agent": f"runpod-python/{version('runpod-python')} (https://runpod.io; [email protected])"
-}
+HEADERS = {"User-Agent": "runpod-python/0.0.0 (https://runpod.io; [email protected])"}
 
 
 def calculate_chunk_size(file_size: int) -> int:

runpod/serverless/work_loop.py

@@ -9,14 +9,16 @@
 import aiohttp
 
 import runpod.serverless.modules.logging as log
-from .modules.heartbeat import start_ping
+from .modules.heartbeat import HeartbeatSender
 from .modules.job import get_job, run_job, send_result
 from .modules.worker_state import REF_COUNT_ZERO, set_job_id
 from .utils import rp_debugger
 
 
 _TIMEOUT = aiohttp.ClientTimeout(total=300, connect=2, sock_connect=2)
 
+heartbeat = HeartbeatSender()
+
 
 def _get_auth_header() -> dict:
     '''
@@ -40,7 +42,7 @@ async def start_worker(config):
 
     async with aiohttp.ClientSession(headers=auth_header, timeout=_TIMEOUT) as session:
 
-        start_ping()
+        heartbeat.start_ping()
 
         while True:
             job = await get_job(session, config)

setup.cfg

@@ -29,7 +29,6 @@ install_requires =
     aiohttp >= 3.8.3
     backoff >= 2.2.1
     boto3 >= 1.26.142
-    importlib-metadata >= 6.6.0
     pillow >= 9.3.0
     py-cpuinfo >= 9.0.0
     python-dotenv >= 0.21.0