fix: change heartbeat back to daemon thread

Author: zhl146

runpod/serverless/modules/heartbeat.py

@@ -1,42 +1,50 @@
 '''enables heartbeats'''
 
-import asyncio
+import time
+import os
+
+import requests
 
 import runpod.serverless.modules.logging as log
 from .worker_state import get_current_job_id, ping_url, ping_interval
 
+# COUNTER = 0
+
 
-async def send_ping(session):
+def heartbeat_ping(session):
     '''
-    sends the http request
+    Pings the heartbeat endpoint
     '''
+    ping_params = None
+
     try:
-        ping_params = None
         job_id = get_current_job_id()
 
         if job_id is not None:
             ping_params = {
                 'job_id': job_id,
             }
-        await session.get(ping_url, params=ping_params,
-                          timeout=int(ping_interval/1000))
 
-        log.debug(f'Heartbeat URL: {ping_url} Interval: {ping_interval}ms')
-        log.debug(f"Heartbeat Params: {ping_params}")
+            session.get(
+                ping_url,
+                params=ping_params,
+                timeout=int(ping_interval/1000)
+            )
 
+        log.info(
+            f'Heartbeat sent to {ping_url} interval: {ping_interval}ms params: {ping_params}')
     except Exception as err:  # pylint: disable=broad-except
-        log.warn(f"Error while sending heartbeat: {err}")
+        log.error(err)
 
 
-async def heartbeat_ping(session):
+def start_heartbeat():
     '''
     manages heartbeat timing
     '''
-    if ping_url is not None:
-        while True:
-
-            asyncio.create_task(
-                send_ping(session)
-            )
 
-            await asyncio.sleep(ping_interval/1000)
+    session = requests.Session()
+    session.headers.update({
+        "Authorization": f"{os.environ.get('RUNPOD_AI_API_KEY')}"})
+    while True:
+        heartbeat_ping(session)
+        time.sleep(ping_interval/1000)

runpod/serverless/work_loop.py

@@ -5,12 +5,12 @@
 
 import os
 import json
-import asyncio
+from threading import Thread
 
 import aiohttp
 
 import runpod.serverless.modules.logging as log
-from .modules.heartbeat import heartbeat_ping
+from .modules.heartbeat import start_heartbeat
 from .modules.job import get_job, run_job, send_result
 from .modules.worker_state import set_job_id
 
@@ -27,7 +27,9 @@ async def start_worker(config):
 
     async with aiohttp.ClientSession(headers=auth_header) as session:
 
-        asyncio.create_task(heartbeat_ping(session))
+        heartbeat_thread = Thread(target=start_heartbeat, daemon=True)
+        heartbeat_thread.daemon = True
+        heartbeat_thread.start()
 
         while True:
             # GET JOB
@@ -49,8 +51,10 @@ async def start_worker(config):
             try:
                 job_data = json.dumps(job_result, ensure_ascii=False)
             except Exception as err:  # pylint: disable=broad-except
-                log.error(f"Error while serializing job result {job['id']}: {err}")
-                job_data = json.dumps({"error": "unable to serialize job output"})
+                log.error(
+                    f"Error while serializing job result {job['id']}: {err}")
+                job_data = json.dumps(
+                    {"error": "unable to serialize job output"})
 
             # SEND RESULTS
             await send_result(session, job_data, job)