fix: lazy-loaded get_heartbeat and get_jobs_progress

Author: deanq

runpod/serverless/modules/rp_fastapi.py

@@ -16,8 +16,8 @@
 from ...version import __version__ as runpod_version
 from .rp_handler import is_generator
 from .rp_job import run_job, run_job_generator
-from .rp_ping import Heartbeat
-from .worker_state import Job, JobsProgress
+from .worker_state import Job, get_jobs_progress
+from .rp_ping import get_heartbeat
 
 RUNPOD_ENDPOINT_ID = os.environ.get("RUNPOD_ENDPOINT_ID", None)
 
@@ -96,8 +96,6 @@
 
 
 # ------------------------------ Initializations ----------------------------- #
-job_list = JobsProgress()
-heartbeat = Heartbeat()
 
 
 # ------------------------------- Input Objects ------------------------------ #
@@ -185,7 +183,7 @@ def __init__(self, config: Dict[str, Any]):
         3. Sets the handler for processing jobs.
         """
         # Start the heartbeat thread.
-        heartbeat.start_ping()
+        get_heartbeat().start_ping()
 
         self.config = config
 
@@ -286,12 +284,12 @@ async def _realtime(self, job: Job):
         Performs model inference on the input data using the provided handler.
         If handler is not provided, returns an error message.
         """
-        job_list.add(job.id)
+        get_jobs_progress().add(job.id)
 
         # Process the job using the provided handler, passing in the job input.
         job_results = await run_job(self.config["handler"], job.__dict__)
 
-        job_list.remove(job.id)
+        get_jobs_progress().remove(job.id)
 
         # Return the results of the job processing.
         return jsonable_encoder(job_results)
@@ -304,7 +302,7 @@ async def _realtime(self, job: Job):
     async def _sim_run(self, job_request: DefaultRequest) -> JobOutput:
         """Development endpoint to simulate run behavior."""
         assigned_job_id = f"test-{uuid.uuid4()}"
-        job_list.add({
+        get_jobs_progress().add({
             "id": assigned_job_id,
             "input": job_request.input,
             "webhook": job_request.webhook
@@ -345,7 +343,7 @@ async def _sim_runsync(self, job_request: DefaultRequest) -> JobOutput:
     # ---------------------------------- stream ---------------------------------- #
     async def _sim_stream(self, job_id: str) -> StreamOutput:
         """Development endpoint to simulate stream behavior."""
-        stashed_job = job_list.get(job_id)
+        stashed_job = get_jobs_progress().get(job_id)
         if stashed_job is None:
             return jsonable_encoder(
                 {"id": job_id, "status": "FAILED", "error": "Job ID not found"}
@@ -367,7 +365,7 @@ async def _sim_stream(self, job_id: str) -> StreamOutput:
                 }
             )
 
-        job_list.remove(job.id)
+        get_jobs_progress().remove(job.id)
 
         if stashed_job.webhook:
             thread = threading.Thread(
@@ -384,7 +382,7 @@ async def _sim_stream(self, job_id: str) -> StreamOutput:
     # ---------------------------------- status ---------------------------------- #
     async def _sim_status(self, job_id: str) -> JobOutput:
         """Development endpoint to simulate status behavior."""
-        stashed_job = job_list.get(job_id)
+        stashed_job = get_jobs_progress().get(job_id)
         if stashed_job is None:
             return jsonable_encoder(
                 {"id": job_id, "status": "FAILED", "error": "Job ID not found"}
@@ -400,7 +398,7 @@ async def _sim_status(self, job_id: str) -> JobOutput:
         else:
             job_output = await run_job(self.config["handler"], job.__dict__)
 
-        job_list.remove(job.id)
+        get_jobs_progress().remove(job.id)
 
         if job_output.get("error", None):
             return jsonable_encoder(

runpod/serverless/modules/rp_job.py

@@ -18,12 +18,11 @@
 from .rp_handler import is_generator
 from .rp_http import send_result, stream_result
 from .rp_tips import check_return_size
-from .worker_state import WORKER_ID, REF_COUNT_ZERO, JobsProgress
+from .worker_state import WORKER_ID, REF_COUNT_ZERO, get_jobs_progress
 
 JOB_GET_URL = str(os.environ.get("RUNPOD_WEBHOOK_GET_JOB")).replace("$ID", WORKER_ID)
 
 log = RunPodLogger()
-job_progress = JobsProgress()
 
 
 def _job_get_url(batch_size: int = 1):
@@ -43,7 +42,7 @@ def _job_get_url(batch_size: int = 1):
     else:
         job_take_url = JOB_GET_URL
 
-    job_in_progress = "1" if job_progress.get_job_list() else "0"
+    job_in_progress = "1" if get_jobs_progress().get_job_list() else "0"
     job_take_url += f"&job_in_progress={job_in_progress}"
 
     log.debug(f"rp_job | get_job: {job_take_url}")

runpod/serverless/modules/rp_ping.py

@@ -11,11 +11,27 @@
 
 from runpod.http_client import SyncClientSession
 from runpod.serverless.modules.rp_logger import RunPodLogger
-from runpod.serverless.modules.worker_state import WORKER_ID, JobsProgress
+from runpod.serverless.modules.worker_state import WORKER_ID, get_jobs_progress
 from runpod.version import __version__ as runpod_version
 
 log = RunPodLogger()
-jobs = JobsProgress()  # Contains the list of jobs that are currently running.
+
+# Lazy loading for Heartbeat instance
+_heartbeat_instance = None
+
+
+def get_heartbeat():
+    """Get the global Heartbeat instance with lazy initialization."""
+    global _heartbeat_instance
+    if _heartbeat_instance is None:
+        _heartbeat_instance = Heartbeat()
+    return _heartbeat_instance
+
+
+def reset_heartbeat():
+    """Reset the lazy-loaded Heartbeat instance (useful for testing)."""
+    global _heartbeat_instance
+    _heartbeat_instance = None
 
 
 class Heartbeat:
@@ -97,7 +113,7 @@ def _send_ping(self):
         """
         Sends a heartbeat to the Runpod server.
         """
-        job_ids = jobs.get_job_list()
+        job_ids = get_jobs_progress().get_job_list()
         ping_params = {"job_id": job_ids, "runpod_version": runpod_version}
 
         try:

runpod/serverless/modules/rp_scale.py

@@ -12,10 +12,9 @@
 from ...http_client import AsyncClientSession, ClientSession, TooManyRequests
 from .rp_job import get_job, handle_job
 from .rp_logger import RunPodLogger
-from .worker_state import JobsProgress, IS_LOCAL_TEST
+from .worker_state import IS_LOCAL_TEST, get_jobs_progress
 
 log = RunPodLogger()
-job_progress = JobsProgress()
 
 
 def _handle_uncaught_exception(exc_type, exc_value, exc_traceback):
@@ -149,7 +148,7 @@ def kill_worker(self):
 
     def current_occupancy(self) -> int:
         current_queue_count = self.jobs_queue.qsize()
-        current_progress_count = job_progress.get_job_count()
+        current_progress_count = get_jobs_progress().get_job_count()
 
         log.debug(
             f"JobScaler.status | concurrency: {self.current_concurrency}; queue: {current_queue_count}; progress: {current_progress_count}"
@@ -188,7 +187,7 @@ async def get_jobs(self, session: ClientSession):
 
                 for job in acquired_jobs:
                     await self.jobs_queue.put(job)
-                    job_progress.add(job)
+                    get_jobs_progress().add(job)
                     log.debug("Job Queued", job["id"])
 
                 log.info(f"Jobs in queue: {self.jobs_queue.qsize()}")
@@ -268,6 +267,6 @@ async def handle_job(self, session: ClientSession, job: dict):
             self.jobs_queue.task_done()
 
             # Job is no longer in progress
-            job_progress.remove(job)
+            get_jobs_progress().remove(job)
 
             log.debug("Finished Job", job["id"])

runpod/serverless/modules/worker_state.py

@@ -15,6 +15,24 @@
 
 log = RunPodLogger()
 
+
+# ----------------------------- Lazy Loading Utilities -------------------------- #
+_jobs_progress_instance = None
+
+
+def get_jobs_progress():
+    """Get the global JobsProgress instance with lazy initialization."""
+    global _jobs_progress_instance
+    if _jobs_progress_instance is None:
+        _jobs_progress_instance = JobsProgress()
+    return _jobs_progress_instance
+
+
+def reset_jobs_progress():
+    """Reset the lazy-loaded JobsProgress instance (useful for testing)."""
+    global _jobs_progress_instance
+    _jobs_progress_instance = None
+
 REF_COUNT_ZERO = time.perf_counter()  # Used for benchmarking with the debugger.
 
 WORKER_ID = os.environ.get("RUNPOD_POD_ID", str(uuid.uuid4()))
@@ -72,6 +90,8 @@ class JobsProgress:
     _shared_data: Optional[Any] = None
     _lock: Optional[Any] = None
     _use_multiprocessing: bool = True
+    _fallback_jobs: list = []
+    _fallback_lock: Optional[threading.Lock] = None
 
     def __new__(cls):
         if cls._instance is None:
@@ -99,10 +119,10 @@ def __repr__(self) -> str:
         return f"<{self.__class__.__name__}>: {self.get_job_list()}"
 
     def clear(self) -> None:
-        if self._use_multiprocessing:
+        if self._use_multiprocessing and self._lock is not None and self._shared_data is not None:
             with self._lock:
                 self._shared_data['jobs'][:] = []
-        else:
+        elif not self._use_multiprocessing and self._fallback_lock is not None:
             with self._fallback_lock:
                 self._fallback_jobs.clear()
 
@@ -119,12 +139,12 @@ def add(self, element: Any):
         else:
             raise TypeError("Only Job objects can be added to JobsProgress.")
 
-        if self._use_multiprocessing:
+        if self._use_multiprocessing and self._lock is not None and self._shared_data is not None:
             with self._lock:
                 job_list = self._shared_data['jobs']
                 if not any(job['id'] == job_dict['id'] for job in job_list):
                     job_list.append(job_dict)
-        else:
+        elif not self._use_multiprocessing and self._fallback_lock is not None:
             with self._fallback_lock:
                 if not any(job['id'] == job_dict['id'] for job in self._fallback_jobs):
                     self._fallback_jobs.append(job_dict)
@@ -144,13 +164,13 @@ def get(self, element: Any) -> Optional[Job]:
         else:
             raise TypeError("Only Job objects can be retrieved from JobsProgress.")
 
-        if self._use_multiprocessing:
+        if self._use_multiprocessing and self._lock is not None and self._shared_data is not None:
             with self._lock:
                 for job_dict in self._shared_data['jobs']:
                     if job_dict['id'] == search_id:
                         log.debug(f"JobsProgress | Retrieved job: {job_dict['id']}")
                         return Job(**job_dict)
-        else:
+        elif not self._use_multiprocessing and self._fallback_lock is not None:
             with self._fallback_lock:
                 for job_dict in self._fallback_jobs:
                     if job_dict['id'] == search_id:
@@ -171,15 +191,15 @@ def remove(self, element: Any):
         else:
             raise TypeError("Only Job objects can be removed from JobsProgress.")
 
-        if self._use_multiprocessing:
+        if self._use_multiprocessing and self._lock is not None and self._shared_data is not None:
             with self._lock:
                 job_list = self._shared_data['jobs']
                 for i, job_dict in enumerate(job_list):
                     if job_dict['id'] == job_id:
                         del job_list[i]
                         log.debug(f"JobsProgress | Removed job: {job_dict['id']}")
                         break
-        else:
+        elif not self._use_multiprocessing and self._fallback_lock is not None:
             with self._fallback_lock:
                 for i, job_dict in enumerate(self._fallback_jobs):
                     if job_dict['id'] == job_id:
@@ -191,10 +211,11 @@ def get_job_list(self) -> Optional[str]:
         """
         Returns the list of job IDs as comma-separated string.
         """
-        if self._use_multiprocessing:
+        job_list = []
+        if self._use_multiprocessing and self._lock is not None and self._shared_data is not None:
             with self._lock:
                 job_list = list(self._shared_data['jobs'])
-        else:
+        elif not self._use_multiprocessing and self._fallback_lock is not None:
             with self._fallback_lock:
                 job_list = list(self._fallback_jobs)
 
@@ -208,19 +229,21 @@ def get_job_count(self) -> int:
         """
         Returns the number of jobs.
         """
-        if self._use_multiprocessing:
+        if self._use_multiprocessing and self._lock is not None and self._shared_data is not None:
             with self._lock:
                 return len(self._shared_data['jobs'])
-        else:
+        elif not self._use_multiprocessing and self._fallback_lock is not None:
             with self._fallback_lock:
                 return len(self._fallback_jobs)
+        return 0
 
     def __iter__(self):
         """Make the class iterable - returns Job objects"""
-        if self._use_multiprocessing:
+        job_dicts = []
+        if self._use_multiprocessing and self._lock is not None and self._shared_data is not None:
             with self._lock:
                 job_dicts = list(self._shared_data['jobs'])
-        else:
+        elif not self._use_multiprocessing and self._fallback_lock is not None:
             with self._fallback_lock:
                 job_dicts = list(self._fallback_jobs)
         return iter(Job(**job_dict) for job_dict in job_dicts)
@@ -240,9 +263,10 @@ def __contains__(self, element: Any) -> bool:
         else:
             return False
 
-        if self._use_multiprocessing:
+        if self._use_multiprocessing and self._lock is not None and self._shared_data is not None:
             with self._lock:
                 return any(job['id'] == search_id for job in self._shared_data['jobs'])
-        else:
+        elif not self._use_multiprocessing and self._fallback_lock is not None:
             with self._fallback_lock:
                 return any(job['id'] == search_id for job in self._fallback_jobs)
+        return False

runpod/serverless/worker.py

@@ -8,9 +8,9 @@
 from typing import Any, Dict
 
 from runpod.serverless.modules import rp_logger, rp_local, rp_ping, rp_scale
+from runpod.serverless.modules.rp_ping import get_heartbeat
 
 log = rp_logger.RunPodLogger()
-heartbeat = rp_ping.Heartbeat()
 
 
 def _is_local(config) -> bool:
@@ -36,7 +36,7 @@ def run_worker(config: Dict[str, Any]) -> None:
         config (Dict[str, Any]): Configuration parameters for the worker.
     """
     # Start pinging RunPod to show that the worker is alive.
-    heartbeat.start_ping()
+    get_heartbeat().start_ping()
 
     # Create a JobScaler responsible for adjusting the concurrency
     job_scaler = rp_scale.JobScaler(config)

tests/test_cli/test_cli_sanity.py

@@ -20,20 +20,18 @@ def cli_runner():
 @pytest.fixture(autouse=True)
 def reset_jobs_progress():
     """Reset JobsProgress state before each test."""
+    yield
+    # Cleanup after test
     try:
-        from runpod.serverless.modules.worker_state import JobsProgress
-        JobsProgress._instance = None
-        yield
-        # Cleanup after test
-        if hasattr(JobsProgress, '_instance') and JobsProgress._instance:
-            try:
-                JobsProgress._instance.clear()
-            except Exception:
-                pass
-        JobsProgress._instance = None
-    except ImportError:
-        # JobsProgress might not be available in all test contexts
-        yield
+        from runpod.serverless.modules.worker_state import reset_jobs_progress, JobsProgress
+        from runpod.serverless.modules.rp_ping import reset_heartbeat
+        reset_jobs_progress()
+        reset_heartbeat()
+        # Also reset the singleton instance
+        if hasattr(JobsProgress, '_instance'):
+            JobsProgress._instance = None
+    except (ImportError, AttributeError):
+        pass
 
 
 class TestCLISanity: