add unin_tests and fix some bugs

Author: wanglei19991004

flagscale/runner/runner_train.py

@@ -798,35 +798,35 @@ def query(self, interval=10, timeout=None):
                 cur_time = time.time()
             logger.info(f"Query timeout reached ({timeout}s)")
 
-<<<<<<< HEAD
-=======
     def _monitor_and_collect_logs(self):
         """Monitor training status and collect logs in the background"""
         self._monitoring_active = True
         try:
             while self._monitoring_active:
                 status = self._query_status()
-                
+
                 if hasattr(self, 'resources') and self.resources:
                     for node_rank, (host, _) in enumerate(self.resources.items()):
                         if not self._monitoring_active:
                             break
-                        self._collect_and_diagnose_single_node(host, node_rank, 
-                                                              status != JobStatus.RUNNING)
+                        self._collect_and_diagnose_single_node(
+                            host, node_rank, status != JobStatus.RUNNING
+                        )
                 else:
                     # Single machine
-                    self._collect_and_diagnose_single_node("localhost", 0, 
-                                                          status != JobStatus.RUNNING)
-                
+                    self._collect_and_diagnose_single_node(
+                        "localhost", 0, status != JobStatus.RUNNING
+                    )
+
                 if status == JobStatus.COMPLETED_OR_IDLE:
                     logger.info("Training completed. Final log collection finished.")
                     break
-                    
+
                 for _ in range(10):
                     if not self._monitoring_active:
                         break
                     time.sleep(1)
-                
+
         except Exception as e:
             logger.error(f"Error in log monitoring: {e}")
         finally:
@@ -836,49 +836,56 @@ def _collect_and_diagnose_single_node(self, host, node_rank, process_completed=F
         """Collect and diagnose logs for a single node"""
         try:
             logging_config = self.config.train.system.logging
-            
+
             log_file = collect_logs(
-                self.config, host, node_rank, 
-                logging_config.diagnostic_dir, 
+                self.config,
+                host,
+                node_rank,
+                logging_config.diagnostic_dir,
                 dryrun=False,
-                process_running=not process_completed
+                process_running=not process_completed,
             )
-            
+
             if log_file and os.path.exists(log_file) and os.path.getsize(log_file) > 0:
                 diagnostic_report = generate_diagnostic_report(
                     self.config, host, node_rank, log_file, return_content=True
                 )
-                
+
                 combined_file = os.path.join(
-                    logging_config.diagnostic_dir, 
-                    f"host_{node_rank}_{host}_combined.log"
+                    logging_config.diagnostic_dir, f"host_{node_rank}_{host}_combined.log"
                 )
-                
+
                 os.makedirs(os.path.dirname(combined_file), exist_ok=True)
-                
+
                 with open(log_file, 'r') as log_f:
                     log_content = log_f.read()
-                
+
                 with open(combined_file, 'a') as f:
-                    f.write(f"\n=== Collection Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ===\n")
+                    f.write(
+                        f"\n=== Collection Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ===\n"
+                    )
                     f.write("=== Log Content ===\n")
                     f.write(log_content)
                     f.write("\n=== Diagnostic Report ===\n")
-                    f.write(diagnostic_report if diagnostic_report else "No diagnostic report generated\n")
-                
+                    f.write(
+                        diagnostic_report
+                        if diagnostic_report
+                        else "No diagnostic report generated\n"
+                    )
+
                 logger.debug(f"Updated combined log for {host} (node {node_rank})")
-                
+
                 try:
                     os.remove(log_file)
                 except OSError as e:
                     logger.warning(f"Could not remove temporary log file {log_file}: {e}")
-            
+
             if process_completed:
                 reset_log_collection(host, node_rank)
-                
+
         except Exception as e:
             logger.error(f"Failed to collect and diagnose logs for {host} (node {node_rank}): {e}")
->>>>>>> ba556ac53b36e1272f8c0ec5b6ee476df9293d64
+
 
 class CloudTrainRunner(RunnerBase):
     def __init__(self, config: DictConfig):

tests/scripts/unit_tests/config.yml

@@ -135,7 +135,12 @@ flagscale:
     - export NVTE_FLASH_ATTN=0
     - export NVTE_FUSED_ATTN=0
     - ulimit -n 65535
+  coverage_fold:
+    flagscale
   subset:
+    elastic:
+      type: batch
+      depth: all
     runner:
       type: batch
       depth: all

tests/unit_tests/runner/elastic/__init__.py

@@ -0,0 +1,178 @@
+import os
+import tempfile
+
+from unittest.mock import MagicMock, mock_open, patch
+
+import pytest
+
+from flagscale.runner.elastic.diagnostic import error_types, generate_diagnostic_report
+
+
+class TestDiagnostic:
+    """Test cases for diagnostic module"""
+
+    @pytest.fixture
+    def mock_config(self):
+        """Mock config object"""
+        return MagicMock()
+
+    @pytest.fixture
+    def sample_log_content(self):
+        """Sample log content with various errors"""
+        return """
+        [INFO] Starting training...
+        [ERROR] CUDA out of memory
+        Traceback (most recent call last):
+          File "train.py", line 100, in <module>
+            model.forward()
+        torch.distributed.elastic.rendezvous.api.RendezvousConnectionError: Connection failed
+        OutOfMemoryError: GPU memory exhausted
+        [ERROR] Training failed
+        """
+
+    def test_error_types_dict_exists(self):
+        """Test that error_types dictionary is properly defined"""
+        assert isinstance(error_types, dict)
+        assert len(error_types) > 0
+
+        expected_keys = ['out of memory', 'rendezvous', 'traceback', 'error', 'cuda']
+        for key in expected_keys:
+            assert key in error_types
+
+    def test_generate_diagnostic_report_empty_file(self, mock_config):
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+            f.write("")
+            temp_path = f.name
+
+        try:
+            report = generate_diagnostic_report(
+                mock_config, "localhost", 0, temp_path, return_content=True
+            )
+
+            assert "Diagnostic Report for localhost (node 0)" in report
+            assert "Log file is empty" in report
+        finally:
+            os.unlink(temp_path)
+
+    def test_generate_diagnostic_report_with_errors(self, mock_config, sample_log_content):
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+            f.write(sample_log_content)
+            temp_path = f.name
+
+        try:
+            report = generate_diagnostic_report(
+                mock_config, "localhost", 0, temp_path, return_content=True
+            )
+
+            assert "Diagnostic Report for localhost (node 0)" in report
+            assert "OutOfMemoryError" in report
+            assert "RendezvousConnectionError" in report
+            assert "CodeError" in report
+            assert "GeneralError" in report
+        finally:
+            os.unlink(temp_path)
+
+    def test_generate_diagnostic_report_nonexistent_file(self, mock_config):
+        report = generate_diagnostic_report(
+            mock_config, "localhost", 0, "/nonexistent/file.log", return_content=True
+        )
+
+        assert "Diagnostic Report for localhost (node 0)" in report
+        assert "Log file is empty or does not exist" in report
+
+    def test_generate_diagnostic_report_no_errors(self, mock_config):
+        content = """
+        [INFO] Training started successfully
+        [INFO] Epoch 1/100 completed
+        [INFO] Training finished successfully
+        """
+
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+            f.write(content)
+            temp_path = f.name
+
+        try:
+            report = generate_diagnostic_report(
+                mock_config, "localhost", 0, temp_path, return_content=True
+            )
+
+            assert "Diagnostic Report for localhost (node 0)" in report
+            assert "No errors or unknown error detected" in report
+        finally:
+            os.unlink(temp_path)
+
+    def test_generate_diagnostic_report_file_output(self, mock_config, sample_log_content):
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+            f.write(sample_log_content)
+            temp_path = f.name
+
+        try:
+            # Test file output mode
+            result_path = generate_diagnostic_report(
+                mock_config, "localhost", 0, temp_path, return_content=False
+            )
+
+            assert result_path is not None
+            assert "diagnostic" in result_path or result_path.endswith('.txt')
+        finally:
+            os.unlink(temp_path)
+
+    @patch('flagscale.runner.elastic.diagnostic.logger')
+    def test_generate_diagnostic_report_read_error(self, mock_logger, mock_config):
+        """Test diagnostic report generation with file read error"""
+        with patch('builtins.open', side_effect=PermissionError("Access denied")):
+            report = generate_diagnostic_report(
+                mock_config, "localhost", 0, "/some/file.log", return_content=True
+            )
+
+            assert "Error reading log file" in report
+            mock_logger.error.assert_called()
+
+    def test_error_detection_case_insensitive(self, mock_config):
+        content = """
+        CUDA OUT OF MEMORY ERROR occurred
+        TRACEBACK: Failed to execute
+        RENDEZVOUS connection failed
+        """
+
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+            f.write(content)
+            temp_path = f.name
+
+        try:
+            report = generate_diagnostic_report(
+                mock_config, "localhost", 0, temp_path, return_content=True
+            )
+
+            # Should detect errors regardless of case
+            assert "OutOfMemoryError" in report
+            assert "CodeError" in report
+            assert "RendezvousError" in report
+        finally:
+            os.unlink(temp_path)
+
+    def test_multiple_error_types_detection(self, mock_config):
+        content = """
+        CUDA error: out of memory
+        ImportError: module not found
+        Permission denied: cannot access file
+        Connection timeout occurred
+        Process was killed
+        """
+
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+            f.write(content)
+            temp_path = f.name
+
+        try:
+            report = generate_diagnostic_report(
+                mock_config, "localhost", 0, temp_path, return_content=True
+            )
+
+            assert "CUDAError" in report
+            assert "ImportError" in report
+            assert "PermissionError" in report
+            assert "TimeoutError" in report
+            assert "ProcessKilled" in report
+        finally:
+            os.unlink(temp_path)