run.py

import os
import sys
from dotenv import load_dotenv
# Load the .env file if it exists
load_dotenv()
os.environ["HF_XET_HIGH_PERFORMANCE"] = os.getenv("HF_XET_HIGH_PERFORMANCE", "1")
os.environ["HF_HUB_DISABLE_XET"] = os.getenv("HF_HUB_DISABLE_XET", "0")
os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
os.environ["OPENCV_FFMPEG_LOGLEVEL"] = "-8"
seed = None
if "SEED" in os.environ:
    try:
        seed = int(os.environ["SEED"])
    except ValueError:
        print(f"Invalid SEED value: {os.environ['SEED']}. SEED must be an integer.")

sys.path.insert(0, os.getcwd())

# The UI launches jobs with no console; keep anything we shell out to (torch
# compiles, HF git downloads) from flashing a console window. Must come before
# any import that might spawn a subprocess.
from toolkit.win_console import suppress_child_consoles
suppress_child_consoles()

# must come before ANY torch or fastai imports
# import toolkit.cuda_malloc

# turn off diffusers telemetry until I can figure out how to make it opt-in
os.environ['DISABLE_TELEMETRY'] = 'YES'

# Set ROCm environment variables for better HIP error handling and performance
# These should be set before importing torch
if os.environ.get("AMD_SERIALIZE_KERNEL") is None:
    os.environ["AMD_SERIALIZE_KERNEL"] = "3"  # Better error reporting for HIP errors
if os.environ.get("TORCH_USE_HIP_DSA") is None:
    os.environ["TORCH_USE_HIP_DSA"] = "1"  # Enable device-side assertions
if os.environ.get("HSA_ENABLE_SDMA") is None:
    os.environ["HSA_ENABLE_SDMA"] = "0"  # Disable SDMA for APU compatibility
if os.environ.get("HSA_XNACK") is None:
    os.environ["HSA_XNACK"] = "1"  # Enable XNACK for unified memory on Strix Halo APU
if os.environ.get("GPU_MAX_HEAP_SIZE") is None:
    os.environ["GPU_MAX_HEAP_SIZE"] = "100"  # Allow full heap for APU
if os.environ.get("GPU_MAX_ALLOC_PERCENT") is None:
    os.environ["GPU_MAX_ALLOC_PERCENT"] = "100"
if os.environ.get("PYTORCH_ROCM_ALLOC_CONF") is None:
    os.environ["PYTORCH_ROCM_ALLOC_CONF"] = "max_split_size_mb:768,garbage_collect=1"  # Better VRAM fragmentation

# Workaround for HIPBLAS errors with quantized models
if os.environ.get("ROCBLAS_USE_HIPBLASLT") is None:
    os.environ["ROCBLAS_USE_HIPBLASLT"] = "0"  # Disable HIPBLASLT to avoid quantized model crashes
if os.environ.get("ROCBLAS_LOG_LEVEL") is None:
    os.environ["ROCBLAS_LOG_LEVEL"] = "0"  # Disable verbose logging

# Set HSA_OVERRIDE_GFX_VERSION and PYTORCH_ROCM_ARCH for ROCm
# Must be before torch import so HIP initializes correctly for gfx1151 (Strix Halo)
if os.environ.get("HSA_OVERRIDE_GFX_VERSION") is None or os.environ.get("PYTORCH_ROCM_ARCH") is None:
    # Only attempt ROCm detection on Linux
    import platform
    is_linux = platform.system() == "Linux"
    has_rocm = False
    if is_linux:
        import shutil
        if shutil.which("rocm-smi") or os.path.isdir("/opt/rocm"):
            has_rocm = True
    if has_rocm:
        # Try to detect gfx arch via rocm-smi
        detected_gfx = None
        try:
            import subprocess
            import re
            out = subprocess.run(["rocm-smi", "--showproductname"], capture_output=True, text=True, timeout=5)
            if out.returncode == 0:
                # Look for gfx pattern or GFX Version
                m = re.search(r"gfx\d+", out.stdout)
                if m:
                    detected_gfx = m.group(0)
                # Alternative: GFX Version line
                if not detected_gfx:
                    m2 = re.search(r"GFX Version:\s*(gfx\d+)", out.stdout)
                    if m2:
                        detected_gfx = m2.group(1)
            # Fallback: rocminfo
            if not detected_gfx:
                out2 = subprocess.run(["rocminfo"], capture_output=True, text=True, timeout=5)
                if out2.returncode == 0:
                    m = re.search(r"gfx\d+", out2.stdout)
                    if m:
                        detected_gfx = m.group(0)
        except Exception:
            pass
        # Map gfx to HSA version: gfx1151 -> 11.5.1, gfx1100 -> 11.0.0, etc.
        def _gfx_to_hsa(gfx: str):
            try:
                digits = gfx.replace("gfx", "")
                while len(digits) < 4:
                    digits += "0"
                major = digits[0:2]
                minor = digits[2]
                patch = digits[3]
                return f"{int(major)}.{minor}.{patch}"
            except Exception:
                return None
        if detected_gfx:
            hsa_ver = _gfx_to_hsa(detected_gfx)
            if hsa_ver and os.environ.get("HSA_OVERRIDE_GFX_VERSION") is None:
                os.environ["HSA_OVERRIDE_GFX_VERSION"] = hsa_ver
            # Map to ROCm nightly dir name: gfx110x -> gfx110X-all
            rocm_arch = detected_gfx
            if detected_gfx.startswith("gfx110"):
                rocm_arch = "gfx110X-all"
            if os.environ.get("PYTORCH_ROCM_ARCH") is None:
                os.environ["PYTORCH_ROCM_ARCH"] = rocm_arch
        else:
            # Fallback for Strix Halo when detection fails but ROCm is present
            if os.environ.get("HSA_OVERRIDE_GFX_VERSION") is None:
                os.environ["HSA_OVERRIDE_GFX_VERSION"] = "11.5.1"
            if os.environ.get("PYTORCH_ROCM_ARCH") is None:
                os.environ["PYTORCH_ROCM_ARCH"] = "gfx1151"

import torch

# check if we have DEBUG_TOOLKIT in env
if os.environ.get("DEBUG_TOOLKIT", "0") == "1":
    torch.autograd.set_detect_anomaly(True)

if seed is not None:
    import random
    import numpy as np
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

import argparse
from toolkit.job import get_job
from toolkit.accelerator import get_accelerator
from toolkit.print import print_acc, setup_log_to_file

accelerator = get_accelerator()


def print_end_message(jobs_completed, jobs_failed):
    if not accelerator.is_main_process:
        return
    failure_string = f"{jobs_failed} failure{'' if jobs_failed == 1 else 's'}" if jobs_failed > 0 else ""
    completed_string = f"{jobs_completed} completed job{'' if jobs_completed == 1 else 's'}"

    print_acc("")
    print_acc("========================================")
    print_acc("Result:")
    if len(completed_string) > 0:
        print_acc(f" - {completed_string}")
    if len(failure_string) > 0:
        print_acc(f" - {failure_string}")
    print_acc("========================================")


def main():
    parser = argparse.ArgumentParser()

    # require at lease one config file
    parser.add_argument(
        'config_file_list',
        nargs='+',
        type=str,
        help='Name of config file (eg: person_v1 for config/person_v1.json/yaml), or full path if it is not in config folder, you can pass multiple config files and run them all sequentially'
    )

    # flag to continue if failed job
    parser.add_argument(
        '-r', '--recover',
        action='store_true',
        help='Continue running additional jobs even if a job fails'
    )

    # flag to continue if failed job
    parser.add_argument(
        '-n', '--name',
        type=str,
        default=None,
        help='Name to replace [name] tag in config file, useful for shared config file'
    )
    
    parser.add_argument(
        '-l', '--log',
        type=str,
        default=None,
        help='Log file to write output to'
    )
    args = parser.parse_args()
    
    if args.log is not None:
        setup_log_to_file(args.log)

    config_file_list = args.config_file_list
    if len(config_file_list) == 0:
        raise Exception("You must provide at least one config file")

    jobs_completed = 0
    jobs_failed = 0

    if accelerator.is_main_process:
        print_acc(f"Running {len(config_file_list)} job{'' if len(config_file_list) == 1 else 's'}")

    for config_file in config_file_list:
        try:
            job = get_job(config_file, args.name)
            job.run()
            job.cleanup()
            jobs_completed += 1
        except Exception as e:
            import traceback
            print_acc(f"Error running job: {e}")
            print_acc(f"Traceback: {traceback.format_exc()}")
            jobs_failed += 1
            try:
                job.process[0].on_error(e)
            except Exception as e2:
                print_acc(f"Error running on_error: {e2}")
            if not args.recover:
                print_end_message(jobs_completed, jobs_failed)
                raise e
        except KeyboardInterrupt as e:
            try:
                job.process[0].on_error(e)
            except Exception as e2:
                print_acc(f"Error running on_error: {e2}")
            if not args.recover:
                print_acc("")
                print_acc("========================================")
                print_acc("Job stopped")
                print_acc("========================================")
                sys.exit(0)


if __name__ == '__main__':
    main()