Implement comprehensive cleanup for local worker upon successful job run and upload

Co-authored-by: nolan1999 <54246789+nolan1999@users.noreply.github.com>

Author: Copilot

cli/main.py

@@ -42,8 +42,11 @@ def main() -> None:
     worker_info = info.sys.collect()
 
     while True:
+        container = None
+        job_id = None
+        path_job = None
+        
         try:
-            container = None
             jobs = api_worker.fetch_jobs(
                 limit=1,
                 cpu_cores=worker_info.sys.cores,
@@ -176,9 +179,36 @@ def main() -> None:
                     container.kill()
             else:
                 raise e
-
-        logger.info(f"Job {job_id} finished")
-        shutil.rmtree(path_job)
+        
+        finally:
+            # Clean up resources after successful job run AND upload
+            if job_id:
+                logger.info(f"Cleaning up job {job_id}")
+                
+                # Clean up Docker container
+                if container:
+                    try:
+                        if container.status == "running":
+                            logger.info(f"Stopping running container for job {job_id}")
+                            container.stop()
+                    except Exception as e:
+                        logger.warning(f"Failed to stop container for job {job_id}: {e}")
+                    
+                    try:
+                        logger.info(f"Removing container for job {job_id}")
+                        container.remove()
+                    except Exception as e:
+                        logger.warning(f"Failed to remove container for job {job_id}: {e}")
+                
+                # Clean up job directory
+                if path_job and path_job.exists():
+                    try:
+                        logger.info(f"Removing job directory {path_job}")
+                        shutil.rmtree(path_job)
+                    except Exception as e:
+                        logger.warning(f"Failed to clean up job directory {path_job}: {e}")
+                
+                logger.info(f"Job {job_id} cleanup completed")
 
 
 if __name__ == "__main__":

tests/unit/test_cleanup.py

@@ -0,0 +1,88 @@
+"""Tests for cleanup functionality in the main job processing loop."""
+
+import shutil
+import tempfile
+from pathlib import Path
+from unittest.mock import Mock, MagicMock
+
+import pytest
+
+
+class TestJobCleanup:
+    """Test cleanup behavior for local worker jobs."""
+
+    def test_container_cleanup_logic(self):
+        """Test the container cleanup logic works correctly."""
+        # Mock a container
+        mock_container = Mock()
+        mock_container.status = "running"
+        
+        job_id = "test_job_123"
+        
+        # Test cleanup logic (simulating the finally block)
+        if job_id:
+            if mock_container:
+                try:
+                    if mock_container.status == "running":
+                        mock_container.stop()
+                    mock_container.remove()
+                except Exception as e:
+                    # Should handle exceptions gracefully
+                    pass
+        
+        # Verify methods were called
+        mock_container.stop.assert_called_once()
+        mock_container.remove.assert_called_once()
+
+    def test_container_cleanup_handles_exceptions(self):
+        """Test that container cleanup handles exceptions gracefully."""
+        # Mock a container that raises exceptions
+        mock_container = Mock()
+        mock_container.status = "running"
+        mock_container.stop.side_effect = Exception("Stop failed")
+        mock_container.remove.side_effect = Exception("Remove failed")
+        
+        job_id = "test_job_456"
+        
+        # Test cleanup logic should handle each step independently
+        if job_id:
+            if mock_container:
+                try:
+                    if mock_container.status == "running":
+                        mock_container.stop()
+                except Exception:
+                    pass
+                try:
+                    mock_container.remove()
+                except Exception:
+                    pass
+        
+        # Verify methods were called despite exceptions
+        mock_container.stop.assert_called_once()
+        mock_container.remove.assert_called_once()
+
+    def test_cleanup_defensive_programming(self):
+        """Test that cleanup handles missing variables gracefully."""
+        # This test verifies that the cleanup code can handle cases where
+        # container, job_id, or path_job might be None or undefined
+        
+        # Create a temporary directory for testing
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path_job = Path(temp_dir) / "test_job"
+            path_job.mkdir()
+            
+            # Test cleanup with missing container (should not crash)
+            container = None
+            job_id = "test_job"
+            
+            # Simulate the cleanup code
+            if job_id:
+                if container:
+                    # This should not execute
+                    assert False, "Should not try to clean up None container"
+                
+                if path_job and path_job.exists():
+                    shutil.rmtree(path_job)
+            
+            # Verify directory was cleaned up
+            assert not path_job.exists()