Setting the file. One moment. - Number
- 7.7
- Position
- 7 of 7
- Type
- Python
- Size
- 9 KB
- Lines
- 246
scripts/teardown.py
Python·246 lines·9 KB
16
Calls the `aws` CLI from the current shell, so it inherits the same AWS context
17(profile, region, SSO session) — no Bash/WSL context-sharing problem on Windows.
18
19Usage:
20 python teardown.py <endpoint-name> [<region>]
21"""
22
23from __future__ import annotations
24
25import json
26import os
27import shutil
28import subprocess
29import sys
30import time
31
32VARIANT_DIM = "sagemaker:variant:DesiredInstanceCount"
33IC_DIM = "sagemaker:inference-component:DesiredCopyCount"
34
35
36def log(msg: str) -> None:
37 print(f"[teardown] {msg}", file=sys.stderr, flush=True)
38
39
40def aws_bin() -> str:
41 exe = shutil.which("aws")
42 if not exe:
43 log("ERROR: the 'aws' CLI was not found on PATH. Install AWS CLI v2.")
44 sys.exit(2)
45 return exe
46
47
48def run_aws(args: list[str]) -> subprocess.CompletedProcess:
49 return subprocess.run([aws_bin(), *args], capture_output=True, text=True)
50
51
52def json_out(proc: subprocess.CompletedProcess, default):
53 if proc.returncode != 0:
54 return default
55 try:
56 return json.loads(proc.stdout) or default
57 except json.JSONDecodeError:
58 return default
59
60
61def resolve_region(arg_region: str | None) -> str:
62 """Region from arg, then env, then the active profile's config."""
63 if arg_region:
64 return arg_region
65 for var in ("AWS_REGION", "AWS_DEFAULT_REGION"):
66 if os.environ.get(var):
67 return os.environ[var]
68 proc = run_aws(["configure", "get", "region"])
69 return proc.stdout.strip() if proc.returncode == 0 else ""
70
71
72def delete_alarms(prefix: str, reg: list[str]) -> None:
73 alarms = json_out(
74 run_aws(["cloudwatch", "describe-alarms", "--alarm-name-prefix", prefix,
75 "--query", "MetricAlarms[*].AlarmName", "--output", "json", *reg]),
76 [],
77 )
78 if alarms:
79 run_aws(["cloudwatch", "delete-alarms", "--alarm-names", *alarms, *reg])
80 log(f"Deleted alarms: {' '.join(alarms)}")
81
82
83def delete_scaling_stack(resource_id: str, dimension: str, reg: list[str]) -> None:
84 """Delete every scaling policy on a target, then deregister the target."""
85 policies = json_out(
86 run_aws(["application-autoscaling", "describe-scaling-policies",
87 "--service-namespace", "sagemaker", "--resource-id", resource_id,
88 "--query", "ScalingPolicies[*].PolicyName", "--output", "json", *reg]),
89 [],
90 )
91 for policy in policies:
92 run_aws(["application-autoscaling", "delete-scaling-policy",
93 "--service-namespace", "sagemaker", "--resource-id", resource_id,
94 "--scalable-dimension", dimension, "--policy-name", policy, *reg])
95 log(f"Deleted autoscaling policy: {policy}")
96
97 targets = json_out(
98 run_aws(["application-autoscaling", "describe-scalable-targets",
99 "--service-namespace", "sagemaker", "--resource-ids", resource_id,
100 "--query", "ScalableTargets[*].ResourceId", "--output", "json", *reg]),
101 [],
102 )
103 if resource_id in targets:
104 run_aws(["application-autoscaling", "deregister-scalable-target",
105 "--service-namespace", "sagemaker", "--resource-id", resource_id,
106 "--scalable-dimension", dimension, *reg])
107 log(f"Deregistered scalable target: {resource_id}")
108
109
110def list_inference_components(endpoint_name: str, reg: list[str]) -> list[str]:
111 return json_out(
112 run_aws(["sagemaker", "list-inference-components",
113 "--endpoint-name-equals", endpoint_name,
114 "--query", "InferenceComponents[*].InferenceComponentName",
115 "--output", "json", *reg]),
116 [],
117 )
118
119
120def component_model_name(ic_name: str, reg: list[str]) -> str:
121 desc = json_out(
122 run_aws(["sagemaker", "describe-inference-component",
123 "--inference-component-name", ic_name, *reg]),
124 {},
125 )
126 return desc.get("Specification", {}).get("ModelName", "")
127
128
129def delete_inference_components(ic_names: list[str], reg: list[str],
130 timeout_seconds: int = 900) -> None:
131 """Delete components and wait until they are gone.
132
133 Retries the delete on every poll, because SageMaker refuses it in two
134 transient states seen in practice:
135 - CREATE_IN_PROGRESS (container still starting)
136 - UPDATE_RC_IN_PROGRESS (a scaling action is changing the copy count)
137 An adapter component must also go before its base component, and the
138 ordering falls out of the same retry loop.
139 """
140 pending = list(ic_names)
141 deadline = time.time() + timeout_seconds
142 announced: set[str] = set()
143 while pending and time.time() < deadline:
144 for ic_name in pending:
145 if ic_name not in announced:
146 log(f"Deleting inference component: {ic_name}")
147 announced.add(ic_name)
148 run_aws(["sagemaker", "delete-inference-component",
149 "--inference-component-name", ic_name, *reg])
150 still_there = [
151 ic_name for ic_name in pending
152 if run_aws(["sagemaker", "describe-inference-component",
153 "--inference-component-name", ic_name, *reg]).returncode == 0
154 ]
155 pending = still_there
156 if pending:
157 time.sleep(15)
158
159 if pending:
160 log(f"WARNING: components still present after {timeout_seconds}s: {' '.join(pending)}")
161 else:
162 log("All inference components deleted")
163
164
165def main() -> int:
166 if len(sys.argv) < 2:
167 log(f"Usage: {os.path.basename(sys.argv[0])} <endpoint-name> [<region>]")
168 return 64
169
170 endpoint_name = sys.argv[1]
171 region = resolve_region(sys.argv[2] if len(sys.argv) > 2 else None)
172 if not region:
173 log("ERROR: no AWS region. Pass region as 2nd arg or set AWS_REGION.")
174 return 1
175
176 reg = ["--region", region]
177 log(f"Tearing down endpoint: {endpoint_name} in {region}")
178
179 # Discover what's attached to this endpoint.
180 config_name = ""
181 model_names: list[str] = []
182 desc = run_aws(["sagemaker", "describe-endpoint", "--endpoint-name", endpoint_name, *reg])
183 endpoint_exists = desc.returncode == 0
184 if endpoint_exists:
185 config_name = json_out(desc, {}).get("EndpointConfigName", "")
186 else:
187 log("Endpoint not found — checking for orphan resources anyway")
188
189 if config_name:
190 cfg = json_out(
191 run_aws(["sagemaker", "describe-endpoint-config",
192 "--endpoint-config-name", config_name, *reg]),
193 {},
194 )
195 # Model-based endpoints name the model on the variant. IC-based ones do not.
196 for variant in cfg.get("ProductionVariants", []):
197 if variant.get("ModelName"):
198 model_names.append(variant["ModelName"])
199
200 ic_names = list_inference_components(endpoint_name, reg)
201 if ic_names:
202 log(f"Inference components on this endpoint: {' '.join(ic_names)}")
203 for ic_name in ic_names:
204 model = component_model_name(ic_name, reg)
205 if model and model not in model_names:
206 model_names.append(model)
207
208 # Alarms — discovered by name prefix. deploy.py / deploy_async.py name them
209 # "<endpoint>-*"; deploy_ic.py names the wake alarm "<component>-*".
210 delete_alarms(f"{endpoint_name}-", reg)
211 for ic_name in ic_names:
212 delete_alarms(f"{ic_name}-", reg)
213
214 # Autoscaling — variant target (model-based) and component targets (IC-based).
215 delete_scaling_stack(f"endpoint/{endpoint_name}/variant/AllTraffic", VARIANT_DIM, reg)
216 for ic_name in ic_names:
217 delete_scaling_stack(f"inference-component/{ic_name}", IC_DIM, reg)
218
219 # Inference components must go before the endpoint.
220 if ic_names:
221 delete_inference_components(ic_names, reg)
222
223 # Endpoint (stops billing)
224 if endpoint_exists:
225 run_aws(["sagemaker", "delete-endpoint", "--endpoint-name", endpoint_name, *reg])
226 log(f"Deleted endpoint: {endpoint_name} (billing stopped)")
227
228 # Endpoint config
229 if config_name and run_aws(
230 ["sagemaker", "describe-endpoint-config", "--endpoint-config-name", config_name, *reg]
231 ).returncode == 0:
232 run_aws(["sagemaker", "delete-endpoint-config", "--endpoint-config-name", config_name, *reg])
233 log(f"Deleted endpoint config: {config_name}")
234
235 # Models
236 for model_name in model_names:
237 if run_aws(["sagemaker", "describe-model", "--model-name", model_name, *reg]).returncode == 0:
238 run_aws(["sagemaker", "delete-model", "--model-name", model_name, *reg])
239 log(f"Deleted model: {model_name}")
240
241 log("Teardown complete. Data capture S3 objects (if any) NOT deleted — manage separately.")
242 return 0
243
244
245if __name__ == "__main__":
246 sys.exit(main())