Setting the file. One moment. 93
Creating Amazon Aurora Db Cluster With Instances
Shieldadvanced
Resilience Program Design
Connecting Lambda To DynamoDB
references/training-jobs/code_templates/recipe_training.py
references/training-jobs/code_templates/recipe_training.py
Python·163 lines·6 KB
WaiterError
13from sagemaker.core import Attribution, set_attribution
14from sagemaker.core.helper.session_helper import Session, get_execution_role
15from sagemaker.core.training.configs import (
16 CheckpointConfig,
17 Compute,
18 InputData,
19 Networking,
20 OutputDataConfig,
21 StoppingCondition,
22)
23from sagemaker.train import ModelTrainer
24
25set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)
26
27REGION = ""
28ROLE_ARN = "[ROLE_ARN]"
29TRAINING_RECIPE = "[RECIPE_NAME_OR_YAML]"
30TRAINING_IMAGE_URI = ""
31TRAINING_S3_URI = "[TRAINING_S3_URI]"
32OUTPUT_S3_URI = "[OUTPUT_S3_URI]"
33OUTPUT_KMS_KEY_ID = ""
34CHECKPOINT_S3_URI = ""
35INSTANCE_TYPE = "[INSTANCE_TYPE]"
36INSTANCE_COUNT = None
37VOLUME_SIZE_GB = None
38TRAINING_PLAN_ARN = ""
39SUBNETS = []
40SECURITY_GROUP_IDS = []
41ENABLE_NETWORK_ISOLATION = False
42ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = False
43MAX_RUNTIME_SECONDS = 86400
44BASE_JOB_NAME = "[BASE_JOB_NAME]"
45RECIPE_OVERRIDES = {}
46WAIT_FOR_COMPLETION = False
47WAIT_POLL_SECONDS = 60
48WAIT_STARTUP_SLACK_SECONDS = 600
49
50
51def unique_job_prefix(base_job_name):
52 """Return a collision-resistant prefix with room for the SDK timestamp suffix."""
53 normalized = re.sub(r"[^A-Za-z0-9-]", "-", base_job_name).strip("-") or "training"
54 return f"{normalized[:20].rstrip('-')}-{uuid.uuid4().hex[:10]}"
55
56
57def resolve_submitted_job(sm_client, job_prefix, timeout_seconds=60):
58 """Resolve this submission through public SageMaker APIs only."""
59 deadline = time.monotonic() + timeout_seconds
60 while time.monotonic() < deadline:
61 summaries = sm_client.list_training_jobs(
62 NameContains=job_prefix,
63 SortBy="CreationTime",
64 SortOrder="Descending",
65 MaxResults=10,
66 )["TrainingJobSummaries"]
67 matches = [
68 summary for summary in summaries if summary["TrainingJobName"].startswith(job_prefix)
69 ]
70 if len(matches) == 1:
71 return sm_client.describe_training_job(
72 TrainingJobName=matches[0]["TrainingJobName"]
73 )
74 if len(matches) > 1:
75 raise RuntimeError(f"Multiple training jobs matched unique prefix {job_prefix!r}.")
76 time.sleep(2)
77 raise TimeoutError(f"Training job with prefix {job_prefix!r} was not visible within 60s.")
78
79
80boto_session = boto3.Session(region_name=REGION or None)
81sm_client = boto_session.client("sagemaker")
82sm_session = Session(boto_session=boto_session, sagemaker_client=sm_client)
83role = ROLE_ARN or get_execution_role()
84job_prefix = unique_job_prefix(BASE_JOB_NAME)
85
86# Cell 3: Build the recipe trainer
87if bool(SUBNETS) != bool(SECURITY_GROUP_IDS):
88 raise ValueError("Provide SUBNETS and SECURITY_GROUP_IDS together, or leave both empty.")
89
90compute_kwargs = {"instance_type": INSTANCE_TYPE}
91if INSTANCE_COUNT is not None:
92 compute_kwargs["instance_count"] = INSTANCE_COUNT
93if VOLUME_SIZE_GB is not None:
94 compute_kwargs["volume_size_in_gb"] = VOLUME_SIZE_GB
95if TRAINING_PLAN_ARN:
96 compute_kwargs["training_plan_arn"] = TRAINING_PLAN_ARN
97
98networking = None
99if SUBNETS or ENABLE_NETWORK_ISOLATION or ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION:
100 networking = Networking(
101 subnets=SUBNETS or None,
102 security_group_ids=SECURITY_GROUP_IDS or None,
103 enable_network_isolation=ENABLE_NETWORK_ISOLATION,
104 enable_inter_container_traffic_encryption=ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
105 )
106
107trainer_kwargs = {
108 "training_recipe": TRAINING_RECIPE,
109 "recipe_overrides": RECIPE_OVERRIDES,
110 "compute": Compute(**compute_kwargs),
111 "networking": networking,
112 "output_data_config": OutputDataConfig(
113 s3_output_path=OUTPUT_S3_URI,
114 kms_key_id=OUTPUT_KMS_KEY_ID or None,
115 ),
116 "checkpoint_config": CheckpointConfig(s3_uri=CHECKPOINT_S3_URI)
117 if CHECKPOINT_S3_URI
118 else None,
119 "stopping_condition": StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_SECONDS),
120 "sagemaker_session": sm_session,
121 "role": role,
122 "base_job_name": job_prefix,
123}
124if TRAINING_IMAGE_URI:
125 trainer_kwargs["training_image"] = TRAINING_IMAGE_URI
126
127trainer = ModelTrainer.from_recipe(**trainer_kwargs)
128
129# Review the fully merged recipe before the approval/launch step.
130print(f"Resolved recipe: {trainer.get_resolved_recipe()}")
131
132# Cell 4: Launch and optionally wait
133input_data = None
134if TRAINING_S3_URI:
135 input_data = [InputData(channel_name="training", data_source=TRAINING_S3_URI)]
136
137trainer.train(input_data_config=input_data, wait=False, logs=False)
138details = resolve_submitted_job(sm_client, job_prefix)
139job_name = details["TrainingJobName"]
140job_arn = details["TrainingJobArn"]
141print(f"Training Job Name: {job_name}")
142print(f"Training Job ARN: {job_arn}")
143
144if WAIT_FOR_COMPLETION:
145 wait_budget_seconds = MAX_RUNTIME_SECONDS + WAIT_STARTUP_SLACK_SECONDS
146 max_attempts = max(
147 1,
148 (wait_budget_seconds + WAIT_POLL_SECONDS - 1) // WAIT_POLL_SECONDS + 1,
149 )
150 try:
151 sm_client.get_waiter("training_job_completed_or_stopped").wait(
152 TrainingJobName=job_name,
153 WaiterConfig={"Delay": WAIT_POLL_SECONDS, "MaxAttempts": max_attempts},
154 )
155 except WaiterError:
156 pass # Resolve the real status and FailureReason below.
157 details = sm_client.describe_training_job(TrainingJobName=job_name)
158 if details["TrainingJobStatus"] != "Completed":
159 raise RuntimeError(
160 f"Training job ended in {details['TrainingJobStatus']}: "
161 f"{details.get('FailureReason', 'no failure reason returned')}"
162 )
163 print(f"Model artifacts: {details.get('ModelArtifacts', {}).get('S3ModelArtifacts', '')}")