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/jumpstart_training.py
references/training-jobs/code_templates/jumpstart_training.py
Python·158 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.jumpstart.configs import JumpStartConfig
16from sagemaker.core.training.configs import (
17 CheckpointConfig,
18 Compute,
19 InputData,
20 Networking,
21 OutputDataConfig,
22 StoppingCondition,
23)
24from sagemaker.train import ModelTrainer
25
26set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)
27
28REGION = ""
29ROLE_ARN = "[ROLE_ARN]"
30MODEL_ID = "[JUMPSTART_MODEL_ID]"
31MODEL_VERSION = ""
32TRAINING_S3_URI = "[TRAINING_S3_URI]"
33OUTPUT_S3_URI = "[OUTPUT_S3_URI]"
34OUTPUT_KMS_KEY_ID = ""
35CHECKPOINT_S3_URI = ""
36BASE_JOB_NAME = "[BASE_JOB_NAME]"
37ACCEPT_EULA = False
38HYPERPARAMETERS = {}
39OVERRIDE_INSTANCE_TYPE = ""
40OVERRIDE_INSTANCE_COUNT = None
41SUBNETS = []
42SECURITY_GROUP_IDS = []
43ENABLE_NETWORK_ISOLATION = False
44ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = False
45MAX_RUNTIME_SECONDS = 86400
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 JumpStart trainer
87if bool(SUBNETS) != bool(SECURITY_GROUP_IDS):
88 raise ValueError("Provide SUBNETS and SECURITY_GROUP_IDS together, or leave both empty.")
89
90jumpstart_config = JumpStartConfig(
91 model_id=MODEL_ID,
92 model_version=MODEL_VERSION or None,
93 accept_eula=ACCEPT_EULA,
94)
95trainer_kwargs = {
96 "jumpstart_config": jumpstart_config,
97 "output_data_config": OutputDataConfig(
98 s3_output_path=OUTPUT_S3_URI,
99 kms_key_id=OUTPUT_KMS_KEY_ID or None,
100 ),
101 "checkpoint_config": CheckpointConfig(s3_uri=CHECKPOINT_S3_URI)
102 if CHECKPOINT_S3_URI
103 else None,
104 "stopping_condition": StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_SECONDS),
105 "hyperparameters": HYPERPARAMETERS,
106 "sagemaker_session": sm_session,
107 "role": role,
108 "base_job_name": job_prefix,
109}
110if OVERRIDE_INSTANCE_TYPE or OVERRIDE_INSTANCE_COUNT is not None:
111 compute_kwargs = {}
112 if OVERRIDE_INSTANCE_TYPE:
113 compute_kwargs["instance_type"] = OVERRIDE_INSTANCE_TYPE
114 if OVERRIDE_INSTANCE_COUNT is not None:
115 compute_kwargs["instance_count"] = OVERRIDE_INSTANCE_COUNT
116 trainer_kwargs["compute"] = Compute(**compute_kwargs)
117if SUBNETS or ENABLE_NETWORK_ISOLATION or ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION:
118 trainer_kwargs["networking"] = Networking(
119 subnets=SUBNETS or None,
120 security_group_ids=SECURITY_GROUP_IDS or None,
121 enable_network_isolation=ENABLE_NETWORK_ISOLATION,
122 enable_inter_container_traffic_encryption=ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
123 )
124
125trainer = ModelTrainer.from_jumpstart_config(**trainer_kwargs)
126
127# Cell 4: Launch and optionally wait
128trainer.train(
129 input_data_config=[InputData(channel_name="training", data_source=TRAINING_S3_URI)],
130 wait=False,
131 logs=False,
132)
133details = resolve_submitted_job(sm_client, job_prefix)
134job_name = details["TrainingJobName"]
135job_arn = details["TrainingJobArn"]
136print(f"Training Job Name: {job_name}")
137print(f"Training Job ARN: {job_arn}")
138
139if WAIT_FOR_COMPLETION:
140 wait_budget_seconds = MAX_RUNTIME_SECONDS + WAIT_STARTUP_SLACK_SECONDS
141 max_attempts = max(
142 1,
143 (wait_budget_seconds + WAIT_POLL_SECONDS - 1) // WAIT_POLL_SECONDS + 1,
144 )
145 try:
146 sm_client.get_waiter("training_job_completed_or_stopped").wait(
147 TrainingJobName=job_name,
148 WaiterConfig={"Delay": WAIT_POLL_SECONDS, "MaxAttempts": max_attempts},
149 )
150 except WaiterError:
151 pass # Resolve the real status and FailureReason below.
152 details = sm_client.describe_training_job(TrainingJobName=job_name)
153 if details["TrainingJobStatus"] != "Completed":
154 raise RuntimeError(
155 f"Training job ended in {details['TrainingJobStatus']}: "
156 f"{details.get('FailureReason', 'no failure reason returned')}"
157 )
158 print(f"Model artifacts: {details.get('ModelArtifacts', {}).get('S3ModelArtifacts', '')}")