#!/usr/bin/env python3
"""Simple launcher that imports our model class before running VLMEvalKit.

GPU isolation strategy:
  - We use torchrun --nproc-per-node=N to launch N processes.
  - Each process pins itself to one GPU via CUDA_VISIBLE_DEVICES.
  - We override LOCAL_WORLD_SIZE=1 so VLMEvalKit's run.py does NOT
    re-slice CUDA_VISIBLE_DEVICES (it only does that when LOCAL_WORLD_SIZE>1).
  - We KEEP RANK and WORLD_SIZE intact so that VLMEvalKit's data sharding
    (in vlmeval/inference.py: range(rank, len(dataset), world_size)) works
    correctly — each process only infers its own shard of the dataset.
"""
import sys
import os

# IMPORTANT: Set up CUDA_VISIBLE_DEVICES BEFORE importing torch
# This ensures proper GPU isolation in multi-GPU settings

# Get GPU allocation info from torchrun
RANK = int(os.environ.get('RANK', 0))
WORLD_SIZE = int(os.environ.get('WORLD_SIZE', 1))
LOCAL_WORLD_SIZE = int(os.environ.get("LOCAL_WORLD_SIZE", 1))
LOCAL_RANK = int(os.environ.get("LOCAL_RANK", 0))

# Set CUDA_VISIBLE_DEVICES based on LOCAL_RANK
if LOCAL_WORLD_SIZE > 1:
    # Each process gets one GPU
    os.environ['CUDA_VISIBLE_DEVICES'] = str(LOCAL_RANK)
    # Override LOCAL_WORLD_SIZE so VLMEvalKit's run.py skips its own
    # GPU re-slicing logic (it checks LOCAL_WORLD_SIZE > 1).
    os.environ['LOCAL_WORLD_SIZE'] = '1'
    # KEEP RANK and WORLD_SIZE so data sharding works:
    #   get_rank_and_world_size() reads these env vars and
    #   infer_data() uses them for: range(rank, len(dataset), world_size)
    print(f'[run_vlmeval] RANK={RANK}/{WORLD_SIZE}, LOCAL_RANK={LOCAL_RANK}, '
          f'CUDA_VISIBLE_DEVICES={os.environ["CUDA_VISIBLE_DEVICES"]}')

# Add script directory and models directory to path
script_dir = os.path.dirname(os.path.abspath(__file__))
dvlm_dir = os.path.dirname(os.path.dirname(script_dir))  # dvlm root
models_dir = os.path.join(dvlm_dir, "models")

for path in [script_dir, models_dir]:
    if path not in sys.path:
        sys.path.insert(0, path)

# NOW import torch and model classes (after CUDA_VISIBLE_DEVICES is set)
import vlmeval_qwen_vl_mdm
import vlmeval_nemotron_ministral_mdm

# Import and run VLMEvalKit
import run

print(f"[run_vlmeval] sys.argv: {sys.argv}")
run.main()
