start_wandb_sweep.py

"""
Start either a new wandb sweep or start agents for an existing sweep within a tmux session.

Command line arguments:
    1. sweep_agent: file_name.yaml or entity/project/sweep_id 
    2. gpus_to_use: all, 0:4, 0,1,2,3 (default: all)
    3. agents_per_gpu: number of agents to start per gpu (default: 2)

"""

import jax
import yaml
import wandb
from os import system
import os
import sys

os.environ['WANDB_DISABLE_SERVICE']= "True"
wandb.login()

try:
    sweep_agent = sys.argv[1]
    if sweep_agent.endswith('.yaml'):
        # create a new sweep
        config_file = sweep_agent
        # check if config file exists
        assert os.path.exists(config_file), "Config file does not exist"
        print(f"Creating sweep with config from {config_file}")
        sweep_config = yaml.load(open(config_file, "r"), Loader=yaml.FullLoader)
        project = sweep_config["project"]
        entity = os.environ.get("WANDB_ENTITY")
        sweep_id = wandb.sweep(sweep_config, project=project)
        print(f"Created sweep with id: {sweep_id}")
    elif len(sweep_agent.split('/')) == 3:
        entity, project, sweep_id = sweep_agent.split('/')
    else: 
        raise ValueError("Invalid sweep agent")
except: 
    raise ValueError("No sweep agent")
    

try:
    gpus_to_use = sys.argv[2]
    if gpus_to_use == "all":
        gpus_to_use = list(range(len(jax.devices())))
    elif ":" in gpus_to_use:
        gpus_to_use = list(range(int(gpus_to_use.split(":")[0]), int(gpus_to_use.split(":")[1])))
    elif "," in gpus_to_use:
        gpus_to_use = list(map(int, gpus_to_use.split(",")))
    
except:
    # assumes all gpus
    gpus_to_use = list(range(len(jax.devices())))

assert len(gpus_to_use) <= len(jax.devices()), "More GPUs requested than available"

try:
    agents_per_gpu = int(sys.argv[3])
except:
    agents_per_gpu = 2

print('-- starting wandb sweep agents --')
print(f'entity: {entity}, project: {project}, sweep_id: {sweep_id}')
print(f'sweep command: wandb agent {entity}/{project}/{sweep_id}')
print(f'gpus_to_use: {gpus_to_use}, agents_per_gpu: {agents_per_gpu}, total_agents: {len(gpus_to_use) * agents_per_gpu}')

def tmux(command):
    system('tmux %s' % command)

def tmux_shell(command):
    tmux('send-keys "%s" "C-m"' % command)


# start sesstion
tmux('new-session -d -s sweep')

for gpu in gpus_to_use:
    for agent in range(agents_per_gpu):

        pane_name = f"{gpu}-{agent}"
        tmux(f'new-window -t sweep -n {pane_name}')
        command = f"CUDA_VISIBLE_DEVICES={gpu} XLA_PYTHON_CLIENT_PREALLOCATE=false wandb agent {entity}/{project}/{sweep_id}"

        tmux_shell(command)

print('view runs with: $ tmux a -t sweep')