''' Unified SageMaker launcher for the Robotics (124...) and AD2 WFM (385...) accounts. All accounts now submit via AWS Batch service jobs onto a SAGEMAKER_TRAINING-backed job queue (fair-share scheduling). Select an account variant with --account: - 'robotics-old' (124224456861, legacy, S3 on tri-ml-sandbox) - 'robotics-new' (385697366450, S3 on tri-ml-sandbox cross-account) - 'ad2-gaia-wfm' (385697366450, S3 on gaia-e2e-wfm-datasets) ''' import argparse import ast from datetime import datetime import os from pathlib import Path import subprocess import time import boto3 from rich import print from rich.table import Table from sagemaker import Session as sm_Session from sagemaker.inputs import TrainingInput from sagemaker.pytorch import PyTorch import yaml # Queue-capable sagemaker. Public GA (sagemaker>=2.249.0) renamed batch_queueing -> aws_batch; the # vendored 2.240.1.dev0 (custom/sagemaker/wheels/) still uses the old name. Support BOTH; the first # arg differs (training_job= vs estimator=), so the estimator is passed POSITIONALLY at submit time. _QueueCls = None try: from sagemaker.aws_batch.training_queue import TrainingQueue as _QueueCls # public >=2.249.0 except ImportError: try: from sagemaker.batch_queueing.queue import Queue as _QueueCls # vendored .dev0 except ImportError: _QueueCls = None _HAS_QUEUE = _QueueCls is not None # NAME = "cosmos-predict2" # legacy name (old ECR repos + S3 paths still use this) NAME = "any4d-cp2" # Some Dockerfiles (e.g. Dockerfile_271-2stage) hardcode their base image in the FROM # line: `FROM 124224456861.dkr.ecr.us-west-2.amazonaws.com/cosmos-predict2-base:latest`. # That ECR lives in the legacy robotics account (124224456861); when launching from any # other account we need an extra docker-login with a profile that has read access to it. # NOTE: this only controls which registry we docker-login to for pull auth -- it does NOT # redirect the hardcoded FROM line. If the base image is replicated elsewhere (e.g. to # 385697366450), update the Dockerfile FROM too and override this env var to match. BASE_IMAGE_REGISTRY = os.environ.get( 'SM_BASE_IMAGE_REGISTRY', '124224456861.dkr.ecr.us-west-2.amazonaws.com', ) BASE_IMAGE_PROFILE = os.environ.get('SM_BASE_IMAGE_PROFILE', 'rob-s3') INSTANCE_MAPPER = { "p4d": "ml.p4d.24xlarge", "p4de": "ml.p4de.24xlarge", "p5": "ml.p5.48xlarge", # 8x H100 80GB "p5en": "ml.p5en.48xlarge", # 8x H200 144GB "p6": "ml.p6-b200.48xlarge", # 8x B200 192GB "g6e": "ml.g6e.48xlarge", "g6e-sm": "ml.g6e.12xlarge", "g5": "ml.g5.48xlarge", "g5-sm": "ml.g5.24xlarge", } # Fair-share Batch queue map. The cv-* aliases (cv / cv-wfm / cv-p5en / cv-wfm-p5en) all resolve to the # cv-wfm-p5en pool (renamed from cv-ml-p5en 2026-05-26; AD2 + CV work share it). QUEUE_MAP = { ('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-wfm', 'cv-p5en', 'cv-wfm-p5en'): ('fss-cv-wfm-p5en-48xlarge-us-west-2', 'p5en'), # ^ Renamed 2026-05-26: was fss-cv-ml-p5en (now DISABLED) -> fss-cv-wfm-p5en (ENABLED). ('cv-spot', 'cv-wfm-spot', 'cv-spot-p5en', 'cv-wfm-spot-p5en'): ('fss-cv-wfm-spot-p5en-48xlarge-us-west-2', 'p5en'), ('cv-spot-p5', 'cv-wfm-spot-p5'): ('fss-cv-wfm-spot-p5-48xlarge-us-west-2', 'p5'), ('cv-spot-p6', 'cv-wfm-spot-p6', 'cv-spot-p6-b200', 'cv-wfm-spot-p6-b200'): ('fss-cv-wfm-spot-p6-b200-48xlarge-us-west-2', 'p6'), # ^ Spot/interruptible cv-wfm queues in robotics-old (124224456861): p5en / p5 / p6-b200. # --spot-instance is auto-enabled for these below (resolved name contains "spot"). ('cam', 'cam-p5'): ('fss-tri-cam-humanoid-p5-48xlarge-us-west-2', 'p5'), } # Per-account defaults ACCOUNT_CONFIGS = { # Legacy robotics path: caller = robs3/rob-s3 (124224456861), FSS queues + SM exec role # in the same 124224456861 account. Image pushed to 124224456861 ECR. Being phased out. 'robotics-old': dict( account_id='124224456861', profile='rob-s3', arn='arn:aws:iam::124224456861:role/SageMaker-SageMakerAllAccess-us-west-2', s3_remote_sync='s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d', s3_checkpoint_base='s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d/sagemaker', s3_bucket='tri-ml-sandbox-16011-us-west-2-datasets', training_plan=None, use_queue=True, max_run=25 * 24 * 60 * 60 - 360, # 25 days (legacy account quota) volume_size=3000, ), # New robotics path (IT-reconfigured 2026-05): caller = rob-sm (RoboticsWFM-BatchOperator # role in 385697366450, same account as ad2-sm). SM exec role + ECR + batch queues all # live in 385697366450. Source / checkpoints still go to the tri-ml-sandbox bucket # (owned by 124224456861) via a cross-account bucket policy. 'robotics-new': dict( account_id='385697366450', profile='rob-sm', arn='arn:aws:iam::385697366450:role/Robotics-WFM-Sagemaker-role-us-west-2', s3_remote_sync='s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d', s3_checkpoint_base='s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d/sagemaker', s3_bucket='tri-ml-sandbox-16011-us-west-2-datasets', training_plan=None, use_queue=True, max_run=25 * 24 * 60 * 60 - 360, # 25 days (new AWS quota L-33A961FD = 28d on 385 acct) volume_size=1000, # 1024 GB account cap ), 'ad2-gaia-wfm': dict( account_id='385697366450', profile='ad2-sm', arn='arn:aws:iam::385697366450:role/AD2-SagemakerAccess-us-west-2', # NOTE(2026-06-24): gaia bucket policy REGRESSED -- the EXEC role AD2-SagemakerAccess lost # ListBucket+PutObject on gaia `any4d/sagemaker/*` (the SM-PLATFORM path). It killed BOTH gaia runs # the same day: s49 (16n) + s51 (2n) failed on profiler-output PutObject (mid-run), and a fresh # s51 resume then failed on checkpoint_s3_uri ListBucket (at startup) -- both AccessDenied. FIX per # report-34 recipe: route SM-platform artifacts (output_s3/profiler/source + checkpoint_s3_uri sync) # to the 385-owned bucket below (PowerUser+exec writable, same-account key). Training DATA + our # checkpoints stay on gaia via the `job.s3_root` hyperparameter (exec role can write `any4d/a4d2/*`). # REVERT both lines to 's3://gaia-e2e-wfm-datasets/any4d[/sagemaker]' once IT re-grants the exec role. s3_remote_sync='s3://sagemaker-us-west-2-385697366450/any4d-sm', s3_checkpoint_base='s3://sagemaker-us-west-2-385697366450/any4d-sm/sagemaker', s3_bucket='gaia-e2e-wfm-datasets', training_plan=None, use_queue=True, # route via Batch onto cv-wfm-p5en (direct CreateTrainingJob is SCP-blocked) max_run=25 * 24 * 60 * 60 - 360, # 25 days (new AWS quota L-33A961FD = 28d on 385 acct) volume_size=1000, # 1024 GB account cap ), } TAGS = [ {"Key": "tri.project", "Value": os.environ.get("TRI_PROJECT", "MM:PJ-0077")}, {"Key": "tri.owner.email", "Value": os.environ.get("TRI_OWNER_EMAIL", "CHANGE.ME@tri.global")}, ] def run_command(command): print(f"[dim]=> {command}[/dim]") subprocess.run(command, shell=True, check=True) def build_anydata_training_inputs(experiment, input_source): """Build SageMaker input channels for AnyData S3 manifest configs.""" 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 resolve_queue(queue_alias): ''' Resolve a queue alias to (fss_queue_name, instance_type_key). ''' for aliases, (queue_name, instance_type) in QUEUE_MAP.items(): if queue_alias in aliases: return queue_name, instance_type all_aliases = [a for aliases in QUEUE_MAP for a in aliases] raise ValueError(f"Invalid queue name: {queue_alias!r}. Valid: {all_aliases}") def build_and_push_image(user, instance_type, version="271", build_type="full", profile="default", region="us-west-2", image_tag="latest"): ''' Build Docker image, push to ECR, return the full image URI. Perms: None: sts:GetCallerIdentity update: + ecr:GetAuthorizationToken, BatchCheckLayerAvailability, InitiateLayerUpload, UploadLayerPart, CompleteLayerUpload, PutImage full: + ecr:DescribeRepositories, CreateRepository ''' print(f"[bold]Building image:[/bold] user={user}, instance_type={instance_type}, " f"version={version}, build_type={build_type}") os.environ["AWS_PROFILE"] = profile account = subprocess.getoutput( f"aws --region {region} --profile {profile} sts get-caller-identity " f"--query Account --output text" ) docker_dir = Path(__file__).parent algorithm_name = f"{user}-{NAME}-{version}" dockerfile = docker_dir / f"Dockerfile_{version}" # Per-agent image TAG so concurrent projects can run in parallel / do not compete. fullname = f"{account}.dkr.ecr.{region}.amazonaws.com/{algorithm_name}:{image_tag}" if build_type is None or build_type == 'None': return fullname login_cmd = (f"aws ecr get-login-password --region {region} --profile {profile} " f"| docker login --username AWS --password-stdin") if build_type == "full": # Extra docker-login for the legacy ECR that hosts the base image (cosmos-predict2-base). # Uses BASE_IMAGE_PROFILE (overridable via SM_BASE_IMAGE_PROFILE env var) which must # have ECR read perms on BASE_IMAGE_REGISTRY. Skipped when the launch profile is # already on that account (otherwise the existing login_cmd at line below covers it). base_login_cmd = ( f"aws ecr get-login-password --region {region} --profile {BASE_IMAGE_PROFILE} " f"| docker login --username AWS --password-stdin {BASE_IMAGE_REGISTRY}" ) if profile != BASE_IMAGE_PROFILE else None commands = [ f"{login_cmd} 763104351884.dkr.ecr.{region}.amazonaws.com", *([base_login_cmd] if base_login_cmd else []), f"docker build -f {dockerfile} --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 " f"--repository-names {algorithm_name} || " f"aws --region {region} --profile {profile} ecr create-repository " f"--repository-name {algorithm_name}"), ] elif build_type == "update": dockerfile_update = docker_dir / "Dockerfile_update" commands = [ f"docker build -f {dockerfile_update} --build-arg BASE_DOCKER={algorithm_name} " f"-t {algorithm_name} .", f"docker tag {algorithm_name} {fullname}", f"{login_cmd} {fullname}", ] else: raise ValueError(f"Unknown build_type: {build_type}") run_command("\n".join(f"{cmd} || exit 1" for cmd in commands)) run_command(f"docker push {fullname}") time.sleep(5) return fullname def parse_args(): parser = argparse.ArgumentParser() # Account selection parser.add_argument("--account", default="robotics-new", choices=list(ACCOUNT_CONFIGS.keys()), help="Target account variant: 'robotics-new' (385697366450, default), " "'robotics-old' (124224456861, legacy), or 'ad2-gaia-wfm' (385697366450)") # Build parser.add_argument("--build-type", default="full") parser.add_argument("--local", action="store_true") parser.add_argument("--user", default=None, help="User name (or set SM_USER env var)") parser.add_argument("--entry_point", default="scripts/train.py") # Experiment parser.add_argument("--config", required=True, help="Config base path") parser.add_argument("--hydra_cfg", default=None, help="Hydra YAML overrides file") parser.add_argument("--experiment", default=None) # AWS (overrides account defaults if provided) parser.add_argument("--region", default="us-west-2") parser.add_argument("--profile", default=None, help="Override account default profile") parser.add_argument("--arn", default=None, help="Override execution role ARN") parser.add_argument("--s3-remote-sync", default=None, help="Override S3 sync root") # Instance parser.add_argument("--instance-count", default=1, type=int) parser.add_argument("--instance-type", default=None, help="Instance type key (e.g. p5en). For robotics, set by --queue.") parser.add_argument("--version", default="271", help="Dockerfile version (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") # Queue (robotics only) parser.add_argument("--fss-identifier", default="default", help="FSS share identifier") parser.add_argument("--priority", default=5, type=int, help="Job priority in queue") parser.add_argument("--queue", default="ml", help="Queue alias (robotics only, see QUEUE_MAP)") parser.add_argument("--name", default=None, help="Job name suffix") # Training plan (ad2 only) parser.add_argument("--training-plan", default=None, help="Training plan ARN (defaults to account config)") parser.add_argument("--no-plan", action="store_true", help="Skip training plan (on-demand)") # Hydra overrides parser.add_argument("overrides", nargs="*", default=[], help="Extra key=value overrides forwarded as hyperparameters") return parser.parse_args() def main(): args = parse_args() # Resolve user from env if not provided if args.user is None: args.user = os.environ.get("SM_USER") assert args.user, "Specify --user or set SM_USER env var" # Load account defaults, allow CLI overrides acfg = ACCOUNT_CONFIGS[args.account] if args.profile is None: args.profile = acfg['profile'] if args.arn is None: args.arn = acfg['arn'] if args.s3_remote_sync is None: args.s3_remote_sync = acfg['s3_remote_sync'] use_queue = acfg['use_queue'] # Resolve instance type if use_queue: # Robotics: queue determines instance type args.queue, args.instance_type = resolve_queue(args.queue) # Spot queues need managed-spot training; auto-enable on the resolved FSS queue name # so any launch path gets it without passing --spot-instance. if "spot" in args.queue and not args.spot_instance: print(f"[yellow]Auto-enabling spot instances for spot queue {args.queue}[/yellow]") args.spot_instance = True if "testing-p6" in args.queue and not args.spot_instance: print("[yellow]Auto-enabling spot instances for testing-p6 queue[/yellow]") args.spot_instance = True else: # AD2: default to p5en, no queue if args.instance_type is None: args.instance_type = 'p5en' args.queue = 'none' # Resolve ARN from env as fallback if args.arn is None: args.arn = os.environ.get("SAGEMAKER_ARN") assert args.arn, "Specify --arn or set SAGEMAKER_ARN" if args.s3_remote_sync is None: args.s3_remote_sync = os.environ.get("S3_REMOTE_SYNC") assert args.s3_remote_sync, "Specify --s3-remote-sync or set S3_REMOTE_SYNC" image_tag = 'gaia' if args.account == 'ad2-gaia-wfm' else 'rob' # Build and push Docker image image_uri = build_and_push_image( args.user, args.instance_type, args.version, build_type=args.build_type, profile=args.profile, region=args.region, image_tag=image_tag, ) # SageMaker session session_kwargs = dict( boto_session=boto3.session.Session( region_name=args.region, profile_name=args.profile, ), ) if acfg['s3_bucket']: session_kwargs['default_bucket'] = acfg['s3_bucket'] sagemaker_session = sm_Session(**session_kwargs) if args.local: from sagemaker.local import LocalSession sagemaker_session = LocalSession() # Job naming base_job_name = args.base_job_name or f"{args.user}-any4d" if args.name is None: now = datetime.now() date_str = f"{now.strftime('%Y-%m-%d-%H-%M-%S')}-{now.microsecond // 1000:03d}" job_name = f"{base_job_name}-{date_str}" else: job_name = f"{base_job_name}--{args.name}".replace('_', '-') # S3 paths output_s3 = f"{args.s3_remote_sync}/sagemaker/{args.user}/{NAME}/{job_name}" checkpoint_s3_uri = ( f"{acfg['s3_checkpoint_base']}/{args.user}/{NAME}/{job_name}" if not args.local else None ) # Build hyperparameters 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: hyperparameters.update(yaml.safe_load(fd)) fit_inputs, input_channel_specs = build_anydata_training_inputs(args.experiment, args.input_source) # Environment variables instance_type = "local_gpu" if args.local else INSTANCE_MAPPER[args.instance_type] 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), "INSTANCE_TYPE": str(args.instance_type), "SM_USE_RESERVED_CAPACITY": "0" if args.spot_instance else "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": args.entry_point, "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:False", # "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", # NOTE(bvh): new (May 2026), saves peak VRAM on DGX, but slower "CUDA_LAUNCH_BLOCKING": "1", "TORCH_USE_CUDA_DSA": "1", "SM_JOB_NAME": job_name, # Fixed at submit so spot restarts reuse the same run dir (see trainer.py prepend). # Honor a pinned RUN_DATETIME from the shell to resume a prior run dir. "RUN_DATETIME": os.environ.get("RUN_DATETIME") or datetime.now().strftime('%m-%d-%H-%M'), "SAGEMAKER": "enabled", "QUEUE": str(args.queue), "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 # Training plan (ad2 only) training_plan = None if not args.no_plan: training_plan = args.training_plan or acfg.get('training_plan') # Print summary max_run = acfg['max_run'] role_name = args.arn.split("/")[-1] table = Table(title=f"SageMaker Job [{args.account}]", show_header=False, title_style="bold cyan", border_style="dim") table.add_column("Key", style="bold") table.add_column("Value") table.add_row("Account", f"{acfg['account_id']} ({args.account})") table.add_row("Profile", args.profile) table.add_row("Execution Role", role_name) table.add_row("Image URI", image_uri) table.add_row("Job name", f"[bold green]{job_name}[/bold green]") table.add_row("Instance", f"{args.instance_count}x {instance_type}") if use_queue: table.add_row("Queue", f"{args.queue} (priority={args.priority})") if training_plan: table.add_row("Training Plan", training_plan.split('/')[-1]) table.add_row("Input Source", args.input_source) table.add_row("FastFile", "enabled" if args.input_source == "fastfile" else "disabled") table.add_row("Input Channels", str(input_channel_specs)) if input_channel_specs: table.add_row("Input Mount Note", "First-time FastFile/File setup can be long; 50M files / 50T may take about 1 hour.") table.add_row("Output S3", output_s3) table.add_row("Checkpoint S3", str(checkpoint_s3_uri)) hp_lines = [f" {k}={v}" for k, v in hyperparameters.items()] table.add_row("Hyperparameters", "\n".join(hp_lines)) print(table) # Create PyTorch estimator estimator_kwargs = dict( entry_point=args.entry_point, sagemaker_session=sagemaker_session, base_job_name=base_job_name, hyperparameters=hyperparameters, role=args.arn, image_uri=image_uri, instance_count=args.instance_count, instance_type=instance_type, train_use_spot_instances=args.spot_instance, output_path=output_s3, job_name=job_name, checkpoint_s3_uri=checkpoint_s3_uri, checkpoint_local_path=None if args.local else "/opt/ml/checkpoints", code_location=output_s3, distribution={"torch_distributed": {"enabled": True}}, max_run=max_run, max_wait=max_run if args.spot_instance else None, environment=environment, keep_alive_period_in_seconds=None if args.spot_instance else 300, train_volume_size=acfg['volume_size'], volume_size=acfg['volume_size'], enable_sagemaker_metrics=True, enable_remote_debug=True, logs=True, ) if args.instance_type == 'p6': estimator_kwargs['disable_profiler'] = True if input_channel_specs: estimator_kwargs['input_mode'] = input_channel_specs[0]['input_mode'] if training_plan: estimator_kwargs['training_plan'] = training_plan if not use_queue: estimator_kwargs['tags'] = TAGS estimator = PyTorch(**estimator_kwargs) # Submit. SDK first stages source -> s3:PutObject (+ GetObject, ListBucket). # Then path depends on use_queue: # True: batch:SubmitServiceJob (+ DescribeJobs, ListJobs) # False: sagemaker:CreateTrainingJob (+ DescribeTrainingJob, StopTrainingJob, iam:PassRole) if use_queue: # This account REQUIRES a Batch queue; a direct submit is SCP-denied. NEVER fall through to # estimator.fit() silently (that was the trap: queuing no-ops -> direct job -> account rejects). if not _HAS_QUEUE: raise RuntimeError( "This account REQUIRES a Batch queue, but no queue-capable sagemaker is installed " "(neither sagemaker.aws_batch.training_queue.TrainingQueue nor " "sagemaker.batch_queueing.queue.Queue importable). A direct SageMaker submit would be " "SCP-denied, so refusing to fall back silently.\nFix (in the launcher env, NOT in docker):\n" " pip install custom/sagemaker/wheels/sagemaker-2.240.1.dev0-py3-none-any.whl\n" " # or public: pip install 'sagemaker>=2.249.0'" ) queue = _QueueCls(args.queue) print(f"[bold green]Starting training job on queue: {queue.queue_name}[/bold green]") tags_dict = {t["Key"]: t["Value"] for t in TAGS} queued_jobs = queue.map( estimator, [fit_inputs], job_names=[job_name], # estimator POSITIONAL (arg name differs across APIs) priority=args.priority, share_identifier=args.fss_identifier, timeout={"attemptDurationSeconds": max_run}, # Pass tags at the Batch level tags=tags_dict, ) print(f"Queued jobs: {queued_jobs}") else: print(f"[bold green]Submitting training job: {job_name}[/bold green]") estimator.fit(inputs=fit_inputs, wait=False, job_name=job_name) print(f"Monitor: aws sagemaker describe-training-job --training-job-name {job_name} " f"--profile {args.profile} --region {args.region}") if __name__ == "__main__": main()