improve code

data.py

@@ -6,7 +6,8 @@ import threading
 import traceback
 import json
 import re
-from typing import List, Tuple, Optional
+import random
+from typing import List, Tuple, Optional, Dict
 
 # NOTE: if caching is an issue, try adding `use_listings_cache=False`
 fs = HfFileSystem()
@@ -56,6 +57,40 @@ KEYS_TO_KEEP = [
     "job_link_nvidia",
 ]
 
+# ============================================================================
+# HELPER FUNCTIONS
+# ============================================================================
+
+def generate_fake_dates(num_days: int = 7) -> List[str]:
+    """Generate fake dates for the last N days."""
+    today = datetime.now()
+    return [(today - timedelta(days=i)).strftime("%Y-%m-%d") for i in range(num_days)]
+
+def parse_json_field(value) -> dict:
+    """Safely parse a JSON field that might be a string or dict."""
+    if isinstance(value, str):
+        try:
+            return json.loads(value)
+        except:
+            return {}
+    return value if isinstance(value, dict) else {}
+
+def extract_date_from_path(path: str, pattern: str) -> Optional[str]:
+    """Extract date from file path using regex pattern."""
+    match = re.search(pattern, path)
+    return match.group(1) if match else None
+
+def get_test_names(tests: list) -> set:
+    """Extract test names from a list of test dictionaries."""
+    return {test.get('line', '') for test in tests}
+
+def safe_extract(row: pd.Series, key: str) -> int:
+    """Safely extract an integer value from a DataFrame row."""
+    return int(row.get(key, 0)) if pd.notna(row.get(key, 0)) else 0
+
+# ============================================================================
+# DATA LOADING FUNCTIONS
+# ============================================================================
 
 def log_dataframe_link(link: str) -> str:
     """
@@ -109,69 +144,37 @@ def read_one_dataframe(json_path: str, device_label: str) -> tuple[pd.DataFrame,
 def get_available_dates() -> List[str]:
     """Get list of available dates from both AMD and NVIDIA datasets."""
     try:
-        # Get AMD dates - the path structure is: YYYY-MM-DD/runs/{run_id}/ci_results_run_models_gpu/model_results.json
+        # Get file lists
         amd_src = "hf://datasets/optimum-amd/transformers_daily_ci/**/runs/**/ci_results_run_models_gpu/model_results.json"
-        files_amd = sorted(fs.glob(amd_src, refresh=True), reverse=True)
-        logger.info(f"Found {len(files_amd)} AMD files")
-        
-        # Get NVIDIA dates - structure is: YYYY-MM-DD/ci_results_run_models_gpu/model_results.json
         nvidia_src = "hf://datasets/hf-internal-testing/transformers_daily_ci/*/ci_results_run_models_gpu/model_results.json"
+        
+        files_amd = sorted(fs.glob(amd_src, refresh=True), reverse=True)
         files_nvidia = sorted(fs.glob(nvidia_src, refresh=True), reverse=True)
-        logger.info(f"Found {len(files_nvidia)} NVIDIA files")
-        
-        # Extract dates from file paths
-        amd_dates = set()
-        for file_path in files_amd:
-            # Pattern to match the date in the AMD path: YYYY-MM-DD/runs/{run_id}/ci_results_run_models_gpu/model_results.json
-            pattern = r'transformers_daily_ci/(\d{4}-\d{2}-\d{2})/runs/[^/]+/ci_results_run_models_gpu/model_results\.json'
-            match = re.search(pattern, file_path)
-            if match:
-                amd_dates.add(match.group(1))
-            else:
-                # Log unmatched paths for debugging
-                logger.debug(f"AMD file path didn't match pattern: {file_path}")
         
-        # Log a few example AMD file paths for debugging
-        if files_amd:
-            logger.info(f"Example AMD file paths: {files_amd[:3]}")
+        logger.info(f"Found {len(files_amd)} AMD files, {len(files_nvidia)} NVIDIA files")
+        
+        # Extract dates using patterns
+        amd_pattern = r'transformers_daily_ci/(\d{4}-\d{2}-\d{2})/runs/[^/]+/ci_results_run_models_gpu/model_results\.json'
+        nvidia_pattern = r'transformers_daily_ci/(\d{4}-\d{2}-\d{2})/ci_results_run_models_gpu/model_results\.json'
+        
+        amd_dates = {extract_date_from_path(f, amd_pattern) for f in files_amd}
+        amd_dates.discard(None)  # Remove None values
         
-        nvidia_dates = set()
-        for file_path in files_nvidia:
-            # Pattern to match the date in the NVIDIA path: YYYY-MM-DD/ci_results_run_models_gpu/model_results.json
-            pattern = r'transformers_daily_ci/(\d{4}-\d{2}-\d{2})/ci_results_run_models_gpu/model_results\.json'
-            match = re.search(pattern, file_path)
-            if match:
-                nvidia_dates.add(match.group(1))
+        nvidia_dates = {extract_date_from_path(f, nvidia_pattern) for f in files_nvidia}
+        nvidia_dates.discard(None)
         
-        logger.info(f"AMD dates: {sorted(amd_dates, reverse=True)[:5]}...")  # Show first 5
-        logger.info(f"NVIDIA dates: {sorted(nvidia_dates, reverse=True)[:5]}...")  # Show first 5
+        logger.info(f"AMD dates: {sorted(amd_dates, reverse=True)[:5]}...")
+        logger.info(f"NVIDIA dates: {sorted(nvidia_dates, reverse=True)[:5]}...")
         
-        # Return intersection of both datasets (dates where both have data)
+        # Return intersection of both datasets
         common_dates = sorted(amd_dates.intersection(nvidia_dates), reverse=True)
         logger.info(f"Common dates: {len(common_dates)} dates where both AMD and NVIDIA have data")
         
-        if common_dates:
-            return common_dates[:30]  # Limit to last 30 days for performance
-        else:
-            # If no real dates available, generate fake dates for the last 7 days
-            logger.warning("No real dates available, generating fake dates for demo purposes")
-            fake_dates = []
-            today = datetime.now()
-            for i in range(7):
-                date = today - timedelta(days=i)
-                fake_dates.append(date.strftime("%Y-%m-%d"))
-            return fake_dates
+        return common_dates[:30] if common_dates else generate_fake_dates()
         
     except Exception as e:
         logger.error(f"Error getting available dates: {e}")
-        # Generate fake dates when there's an error
-        logger.info("Generating fake dates due to error")
-        fake_dates = []
-        today = datetime.now()
-        for i in range(7):
-            date = today - timedelta(days=i)
-            fake_dates.append(date.strftime("%Y-%m-%d"))
-        return fake_dates
+        return generate_fake_dates()
 
 
 def get_data_for_date(target_date: str) -> tuple[pd.DataFrame, str]:
@@ -242,37 +245,30 @@ def get_historical_data(start_date: str, end_date: str, sample_data = False) ->
     """Get historical data for a date range."""
     if sample_data:
         return get_fake_historical_data(start_date, end_date)
+    
     try:
         start_dt = datetime.strptime(start_date, "%Y-%m-%d")
         end_dt = datetime.strptime(end_date, "%Y-%m-%d")
-        
         historical_data = []
-        current_dt = start_dt
         
+        # Load data for each day in range
+        current_dt = start_dt
         while current_dt <= end_dt:
             date_str = current_dt.strftime("%Y-%m-%d")
             try:
                 df, _ = get_data_for_date(date_str)
-                # Only add non-empty dataframes
                 if not df.empty:
                     df['date'] = date_str
                     historical_data.append(df)
                     logger.info(f"Loaded data for {date_str}")
-                else:
-                    logger.warning(f"No data available for {date_str}")
             except Exception as e:
                 logger.warning(f"Could not load data for {date_str}: {e}")
-            
             current_dt += timedelta(days=1)
         
-        # Combine all dataframes
-        combined_df = pd.concat(historical_data, ignore_index=False)
-        return combined_df
+        return pd.concat(historical_data, ignore_index=False) if historical_data else pd.DataFrame()
         
     except Exception as e:
         logger.error(f"Error getting historical data: {e}")
-        # Fall back to fake data when there's an error
-        logger.info("Falling back to fake historical data due to error")
         return get_fake_historical_data(start_date, end_date)
 
 
@@ -326,49 +322,36 @@ def get_fake_historical_data(start_date: str, end_date: str) -> pd.DataFrame:
     try:
         start_dt = datetime.strptime(start_date, "%Y-%m-%d")
         end_dt = datetime.strptime(end_date, "%Y-%m-%d")
-        
-        # Generate fake data for each date in the range
-        historical_data = []
-        current_dt = start_dt
-        
-        # Get base sample data to use as template
         sample_df, _ = get_sample_data()
+        historical_data = []
         
+        # Generate data for each date
+        current_dt = start_dt
         while current_dt <= end_dt:
-            date_str = current_dt.strftime("%Y-%m-%d")
-            
-            # Create a copy of sample data for this date with some random variations
             date_df = sample_df.copy()
-            date_df['date'] = date_str
+            date_df['date'] = current_dt.strftime("%Y-%m-%d")
             
-            # Add some random variation to make it look more realistic
-            import random
+            # Add random variations to make it realistic
             for idx in date_df.index:
-                # Vary the success/failure counts slightly (±20%)
+                # Vary success/skipped counts (±20%)
                 for col in ['success_amd', 'success_nvidia', 'skipped_amd', 'skipped_nvidia']:
-                    if col in date_df.columns:
-                        original_val = date_df.loc[idx, col]
-                        if pd.notna(original_val) and original_val > 0:
-                            variation = random.uniform(0.8, 1.2)
-                            date_df.loc[idx, col] = max(0, int(original_val * variation))
+                    if col in date_df.columns and pd.notna(date_df.loc[idx, col]):
+                        val = date_df.loc[idx, col]
+                        if val > 0:
+                            date_df.loc[idx, col] = max(0, int(val * random.uniform(0.8, 1.2)))
                 
-                # Vary failure counts more dramatically to show trends
+                # Vary failure counts more dramatically (±50-100%)
                 for col in ['failed_multi_no_amd', 'failed_multi_no_nvidia', 'failed_single_no_amd', 'failed_single_no_nvidia']:
-                    if col in date_df.columns:
-                        original_val = date_df.loc[idx, col]
-                        if pd.notna(original_val):
-                            # Sometimes have more failures, sometimes fewer
-                            variation = random.uniform(0.5, 2.0)
-                            date_df.loc[idx, col] = max(0, int(original_val * variation))
+                    if col in date_df.columns and pd.notna(date_df.loc[idx, col]):
+                        val = date_df.loc[idx, col]
+                        date_df.loc[idx, col] = max(0, int(val * random.uniform(0.5, 2.0)))
             
             historical_data.append(date_df)
             current_dt += timedelta(days=1)
         
         if not historical_data:
-            logger.warning("No fake historical data generated")
             return pd.DataFrame()
         
-        # Combine all dataframes
         combined_df = pd.concat(historical_data, ignore_index=False)
         logger.info(f"Generated fake historical data: {len(combined_df)} records from {start_date} to {end_date}")
         return combined_df
@@ -377,53 +360,23 @@ def get_fake_historical_data(start_date: str, end_date: str) -> pd.DataFrame:
         logger.error(f"Error generating fake historical data: {e}")
         return pd.DataFrame()
 
-def safe_extract(row: pd.DataFrame, key: str) -> int:
-    return int(row.get(key, 0)) if pd.notna(row.get(key, 0)) else 0
-
-
 def find_failure_first_seen(historical_df: pd.DataFrame, model_name: str, test_name: str, device: str, gpu_type: str) -> Optional[str]:
-    """
-    Find the first date when a specific test failure appeared in historical data.
-    """
+    """Find the first date when a specific test failure appeared in historical data."""
     if historical_df.empty:
         return None
     
     try:
-        # Normalize model name to match DataFrame index
-        model_name_lower = model_name.lower()
-        
-        # Filter historical data for this model
-        model_data = historical_df[historical_df.index == model_name_lower].copy()
-        
+        model_data = historical_df[historical_df.index == model_name.lower()].copy()
         if model_data.empty:
             return None
         
-        # Sort by date (oldest first)
-        model_data = model_data.sort_values('date')
-        
-        # Check each date for this failure
-        for idx, row in model_data.iterrows():
-            failures = row.get(f'failures_{device}', None)
-            
-            if failures is None or pd.isna(failures):
-                continue
-            
-            # Handle case where failures might be a string (JSON)
-            if isinstance(failures, str):
-                try:
-                    import json
-                    failures = json.loads(failures)
-                except:
-                    continue
-            
-            # Check if this test appears in the failures for this gpu_type
+        # Check each date (oldest first) for this failure
+        for _, row in model_data.sort_values('date').iterrows():
+            failures = parse_json_field(row.get(f'failures_{device}'))
             if gpu_type in failures:
                 for test in failures[gpu_type]:
-                    test_line = test.get('line', '')
-                    if test_line == test_name:
-                        # Found the first occurrence
-                        return row.get('date', None)
-        
+                    if test.get('line', '') == test_name:
+                        return row.get('date')
         return None
         
     except Exception as e:
@@ -431,148 +384,89 @@ def find_failure_first_seen(historical_df: pd.DataFrame, model_name: str, test_n
         return None
 
 
+def _find_device_regressions(model_name: str, current_failures: dict, yesterday_failures: dict, device: str) -> list[dict]:
+    """Helper to find regressions for a specific device."""
+    regressions = []
+    for gpu_type in ['single', 'multi']:
+        current_tests = get_test_names(current_failures.get(gpu_type, []))
+        yesterday_tests = get_test_names(yesterday_failures.get(gpu_type, []))
+        
+        # Find NEW failures: failing NOW but NOT yesterday
+        new_tests = current_tests - yesterday_tests
+        for test_name in new_tests:
+            if test_name:  # Skip empty names
+                regressions.append({
+                    'model': model_name,
+                    'test': test_name.split('::')[-1],  # Short name
+                    'test_full': test_name,  # Full name
+                    'device': device,
+                    'gpu_type': gpu_type
+                })
+    return regressions
+
 def find_new_regressions(current_df: pd.DataFrame, historical_df: pd.DataFrame) -> list[dict]:
-    """
-    Compare CURRENT failures against PREVIOUS day's failures to find NEW regressions.
-    
-    A regression is a test that:
-    - Is failing in the CURRENT/LATEST run (current_df)
-    - Was NOT failing in the PREVIOUS run (yesterday in historical_df)
-    """
+    """Compare current failures against previous day's failures to find new regressions."""
     if current_df.empty or historical_df.empty:
         return []
     
-    new_regressions = []
-    
-    # Get the most recent date from historical data (this is "yesterday")
+    # Get yesterday's data
     available_dates = sorted(historical_df['date'].unique(), reverse=True)
-    if len(available_dates) < 1:
-        # No history to compare against
+    if not available_dates:
         return []
     
-    yesterday_date = available_dates[0]
-    yesterday_data = historical_df[historical_df['date'] == yesterday_date]
+    yesterday_data = historical_df[historical_df['date'] == available_dates[0]]
+    new_regressions = []
     
-    # For each model in current data, compare against yesterday
+    # For each model, compare current vs yesterday
     for model_name in current_df.index:
-        model_name_lower = model_name.lower()
-        
-        # Get CURRENT failures from current_df
         current_row = current_df.loc[model_name]
+        yesterday_row = yesterday_data[yesterday_data.index == model_name.lower()]
         
-        # Get YESTERDAY's failures from historical_df
-        yesterday_row = yesterday_data[yesterday_data.index == model_name_lower]
-        yesterday_failures_amd = {}
-        yesterday_failures_nvidia = {}
+        # Parse current failures
+        current_amd = parse_json_field(current_row.get('failures_amd', {}))
+        current_nvidia = parse_json_field(current_row.get('failures_nvidia', {}))
         
+        # Parse yesterday failures
+        yesterday_amd = {}
+        yesterday_nvidia = {}
         if not yesterday_row.empty:
             yesterday_row = yesterday_row.iloc[0]
-            yesterday_failures_amd = yesterday_row.get('failures_amd', {})
-            yesterday_failures_nvidia = yesterday_row.get('failures_nvidia', {})
-            
-            # Handle string/dict conversion
-            if isinstance(yesterday_failures_amd, str):
-                try:
-                    yesterday_failures_amd = json.loads(yesterday_failures_amd)
-                except:
-                    yesterday_failures_amd = {}
-            if isinstance(yesterday_failures_nvidia, str):
-                try:
-                    yesterday_failures_nvidia = json.loads(yesterday_failures_nvidia)
-                except:
-                    yesterday_failures_nvidia = {}
-        
-        # Get CURRENT failures
-        current_failures_amd = current_row.get('failures_amd', {})
-        current_failures_nvidia = current_row.get('failures_nvidia', {})
-        
-        # Handle string/dict conversion
-        if isinstance(current_failures_amd, str):
-            try:
-                current_failures_amd = json.loads(current_failures_amd)
-            except:
-                current_failures_amd = {}
-        if isinstance(current_failures_nvidia, str):
-            try:
-                current_failures_nvidia = json.loads(current_failures_nvidia)
-            except:
-                current_failures_nvidia = {}
-        
-        # Check AMD failures - find tests failing NOW but NOT yesterday
-        for gpu_type in ['single', 'multi']:
-            current_tests = current_failures_amd.get(gpu_type, [])
-            yesterday_tests = yesterday_failures_amd.get(gpu_type, [])
-            
-            # Get test names
-            current_test_names = {test.get('line', '') for test in current_tests}
-            yesterday_test_names = {test.get('line', '') for test in yesterday_tests}
-            
-            # Find NEW failures: failing NOW but NOT yesterday
-            new_tests = current_test_names - yesterday_test_names
-            for test_name in new_tests:
-                if test_name:  # Skip empty names
-                    new_regressions.append({
-                        'model': model_name,
-                        'test': test_name.split('::')[-1],  # Short name
-                        'test_full': test_name,  # Full name
-                        'device': 'amd',
-                        'gpu_type': gpu_type
-                    })
-        
-        # Check NVIDIA failures - find tests failing NOW but NOT yesterday
-        for gpu_type in ['single', 'multi']:
-            current_tests = current_failures_nvidia.get(gpu_type, [])
-            yesterday_tests = yesterday_failures_nvidia.get(gpu_type, [])
-            
-            # Get test names
-            current_test_names = {test.get('line', '') for test in current_tests}
-            yesterday_test_names = {test.get('line', '') for test in yesterday_tests}
-            
-            # Find NEW failures: failing NOW but NOT yesterday
-            new_tests = current_test_names - yesterday_test_names
-            for test_name in new_tests:
-                if test_name:  # Skip empty names
-                    new_regressions.append({
-                        'model': model_name,
-                        'test': test_name.split('::')[-1],  # Short name
-                        'test_full': test_name,  # Full name
-                        'device': 'nvidia',
-                        'gpu_type': gpu_type
-                    })
+            yesterday_amd = parse_json_field(yesterday_row.get('failures_amd', {}))
+            yesterday_nvidia = parse_json_field(yesterday_row.get('failures_nvidia', {}))
+        
+        # Find regressions for both devices
+        new_regressions.extend(_find_device_regressions(model_name, current_amd, yesterday_amd, 'amd'))
+        new_regressions.extend(_find_device_regressions(model_name, current_nvidia, yesterday_nvidia, 'nvidia'))
     
     return new_regressions
 
 
 def extract_model_data(row: pd.Series) -> tuple[dict[str, int], dict[str, int], int, int, int, int]:
     """Extract and process model data from DataFrame row."""
-    # Handle missing values and get counts directly from dataframe
-    success_nvidia = safe_extract(row, "success_nvidia")
-    success_amd = safe_extract(row, "success_amd")
-
-    skipped_nvidia = safe_extract(row, "skipped_nvidia")
-    skipped_amd = safe_extract(row, "skipped_amd")
+    # Extract all counts
+    counts = {key: safe_extract(row, key) for key in [
+        'success_amd', 'success_nvidia', 'skipped_amd', 'skipped_nvidia',
+        'failed_multi_no_amd', 'failed_multi_no_nvidia', 
+        'failed_single_no_amd', 'failed_single_no_nvidia'
+    ]}
     
-    failed_multi_amd = safe_extract(row, 'failed_multi_no_amd')
-    failed_multi_nvidia = safe_extract(row, 'failed_multi_no_nvidia')
-    failed_single_amd = safe_extract(row, 'failed_single_no_amd')
-    failed_single_nvidia = safe_extract(row, 'failed_single_no_nvidia')
-    # Calculate total failures
-    total_failed_amd = failed_multi_amd + failed_single_amd
-    total_failed_nvidia = failed_multi_nvidia + failed_single_nvidia
-    # Create stats dictionaries directly from dataframe values
+    # Create stats dictionaries
     amd_stats = {
-        'passed': success_amd,
-        'failed': total_failed_amd,
-        'skipped': skipped_amd,
-        'error': 0     # Not available in this dataset
+        'passed': counts['success_amd'],
+        'failed': counts['failed_multi_no_amd'] + counts['failed_single_no_amd'],
+        'skipped': counts['skipped_amd'],
+        'error': 0
     }
     nvidia_stats = {
-        'passed': success_nvidia,
-        'failed': total_failed_nvidia,
-        'skipped': skipped_nvidia,
-        'error': 0     # Not available in this dataset
+        'passed': counts['success_nvidia'],
+        'failed': counts['failed_multi_no_nvidia'] + counts['failed_single_no_nvidia'],
+        'skipped': counts['skipped_nvidia'],
+        'error': 0
     }
-    return amd_stats, nvidia_stats, failed_multi_amd, failed_single_amd, failed_multi_nvidia, failed_single_nvidia
+    
+    return (amd_stats, nvidia_stats, counts['failed_multi_no_amd'], 
+            counts['failed_single_no_amd'], counts['failed_multi_no_nvidia'], 
+            counts['failed_single_no_nvidia'])
 
 
 
@@ -643,19 +537,11 @@ class CIResults:
         """Load all available historical data at startup."""
         try:
             if not self.available_dates:
-                # Generate fake dates when no real dates are available
-                fake_dates = []
-                today = datetime.now()
-                for i in range(7):
-                    date = today - timedelta(days=i)
-                    fake_dates.append(date.strftime("%Y-%m-%d"))
-                self.available_dates = fake_dates
+                self.available_dates = generate_fake_dates()
                 logger.info(f"No available dates found, generated {len(self.available_dates)} sample dates.")
-                
-            logger.info(f"Loading all historical data for {len(self.available_dates)} dates...")
-            start_date = self.available_dates[-1]  # Oldest date
-            end_date = self.available_dates[0]     # Newest date
             
+            logger.info(f"Loading all historical data for {len(self.available_dates)} dates...")
+            start_date, end_date = self.available_dates[-1], self.available_dates[0]
             self.all_historical_data = get_historical_data(start_date, end_date, self.sample_data)
             logger.info(f"All historical data loaded: {len(self.all_historical_data)} records")
         except Exception as e:
@@ -672,17 +558,15 @@ class CIResults:
                 self.historical_df = pd.DataFrame()
                 return
             
-            # Filter the pre-loaded data by date range
+            # Filter by date range
             start_dt = datetime.strptime(start_date, "%Y-%m-%d")
             end_dt = datetime.strptime(end_date, "%Y-%m-%d")
             
-            # Filter data within the date range
-            filtered_data = []
-            for date_str in self.all_historical_data['date'].unique():
-                date_dt = datetime.strptime(date_str, "%Y-%m-%d")
-                if start_dt <= date_dt <= end_dt:
-                    date_data = self.all_historical_data[self.all_historical_data['date'] == date_str]
-                    filtered_data.append(date_data)
+            filtered_data = [
+                self.all_historical_data[self.all_historical_data['date'] == date_str]
+                for date_str in self.all_historical_data['date'].unique()
+                if start_dt <= datetime.strptime(date_str, "%Y-%m-%d") <= end_dt
+            ]
             
             if filtered_data:
                 self.historical_df = pd.concat(filtered_data, ignore_index=False)