fix: JobsProgress uses _MultiprocessingStorage with _ThreadSafeStorage as fallback

Author: deanq

runpod/serverless/modules/rp_scale.py

@@ -101,7 +101,7 @@ def start(self):
             signal.signal(signal.SIGTERM, self.handle_shutdown)
             signal.signal(signal.SIGINT, self.handle_shutdown)
         except ValueError:
-            log.warning("Signal handling is only supported in the main thread.")
+            log.warn("Signal handling is only supported in the main thread.")
 
         # Start the main loop
         # Run forever until the worker is signalled to shut down.

runpod/serverless/modules/worker_state.py

@@ -5,9 +5,10 @@
 import os
 import time
 import uuid
+import threading
 from multiprocessing import Manager
 from multiprocessing.managers import SyncManager
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, List
 
 from .rp_logger import RunPodLogger
 
@@ -63,45 +64,148 @@ def __str__(self) -> str:
 # ---------------------------------------------------------------------------- #
 #                                    Tracker                                   #
 # ---------------------------------------------------------------------------- #
+
+class _JobStorage:
+    """Abstract storage backend for jobs."""
+    
+    def add_job(self, job_dict: Dict[str, Any]) -> None:
+        raise NotImplementedError
+    
+    def remove_job(self, job_id: str) -> None:
+        raise NotImplementedError
+        
+    def get_job(self, job_id: str) -> Optional[Dict[str, Any]]:
+        raise NotImplementedError
+        
+    def get_all_jobs(self) -> List[Dict[str, Any]]:
+        raise NotImplementedError
+        
+    def clear_jobs(self) -> None:
+        raise NotImplementedError
+    
+    def job_exists(self, job_id: str) -> bool:
+        raise NotImplementedError
+
+
+class _MultiprocessingStorage(_JobStorage):
+    """Multiprocessing-based storage for GIL-free operation."""
+    
+    def __init__(self):
+        self._manager = Manager()
+        self._shared_data = self._manager.dict()
+        self._shared_data['jobs'] = self._manager.list()
+        self._lock = self._manager.Lock()
+        
+    def add_job(self, job_dict: Dict[str, Any]) -> 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)
+    
+    def remove_job(self, job_id: str) -> None:
+        with self._lock:
+            job_list = self._shared_data['jobs']
+            for i, job in enumerate(job_list):
+                if job['id'] == job_id:
+                    del job_list[i]
+                    break
+    
+    def get_job(self, job_id: str) -> Optional[Dict[str, Any]]:
+        with self._lock:
+            for job in self._shared_data['jobs']:
+                if job['id'] == job_id:
+                    return dict(job)
+        return None
+    
+    def get_all_jobs(self) -> List[Dict[str, Any]]:
+        with self._lock:
+            return [dict(job) for job in self._shared_data['jobs']]
+    
+    def clear_jobs(self) -> None:
+        with self._lock:
+            self._shared_data['jobs'][:] = []
+    
+    def job_exists(self, job_id: str) -> bool:
+        with self._lock:
+            return any(job['id'] == job_id for job in self._shared_data['jobs'])
+
+
+class _ThreadSafeStorage(_JobStorage):
+    """Thread-safe storage fallback when multiprocessing fails."""
+    
+    def __init__(self):
+        self._jobs: List[Dict[str, Any]] = []
+        self._lock = threading.Lock()
+    
+    def add_job(self, job_dict: Dict[str, Any]) -> None:
+        with self._lock:
+            if not any(job['id'] == job_dict['id'] for job in self._jobs):
+                self._jobs.append(job_dict)
+    
+    def remove_job(self, job_id: str) -> None:
+        with self._lock:
+            for i, job in enumerate(self._jobs):
+                if job['id'] == job_id:
+                    del self._jobs[i]
+                    break
+    
+    def get_job(self, job_id: str) -> Optional[Dict[str, Any]]:
+        with self._lock:
+            for job in self._jobs:
+                if job['id'] == job_id:
+                    return job.copy()
+        return None
+    
+    def get_all_jobs(self) -> List[Dict[str, Any]]:
+        with self._lock:
+            return self._jobs.copy()
+    
+    def clear_jobs(self) -> None:
+        with self._lock:
+            self._jobs.clear()
+    
+    def job_exists(self, job_id: str) -> bool:
+        with self._lock:
+            return any(job['id'] == job_id for job in self._jobs)
+
+
 class JobsProgress:
-    """Track the state of current jobs in progress using shared memory."""
+    """Track the state of current jobs in progress using shared memory or thread-safe fallback."""
 
     _instance: Optional['JobsProgress'] = None
-    # Singleton
+    _storage: Optional[_JobStorage] = None
+
     def __new__(cls):
         if cls._instance is None:
-            cls._instance = super().__new__(cls)
+            cls._instance = object.__new__(cls)
+            cls._instance._storage = None
         return cls._instance
 
     def __init__(self):
-        if not hasattr(self, '_initialized'):
-            self._manager: Optional[SyncManager] = None
-            self._shared_data: Optional[Any] = None
-            self._lock: Optional[Any] = None
-            self._initialized = True
-
+        pass
+    
     def _ensure_initialized(self):
-        """Initialize the multiprocessing manager and shared data structures only when needed."""
-        if self._manager is None:
-            self._manager = Manager()
-            self._shared_data = self._manager.dict()
-            self._shared_data['jobs'] = self._manager.list()
-            self._lock = self._manager.Lock()
+        """Lazily initialize storage backend."""
+        if self._storage is None:
+            try:
+                self._storage = _MultiprocessingStorage()
+                log.debug("JobsProgress | Using multiprocessing for GIL-free operation")
+            except Exception as e:
+                log.warn(f"JobsProgress | Multiprocessing failed ({e}), falling back to thread-safe mode")
+                self._storage = _ThreadSafeStorage()
 
     def __repr__(self) -> str:
         return f"<{self.__class__.__name__}>: {self.get_job_list()}"
 
     def clear(self) -> None:
         self._ensure_initialized()
-        with self._lock:
-            self._shared_data['jobs'][:] = []
+        self._storage.clear_jobs()
 
     def add(self, element: Any):
         """
         Adds a Job object to the set.
         """
         self._ensure_initialized()
-        
         if isinstance(element, str):
             job_dict = {'id': element}
         elif isinstance(element, dict):
@@ -111,16 +215,8 @@ def add(self, element: Any):
         else:
             raise TypeError("Only Job objects can be added to JobsProgress.")
 
-        with self._lock:
-            # Check if job already exists
-            job_list = self._shared_data['jobs']
-            for existing_job in job_list:
-                if existing_job['id'] == job_dict['id']:
-                    return  # Job already exists
-            
-            # Add new job
-            job_list.append(job_dict)
-            log.debug(f"JobsProgress | Added job: {job_dict['id']}")
+        self._storage.add_job(job_dict)
+        log.debug(f"JobsProgress | Added job: {job_dict['id']}")
 
     def get(self, element: Any) -> Optional[Job]:
         """
@@ -129,28 +225,24 @@ def get(self, element: Any) -> Optional[Job]:
         If the element is a string, searches for Job with that id.
         """
         self._ensure_initialized()
-        
         if isinstance(element, str):
             search_id = element
         elif isinstance(element, Job):
             search_id = element.id
         else:
             raise TypeError("Only Job objects can be retrieved from JobsProgress.")
 
-        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)
-        
+        job_dict = self._storage.get_job(search_id)
+        if job_dict:
+            log.debug(f"JobsProgress | Retrieved job: {job_dict['id']}")
+            return Job(**job_dict)
         return None
 
     def remove(self, element: Any):
         """
         Removes a Job object from the set.
         """
         self._ensure_initialized()
-        
         if isinstance(element, str):
             job_id = element
         elif isinstance(element, dict):
@@ -160,23 +252,15 @@ def remove(self, element: Any):
         else:
             raise TypeError("Only Job objects can be removed from JobsProgress.")
 
-        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
+        self._storage.remove_job(job_id)
+        log.debug(f"JobsProgress | Removed job: {job_id}")
 
     def get_job_list(self) -> Optional[str]:
         """
         Returns the list of job IDs as comma-separated string.
         """
-        if self._manager is None:
-            return None
-            
-        with self._lock:
-            job_list = list(self._shared_data['jobs'])
+        self._ensure_initialized()
+        job_list = self._storage.get_all_jobs()
 
         if not job_list:
             return None
@@ -188,22 +272,13 @@ def get_job_count(self) -> int:
         """
         Returns the number of jobs.
         """
-        if self._manager is None:
-            return 0
-            
-        with self._lock:
-            return len(self._shared_data['jobs'])
+        self._ensure_initialized()
+        return len(self._storage.get_all_jobs())
 
     def __iter__(self):
         """Make the class iterable - returns Job objects"""
-        if self._manager is None:
-            return iter([])
-            
-        with self._lock:
-            # Create a snapshot of jobs to avoid holding lock during iteration
-            job_dicts = list(self._shared_data['jobs'])
-        
-        # Return an iterator of Job objects
+        self._ensure_initialized()
+        job_dicts = self._storage.get_all_jobs()
         return iter(Job(**job_dict) for job_dict in job_dicts)
 
     def __len__(self):
@@ -212,9 +287,7 @@ def __len__(self):
 
     def __contains__(self, element: Any) -> bool:
         """Support 'in' operator"""
-        if self._manager is None:
-            return False
-            
+        self._ensure_initialized()
         if isinstance(element, str):
             search_id = element
         elif isinstance(element, Job):
@@ -224,8 +297,4 @@ def __contains__(self, element: Any) -> bool:
         else:
             return False
 
-        with self._lock:
-            for job_dict in self._shared_data['jobs']:
-                if job_dict['id'] == search_id:
-                    return True
-        return False
+        return self._storage.job_exists(search_id)

tests/test_cli/test_cli_sanity.py

@@ -84,23 +84,32 @@ def test_import_safety(self):
         try:
             from runpod.serverless.modules.worker_state import JobsProgress
             jobs = JobsProgress()
-            # Ensure lazy initialization is working
-            self.assertIsNone(jobs._manager, 
-                            "Manager should not be created until first use")
+            # Ensure lazy initialization is working - storage should be None until first use
+            self.assertIsNone(jobs._storage, 
+                            "Storage backend should not be created until first use")
             self.assertTrue(True, "JobsProgress import and instantiation successful")
         except Exception as e:
             self.fail(f"Failed to import/instantiate JobsProgress: {e}")
 
-        # Test that read-only operations work efficiently
+        # Test that operations trigger proper initialization
         try:
             from runpod.serverless.modules.worker_state import JobsProgress
             jobs = JobsProgress()
-            count = jobs.get_job_count()  # Should work without heavy initialization
+            # Initially no storage backend
+            self.assertIsNone(jobs._storage, "Storage should be None initially")
+            
+            # First operation should initialize storage backend
+            count = jobs.get_job_count()
             self.assertEqual(count, 0)
-            self.assertIsNone(jobs._manager, 
-                            "Manager should not be created for read-only operations")
+            self.assertIsNotNone(jobs._storage, 
+                               "Storage backend should be created after first operation")
+            
+            # Verify storage backend is one of the expected types
+            from runpod.serverless.modules.worker_state import _MultiprocessingStorage, _ThreadSafeStorage
+            self.assertIsInstance(jobs._storage, (_MultiprocessingStorage, _ThreadSafeStorage),
+                                "Storage should be either multiprocessing or thread-safe backend")
         except Exception as e:
-            self.fail(f"Read-only operations failed: {e}")
+            self.fail(f"Storage initialization failed: {e}")
 
     def test_cli_entry_point_import(self):
         """