Fixing worker-state initialization bug

Author: justinwlin

runpod/serverless/modules/worker_state.py

@@ -67,36 +67,41 @@ class JobsProgress:
     """Track the state of current jobs in progress using shared memory."""
 
     _instance: Optional['JobsProgress'] = None
-    _manager: SyncManager
-    _shared_data: Any
-    _lock: Any
-
+    # Singleton
     def __new__(cls):
         if cls._instance is None:
-            instance = object.__new__(cls)
-            # Initialize instance variables
-            instance._manager = Manager()
-            instance._shared_data = instance._manager.dict()
-            instance._shared_data['jobs'] = instance._manager.list()
-            instance._lock = instance._manager.Lock()
-            cls._instance = instance
+            cls._instance = super().__new__(cls)
         return cls._instance
 
     def __init__(self):
-        # Everything is already initialized in __new__
-        pass
+        if not hasattr(self, '_initialized'):
+            self._manager: Optional[SyncManager] = None
+            self._shared_data: Optional[Any] = None
+            self._lock: Optional[Any] = None
+            self._initialized = True
+
+    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()
 
     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'][:] = []
 
     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):
@@ -123,6 +128,8 @@ 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):
@@ -142,6 +149,8 @@ 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):
@@ -153,7 +162,6 @@ def remove(self, element: Any):
 
         with self._lock:
             job_list = self._shared_data['jobs']
-            # Find and remove the job
             for i, job_dict in enumerate(job_list):
                 if job_dict['id'] == job_id:
                     del job_list[i]
@@ -164,6 +172,9 @@ 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'])
 
@@ -177,11 +188,17 @@ 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'])
 
     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'])
@@ -195,14 +212,15 @@ def __len__(self):
 
     def __contains__(self, element: Any) -> bool:
         """Support 'in' operator"""
+        if self._manager is None:
+            return False
+            
         if isinstance(element, str):
             search_id = element
         elif isinstance(element, Job):
             search_id = element.id
-        elif isinstance(element, dict):
-            search_id = element.get('id')
         else:
-            return False
+            raise TypeError("Only Job objects can be checked in JobsProgress.")
 
         with self._lock:
             for job_dict in self._shared_data['jobs']:

tests/test_cli/test_cli_sanity.py

@@ -0,0 +1,117 @@
+"""
+CLI Sanity Checks
+
+These tests ensure that basic CLI operations work correctly and efficiently.
+"""
+
+import subprocess
+import sys
+import unittest
+from click.testing import CliRunner
+
+from runpod.cli.entry import runpod_cli
+
+
+class TestCLISanity(unittest.TestCase):
+    """Test basic CLI functionality and import safety"""
+
+    def test_help_command_works(self):
+        """
+        Test that --help commands work correctly for all CLI commands.
+        """
+        runner = CliRunner()
+        
+        # Test main help
+        result = runner.invoke(runpod_cli, ["--help"])
+        self.assertEqual(result.exit_code, 0, f"Main --help failed: {result.output}")
+        self.assertIn("A collection of CLI functions for RunPod", result.output)
+        
+        # Test subcommand help
+        result = runner.invoke(runpod_cli, ["pod", "--help"])
+        self.assertEqual(result.exit_code, 0, f"Pod --help failed: {result.output}")
+        self.assertIn("Manage and interact with pods", result.output)
+        
+        result = runner.invoke(runpod_cli, ["config", "--help"])
+        self.assertEqual(result.exit_code, 0, f"Config --help failed: {result.output}")
+        
+        result = runner.invoke(runpod_cli, ["project", "--help"])
+        self.assertEqual(result.exit_code, 0, f"Project --help failed: {result.output}")
+        
+        result = runner.invoke(runpod_cli, ["ssh", "--help"])
+        self.assertEqual(result.exit_code, 0, f"SSH --help failed: {result.output}")
+        
+        result = runner.invoke(runpod_cli, ["exec", "--help"])
+        self.assertEqual(result.exit_code, 0, f"Exec --help failed: {result.output}")
+
+    def test_help_command_subprocess(self):
+        """
+        Test --help commands using subprocess to ensure they work in real-world usage.
+        """
+        # Test main help using the installed runpod command
+        result = subprocess.run(
+            ["runpod", "--help"],
+            capture_output=True,
+            text=True,
+            timeout=10  # Prevent hanging
+        )
+        self.assertEqual(result.returncode, 0, 
+                        f"Subprocess --help failed: {result.stderr}")
+        self.assertIn("A collection of CLI functions for RunPod", result.stdout)
+        
+        # Test pod help
+        result = subprocess.run(
+            ["runpod", "pod", "--help"],
+            capture_output=True,
+            text=True,
+            timeout=10
+        )
+        self.assertEqual(result.returncode, 0, 
+                        f"Subprocess pod --help failed: {result.stderr}")
+        self.assertIn("Manage and interact with pods", result.stdout)
+
+    def test_import_safety(self):
+        """
+        Test that importing runpod modules works correctly.
+        """
+        # Test importing main package
+        try:
+            import runpod
+            self.assertTrue(True, "Main runpod import successful")
+        except Exception as e:
+            self.fail(f"Failed to import runpod: {e}")
+        
+        # Test importing serverless modules
+        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")
+            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
+        try:
+            from runpod.serverless.modules.worker_state import JobsProgress
+            jobs = JobsProgress()
+            count = jobs.get_job_count()  # Should work without heavy initialization
+            self.assertEqual(count, 0)
+            self.assertIsNone(jobs._manager, 
+                            "Manager should not be created for read-only operations")
+        except Exception as e:
+            self.fail(f"Read-only operations failed: {e}")
+
+    def test_cli_entry_point_import(self):
+        """
+        Test that the CLI entry point can be imported without issues.
+        """
+        try:
+            from runpod.cli.entry import runpod_cli
+            self.assertTrue(callable(runpod_cli), "runpod_cli should be callable")
+        except Exception as e:
+            self.fail(f"Failed to import CLI entry point: {e}")
+
+
+if __name__ == "__main__":
+    unittest.main()