SignerX/SignVerse-2M

1#!/usr/bin/env python32from __future__ import annotations3 4import argparse5import csv6import json7import os8import re9import subprocess10from collections import Counter11from pathlib import Path12 13from utils.stats_npz import processed_complete_counter_path, read_processed_complete_counter14 15 16VIDEO_EXTS = {'.mp4', '.mkv', '.webm', '.mov'}17ARRAY_RANGE_RE = re.compile(r'^(\d+)_\[(.+)\]$')18PROCESSED_REQUIRED_COLUMNS = {19    'video_id',20    'download_status',21    'process_status',22    'upload_status',23    'archive_name',24}25GPU_PARTITIONS = ['gpu', 'gpu-redhat', 'cgpu']26DEFAULT_VIDEOS_PER_DWPOSE_JOB = 2027 28 29def run_command(cmd: list[str]) -> str:30    try:31        proc = subprocess.run(cmd, check=False, capture_output=True, text=True)32    except OSError:33        return ''34    return (proc.stdout or '').strip()35 36 37def count_claims(directory: Path) -> int:38    if not directory.exists():39        return 040    return sum(1 for _ in directory.glob('*.claim'))41 42 43def aggregate_claims_by_job_key(directory: Path) -> dict[str, int]:44    counts: Counter[str] = Counter()45    if not directory.exists():46        return {}47    for path in directory.glob('*.claim'):48        try:49            lines = path.read_text(encoding='utf-8').splitlines()50        except OSError:51            continue52        job_key = ''53        for line in lines:54            if line.startswith('job_key='):55                job_key = line.split('=', 1)[1].strip()56                break57        if job_key:58            counts[job_key] += 159    return dict(counts)60 61 62def read_videos_per_dwpose_job(root_dir: Path) -> int:63    worker = root_dir / 'slurm' / 'process_dwpose_array.slurm'64    if not worker.exists():65        return DEFAULT_VIDEOS_PER_DWPOSE_JOB66    try:67        for line in worker.read_text().splitlines():68            if line.startswith('VIDEOS_PER_JOB='):69                m = re.search(r'\$\{VIDEOS_PER_JOB:-([0-9]+)\}', line)70                if m:71                    return int(m.group(1))72    except Exception:73        pass74    return DEFAULT_VIDEOS_PER_DWPOSE_JOB75 76 77def sum_file_sizes(paths: list[Path]) -> int:78    total = 079    for path in paths:80        try:81            total += path.stat().st_size82        except FileNotFoundError:83            continue84    return total85 86 87def count_uploaded(progress_path: Path) -> tuple[int, int]:88    if not progress_path.exists():89        return 0, 090    try:91        data = json.loads(progress_path.read_text())92    except Exception:93        return 0, 094    archives = data.get('archives', {})95    uploaded_folders = data.get('uploaded_folders', {})96    return len(archives), len(uploaded_folders)97 98 99def expand_task_count(jobid_token: str) -> int:100    m = ARRAY_RANGE_RE.match(jobid_token)101    if not m:102        return 1103    body = m.group(2)104    if '%' in body:105        body = body.split('%', 1)[0]106    total = 0107    for part in body.split(','):108        part = part.strip()109        if not part:110            continue111        if '-' in part:112            a, b = part.split('-', 1)113            try:114                total += int(b) - int(a) + 1115            except ValueError:116                total += 1117        else:118            total += 1119    return max(total, 1)120 121 122def queue_status(username: str) -> dict[str, object]:123    output = run_command(['squeue', '-u', username, '-h', '-o', '%i|%j|%T|%P'])124    job_counts: Counter[str] = Counter()125    partition_counts: Counter[str] = Counter()126    active_tasks_by_partition: Counter[str] = Counter()127    running_dwpose = 0128    running_download = 0129    pending_download = 0130    pending_dwpose = 0131    if output:132        for line in output.splitlines():133            parts = line.split('|')134            if len(parts) != 4:135                continue136            jobid_token, job, state, partition = parts137            count = expand_task_count(jobid_token)138            job_counts[f'{job}|{state}'] += count139            partition_counts[f'{job}|{partition}|{state}'] += count140            if state in {'RUNNING', 'PENDING', 'CONFIGURING'}:141                active_tasks_by_partition[partition] += count142            if job == 'dwpose' and state == 'RUNNING':143                running_dwpose += count144            if job == 'download' and state == 'RUNNING':145                running_download += count146            if job == 'download' and state in {'PENDING', 'CONFIGURING'}:147                pending_download += count148            if job == 'dwpose' and state in {'PENDING', 'CONFIGURING'}:149                pending_dwpose += count150    total_download = running_download + pending_download151    return {152        'running_dwpose': running_dwpose,153        'running_download': running_download,154        'pending_dwpose_jobs': pending_dwpose,155        'pending_download_jobs': pending_download,156        'total_download_jobs': total_download,157        'job_state_counts': dict(job_counts),158        'job_partition_state_counts': dict(partition_counts),159        'active_tasks_by_partition': dict(active_tasks_by_partition),160    }161 162 163def gpu_partition_capacity(partitions: list[str], active_tasks_by_partition: dict[str, int]) -> list[dict[str, object]]:164    qos_limit_by_part: dict[str, int] = {}165    qos_output = run_command(['sacctmgr', 'show', 'qos', 'format=Name,MaxSubmitPU', '-P'])166    if qos_output:167        for line in qos_output.splitlines():168            if not line.strip() or '|' not in line:169                continue170            name, max_submit = line.split('|', 1)171            name = name.strip()172            max_submit = max_submit.strip()173            if name in partitions and max_submit:174                try:175                    qos_limit_by_part[name] = int(max_submit)176                except ValueError:177                    pass178 179    rows: list[dict[str, object]] = []180    for partition in partitions:181        free_gpus = 0182        nodes_output = run_command(['sinfo', '-h', '-N', '-p', partition, '-o', '%N'])183        nodes = [line.strip() for line in nodes_output.splitlines() if line.strip()]184        for node in nodes:185            node_line = run_command(['scontrol', 'show', 'node', node, '-o'])186            if not node_line:187                continue188            state_m = re.search(r'\bState=([^ ]+)', node_line)189            state = state_m.group(1).lower() if state_m else ''190            if any(flag in state for flag in ('drain', 'drained', 'down', 'fail', 'inval')):191                continue192            cfg_m = re.search(r'\bCfgTRES=.*?(?:,|^)gres/gpu=(\d+)', node_line)193            alloc_m = re.search(r'\bAllocTRES=.*?(?:,|^)gres/gpu=(\d+)', node_line)194            total = int(cfg_m.group(1)) if cfg_m else 0195            used = int(alloc_m.group(1)) if alloc_m else 0196            free = total - used197            if free > 0:198                free_gpus += free199        active_tasks = int(active_tasks_by_partition.get(partition, 0))200        qos_limit = qos_limit_by_part.get(partition)201        submit_slots = free_gpus202        if qos_limit is not None:203            submit_slots = min(submit_slots, max(0, qos_limit - active_tasks))204        rows.append({205            'partition': partition,206            'free_gpus': free_gpus,207            'active_tasks': active_tasks,208            'qos_limit': qos_limit,209            'submit_slots': submit_slots,210        })211    return rows212 213 214def filesystem_avail_bytes(path: Path) -> int:215    try:216        proc = subprocess.run(['df', '-B1', str(path)], check=False, capture_output=True, text=True)217        lines = (proc.stdout or '').splitlines()218        if len(lines) < 2:219            return 0220        fields = lines[1].split()221        if len(fields) < 4:222            return 0223        return int(fields[3])224    except Exception:225        return 0226 227 228def human_bytes(num: int) -> str:229    value = float(num)230    for unit in ['B', 'KB', 'MB', 'GB', 'TB', 'PB']:231        if value < 1024.0:232            return f'{value:.1f}{unit}'233        value /= 1024.0234    return f'{value:.1f}EB'235 236 237def read_source_manifest_count(path: Path) -> int:238    if not path.exists():239        return 0240    count = 0241    with path.open('r', encoding='utf-8-sig', newline='') as f:242        reader = csv.reader(f)243        for row in reader:244            if not row:245                continue246            if not (row[0] or '').strip():247                continue248            count += 1249    return count250 251 252def fast_count_complete(dataset_dir: Path, scratch_dataset_dir: Path) -> int:253    roots = [str(p) for p in (dataset_dir, scratch_dataset_dir) if p.exists()]254    if not roots:255        return 0256    cmd = ['find', *roots, '-path', '*/npz/.complete']257    try:258        proc = subprocess.run(cmd, check=False, capture_output=True, text=True)259    except OSError:260        return 0261    if proc.returncode not in (0, 1):262        return 0263    out = proc.stdout264    if not out:265        return 0266    return sum(1 for line in out.splitlines() if line.strip())267 268 269def read_processed_progress(path: Path) -> dict[str, object]:270    result = {271        'csv_exists': path.exists(),272        'csv_ok': False,273        'csv_error': '',274        'processed_rows': 0,275        'download_ok_rows': 0,276        'download_skipped_rows': 0,277        'download_running_rows': 0,278        'download_pending_rows': 0,279        'process_ok_rows': 0,280        'process_running_rows': 0,281        'upload_uploaded_rows': 0,282    }283    if not path.exists():284        result['csv_error'] = 'missing'285        return result286    try:287        with path.open('r', encoding='utf-8-sig', newline='') as f:288            reader = csv.DictReader(f)289            fieldnames = list(reader.fieldnames or [])290            missing = sorted(PROCESSED_REQUIRED_COLUMNS - set(fieldnames))291            if missing:292                result['csv_error'] = f'missing_columns:{",".join(missing)}'293                return result294            rows = list(reader)295        result['processed_rows'] = len(rows)296        for row in rows:297            d = (row.get('download_status') or '').strip()298            p = (row.get('process_status') or '').strip()299            u = (row.get('upload_status') or '').strip()300            if d == 'ok':301                result['download_ok_rows'] += 1302            elif d == 'skipped':303                result['download_skipped_rows'] += 1304            elif d == 'running':305                result['download_running_rows'] += 1306            else:307                result['download_pending_rows'] += 1308            if p == 'ok':309                result['process_ok_rows'] += 1310            elif p == 'running':311                result['process_running_rows'] += 1312            if u == 'uploaded':313                result['upload_uploaded_rows'] += 1314        result['csv_ok'] = True315        return result316    except Exception as exc:317        result['csv_error'] = str(exc)318        return result319 320 321def run_sync(runtime_root: Path) -> str:322    sync_script = Path('/cache/home/sf895/SignVerse-2M/scripts/sync_processed_csv_from_runtime.py')323    if not sync_script.exists():324        return 'missing_sync_script'325    cmd = [326        'python3', str(sync_script),327        '--source-metadata-csv', str(runtime_root / 'SignVerse-2M-metadata_ori.csv'),328        '--output-metadata-csv', str(runtime_root / 'SignVerse-2M-metadata_processed.csv'),329        '--raw-video-dir', str(runtime_root / 'raw_video'),330        '--scratch-raw-video-dir', str(Path(f'/scratch/{os.environ.get("USER", "sf895")}/SignVerse-2M-runtime/raw_video')),331        '--raw-caption-dir', str(runtime_root / 'raw_caption'),332        '--raw-metadata-dir', str(runtime_root / 'raw_metadata'),333        '--dataset-dir', str(runtime_root / 'dataset'),334        '--scratch-dataset-dir', str(Path(f'/scratch/{os.environ.get("USER", "sf895")}/SignVerse-2M-runtime/dataset')),335        '--progress-path', str(runtime_root / 'archive_upload_progress.json'),336        '--status-journal-path', str(runtime_root / 'upload_status_journal.jsonl'),337    ]338    try:339        proc = subprocess.run(cmd, check=False, capture_output=True, text=True)340    except OSError as exc:341        return f'error:{exc}'342    if proc.returncode == 0:343        return (proc.stdout or '').strip() or 'ok'344    err = (proc.stderr or proc.stdout or '').strip()345    return f'failed:{err}'346 347 348def main() -> None:349    parser = argparse.ArgumentParser(description='Report SignVerse runtime status.')350    parser.add_argument('--runtime-root', default='/home/sf895/SignVerse-2M-runtime')351    parser.add_argument('--username', default='sf895')352    parser.add_argument('--no-sync', action='store_true')353    parser.add_argument('--json', action='store_true')354    parser.add_argument('--include-partitions', action='store_true')355    parser.add_argument('--scan-csv', action='store_true')356    parser.add_argument('--scan-complete', action='store_true')357    parser.add_argument('--scan-runtime-size', action='store_true')358    args = parser.parse_args()359 360    runtime_root = Path(args.runtime_root)361    root_dir = Path('/cache/home/sf895/SignVerse-2M')362    raw_dir = runtime_root / 'raw_video'363    scratch_raw_dir = Path(f'/scratch/{os.environ.get("USER", "sf895")}/SignVerse-2M-runtime/raw_video')364    dataset_dir = runtime_root / 'dataset'365    scratch_dataset_dir = Path(f'/scratch/{os.environ.get("USER", "sf895")}/SignVerse-2M-runtime/dataset')366    claims_dir = runtime_root / 'slurm' / 'state' / 'claims'367    download_claims_dir = runtime_root / 'slurm' / 'state' / 'download_claims'368    progress_path = runtime_root / 'archive_upload_progress.json'369    source_csv = runtime_root / 'SignVerse-2M-metadata_ori.csv'370    processed_csv = runtime_root / 'SignVerse-2M-metadata_processed.csv'371 372    sync_result = 'skipped'373    if not args.no_sync:374        sync_result = run_sync(runtime_root)375 376    raw_complete: dict[str, Path] = {}377    raw_temp: list[Path] = []378    for current_raw_dir in [raw_dir, scratch_raw_dir]:379        if not current_raw_dir.exists():380            continue381        for path in current_raw_dir.iterdir():382            if not path.is_file():383                continue384            if path.suffix.lower() in VIDEO_EXTS:385                raw_complete.setdefault(path.stem, path)386            else:387                raw_temp.append(path)388 389    raw_size = sum_file_sizes(list(raw_complete.values()))390    runtime_size = 0391    if args.scan_runtime_size and runtime_root.exists():392        proc = subprocess.run(['du', '-sb', str(runtime_root)], check=False, capture_output=True, text=True)393        if proc.returncode == 0 and proc.stdout.strip():394            try:395                runtime_size = int(proc.stdout.split()[0])396            except Exception:397                runtime_size = 0398 399    source_rows = read_source_manifest_count(source_csv)400    if args.scan_csv:401        progress = read_processed_progress(processed_csv)402    else:403        progress = {404            'csv_exists': processed_csv.exists(),405            'csv_ok': False,406            'csv_error': 'skipped',407            'processed_rows': 0,408            'download_ok_rows': 0,409            'download_skipped_rows': 0,410            'download_running_rows': 0,411            'download_pending_rows': 0,412            'process_ok_rows': 0,413            'process_running_rows': 0,414            'upload_uploaded_rows': 0,415        }416    videos_per_dwpose_job = read_videos_per_dwpose_job(root_dir)417 418    payload = {419        'sync_result': sync_result,420        'download_normal': len(raw_temp) == 0,421        'raw_videos': len(raw_complete),422        'raw_temp_files': len(raw_temp),423        'sent_to_gpu': count_claims(claims_dir),424        'processed_complete': read_processed_complete_counter(processed_complete_counter_path(runtime_root / 'stats.npz')),425        'active_downloads': count_claims(download_claims_dir),426        'uploaded_archives': 0,427        'uploaded_folders': 0,428        'raw_size_bytes': raw_size,429        'runtime_size_bytes': runtime_size,430        'filesystem_avail_bytes': filesystem_avail_bytes(runtime_root),431        'source_rows': source_rows,432        'csv_exists': progress['csv_exists'],433        'csv_ok': progress['csv_ok'],434        'csv_error': progress['csv_error'],435        'processed_rows': progress['processed_rows'],436        'download_ok_rows': progress['download_ok_rows'],437        'download_skipped_rows': progress['download_skipped_rows'],438        'download_running_rows': progress['download_running_rows'],439        'download_pending_rows': progress['download_pending_rows'],440        'process_ok_rows': progress['process_ok_rows'],441        'process_running_rows': progress['process_running_rows'],442        'upload_uploaded_rows': progress['upload_uploaded_rows'],443    }444    uploaded_archives, uploaded_folders = count_uploaded(progress_path)445    payload['uploaded_archives'] = uploaded_archives446    payload['uploaded_folders'] = uploaded_folders447    payload.update(queue_status(args.username))448    process_claims_by_job = aggregate_claims_by_job_key(claims_dir)449    payload['videos_per_dwpose_job'] = videos_per_dwpose_job450    payload['process_claim_job_keys'] = len(process_claims_by_job)451    payload['process_claim_videos_actual'] = sum(process_claims_by_job.values())452    payload['process_claim_videos_max_per_job'] = max(process_claims_by_job.values(), default=0)453    payload['running_dwpose_jobs'] = payload['running_dwpose']454    payload['running_dwpose_videos_estimated'] = payload['running_dwpose'] * videos_per_dwpose_job455    payload['pending_dwpose_videos_estimated'] = payload['pending_dwpose_jobs'] * videos_per_dwpose_job456    payload['total_dwpose_jobs'] = payload['running_dwpose'] + payload['pending_dwpose_jobs']457    payload['total_dwpose_videos_estimated'] = payload['running_dwpose_videos_estimated'] + payload['pending_dwpose_videos_estimated']458    if args.include_partitions:459        payload['gpu_partition_capacity'] = gpu_partition_capacity(GPU_PARTITIONS, payload.get('active_tasks_by_partition', {}))460    else:461        payload['gpu_partition_capacity'] = []462        payload['job_partition_state_counts'] = {}463        payload['active_tasks_by_partition'] = {}464    payload['csv_row_match'] = (payload['processed_rows'] == payload['source_rows']) if payload['csv_ok'] else False465 466    if args.json:467        print(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True))468        return469 470    print(f"sync_result={payload['sync_result']}")471    print(f"download_normal={payload['download_normal']}")472    print(f"raw_videos={payload['raw_videos']}")473    print(f"raw_temp_files={payload['raw_temp_files']}")474    print(f"sent_to_gpu={payload['sent_to_gpu']}")475    print(f"running_dwpose={payload['running_dwpose']}")476    print(f"running_dwpose_jobs={payload['running_dwpose_jobs']}")477    print(f"pending_dwpose_jobs={payload['pending_dwpose_jobs']}")478    print(f"total_dwpose_jobs={payload['total_dwpose_jobs']}")479    print(f"videos_per_dwpose_job={payload['videos_per_dwpose_job']}")480    print(f"process_claim_job_keys={payload['process_claim_job_keys']}")481    print(f"process_claim_videos_actual={payload['process_claim_videos_actual']}")482    print(f"process_claim_videos_max_per_job={payload['process_claim_videos_max_per_job']}")483    print(f"running_dwpose_videos_estimated={payload['running_dwpose_videos_estimated']}")484    print(f"pending_dwpose_videos_estimated={payload['pending_dwpose_videos_estimated']}")485    print(f"total_dwpose_videos_estimated={payload['total_dwpose_videos_estimated']}")486    print(f"processed_complete={payload['processed_complete']}")487    print(f"active_downloads={payload['active_downloads']}")488    print(f"running_download_jobs={payload['running_download']}")489    print(f"pending_download_jobs={payload['pending_download_jobs']}")490    print(f"total_download_jobs={payload['total_download_jobs']}")491    print(f"uploaded_archives={payload['uploaded_archives']}")492    print(f"uploaded_folders={payload['uploaded_folders']}")493    print(f"source_rows={payload['source_rows']}")494    print(f"processed_rows={payload['processed_rows']}")495    print(f"csv_ok={payload['csv_ok']}")496    print(f"csv_row_match={payload['csv_row_match']}")497    print(f"csv_error={payload['csv_error']}")498    print(f"download_ok_rows={payload['download_ok_rows']}")499    print(f"download_skipped_rows={payload['download_skipped_rows']}")500    print(f"download_running_rows={payload['download_running_rows']}")501    print(f"download_pending_rows={payload['download_pending_rows']}")502    print(f"process_ok_rows={payload['process_ok_rows']}")503    print(f"process_running_rows={payload['process_running_rows']}")504    print(f"upload_uploaded_rows={payload['upload_uploaded_rows']}")505    for key in sorted(payload.get('job_partition_state_counts', {})):506        print(f"job_partition_state[{key}]={payload['job_partition_state_counts'][key]}")507    for row in payload.get('gpu_partition_capacity', []):508        qos_limit = row['qos_limit'] if row['qos_limit'] is not None else 'na'509        print(510            f"gpu_partition[{row['partition']}]=free_gpus={row['free_gpus']},"511            f"active_tasks={row['active_tasks']},qos_limit={qos_limit},submit_slots={row['submit_slots']}"512        )513    print(f"raw_size={human_bytes(payload['raw_size_bytes'])}")514    print(f"runtime_size={human_bytes(payload['runtime_size_bytes'])}")515    print(f"filesystem_avail={human_bytes(payload['filesystem_avail_bytes'])}")516 517 518if __name__ == '__main__':519    main()520