def download_directory_cached()

in lm_human_preferences/utils/gcs.py [0:0]

• 17 lines of code
• 5 McCabe index (conditional complexity)

def download_directory_cached(url, comm=None):
    """ Given a GCS path url, caches the contents locally.
    WARNING: only use this function if contents under the path won't change!
    """
    cache_dir = '/tmp/gcs-cache'
    bucket_name, path = parse_url(url)
    is_master = not comm or comm.Get_rank() == 0
    local_path = os.path.join(cache_dir, bucket_name, path)

    sentinel = os.path.join(local_path, 'SYNCED')
    if is_master:
        if not os.path.exists(local_path):
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            cmd = 'gsutil', '-m', 'cp', '-r', url, os.path.dirname(local_path) + '/'
            print(' '.join(cmd))
            subprocess.check_call(cmd)
            open(sentinel, 'a').close()
    else:
        while not os.path.exists(sentinel):
            time.sleep(1)
    return local_path