import argparse import ast import time import os import subprocess from datetime import datetime import os from pathlib import Path import subprocess import sys import time import yaml import boto3 import yaml from sagemaker import Session as sm_Session from sagemaker.pytorch import PyTorch from sagemaker.inputs import TrainingInput try: from sagemaker.batch_queueing.queue import Queue print(f"Loading SageMaker batch queueing.") is_sm_queue = True except Exception as e: print(f"Could not load SageMaker batch queueing: {e}.") # os.system("bash scripts/setup_sm_batch.sh") # os.system("pip install 'sagemaker>=3.0.0,<4.0.0'") print(f"Please run: pip install 'sagemaker>=3.0.0,<4.0.0'") sys.exit(1) from sagemaker.batch_queueing.queue import Queue print(f"Loading SageMaker batch queueing.") is_sm_queue = True # except: # is_sm_queue = False is_sm_queue = True NAME = "cosmos-predict2" INSTANCE_MAPPER = { "p4d": "ml.p4d.24xlarge", "p4de": "ml.p4de.24xlarge", "p5": "ml.p5.48xlarge", "p5en": "ml.p5en.48xlarge", # "p6": "ml.p6.48xlarge", "p6": "ml.p6-b200.48xlarge", "g6e": "ml.g6e.48xlarge", "g6e-small": "ml.g6e.12xlarge", "g5": "ml.g5.48xlarge", "g5-small": "ml.g5.24xlarge" } # from Sedrick: # INSTANCE_MAPPER = { # "p4de": "ml.p4de.24xlarge", # "p5": "ml.p5.48xlarge", # "p5en": "ml.p5en.48xlarge", # "p6": "ml.p6-b200.48xlarge", # } # Map simple queue aliases to the actual FSS queue names and set instance_type appropriately QUEUE_MAP = { # [aliases] : (fss-queue-name, instance_type) ('ml', 'ml-p5',): ('fss-ml-p5-48xlarge-us-west-2', 'p5'), ('testing', 'testing-p5',): ('fss-testing-p5-48xlarge-us-west-2', 'p5'), ('testing-p6',): ('fss-testing-p6-b200-48xlarge-us-west-2', 'p6'), ('vla', 'vla-p5',): ('fss-vla-p5-48xlarge-us-west-2', 'p5'), ('vla-p5en',): ('fss-vla-p5en-48xlarge-us-west-2', 'p5en'), ('cv', 'cv-ml', 'cv-p5en', 'cv-ml-p5en', ): ('fss-cv-ml-p5en-48xlarge-us-west-2', 'p5en'), ('cam', 'cam-p5',): ('fss-tri-cam-humanoid-p5-48xlarge-us-west-2', 'p5'), } def run_command(command): print(f"=> {command}") subprocess.run(command, shell=True, check=True) def _build_anydata_training_inputs(experiment, input_source): if input_source not in ("fastfile", "file"): return None, [] if not experiment: raise ValueError(f"input-source={input_source} requires an AnyData experiment") stem = experiment[len("any4d_"):] if experiment.startswith("any4d_") else experiment matches = sorted(Path("custom/experiment").rglob(f"{stem}.py")) if len(matches) != 1: raise ValueError( f"Expected one experiment file for {experiment!r}, found {len(matches)}: " f"{[str(m) for m in matches]}" ) tree = ast.parse(matches[0].read_text()) config_paths = sorted({ node.value for node in ast.walk(tree) if isinstance(node, ast.Constant) and isinstance(node.value, str) and node.value.startswith("custom/config/anydata/") and node.value.endswith((".yaml", ".yml")) }) if not config_paths: raise ValueError( f"input-source={input_source} requires anydata YAML paths in experiment={experiment!r}" ) channel_roots = set() for config_path in config_paths: with open(config_path, "r") as fp: config = yaml.safe_load(fp)["dataset"] path_prefix = str(config.get("path_prefix", "")).rstrip("/") for path_entry in config["path"]: uri = str(path_entry[0] if isinstance(path_entry, list) else path_entry) if not uri.startswith("s3://") and path_prefix.startswith("s3://"): uri = f"{path_prefix}/{uri.lstrip('/')}" if uri.startswith("s3://"): channel_roots.add(os.path.dirname(uri).rstrip("/")) input_mode = "FastFile" if input_source == "fastfile" else "File" distribution = "FullyReplicated" if input_source == "fastfile" else "ShardedByS3Key" channel_specs = [ { "channel": f"anydata{i:03d}", "s3_root": channel_root, "input_mode": input_mode, } for i, channel_root in enumerate(sorted(channel_roots)) ] inputs = { spec["channel"]: TrainingInput( s3_data=spec["s3_root"], input_mode=input_mode, s3_data_type="S3Prefix", distribution=distribution, ) for spec in channel_specs } return inputs, channel_specs def get_image(user, instance_type, version="271", build_type="full", profile="default", region="us-east-1"): print(f"Building image for user {user}, instance_type {instance_type}, version {version}, build_type {build_type}") 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}" 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 is None: 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("--config", help="config base", required=True) parser.add_argument("--hydra_cfg", help="hydra groups to override with", type=str, default=None) parser.add_argument("--experiment", help="which experiment to run", type=str, default=None) # 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('--input-source', choices=['s3', 'fastfile', 'file', '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='ml', help="Job queue") parser.add_argument("--name", type=str, default=None) parser.add_argument("overrides", nargs="*", default=[], help="Extra key=value overrides forwarded as hyperparameters (e.g. dataloader_train.batch_size=2 trainer.grad_accum_iter=1)") args = parser.parse_args() # Find the matching FSS queue found = False for aliases, (queue_name, instance_type) in QUEUE_MAP.items(): if args.queue in aliases: args.queue = queue_name args.instance_type = instance_type found = True break if not found: raise ValueError(f'Invalid queue name {args.queue}') # Auto-enable spot instances for testing-p6 queue (uses spot, not reserved capacity) if "testing-p6" in args.queue and not args.spot_instance: print(f"Auto-enabling spot instances for testing-p6 queue: {args.queue}") args.spot_instance = True 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"] image_uri = get_image( args.user, args.instance_type, args.version, region=args.region, build_type=args.build_type, profile=args.profile, ) # (g5, g6 series don't go into queue) if args.instance_type.startswith("g"): is_sm_queue = False else: is_sm_queue = True ########## # 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() # Format example: 2023-03-03-10-14-02-324 now_ms_str = f"{now.microsecond // 1000:03d}" date_str = f"{now.strftime('%Y-%m-%d-%H-%M-%S')}-{now_ms_str}" job_name = "-".join([base, date_str]) return job_name if args.name is None: job_name = get_job_name(base_job_name) else: job_name = f"{base_job_name}--{args.name}".replace('_', '-') 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 = 25 * 24 * 60 * 60 max_wait = 25 * 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 fit_inputs, input_channel_specs = _build_anydata_training_inputs(args.experiment, args.input_source) hyperparameters = { "config": args.config, } if args.experiment is not None: hyperparameters["experiment"] = args.experiment for ov in args.overrides: if "=" not in ov: raise ValueError(f"Override must be key=value, got: {ov}") key, value = ov.split("=", 1) hyperparameters[key] = value 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) 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": "0" if "testing-p6" in args.queue else "1", # NOTE(bvh): ^ testing-p6 queue uses spot (non-reserved) "FI_EFA_FORK_SAFE": "1", "NVTE_FUSED_ATTN": "0", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "AWS_S3_USE_CRT": "1", # "NCCL_DEBUG": "INFO", "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), "NCCL_DEBUG": "TRACE", "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", "NCCL_DEBUG_SUBSYS": "ALL", "NCCL_DEBUG_FILE": "/opt/ml/output/nccl_%h_%p.log", "ANY4D_SM_INPUT_SOURCE": args.input_source, } if input_channel_specs: environment["SM_FFM_WITH_MOUNTPOINT_ENABLED"] = "true" environment["ANY4D_SM_INPUT_CHANNEL_COUNT"] = str(len(input_channel_specs)) for i, spec in enumerate(input_channel_specs): environment[f"ANY4D_SM_INPUT_CHANNEL_{i:03d}"] = spec["channel"] environment[f"ANY4D_SM_INPUT_S3_ROOT_{i:03d}"] = spec["s3_root"] 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}') print(f'Queue: {args.queue}') print(f'Input source: {args.input_source}') print(f'FastFile mode: {"enabled" if args.input_source == "fastfile" else "disabled"}') print(f'Input channels: {input_channel_specs}') if input_channel_specs: print('Input mount note: First-time FastFile/File setup can be long; 50M files / 50T may take about 1 hour.') print('#############################################################') print() print() estimator_kwargs = dict( 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=keep_alive_period_in_seconds, tags=tags, # subnets=subnets[region], # security_group_ids=security_group_ids[region], train_volume_size=3000, volume_size=3000, # profiler_config=profiler_config, # rules=profiler_rules, enable_sagemaker_metrics=True, enable_remote_debug=True, logs=True, ) # NOTE(bvh): this is needed temporarily (I think) # https://tri-internal.slack.com/archives/C06UKGN6HC6/p1772143811777889 if args.instance_type == 'p6': estimator_kwargs['disable_profiler'] = True if input_channel_specs: estimator_kwargs['input_mode'] = input_channel_specs[0]['input_mode'] estimator = PyTorch(**estimator_kwargs) 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=[fit_inputs], job_names=[job_name], priority=args.priority, share_identifier=args.fss_identifier, timeout={"attemptDurationSeconds": max_run}, ) print(f"Queued jobs: {queued_jobs}") else: estimator.fit(inputs=fit_inputs) if __name__ == "__main__": main()