utils/dist.py

import os
import torch as T
import re
from tqdm import tqdm
from datetime import timedelta

import requests
import hashlib

from io import BytesIO  

def rank0():
    rank = os.environ.get('RANK')
    if rank is None or rank == '0':
        return True
    else:
        return False
    
def local0():
    local_rank = os.environ.get('LOCAL_RANK')
    if local_rank is None or local_rank == '0':
        return True
    else:
        return False
class tqdm0(tqdm):
    def __init__(self, *args, **kwargs):
        total = kwargs.get('total', None)
        if total is None and len(args) > 0:
            try:
                total = len(args[0])
            except TypeError:
                pass
        if total is not None:
            kwargs['miniters'] = max(1, total // 20)
        super().__init__(*args, **kwargs, disable=not rank0(), bar_format='{bar}| {n_fmt}/{total_fmt} [{rate_fmt}{postfix}]')
    
def print0(*args, **kwargs):
    if rank0():
        print(*args, **kwargs)

_PRINTED_IDS = set()

def printonce(*args, id=None, **kwargs):
    if id is None:
        id = ' '.join(map(str, args))
    
    if id not in _PRINTED_IDS:
        print(*args, **kwargs)
        _PRINTED_IDS.add(id)

def print0once(*args, **kwargs):
    if rank0(): 
        printonce(*args, **kwargs)

def init_dist():
    if T.distributed.is_initialized():
        print0('Distributed already initialized')
        rank = T.distributed.get_rank()
        local_rank = int(os.environ.get('LOCAL_RANK', 0))
        world_size = T.distributed.get_world_size()
    else:
        try:
            rank = int(os.environ['RANK'])
            local_rank = int(os.environ['LOCAL_RANK'])
            world_size = int(os.environ['WORLD_SIZE'])
            device = f'cuda:{local_rank}'
            T.cuda.set_device(device)
            T.distributed.init_process_group(backend='nccl', timeout=timedelta(minutes=30), rank=rank, world_size=world_size, device_id=T.device(device))
            print(f'Rank {rank} of {world_size}.')
        except Exception as e:
            print0once(f'Not initializing distributed env: {e}')
            rank = 0
            local_rank = 0
            world_size = 1
    return rank, local_rank, world_size

def load_ckpt(load_from_location, expected_hash=None):
    if local0():
        os.makedirs('ckpt', exist_ok=True)
        url = f"https://ckpt.si.inc/hertz-dev/{load_from_location}.pt"
        save_path = f"ckpt/{load_from_location}.pt"
        if not os.path.exists(save_path):
            response = requests.get(url, stream=True)
            total_size = int(response.headers.get('content-length', 0))
            with open(save_path, 'wb') as f, tqdm(total=total_size, desc=f'Downloading {load_from_location}.pt', unit='GB', unit_scale=1/(1024*1024*1024)) as pbar:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
                    pbar.update(len(chunk))
        if expected_hash is not None:
            with open(save_path, 'rb') as f:
                file_hash = hashlib.md5(f.read()).hexdigest()
            if file_hash != expected_hash:
                print(f'Hash mismatch for {save_path}. Expected {expected_hash} but got {file_hash}. Deleting checkpoint and trying again.')
                os.remove(save_path)
                return load_ckpt(load_from_location, expected_hash)
    if T.distributed.is_initialized():
        T.distributed.barrier() # so that ranks don't try to load checkpoint before it's finished downloading
    loaded = T.load(f"ckpt/{load_from_location}.pt", weights_only=False, map_location='cpu')    
    return loaded