272 lines
9.6 KiB
Python
272 lines
9.6 KiB
Python
|
"""
|
||
|
|
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()
|