iDRAC_Info/data/scripts/collect_gpu.py

import os
import re
import sys
import yaml
import argparse
import subprocess
import logging
import fnmatch
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from dotenv import load_dotenv, find_dotenv

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [INFO] root: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

# .env 로드
load_dotenv(find_dotenv())

IDRAC_USER = os.getenv("IDRAC_USER")
IDRAC_PASS = os.getenv("IDRAC_PASS")

# 경로 설정
BASE_DIR = Path(__file__).resolve().parent
PROFILES_FILE = BASE_DIR / "profiles" / "gpu_profiles.yaml"
OUTPUT_DIR = BASE_DIR.parent / "temp" / "staging"

def load_profile(profile_name="default"):
    """Load specific profile from gpu_profiles.yaml"""
    if not PROFILES_FILE.exists():
        logging.warning(f"Profile file not found at {PROFILES_FILE}. Using defaults.")
        return {}

    try:
        with open(PROFILES_FILE, 'r', encoding='utf-8') as f:
            data = yaml.safe_load(f)
            profiles = data.get('profiles', {})
            return profiles.get(profile_name, {})
    except Exception as e:
        logging.error(f"Error loading profile: {e}")
        return {}

def resolve_output_dir() -> Path:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    return OUTPUT_DIR

def parse_gpu_serials_from_hwinventory(hwinv_text: str, target_slots=None) -> dict:
    """
    Parse GPU serials and filter based on target_slots patterns.
    target_slots: list of glob patterns (e.g. ["Video.Slot.*"])
    """
    results = {}
    blocks = re.split(r"\n\s*\n", hwinv_text, flags=re.MULTILINE)

    for block in blocks:
        # Check if it looks like a Video/GPU block
        if not re.search(r"\[?InstanceID:\s*Video\.Slot\.", block):
            continue

        # Extract FQDD
        fqdd = None
        m_fqdd = re.search(r"^FQDD\s*=\s*([^\r\n]+)", block, flags=re.MULTILINE)
        if m_fqdd:
            fqdd = m_fqdd.group(1).strip()
        else:
            m_hdr = re.search(r"\[InstanceID:\s*(Video\.Slot\.[^\]\r\n]+)\]", block)
            if m_hdr:
                fqdd = m_hdr.group(1).strip()

        if not fqdd:
            continue

        # Filtering logic
        if target_slots:
            matched = False
            for pattern in target_slots:
                # fnmatch supports simple unix shell style wildcards (*, ?, [], etc)
                if fnmatch.fnmatch(fqdd, pattern):
                    matched = True
                    break
            if not matched:
                continue

        # Extract SerialNumber
        m_sn = re.search(r"^SerialNumber\s*=\s*([^\r\n]+)", block, flags=re.MULTILINE)
        serial = m_sn.group(1).strip() if m_sn else "Not Found"

        results[fqdd] = serial

    return results

def get_slot_index(fqdd: str):
    """
    Extract logic for sorting.
    Format usually: Video.Slot.X-Y -> returns (X, Y) tuple
    """
    m = re.search(r"Video\.Slot\.(\d+)-(\d+)", fqdd)
    if m:
        return (int(m.group(1)), int(m.group(2)))
    m_simple = re.search(r"Video\.Slot\.(\d+)", fqdd)
    if m_simple:
        return (int(m_simple.group(1)), 0)
    return (float('inf'), float('inf'))

def fetch_idrac_info(idrac_ip: str, output_dir: Path, gpu_settings: dict) -> None:
    try:
        # 1. Get SVC Tag
        cmd_getsysinfo = [
            "racadm", "-r", idrac_ip, "-u", IDRAC_USER or "", "-p", IDRAC_PASS or "", "getsysinfo"
        ]
        # execute without shell=True for security (if possible), but subprocess.getoutput uses shell
        # Here we use subprocess.run for better control if needed, but sticking to logic similar to original for now
        # Ideally: result = subprocess.run(cmd_getsysinfo, capture_output=True, text=True)
        # But let's stick to getoutput wrapper logic or use run compatible way
        
        # Using subprocess.run is safer and modern
        res_sys = subprocess.run(cmd_getsysinfo, capture_output=True, text=True, encoding='utf-8', errors='ignore')
        getsysinfo = res_sys.stdout
        
        svc_tag_match = re.search(r"SVC Tag\s*=\s*(\S+)", getsysinfo)
        svc_tag = svc_tag_match.group(1) if svc_tag_match else None

        if not svc_tag:
            logging.error(f"Failed to retrieve SVC Tag for IP: {idrac_ip}")
            return

        # 2. Get Hardware Inventory
        cmd_hwinv = [
            "racadm", "-r", idrac_ip, "-u", IDRAC_USER or "", "-p", IDRAC_PASS or "", "hwinventory"
        ]
        res_inv = subprocess.run(cmd_hwinv, capture_output=True, text=True, encoding='utf-8', errors='ignore')
        hwinv_text = res_inv.stdout

        # 3. Parse and Filter
        target_slots = gpu_settings.get('target_slots', ["Video.Slot.*"])
        gpu_map = parse_gpu_serials_from_hwinventory(hwinv_text, target_slots)

        # 4. Sort
        sort_order = gpu_settings.get('sort_order', 'slot_index')
        
        if sort_order == 'slot_index':
            sorted_keys = sorted(gpu_map.keys(), key=lambda k: (get_slot_index(k), k))
        elif sort_order == 'serial':
            sorted_keys = sorted(gpu_map.keys(), key=lambda k: gpu_map[k])
        else:
            sorted_keys = sorted(gpu_map.keys())

        # 5. Write Output (Text format)
        output_file = output_dir / f"{svc_tag}.txt"
        with output_file.open("w", encoding="utf-8", newline="\n") as f:
            f.write(f"{svc_tag}\n")

            if not gpu_map:
                f.write("No matching GPU inventory found.\n")
                return

            for fqdd in sorted_keys:
                serial = gpu_map[fqdd]
                f.write(f"{fqdd}: {serial}\n")

            # Serials only line
            serials_only = [gpu_map[k] for k in sorted_keys if gpu_map[k] != "Not Found"]
            if serials_only:
                f.write(f"SERIALS: {';'.join(serials_only)}\n")

        logging.info(f"Saved {output_file}")

    except Exception as e:
        logging.error(f"Error processing IP {idrac_ip}: {e}")

def main():
    parser = argparse.ArgumentParser(description="Collect GPU Serial Info")
    parser.add_argument("ip_file", help="File containing list of IP addresses")
    parser.add_argument("--profile", default="default", help="Profile name from gpu_profiles.yaml")
    
    args = parser.parse_args()
    
    ip_path = Path(args.ip_file)
    if not ip_path.is_file():
        logging.error(f"IP file {args.ip_file} does not exist.")
        return

    # Load Profile
    profile = load_profile(args.profile)
    logging.info(f"Loaded profile: {args.profile}")
    
    gpu_settings = profile.get('gpu_settings', {})
    
    output_dir = resolve_output_dir()
    
    with ip_path.open("r", encoding="utf-8") as file:
        ip_addresses = [line.strip() for line in file if line.strip()]

    with ThreadPoolExecutor(max_workers=50) as executor:
        future_to_ip = {
            executor.submit(fetch_idrac_info, ip, output_dir, gpu_settings): ip 
            for ip in ip_addresses
        }
        for future in as_completed(future_to_ip):
            ip = future_to_ip[future]
            try:
                future.result()
            except Exception as e:
                logging.error(f"Error on {ip}: {e}")

if __name__ == "__main__":
    main()