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/custom_training.py
references/training-jobs/code_templates/custom_training.py
Python·175 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 SourceCode,
22 StoppingCondition,
23)
24from sagemaker.train import ModelTrainer
25
26set_attribution(Attribution.SAGEMAKER_AGENT_PLUGIN)
27
28REGION = ""
29ROLE_ARN = "[ROLE_ARN]"
30TRAINING_IMAGE = "[TRAINING_IMAGE_URI]"
31ALGORITHM_NAME = ""
32SOURCE_DIR = "[SOURCE_DIR]"
33ENTRY_SCRIPT = "[ENTRY_SCRIPT]"
34REQUIREMENTS = "[REQUIREMENTS]"
35TRAINING_S3_URI = "[TRAINING_S3_URI]"
36OUTPUT_S3_URI = "[OUTPUT_S3_URI]"
37OUTPUT_KMS_KEY_ID = ""
38INSTANCE_TYPE = "[INSTANCE_TYPE]"
39INSTANCE_COUNT = 1
40VOLUME_SIZE_GB = 30
41MAX_RUNTIME_SECONDS = 86400
42USE_SPOT = False
43MAX_WAIT_SECONDS = 172800
44CHECKPOINT_S3_URI = ""
45SUBNETS = []
46SECURITY_GROUP_IDS = []
47ENABLE_NETWORK_ISOLATION = False
48ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = False
49HYPERPARAMETERS = {}
50BASE_JOB_NAME = "[BASE_JOB_NAME]"
51WAIT_FOR_COMPLETION = False
52WAIT_POLL_SECONDS = 60
53WAIT_STARTUP_SLACK_SECONDS = 600
54
55
56def unique_job_prefix(base_job_name):
57 """Return a collision-resistant prefix with room for the SDK timestamp suffix."""
58 normalized = re.sub(r"[^A-Za-z0-9-]", "-", base_job_name).strip("-") or "training"
59 return f"{normalized[:20].rstrip('-')}-{uuid.uuid4().hex[:10]}"
60
61
62def resolve_submitted_job(sm_client, job_prefix, timeout_seconds=60):
63 """Resolve this submission through public SageMaker APIs only."""
64 deadline = time.monotonic() + timeout_seconds
65 while time.monotonic() < deadline:
66 summaries = sm_client.list_training_jobs(
67 NameContains=job_prefix,
68 SortBy="CreationTime",
69 SortOrder="Descending",
70 MaxResults=10,
71 )["TrainingJobSummaries"]
72 matches = [
73 summary for summary in summaries if summary["TrainingJobName"].startswith(job_prefix)
74 ]
75 if len(matches) == 1:
76 return sm_client.describe_training_job(
77 TrainingJobName=matches[0]["TrainingJobName"]
78 )
79 if len(matches) > 1:
80 raise RuntimeError(f"Multiple training jobs matched unique prefix {job_prefix!r}.")
81 time.sleep(2)
82 raise TimeoutError(f"Training job with prefix {job_prefix!r} was not visible within 60s.")
83
84
85boto_session = boto3.Session(region_name=REGION or None)
86sm_client = boto_session.client("sagemaker")
87sm_session = Session(boto_session=boto_session, sagemaker_client=sm_client)
88role = ROLE_ARN or get_execution_role()
89job_prefix = unique_job_prefix(BASE_JOB_NAME)
90
91# Cell 3: Build the custom trainer
92if bool(TRAINING_IMAGE) == bool(ALGORITHM_NAME):
93 raise ValueError("Provide exactly one of TRAINING_IMAGE or ALGORITHM_NAME.")
94if bool(SOURCE_DIR) != bool(ENTRY_SCRIPT):
95 raise ValueError("Provide SOURCE_DIR and ENTRY_SCRIPT together, or leave both empty.")
96if USE_SPOT and not CHECKPOINT_S3_URI:
97 raise ValueError("CHECKPOINT_S3_URI is required when USE_SPOT=True.")
98if bool(SUBNETS) != bool(SECURITY_GROUP_IDS):
99 raise ValueError("Provide SUBNETS and SECURITY_GROUP_IDS together, or leave both empty.")
100
101trainer_kwargs = {
102 "compute": Compute(
103 instance_type=INSTANCE_TYPE,
104 instance_count=INSTANCE_COUNT,
105 volume_size_in_gb=VOLUME_SIZE_GB,
106 enable_managed_spot_training=USE_SPOT,
107 ),
108 "output_data_config": OutputDataConfig(
109 s3_output_path=OUTPUT_S3_URI,
110 kms_key_id=OUTPUT_KMS_KEY_ID or None,
111 ),
112 "stopping_condition": StoppingCondition(
113 max_runtime_in_seconds=MAX_RUNTIME_SECONDS,
114 max_wait_time_in_seconds=MAX_WAIT_SECONDS if USE_SPOT else None,
115 ),
116 "hyperparameters": HYPERPARAMETERS,
117 "sagemaker_session": sm_session,
118 "role": role,
119 "base_job_name": job_prefix,
120}
121if TRAINING_IMAGE:
122 trainer_kwargs["training_image"] = TRAINING_IMAGE
123else:
124 trainer_kwargs["algorithm_name"] = ALGORITHM_NAME
125if SOURCE_DIR:
126 trainer_kwargs["source_code"] = SourceCode(
127 source_dir=SOURCE_DIR,
128 entry_script=ENTRY_SCRIPT,
129 requirements=REQUIREMENTS or None,
130 )
131if CHECKPOINT_S3_URI:
132 trainer_kwargs["checkpoint_config"] = CheckpointConfig(s3_uri=CHECKPOINT_S3_URI)
133if SUBNETS or ENABLE_NETWORK_ISOLATION or ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION:
134 trainer_kwargs["networking"] = Networking(
135 subnets=SUBNETS or None,
136 security_group_ids=SECURITY_GROUP_IDS or None,
137 enable_network_isolation=ENABLE_NETWORK_ISOLATION,
138 enable_inter_container_traffic_encryption=ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION,
139 )
140
141trainer = ModelTrainer(**trainer_kwargs)
142
143# Cell 4: Launch and optionally wait
144trainer.train(
145 input_data_config=[InputData(channel_name="training", data_source=TRAINING_S3_URI)],
146 wait=False,
147 logs=False,
148)
149details = resolve_submitted_job(sm_client, job_prefix)
150job_name = details["TrainingJobName"]
151job_arn = details["TrainingJobArn"]
152print(f"Training Job Name: {job_name}")
153print(f"Training Job ARN: {job_arn}")
154
155if WAIT_FOR_COMPLETION:
156 configured_wait_seconds = MAX_WAIT_SECONDS if USE_SPOT else MAX_RUNTIME_SECONDS
157 wait_budget_seconds = configured_wait_seconds + WAIT_STARTUP_SLACK_SECONDS
158 max_attempts = max(
159 1,
160 (wait_budget_seconds + WAIT_POLL_SECONDS - 1) // WAIT_POLL_SECONDS + 1,
161 )
162 try:
163 sm_client.get_waiter("training_job_completed_or_stopped").wait(
164 TrainingJobName=job_name,
165 WaiterConfig={"Delay": WAIT_POLL_SECONDS, "MaxAttempts": max_attempts},
166 )
167 except WaiterError:
168 pass # Resolve the real status and FailureReason below.
169 details = sm_client.describe_training_job(TrainingJobName=job_name)
170 if details["TrainingJobStatus"] != "Completed":
171 raise RuntimeError(
172 f"Training job ended in {details['TrainingJobStatus']}: "
173 f"{details.get('FailureReason', 'no failure reason returned')}"
174 )
175 print(f"Model artifacts: {details.get('ModelArtifacts', {}).get('S3ModelArtifacts', '')}")