third_party/diffusion_policy/ray_train_multirun.py

"""
Start local ray cluster
(robodiff)$ export CUDA_VISIBLE_DEVICES=0,1,2 # select GPUs to be managed by the ray cluster
(robodiff)$ ray start --head --num-gpus=3

Training:
python ray_train_multirun.py --config-name=train_diffusion_unet_lowdim_workspace --seeds=42,43,44 --monitor_key=test/mean_score -- logger.mode=online training.eval_first=True
"""
import os
import ray
import click
import hydra
import yaml
import wandb
import pathlib
import collections
from pprint import pprint
from omegaconf import OmegaConf
from ray_exec import worker_fn
from ray.util.placement_group import (
    placement_group,
)
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

OmegaConf.register_new_resolver("eval", eval, replace=True)

@click.command()
@click.option('--config-name', '-cn', required=True, type=str)
@click.option('--config-dir', '-cd', default=None, type=str)
@click.option('--seeds', '-s', default='42,43,44', type=str)
@click.option('--monitor_key', '-k', multiple=True, default=['test/mean_score'])
@click.option('--ray_address', '-ra', default='auto')
@click.option('--num_cpus', '-nc', default=7, type=float)
@click.option('--num_gpus', '-ng', default=1, type=float)
@click.option('--max_retries', '-mr', default=0, type=int)
@click.option('--monitor_max_retires', default=3, type=int)
@click.option('--data_src', '-d', default='./data', type=str)
@click.option('--unbuffer_python', '-u', is_flag=True, default=False)
@click.option('--single_node', '-sn', is_flag=True, default=False, help='run all experiments on a single machine')
@click.argument('command_args', nargs=-1, type=str)
def main(config_name, config_dir, seeds, monitor_key, ray_address, 
    num_cpus, num_gpus, max_retries, monitor_max_retires,
    data_src, unbuffer_python, 
    single_node, command_args):
    # parse args
    seeds = [int(x) for x in seeds.split(',')]
    # expand path
    if data_src is not None:
        data_src = os.path.abspath(os.path.expanduser(data_src))

    # initialize hydra
    if config_dir is None:
        config_path_abs = pathlib.Path(__file__).parent.joinpath(
            'diffusion_policy','config')
        config_path_rel = str(config_path_abs.relative_to(pathlib.Path.cwd()))
    else:
        config_path_rel = config_dir

    run_command_args = list()
    monitor_command_args = list()
    with hydra.initialize(
        version_base=None, 
        config_path=config_path_rel):

        # generate raw config
        cfg = hydra.compose(
            config_name=config_name, 
            overrides=command_args)
        OmegaConf.resolve(cfg)
    
        # manually create output dir
        output_dir = pathlib.Path(cfg.multi_run.run_dir)
        output_dir.mkdir(parents=True, exist_ok=False)
        config_path = output_dir.joinpath('config.yaml')
        print(output_dir)

        # save current config
        yaml.dump(OmegaConf.to_container(cfg, resolve=True), 
            config_path.open('w'), default_flow_style=False)

        # wandb
        wandb_group_id = wandb.util.generate_id()
        name_base = cfg.multi_run.wandb_name_base

        # create monitor command args
        monitor_command_args = [
            'python',
            'multirun_metrics.py',
            '--input', str(output_dir),
            '--use_wandb',
            '--project', 'diffusion_policy_metrics',
            '--group', wandb_group_id
        ]
        for k in monitor_key:
            monitor_command_args.extend([
                '--key', k
            ])

        # generate command args
        run_command_args = list()
        for i, seed in enumerate(seeds):
            test_start_seed = (seed + 1) * 100000
            this_output_dir = output_dir.joinpath(f'train_{i}')
            this_output_dir.mkdir()
            wandb_name = name_base + f'_train_{i}'
            wandb_run_id = wandb_group_id + f'_train_{i}'

            this_command_args = [
                'python',
                'train.py',
                '--config-name='+config_name,
                '--config-dir='+config_path_rel
            ]

            this_command_args.extend(command_args)
            this_command_args.extend([
                f'training.seed={seed}',
                f'task.env_runner.test_start_seed={test_start_seed}',
                f'logging.name={wandb_name}',
                f'logging.id={wandb_run_id}',
                f'logging.group={wandb_group_id}',
                f'hydra.run.dir={this_output_dir}'
            ])
            run_command_args.append(this_command_args)

    # init ray
    root_dir = os.path.dirname(__file__)
    runtime_env = {
        'working_dir': root_dir,
        'excludes': ['.git'],
        'pip': ['dm-control==1.0.9']
    }
    ray.init(
        address=ray_address, 
        runtime_env=runtime_env
    )
    
    # create resources for train
    train_resources = dict()

    train_bundle = dict(train_resources)
    train_bundle['CPU'] = num_cpus
    train_bundle['GPU'] = num_gpus

    # create resources for monitor
    monitor_resources = dict()
    monitor_resources['CPU'] = 1
    
    monitor_bundle = dict(monitor_resources)

    # aggregate bundle
    bundle = collections.defaultdict(lambda:0)
    n_train_bundles = 1
    if single_node:
        n_train_bundles = len(seeds)
    for _ in range(n_train_bundles):
        for k, v in train_bundle.items():
            bundle[k] += v
    for k, v in monitor_bundle.items():
        bundle[k] += v
    bundle = dict(bundle)

    # create placement group
    print("Creating placement group with resources:")
    pprint(bundle)
    pg = placement_group([bundle])

    # run
    task_name_map = dict()
    task_refs = list()
    for i, this_command_args in enumerate(run_command_args):
        if single_node or i == (len(run_command_args) - 1):
            print(f'Training worker {i} with placement group.')
            ray.get(pg.ready())
            print("Placement Group created!")
            worker_ray = ray.remote(worker_fn).options(
                num_cpus=num_cpus,
                num_gpus=num_gpus,
                max_retries=max_retries,
                resources=train_resources,
                retry_exceptions=True,
                scheduling_strategy=PlacementGroupSchedulingStrategy(
                    placement_group=pg)
            )
        else:
            print(f'Training worker {i} without placement group.')
            worker_ray = ray.remote(worker_fn).options(
                num_cpus=num_cpus,
                num_gpus=num_gpus,
                max_retries=max_retries,
                resources=train_resources,
                retry_exceptions=True,
            )
        task_ref = worker_ray.remote(
            this_command_args, data_src, unbuffer_python)
        task_refs.append(task_ref)
        task_name_map[task_ref] = f'train_{i}'

    # monitor worker is always packed on the same node
    # as training worker 0
    ray.get(pg.ready())
    monitor_worker_ray = ray.remote(worker_fn).options(
        num_cpus=1, 
        num_gpus=0,
        max_retries=monitor_max_retires,
        # resources=monitor_resources,
        retry_exceptions=True,
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=pg)
        )
    monitor_ref = monitor_worker_ray.remote(
            monitor_command_args, data_src, unbuffer_python)
    task_name_map[monitor_ref] = 'metrics'

    try:
        # normal case
        ready_refs = list()
        rest_refs = task_refs
        while len(ready_refs) < len(task_refs):
            this_ready_refs, rest_refs = ray.wait(rest_refs, 
                num_returns=1, timeout=None, fetch_local=True)
            cancel_other_tasks = False
            for ref in this_ready_refs:
                task_name = task_name_map[ref]
                try:
                    result = ray.get(ref)
                    print(f"Task {task_name} finished with result: {result}")
                except KeyboardInterrupt as e:
                    # skip to outer try catch
                    raise KeyboardInterrupt
                except Exception as e:
                    print(f"Task {task_name} raised exception: {e}")
                    this_cancel_other_tasks = True
                    if isinstance(e, ray.exceptions.RayTaskError):
                        if isinstance(e.cause, ray.exceptions.TaskCancelledError):
                            this_cancel_other_tasks = False
                    cancel_other_tasks = cancel_other_tasks or this_cancel_other_tasks
                ready_refs.append(ref)
            if cancel_other_tasks:
                print('Exception! Cancelling all other tasks.')
                # cancel all other refs
                for _ref in rest_refs:
                    ray.cancel(_ref, force=False)
        print("Training tasks done.")
        ray.cancel(monitor_ref, force=False)
    except KeyboardInterrupt:
        print('KeyboardInterrupt received in the driver.')
        # a KeyboardInterrupt will be raised in worker
        _ = [ray.cancel(x, force=False) for x in task_refs + [monitor_ref]]
        print('KeyboardInterrupt sent to workers.')
    except Exception as e:
        # worker will be terminated
        _ = [ray.cancel(x, force=True) for x in task_refs + [monitor_ref]]
        raise e

    for ref in task_refs + [monitor_ref]:
        task_name = task_name_map[ref]
        try:
            result = ray.get(ref)
            print(f"Task {task_name} finished with result: {result}")
        except KeyboardInterrupt as e:
            # force kill everything.
            print("Force killing all workers")
            _ = [ray.cancel(x, force=True) for x in task_refs]
            ray.cancel(monitor_ref, force=True)
        except Exception as e:
            print(f"Task {task_name} raised exception: {e}")


if __name__ == "__main__":
    main()