Move more cache functions over

Author: ZainRizvi

tools/analytics/org/analyze_runner_usage.py

@@ -55,7 +55,6 @@
 """
 
 import argparse
-import json
 import logging
 import os
 from collections import defaultdict
@@ -65,10 +64,9 @@
 
 import requests
 import yaml
+from cache_manager import get_cache_stats, make_cached_request
 from dotenv import load_dotenv
 
-from cache_manager import CACHE_DIR, CacheManager
-
 
 load_dotenv()
 
@@ -141,58 +139,9 @@
     "linux.g5.4xlarge.nvidia.cpu",  # a nonexistent label used by a repo
 ]
 
-HEADERS = {
-    "Authorization": f"Bearer {GITHUB_TOKEN}",
-    "Accept": "application/vnd.github+json",
-}
-
 BASE_URL = "https://api.github.com"
 WORKFLOW_RUN_LOOKBACK = (datetime.utcnow() - timedelta(days=180)).isoformat() + "Z"
 
-# Global cache manager instance
-cache_manager = CacheManager()
-
-
-def make_cached_request(
-    url: str, headers: Optional[Dict[str, str]] = None
-) -> Optional[Dict]:
-    """
-    Make an HTTP request with caching. Returns the JSON response if successful.
-
-    Args:
-        url: The URL to request
-        headers: Optional headers for the request
-
-    Returns:
-        JSON response data if successful, None if failed
-    """
-    # Check cache first
-    cached_response = cache_manager.get(url)
-    if cached_response:
-        logging.info(f"[make_cached_request] Using cached response for: {url}")
-        return cached_response
-
-    # Make actual HTTP request
-    logging.info(f"[make_cached_request] Making HTTP request to: {url}")
-    try:
-        response = requests.get(url, headers=headers or HEADERS)
-        response.raise_for_status()
-        data = response.json()
-
-        # Cache successful response
-        cache_manager.set(url, data)
-        logging.info(f"[make_cached_request] Successfully cached response for: {url}")
-        return data
-
-    except requests.exceptions.RequestException as e:
-        logging.error(f"[make_cached_request] HTTP request failed for {url}: {e}")
-        return None
-    except json.JSONDecodeError as e:
-        logging.error(
-            f"[make_cached_request] Failed to parse JSON response for {url}: {e}"
-        )
-        return None
-
 
 def get_repos(org: str) -> List[str]:
     logging.info(f"[get_repos] Start fetching repositories for org: {org}")
@@ -445,32 +394,6 @@ def save_to_yaml(data: Dict, filename: str = "runner_labels_summary.yml"):
     logging.info(f"[save_to_yaml] Data successfully saved to {filename}")
 
 
-def clear_cache():
-    """Clear all cached data."""
-    import shutil
-
-    if CACHE_DIR.exists():
-        shutil.rmtree(CACHE_DIR)
-        CACHE_DIR.mkdir(exist_ok=True)
-        logging.info(f"[clear_cache] Cleared cache directory: {CACHE_DIR}")
-    else:
-        logging.info(f"[clear_cache] Cache directory does not exist: {CACHE_DIR}")
-
-
-def get_cache_stats():
-    """Get statistics about the cache."""
-    if not CACHE_DIR.exists():
-        return {"total_files": 0, "total_size_mb": 0}
-
-    cache_files = list(CACHE_DIR.glob("*.json"))
-    total_size = sum(f.stat().st_size for f in cache_files)
-
-    return {
-        "total_files": len(cache_files),
-        "total_size_mb": round(total_size / (1024 * 1024), 2),
-    }
-
-
 def download_scale_config(url: str, dest: str = "scale-config.yml") -> bool:
     """Download scale-config.yml from the given URL if it does not exist locally."""
     if os.path.exists(dest):

tools/analytics/org/cache_manager.py

@@ -4,6 +4,8 @@
 from pathlib import Path
 from typing import Dict, Optional
 
+import requests
+
 
 # Cache configuration
 CACHE_DIR = Path("cache")
@@ -95,4 +97,81 @@ def set(self, url: str, data: Dict) -> None:
                 json.dump(data, f, indent=2)
             logging.debug(f"[CacheManager] Cached response for URL: {url}")
         except IOError as e:
-            logging.error(f"[CacheManager] Failed to write cache for {url}: {e}")
+            logging.error(f"[CacheManager] Failed to write cache for {url}: {e}")
+
+
+# Global cache manager instance
+cache_manager = CacheManager()
+
+
+def get_cache_stats():
+    """Get statistics about the cache."""
+    if not CACHE_DIR.exists():
+        return {"total_files": 0, "total_size_mb": 0}
+
+    cache_files = list(CACHE_DIR.glob("*.json"))
+    total_size = sum(f.stat().st_size for f in cache_files)
+
+    return {
+        "total_files": len(cache_files),
+        "total_size_mb": round(total_size / (1024 * 1024), 2),
+    }
+
+
+def clear_cache():
+    """Clear all cached data."""
+    import shutil
+
+    if CACHE_DIR.exists():
+        shutil.rmtree(CACHE_DIR)
+        CACHE_DIR.mkdir(exist_ok=True)
+        logging.info(f"[clear_cache] Cleared cache directory: {CACHE_DIR}")
+    else:
+        logging.info(f"[clear_cache] Cache directory does not exist: {CACHE_DIR}")
+
+
+HEADERS = {
+    "Authorization": f"Bearer {GITHUB_TOKEN}",
+    "Accept": "application/vnd.github+json",
+}
+
+
+def make_cached_request(
+    url: str, headers: Optional[Dict[str, str]] = None
+) -> Optional[Dict]:
+    """
+    Make an HTTP request with caching. Returns the JSON response if successful.
+
+    Args:
+        url: The URL to request
+        headers: Optional headers for the request
+
+    Returns:
+        JSON response data if successful, None if failed
+    """
+    # Check cache first
+    cached_response = cache_manager.get(url)
+    if cached_response:
+        logging.info(f"[make_cached_request] Using cached response for: {url}")
+        return cached_response
+
+    # Make actual HTTP request
+    logging.info(f"[make_cached_request] Making HTTP request to: {url}")
+    try:
+        response = requests.get(url, headers=headers or HEADERS)
+        response.raise_for_status()
+        data = response.json()
+
+        # Cache successful response
+        cache_manager.set(url, data)
+        logging.info(f"[make_cached_request] Successfully cached response for: {url}")
+        return data
+
+    except requests.exceptions.RequestException as e:
+        logging.error(f"[make_cached_request] HTTP request failed for {url}: {e}")
+        return None
+    except json.JSONDecodeError as e:
+        logging.error(
+            f"[make_cached_request] Failed to parse JSON response for {url}: {e}"
+        )
+        return None