GPU Compute Capability Scraper

get_compute_capability.py

from bs4 import BeautifulSoup
import requests

def get_gpu_info(url):
    """Parses the given URL and extracts GPU, compute capability, and category information."""
    try:
        response = requests.get(url, allow_redirects=True)
        response.raise_for_status()
        soup = BeautifulSoup(response.content, 'html.parser')

        all_gpu_info = []

        # Look for tables with GPU information
        tables = soup.find_all('table')
        
        for table in tables:
            headers = [th.text.strip() for th in table.find_all('th')]
            
            # Check if this is the compute capability table
            if 'Compute Capability' in headers:
                # Find the index of each column
                cc_idx = headers.index('Compute Capability')
                category_indices = {i: headers[i] for i in range(len(headers)) if i != cc_idx}
                
                # Parse each row
                for row in table.find_all('tr')[1:]:  # Skip header row
                    cells = row.find_all('td')
                    if len(cells) <= cc_idx:
                        continue
                    
                    compute_capability = cells[cc_idx].text.strip()
                    
                    # Extract GPUs from each category column
                    for idx, category in category_indices.items():
                        if idx < len(cells):
                            cell = cells[idx]
                            
                            # Try to find individual GPU names by looking at HTML structure
                            # Check for <br> tags or other line separators
                            gpu_names = []
                            
                            # Replace <br> tags with a delimiter
                            for br in cell.find_all('br'):
                                br.replace_with('\n')
                            
                            # Get text and split by newlines
                            cell_text = cell.get_text(separator='\n')
                            gpu_names = [name.strip() for name in cell_text.split('\n') if name.strip()]
                            
                            for gpu_name in gpu_names:
                                if gpu_name:  # Only add non-empty names
                                    all_gpu_info.append({
                                        'GPU': gpu_name,
                                        'Compute Capability': compute_capability,
                                        'Category': category
                                    })

        return all_gpu_info

    except requests.exceptions.RequestException as e:
        print(f"Error fetching URL: {e}")
        return []
    except Exception as e:
        print(f"An error occurred: {e}")
        return []

def sort_gpus(gpu_list):
    """Sorts the GPU list by compute capability (descending)."""
    def compute_capability_key(gpu):
        cc = gpu.get('Compute Capability', '0.0') # Default to 0.0 if not found
        try:
            return float(cc) #Convert to float for numerical sorting
        except ValueError: # Handle cases like "5.0+"
            try:
                return float(cc[:-1]) # Try removing the last character (+)
            except ValueError:
                return 0.0 # Default to 0 if neither works
    return sorted(gpu_list, key=compute_capability_key, reverse=True)

def print_markdown(gpu_list):
    """Prints the GPU information in Markdown format."""
    if not gpu_list:
        print("No GPU information found.")
        return

    print("| GPU | Category | Compute Capability |")
    print("|:---:|:---:|:---:|")  # Markdown table header separator

    for gpu in gpu_list:
        gpu_name = gpu.get('GPU', 'N/A').replace("|", "\\|") # Escape pipes in GPU names
        category = gpu.get('Category', 'N/A').replace("|", "\\|") # Escape pipes in category names
        cc = gpu.get('Compute Capability', 'N/A').replace("|", "\\|") # Escape pipes in CC
        print(f"| {gpu_name} | {category} | {cc} |")

# Get GPU information from both webpages
latest_gpu_info = get_gpu_info('https://developer.nvidia.com/cuda/gpus')
legacy_gpu_info = get_gpu_info('https://developer.nvidia.com/cuda/gpus/legacy')

# Combine the GPU information from both sources
combined_gpu_info = latest_gpu_info + legacy_gpu_info

# Sort the GPU information
sorted_gpu_info = sort_gpus(combined_gpu_info)

# Print the sorted information in Markdown format
print_markdown(sorted_gpu_info)