import argparse import time import os import re import subprocess from datetime import datetime from pathlib import Path import boto3 import yaml from sagemaker import Session as sm_Session from sagemaker.pytorch import PyTorch try: from sagemaker.aws_batch.training_queue import TrainingQueue as Queue print("Loading SageMaker AWS Batch training queue.") is_sm_queue = True except Exception as e: print(f"Could not load SageMaker AWS Batch training queue: {e}.") is_sm_queue = False NAME = "anydata" INSTANCE_MAPPER = { "p4d": "ml.p4d.24xlarge", "p4de": "ml.p4de.24xlarge", "p5": "ml.p5.48xlarge", "p5en": "ml.p5en.48xlarge", "g6e": "ml.g6e.48xlarge", "g6e-small": "ml.g6e.12xlarge", "g5": "ml.g5.48xlarge", "g5-small": "ml.g5.24xlarge", "r7i": "ml.r7i.48xlarge", "r7i-small": "ml.r7i.12xlarge", } def run_command(command): print(f"=> {command}") subprocess.run(command, shell=True, check=True) def get_image(user, instance_type, version="280", build_type="full", profile="default", region="us-east-1"): os.environ["AWS_PROFILE"] = f"{profile}" account = subprocess.getoutput( f"aws --region {region} --profile {profile} sts get-caller-identity --query Account --output text" ) docker_dir = Path(__file__).parent if instance_type in INSTANCE_MAPPER.keys(): algorithm_name = f"{user}-{NAME}-{version}" dockerfile_base = docker_dir / f"Dockerfile_{version}_data" dockerfile_update = docker_dir / "Dockerfile_update" else: raise ValueError(f"Unknown instance_type: {instance_type}") fullname = f"{account}.dkr.ecr.{region}.amazonaws.com/{algorithm_name}:latest" if build_type in (None, "none", "skip", "reuse"): print("Reusing existing ECR image (skip build/push).") return fullname login_cmd = f"aws ecr get-login-password --region {region} --profile {profile} | docker login --username AWS --password-stdin" # NOTE (Dian): no "Dockerfile_update" build is needed since the code copy happens late in the SM dockerfile, # therefore docker will only update top layers in case of any code changes. the update method might # lead to max depth reached issues due to recursive basing. print("Building container") if build_type == "full": print("Building container") commands = [ # Log in to Sagemaker account to get image. f"{login_cmd} 763104351884.dkr.ecr.{region}.amazonaws.com", f"docker build -f {dockerfile_base} --build-arg AWS_REGION={region} -t {algorithm_name} .", f"docker tag {algorithm_name} {fullname}", f"{login_cmd} {fullname}", ( f"aws --region {region} --profile {profile} ecr describe-repositories --repository-names {algorithm_name} || " f"aws --region {region} --profile {profile} ecr create-repository --repository-name {algorithm_name}" ), ] elif build_type == "update": print("Updating container") commands = [ f"docker build -f {dockerfile_update} --build-arg BASE_DOCKER={algorithm_name} -t {algorithm_name} .", f"docker tag {algorithm_name} {fullname}", f"{login_cmd} {fullname}", ] else: raise ValueError(f"Unknown build_type: {build_type}") # Create command, making sure to exit if any part breaks. command = "\n".join([f"{x} || exit 1" for x in commands]) run_command(command) run_command(f"docker push {fullname}") print("Sleeping for 5 seconds to ensure push succeeded") time.sleep(5) return fullname def main(): # Use first line of file docstring as description if it exists. parser = argparse.ArgumentParser() parser.add_argument("--build-type", default="full") parser.add_argument("--local", action="store_true") parser.add_argument("--user", required=True, help="User name") parser.add_argument("--entry_point", type=str, default="scripts/train.py") parser.add_argument("--args_file", help="YAML file with hyperparameters", type=str, required=True) # AWS profile args parser.add_argument("--region", default="us-west-2", help="AWS region") parser.add_argument("--profile", default="default", help="AWS profile to use") parser.add_argument("--arn", default=None, help="If None, reads from SAGEMAKER_ARN env var") parser.add_argument( "--s3-remote-sync", default=None, help="S3 path to sync to. If none, reads from S3_REMOTE_SYNC env var" ) # Instance args parser.add_argument("--instance-count", default=1, type=int, help="Number of instances") parser.add_argument("--instance-type", default="p4de", choices=list(INSTANCE_MAPPER.keys())) parser.add_argument("--version", default="271", type=str, help="Choose from: (271, 271-2stage)") parser.add_argument("--spot-instance", action="store_true") parser.add_argument('--base-job-name', type=str) parser.add_argument("--split_json", type=str, help="Path to split JSON file", required=False) parser.add_argument('--input-source', choices=['s3', 'lustre', 'local'], default='s3') # Jobs Queue parser.add_argument("--fss-identifier", default="default", help="Share identifier for FSS queue") parser.add_argument("--priority", default=5, type=int, help="Priority of the job") parser.add_argument("--queue", type=str, default=None, help="Job queue") parser.add_argument("--name", type=str, default=None) args = parser.parse_args() if args.queue == 'mvdp': args.queue = "fss-mvdp-p4de-24xlarge-us-west-2" args.instance_type = 'p4de' elif args.queue == 'vla': args.queue = "fss-vla-p4de-24xlarge-us-west-2" args.instance_type = 'p4de' elif args.queue == 'dp': args.queue = "fss-dp-p4de-24xlarge-us-west-2" args.instance_type = 'p4de' elif args.queue == 'ml': args.queue = "fss-ml-p4de-24xlarge-us-west-2" args.instance_type = 'p4de' elif args.queue == 'mlp5': args.queue = f"fss-ml-p5-48xlarge-us-west-2" args.instance_type = 'p5' elif args.queue == 'cv-p5en': args.queue = f"fss-cv-ml-p5en-48xlarge-us-west-2" args.instance_type = 'p5en' else: # raise ValueError(f'Invalid queue name {args.queue}') print(f"not using queue") main_after_setup_move(args) def main_after_setup_move(args): if args.arn is None: assert "SAGEMAKER_ARN" in os.environ, "Please specify --arn or set the SAGEMAKER_ARN environment variable" args.arn = os.environ["SAGEMAKER_ARN"] if args.s3_remote_sync is None: assert ( "S3_REMOTE_SYNC" in os.environ ), "Please specify --s3-remote-sync or set the S3_REMOTE_SYNC environment variable" args.s3_remote_sync = os.environ["S3_REMOTE_SYNC"] normalized_build_type = ( None if str(args.build_type).lower() in {"none", "skip", "reuse"} else args.build_type ) image_uri = get_image( args.user, args.instance_type, args.version, region=args.region, build_type=normalized_build_type, profile=args.profile, ) # Use queueing only when an actual queue name is provided. queue_name = str(args.queue).strip().lower() if args.queue is not None else "" is_sm_queue = queue_name not in {"", "none", "null"} ########## # Create session and make sure of account and region ########## sagemaker_session = sm_Session( boto_session=boto3.session.Session( region_name=args.region, profile_name=args.profile ) ) if args.local: from sagemaker.local import LocalSession sagemaker_session = LocalSession() role = args.arn # provide a pre-existing role ARN as an alternative to creating a new role role_name = role.split(["/"][-1]) # client = boto3.client("sts", config=boto3_config) # account = client.get_caller_identity()["Account"] account = '124224456861' # client.get_caller_identity()["Account"] # account = subprocess.getoutput( # f"aws --region {args.region} --profile {args.profile} sts get-caller-identity --query Account --output text" # ) # session = boto3.session.Session() session = boto3.session.Session(region_name=args.region) region = session.region_name ########## # Configure the training ########## base_job_name = args.base_job_name # f"{args.user.replace('.', '-')}-{NAME}" def get_job_name(base): now = datetime.now() # Compact timestamp keeps SageMaker training job names below 63 chars. date_str = now.strftime("%y%m%d-%H%M%S") max_base_len = 63 - len(date_str) - 1 return "-".join([base[:max_base_len].rstrip("-"), date_str]) if args.name is None: job_name = get_job_name(base_job_name) else: job_suffix = args.name user_prefix = f"{args.user}-" if job_suffix.startswith(user_prefix): job_suffix = job_suffix[len(user_prefix):] named_base = f"{base_job_name}-{job_suffix}".replace('_', '-') job_name = get_job_name(named_base) output_root = f"{args.s3_remote_sync}/sagemaker/{args.user}/{NAME}/" output_s3 = os.path.join(output_root, job_name) tags = [ { "Key": "tri.project", "Value": "MM:PJ-0077", }, { "Key": "tri.owner.email", "Value": "fangzhou.cheng.ctr@tri.global", }, ] max_run = 14 * 24 * 60 * 60 max_wait = 5 * 24 * 60 * 60 if args.spot_instance else None keep_alive_period_in_seconds = 300 if not args.spot_instance else None entry_point = args.entry_point instance_type = "local_gpu" if args.local else INSTANCE_MAPPER[args.instance_type] instance_count = args.instance_count train_use_spot_instances = args.spot_instance checkpoint_s3_uri = os.path.join( f's3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/{args.user}/{NAME}', job_name) checkpoint_s3_uri = None if args.local else checkpoint_s3_uri checkpoint_local_path = "/opt/ml/checkpoints" checkpoint_local_path = None if args.local else checkpoint_local_path # hyperparameters = { # "config": args.config, # } # if args.experiment is not None: # hyperparameters["experiment"] = args.experiment # if args.hydra_cfg is not None: # with open(args.hydra_cfg, 'r') as fd: # hydra_config = yaml.safe_load(fd) # hyperparameters.update(hydra_config) if args.args_file is not None: with open(args.args_file, "r") as fd: hyperparameters = yaml.safe_load(fd) if args.split_json: # For split fanout jobs, use this split file as the download source. # Do not pass split_json through to web_sagemaker.py because that rewrites # the dataset recipe key (e.g., DTU/all -> DTU/sm_split_0). hyperparameters["s3_download_data"] = args.split_json match = re.search(r"_(\d+)\.json$", os.path.basename(args.split_json)) if match: # Keep fanout input/output indices aligned: # sm_split_0.json -> split_all_0.json. hyperparameters["split_output_index"] = int(match.group(1)) distribution={ "torch_distributed": { "enabled": True, } } environment = { "WANDB_API_KEY": os.environ.get('WANDB_API_KEY', ''), "WANDB_ENTITY": os.environ.get('WANDB_ENTITY', ''), "WANDB__SERVICE_WAIT": "300", "HF_TOKEN": os.environ.get("HF_TOKEN", ""), "HF_HOME": "/tmp", "INSTANCE_COUNT": str(args.instance_count), "SM_USE_RESERVED_CAPACITY": "1", "FI_EFA_FORK_SAFE": "1", "NVTE_FUSED_ATTN": "0", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "AWS_S3_USE_CRT": "1", "TOKENIZERS_PARALLELISM": "false", "SAGEMAKER_PROGRAM": entry_point, "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:False", "CUDA_LAUNCH_BLOCKING": "1", "TORCH_USE_CUDA_DSA": "1", "SM_JOB_NAME": job_name, "SAGEMAKER": "enabled", "QUEUE": str(args.queue), "INSTANCE_TYPE": str(args.instance_type), "INSTANCE_COUNT": str(args.instance_count), "LOCAL_PATH": "/opt/ml/input/data", # where SM mounts input "S3_PATH": "s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_datasets/processed", "NCCL_DEBUG": "INFO", "NCCL_DEBUG_SUBSYS": "INIT", "NCCL_DEBUG_FILE": "/tmp/nccl_%h_%p.log", "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", # enable async error reporting "NCCL_TIMEOUT": "300", # 180 seconds timeout for stuck collectives } if args.name is not None: environment['JOB_NAME'] = args.name security_group_ids = { 'us-east-1': [ 'sg-0afb9fb0e79a54061' ], 'us-west-2': [ 'sg-029d1d476bc087e31', ], } subnets = { 'us-east-1': [ 'subnet-07bf42d7c9cb929e4', 'subnet-0e260ba29726b9fbb', ], 'us-west-2': [ 'subnet-0610f766a4cd5cdae', 'subnet-029adfb9e225d68f8', 'subnet-01cc1bfeaf20155b5', ] } print() print() print('#############################################################') print(f'SageMaker Execution Role: {role}') print(f'The name of the Execution role: {role_name[-1]}') print(f'SM Queue: {is_sm_queue}-{args.priority}-{args.fss_identifier}') print(f'AWS region: {region}') print(f'AWS profile: {args.profile}') print(f'AWS account: {account}') print(f'Entry point: {entry_point}') print(f'Image uri: {image_uri}') print(f'Job name: {job_name}') print(f'Configuration file: {hyperparameters}') print(f'Instance count: {instance_count}') print(f'Instance type: {instance_type}') volume_size = 7000 if args.instance_type == "p5en" else 3000 print(f'Queue: {args.queue}') print(f'Volume size (GB): {volume_size}') print('#############################################################') print() print() print(f"launch_sm hyperparameters: {hyperparameters}") estimator = PyTorch( entry_point=entry_point, sagemaker_session=sagemaker_session, base_job_name=base_job_name, hyperparameters=hyperparameters, role=role, image_uri=image_uri, instance_count=instance_count, instance_type=instance_type, train_use_spot_instances=train_use_spot_instances, output_path=output_s3, job_name=job_name, checkpoint_s3_uri=checkpoint_s3_uri, checkpoint_local_path=checkpoint_local_path, code_location=output_s3, distribution=distribution, max_run=max_run, max_wait=max_wait, # debugger_hook_config=True, environment=environment, keep_alive_period_in_seconds=0, tags=tags, # subnets=subnets[region], # security_group_ids=security_group_ids[region], volume_size=volume_size, # profiler_config=profiler_config, # rules=profiler_rules, enable_sagemaker_metrics=True, enable_remote_debug=True, logs=True, ) if is_sm_queue: # queue_name = "fss-ml-p5-48xlarge-us-west-2" queue = Queue(args.queue) print(f"Starting training job on queue: {queue.queue_name}") queued_jobs = queue.map( estimator, inputs=[None], job_names=[job_name], priority=args.priority, share_identifier=args.fss_identifier, timeout={"attemptDurationSeconds": max_run}, ) print(f"Queued jobs: {queued_jobs}") else: print(f"Submitting job asynchronously (wait=False): {job_name}") estimator.fit(wait=False, job_name=job_name) if __name__ == "__main__": main()