100 lines
4.3 KiB
Python
100 lines
4.3 KiB
Python
|
|
import yaml
|
|
import os
|
|
import argparse
|
|
|
|
|
|
def generate_yaml(
|
|
gpu_memory=40, # 40G or 80G?
|
|
gpu_num=1,
|
|
cpu_num_per_gpu=6,
|
|
memory_per_gpu=30,
|
|
replicas=1,
|
|
project_name='magic123',
|
|
project_support_alias='img2mesh',
|
|
pre_run_event='mkdir -p /fsx/code && ln -s /nfs/code/gqian /fsx/code/ && cd /fsx/code/gqian/img2mesh',
|
|
command="while :; do sleep 1000; done",
|
|
job_name='debug',
|
|
force_node=False,
|
|
**kwargs
|
|
):
|
|
data = {
|
|
'docker_image': '440036398022.dkr.ecr.us-west-2.amazonaws.com/facecraft-ml:efa',
|
|
'project_name': project_name,
|
|
'project_support_alias': project_support_alias,
|
|
'team': 'creative_vision',
|
|
#'fsx': 'fs-0b933bba2f17fe699', # 100T genai filesystem
|
|
'fsx': 'fs-056caaa56fa5cc5f3', # 2T personal filesystem of gqian
|
|
'gpu_type': 'nvidia-tesla-a100',
|
|
'gpu_num': gpu_num,
|
|
'cpu_num': str(int(cpu_num_per_gpu * gpu_num)),
|
|
'memory': str(int(memory_per_gpu * gpu_num)),
|
|
'gpu_memory': str(int(gpu_memory)),
|
|
'pytorchjob': {
|
|
'replicas': replicas
|
|
},
|
|
'efa': True,
|
|
'script': {
|
|
'pre_run_event': str(pre_run_event),
|
|
'command': str(command),
|
|
'jobs': [
|
|
{'name': job_name}
|
|
]
|
|
}
|
|
}
|
|
if gpu_num == 1 and not force_node:
|
|
gpu_yaml = {
|
|
'custom_node_labels': {
|
|
'use_case': 'p4d_debug',
|
|
},
|
|
}
|
|
else:
|
|
gpu_yaml = {
|
|
'custom_node_labels': {
|
|
'snap.com/spine': 'unknown',
|
|
'snap.com/region': 'us-west-2c',
|
|
},
|
|
}
|
|
data.update(gpu_yaml)
|
|
return data
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Script to generate yaml file for AWS Pytorch Job")
|
|
parser.add_argument('--yaml_folder', type=str, default='./', help='path to save the yaml folder')
|
|
parser.add_argument('--gpu_memory', type=int, default=40, help='GPU memory (in GB)')
|
|
parser.add_argument('--gpu_num', type=int, default=1, help='Number of GPUs')
|
|
parser.add_argument('--cpu_num_per_gpu', type=int, default=6, help='Number of CPUs per GPU')
|
|
parser.add_argument('--memory_per_gpu', type=float, default=30.0, help='Memory per GPU')
|
|
parser.add_argument('--replicas', type=int, default=1, help='Number of replicas')
|
|
parser.add_argument('--project_name', type=str, default='magic123', help='Project name')
|
|
parser.add_argument('--project_support_alias', type=str, default='img2mesh', help='Project support alias')
|
|
#parser.add_argument('--pre_run_event', type=str, default='export PATH=/nfs/code/gqian/miniconda3/bin:$PATH && conda init bash && source ~/.bashrc && cd /nfs/code/gqian/img2mesh ', help='Pre-run event command')
|
|
parser.add_argument('--pre_run_event', type=str, default='cd /nfs/code/gqian/img2mesh', help='Pre-run event command')
|
|
parser.add_argument('--command', type=str, default='while :; do sleep 1000; done', help='Command')
|
|
parser.add_argument('--job_name', type=str, default='debug', help='Job name')
|
|
parser.add_argument('--force_node', action='store_true',
|
|
help="use normal cluster not debug cluster")
|
|
args, unknown = parser.parse_known_args()
|
|
args.job_name = args.job_name.replace('_', '-').replace('.', '-')[:51] # do not support _ and . in job name, and max length is limited (around 70)
|
|
|
|
# "bash scripts/magic123/run_single_bothpriors.sh 0 r256 data/nerf4/drums rgba.png --h 300 --w 300"
|
|
data = generate_yaml(**vars(args))
|
|
yaml_str = yaml.safe_dump(data)
|
|
|
|
# Write the YAML content to a file
|
|
os.makedirs(args.yaml_folder, exist_ok=True)
|
|
yaml_path = os.path.join(args.yaml_folder, f'{args.job_name}.yaml')
|
|
with open(yaml_path, 'w') as file:
|
|
file.write(yaml_str)
|
|
print(f'YAML file saved to {yaml_path}')
|
|
|
|
# launch the job using snap_rutls
|
|
os.system(f'yes yes | snap_rutils cluster run {yaml_path} -s')
|
|
|
|
# show the job status
|
|
os.system(f'kubectl get pods | grep {args.job_name} ')
|
|
|
|
# show the job logs
|
|
os.system(f'kubectl logs {args.job_name}-worker-0')
|