Coverage for lambda / analytics-cleanup / handler.py: 100.00%
293 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Analytics stack cleanup handler.
3Runs as a CloudFormation custom resource on stack deletion. Removes all
4SageMaker apps, spaces and user profiles from the Studio domain (waiting
5for each to fully drain) and all EFS access points from the Studio file
6system so CloudFormation can delete the domain and EFS cleanly.
8If draining apps/spaces/user-profiles fails, the handler raises — this
9fails the custom resource and stops CloudFormation before it tries (and
10fails) to delete the domain. EFS/security-group cleanup errors are
11logged but non-fatal.
13Environment variables:
14 DOMAIN_ID: SageMaker Studio domain ID
15 EFS_ID: EFS file system ID
16 REGION: AWS region
17"""
19from __future__ import annotations
21import logging
22import os
23import time
24from typing import Any
26import boto3
27from botocore.exceptions import ClientError
29# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
30# Generated at (UTC): 2026-09-01T14:42:56Z
31# Generated from Git commit: 89b000378ed5a912a38c06f4feab2b029936ebcc
32# Flowchart(s) generated from this file:
33# * ``handler`` -> ``diagrams/code_diagrams/lambda/analytics-cleanup/handler.handler.html``
34# (PNG: ``diagrams/code_diagrams/lambda/analytics-cleanup/handler.handler.png``)
35# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
36# <pyflowchart-code-diagram> END
39logger = logging.getLogger(__name__)
40logger.setLevel(logging.INFO)
42# ---------------------------------------------------------------------------
43# Tunables
44# ---------------------------------------------------------------------------
45#
46# These constants govern how aggressively we poll SageMaker/EFS while waiting
47# for asynchronous delete operations to drain. The defaults are sized for the
48# worst case we've observed in production (a domain with a handful of users,
49# one JupyterLab app each). If you see timeouts in CloudWatch, raise the
50# ``*_WAIT_SECONDS`` values — and correspondingly bump ``timeout=`` on the
51# ``CleanupFunction`` in ``gco/stacks/analytics_stack.py`` so the Lambda
52# doesn't exceed its own execution timeout.
53#
54# Lower the values if you want the custom resource to fail fast during
55# local testing, but remember that CloudFormation will then hit the domain
56# delete before the async drains finish and fail with
57# ``Unable to delete Domain ... because UserProfile(s) are associated``.
59# Interval between list-and-check iterations of every drain wait loop.
60# Small enough to keep typical stacks responsive, large enough to avoid
61# throttling SageMaker/EFS control-plane APIs.
62DRAIN_POLL_INTERVAL_SECONDS = 5
64# Brief pause after issuing ``delete_space`` / ``delete_user_profile`` to
65# avoid ThrottlingException when a domain has dozens of resources.
66DELETE_PACE_SECONDS = 1
68# Maximum time to wait for SageMaker apps to reach ``Deleted``/``Failed``.
69# Apps typically go terminal within 30-60 seconds; 2 minutes leaves plenty
70# of headroom without risking a Lambda timeout.
71APP_DELETE_WAIT_SECONDS = 120
73# Maximum time to wait for SageMaker spaces to disappear from ``list_spaces``.
74# Spaces drain faster than user profiles but may be gated by app deletion.
75SPACE_DELETE_WAIT_SECONDS = 180
77# Maximum time to wait for SageMaker user profiles to disappear from
78# ``list_user_profiles``. This is the critical wait — if CloudFormation
79# reaches ``DeleteDomain`` with any profiles still lingering, the whole
80# stack fails with a UserProfile-in-use error and has to be manually
81# unstuck.
82USER_PROFILE_DELETE_WAIT_SECONDS = 180
84# Maximum time to wait for SageMaker-managed EFS mount targets to drain.
85# Mount target deletion is usually sub-30s but can stall briefly while
86# the ENIs are detached.
87MOUNT_TARGET_DELETE_WAIT_SECONDS = 120
89# SageMaker-managed NFS security group retry behaviour. Right after
90# ``DeleteDomain``, SageMaker's control plane briefly holds an internal
91# reference on one of the two NFS SGs (typically the outbound one),
92# causing ``delete_security_group`` to fail with ``DependencyViolation``
93# for ~30-60 seconds before clearing. We retry that many times, pausing
94# between attempts; if the reference is still there on the final try we
95# promote the failure to a critical error so CloudFormation stops
96# before the VPC delete hits the same dependency and fails the stack.
97SG_DELETE_MAX_ATTEMPTS = 4
98SG_DELETE_RETRY_BACKOFF_SECONDS = 15
101def _poll_iterations(total_wait_seconds: int) -> int:
102 """Return the number of iterations for a drain loop with
103 ``DRAIN_POLL_INTERVAL_SECONDS`` between polls."""
104 return max(1, total_wait_seconds // DRAIN_POLL_INTERVAL_SECONDS)
107def handler(event: dict[str, Any], context: object) -> dict[str, Any]:
108 """CloudFormation custom resource handler."""
109 request_type = event.get("RequestType", "")
110 physical_id = event.get("PhysicalResourceId", "analytics-cleanup")
112 if request_type != "Delete":
113 logger.info("RequestType=%s, nothing to do", request_type)
114 return {"Status": "SUCCESS", "PhysicalResourceId": physical_id}
116 region = os.environ["REGION"]
117 domain_id = os.environ["DOMAIN_ID"]
118 efs_id = os.environ.get("EFS_ID", "")
119 vpc_id = os.environ.get("VPC_ID", "")
121 errors: list[str] = []
122 # Errors from deleting apps/spaces/user-profiles block the domain
123 # delete — if we return SUCCESS with these still present, CloudFormation
124 # will immediately fail on ``AWS::SageMaker::Domain`` with a
125 # ``Unable to delete Domain ... because UserProfile(s) are associated
126 # with it`` error. We track these separately so we can fail the custom
127 # resource and give the operator a useful log pointer instead.
128 critical_errors: list[str] = []
130 # Delete all apps (must be deleted before spaces/profiles)
131 app_errors = _delete_apps(region, domain_id)
132 errors.extend(app_errors)
133 critical_errors.extend(app_errors)
135 # Delete all spaces
136 space_errors = _delete_spaces(region, domain_id)
137 errors.extend(space_errors)
138 critical_errors.extend(space_errors)
140 # Delete all user profiles from the domain
141 profile_errors = _delete_user_profiles(region, domain_id)
142 errors.extend(profile_errors)
143 critical_errors.extend(profile_errors)
145 # Remove EFS resource policies that trigger the intersection
146 # authorization model. Both the CDK-managed EFS and the SageMaker-
147 # managed EFS have resource policies that block DescribeMountTargets.
148 if efs_id:
149 _delete_efs_resource_policy(region, efs_id)
150 sm_efs_id = _get_sagemaker_home_efs_id(region, domain_id)
151 if sm_efs_id:
152 _delete_efs_resource_policy(region, sm_efs_id)
154 # Delete SageMaker-managed EFS (created internally by the domain)
155 errors.extend(_delete_sagemaker_managed_efs(region, domain_id))
157 # Delete SageMaker-managed NFS security groups
158 if vpc_id:
159 errors.extend(_delete_sagemaker_security_groups(region, domain_id, vpc_id))
161 if errors:
162 logger.warning("Cleanup completed with %d errors: %s", len(errors), errors)
163 else:
164 logger.info("Cleanup completed successfully")
166 # If apps/spaces/user-profiles still have unresolved errors, fail the
167 # custom resource. CloudFormation will stop before attempting to
168 # delete the domain, surface the error to the operator, and the stack
169 # stays in a retriable state. EFS/security-group errors are logged
170 # but non-fatal — they don't block the domain delete and are cleaned
171 # up best-effort on the next attempt.
172 if critical_errors:
173 raise RuntimeError(
174 "Analytics cleanup failed to fully drain the SageMaker domain "
175 f"({len(critical_errors)} error(s)). See CloudWatch Logs for "
176 f"details: {critical_errors}"
177 )
179 return {"Status": "SUCCESS", "PhysicalResourceId": physical_id}
182def _delete_apps(region: str, domain_id: str) -> list[str]:
183 """Delete all apps in the domain. Apps must be deleted before spaces/profiles."""
184 errors: list[str] = []
185 sm = boto3.client("sagemaker", region_name=region)
187 try:
188 paginator = sm.get_paginator("list_apps")
189 for page in paginator.paginate(DomainIdEquals=domain_id):
190 for app in page.get("Apps", []):
191 if app.get("Status") in ("Deleted", "Failed"):
192 continue
193 app_name = app["AppName"]
194 app_type = app["AppType"]
195 space_name = app.get("SpaceName")
196 user_profile = app.get("UserProfileName")
197 try:
198 kwargs = {
199 "DomainId": domain_id,
200 "AppType": app_type,
201 "AppName": app_name,
202 }
203 if space_name:
204 kwargs["SpaceName"] = space_name
205 if user_profile:
206 kwargs["UserProfileName"] = user_profile
207 sm.delete_app(**kwargs)
208 logger.info("Deleted app: %s (%s)", app_name, app_type)
209 except ClientError as e:
210 if "does not exist" not in str(e):
211 msg = f"Failed to delete app {app_name}: {e}"
212 logger.error(msg)
213 errors.append(msg)
215 # Wait for apps to finish deleting. A timeout must be surfaced to the
216 # custom resource handler; otherwise CloudFormation immediately moves
217 # on to space/profile/domain deletion while apps are still attached.
218 active: list[str] = []
219 for _ in range(_poll_iterations(APP_DELETE_WAIT_SECONDS)):
220 time.sleep(DRAIN_POLL_INTERVAL_SECONDS)
221 active = []
222 for page in paginator.paginate(DomainIdEquals=domain_id):
223 for app in page.get("Apps", []):
224 if app.get("Status") not in ("Deleted", "Failed"):
225 active.append(app["AppName"])
226 if not active:
227 break
228 logger.info("Waiting for %d app(s) to delete: %s", len(active), active)
229 else:
230 msg = f"Timed out waiting for apps to delete in {domain_id}: {active}"
231 logger.error(msg)
232 errors.append(msg)
234 except ClientError as e:
235 msg = f"Failed to list apps: {e}"
236 logger.error(msg)
237 errors.append(msg)
239 return errors
242def _delete_spaces(region: str, domain_id: str) -> list[str]:
243 """Delete all spaces in the domain and wait for them to be gone.
245 Spaces must be fully removed before user profiles can be deleted, and
246 user profiles must be fully removed before the domain can be deleted.
247 ``delete_space`` is asynchronous, so after issuing the deletes we poll
248 ``list_spaces`` until it returns empty (or a timeout elapses).
249 """
250 errors: list[str] = []
251 sm = boto3.client("sagemaker", region_name=region)
253 try:
254 paginator = sm.get_paginator("list_spaces")
255 for page in paginator.paginate(DomainIdEquals=domain_id):
256 for space in page.get("Spaces", []):
257 space_name = space["SpaceName"]
258 # Skip spaces already being deleted; the wait loop below
259 # will still account for them.
260 if space.get("Status") == "Deleting":
261 continue
262 try:
263 sm.delete_space(DomainId=domain_id, SpaceName=space_name)
264 logger.info("Deleted space: %s", space_name)
265 time.sleep(DELETE_PACE_SECONDS)
266 except ClientError as e:
267 if "does not exist" not in str(e):
268 msg = f"Failed to delete space {space_name}: {e}"
269 logger.error(msg)
270 errors.append(msg)
272 # Wait for spaces to finish deleting.
273 for _ in range(_poll_iterations(SPACE_DELETE_WAIT_SECONDS)):
274 remaining: list[str] = []
275 for page in paginator.paginate(DomainIdEquals=domain_id):
276 for space in page.get("Spaces", []):
277 remaining.append(space["SpaceName"])
278 if not remaining:
279 break
280 logger.info("Waiting for %d space(s) to delete: %s", len(remaining), remaining)
281 time.sleep(DRAIN_POLL_INTERVAL_SECONDS)
282 else:
283 msg = f"Timed out waiting for spaces to delete in {domain_id}: {remaining}"
284 logger.error(msg)
285 errors.append(msg)
286 except ClientError as e:
287 msg = f"Failed to list spaces: {e}"
288 logger.error(msg)
289 errors.append(msg)
291 return errors
294def _delete_user_profiles(region: str, domain_id: str) -> list[str]:
295 """Delete all user profiles in the domain and wait for them to be gone.
297 ``delete_user_profile`` is asynchronous — it puts the profile into
298 ``Deleting`` state and returns immediately. If we don't wait for the
299 list to drain, CloudFormation will race ahead to delete the domain
300 and fail with ``Unable to delete Domain ... because UserProfile(s)
301 are associated with it``. This function polls ``list_user_profiles``
302 until it's empty (or a timeout elapses).
303 """
304 errors: list[str] = []
305 sm = boto3.client("sagemaker", region_name=region)
307 try:
308 paginator = sm.get_paginator("list_user_profiles")
309 for page in paginator.paginate(DomainIdEquals=domain_id):
310 for profile in page.get("UserProfiles", []):
311 name = profile["UserProfileName"]
312 # Skip profiles already being deleted; the wait loop below
313 # will still account for them.
314 if profile.get("Status") == "Deleting":
315 continue
316 try:
317 sm.delete_user_profile(DomainId=domain_id, UserProfileName=name)
318 logger.info("Deleted user profile: %s", name)
319 # Brief pause to avoid throttling
320 time.sleep(DELETE_PACE_SECONDS)
321 except ClientError as e:
322 msg = f"Failed to delete profile {name}: {e}"
323 logger.error(msg)
324 errors.append(msg)
326 # Wait for profiles to finish deleting. Without this, CloudFormation
327 # will race ahead to delete the domain while profiles are still in
328 # ``Deleting`` state and fail the stack.
329 remaining: list[str] = []
330 for _ in range(_poll_iterations(USER_PROFILE_DELETE_WAIT_SECONDS)):
331 remaining = []
332 for page in paginator.paginate(DomainIdEquals=domain_id):
333 for profile in page.get("UserProfiles", []):
334 remaining.append(profile["UserProfileName"])
335 if not remaining:
336 break
337 logger.info(
338 "Waiting for %d user profile(s) to delete: %s",
339 len(remaining),
340 remaining,
341 )
342 time.sleep(DRAIN_POLL_INTERVAL_SECONDS)
343 else:
344 msg = f"Timed out waiting for user profiles to delete in {domain_id}: {remaining}"
345 logger.error(msg)
346 errors.append(msg)
347 except ClientError as e:
348 msg = f"Failed to list user profiles: {e}"
349 logger.error(msg)
350 errors.append(msg)
352 return errors
355def _delete_access_points(region: str, efs_id: str) -> list[str]:
356 """Delete all access points on the file system. Returns a list of error messages."""
357 errors: list[str] = []
358 efs = boto3.client("efs", region_name=region)
360 try:
361 paginator = efs.get_paginator("describe_access_points")
362 for page in paginator.paginate(FileSystemId=efs_id):
363 for ap in page.get("AccessPoints", []):
364 ap_id = ap["AccessPointId"]
365 try:
366 efs.delete_access_point(AccessPointId=ap_id)
367 logger.info("Deleted access point: %s", ap_id)
368 except ClientError as e:
369 msg = f"Failed to delete access point {ap_id}: {e}"
370 logger.error(msg)
371 errors.append(msg)
372 except ClientError as e:
373 msg = f"Failed to list access points: {e}"
374 logger.error(msg)
375 errors.append(msg)
377 return errors
380def _delete_sagemaker_security_groups(region: str, domain_id: str, vpc_id: str) -> list[str]:
381 """Delete SageMaker-managed security groups for the domain.
383 SageMaker creates security groups named
384 ``security-group-for-outbound-nfs-<domain-id>`` and
385 ``security-group-for-inbound-nfs-<domain-id>`` when the domain uses
386 a custom EFS. These are tagged "[DO NOT DELETE]" but must be removed
387 for the VPC to be deletable.
389 The two SGs cross-reference each other (outbound rules on one point
390 to the other), creating a circular dependency. We must revoke all
391 ingress/egress rules before deleting.
393 Right after ``DeleteDomain``, SageMaker's control plane briefly
394 retains an internal reference on one of the NFS SGs — typically the
395 outbound one — causing ``delete_security_group`` to fail with
396 ``DependencyViolation``. The reference reliably clears within 30-60s.
397 We retry the delete ``SG_DELETE_MAX_ATTEMPTS`` times with
398 ``SG_DELETE_RETRY_BACKOFF_SECONDS`` between attempts, and only
399 surface an error if the SG is still undeletable after the final try.
400 """
401 errors: list[str] = []
402 ec2 = boto3.client("ec2", region_name=region)
404 try:
405 response = ec2.describe_security_groups(
406 Filters=[
407 {"Name": "vpc-id", "Values": [vpc_id]},
408 {"Name": "group-name", "Values": [f"*{domain_id}*"]},
409 ]
410 )
411 sgs = response.get("SecurityGroups", [])
413 # First pass: revoke all rules to break cross-references.
414 for sg in sgs:
415 sg_id = sg["GroupId"]
416 try:
417 if sg.get("IpPermissions"):
418 ec2.revoke_security_group_ingress(
419 GroupId=sg_id, IpPermissions=sg["IpPermissions"]
420 )
421 if sg.get("IpPermissionsEgress"):
422 ec2.revoke_security_group_egress(
423 GroupId=sg_id, IpPermissions=sg["IpPermissionsEgress"]
424 )
425 except ClientError as e:
426 logger.warning("Failed to revoke rules on %s: %s", sg_id, e)
428 # Second pass: delete the security groups, retrying on
429 # ``DependencyViolation`` for ones where SageMaker still holds
430 # a transient reference post-DeleteDomain.
431 pending = [(sg["GroupId"], sg.get("GroupName", "")) for sg in sgs]
432 for attempt in range(1, SG_DELETE_MAX_ATTEMPTS + 1):
433 still_pending: list[tuple[str, str]] = []
434 for sg_id, sg_name in pending:
435 try:
436 ec2.delete_security_group(GroupId=sg_id)
437 logger.info("Deleted SageMaker security group: %s (%s)", sg_id, sg_name)
438 except ClientError as e:
439 code = e.response.get("Error", {}).get("Code", "")
440 # ``InvalidGroup.NotFound`` means some other actor
441 # already deleted it — treat as success.
442 if code == "InvalidGroup.NotFound":
443 logger.info("SG %s (%s) already deleted", sg_id, sg_name)
444 continue
445 if code == "DependencyViolation":
446 logger.warning(
447 "SG %s (%s) has a dependent object (attempt %d/%d); will retry",
448 sg_id,
449 sg_name,
450 attempt,
451 SG_DELETE_MAX_ATTEMPTS,
452 )
453 still_pending.append((sg_id, sg_name))
454 continue
455 msg = f"Failed to delete security group {sg_id}: {e}"
456 logger.error(msg)
457 errors.append(msg)
458 if not still_pending:
459 break
460 pending = still_pending
461 if attempt < SG_DELETE_MAX_ATTEMPTS:
462 time.sleep(SG_DELETE_RETRY_BACKOFF_SECONDS)
463 else:
464 # Exhausted retries with SGs still undeletable.
465 for sg_id, sg_name in pending:
466 msg = (
467 f"Failed to delete security group {sg_id} ({sg_name}) "
468 f"after {SG_DELETE_MAX_ATTEMPTS} attempts: "
469 "DependencyViolation did not clear. CloudFormation "
470 "will fail to delete the VPC; manually delete the "
471 "SG after its dependent object releases."
472 )
473 logger.error(msg)
474 errors.append(msg)
475 except ClientError as e:
476 msg = f"Failed to list SageMaker security groups: {e}"
477 logger.error(msg)
478 errors.append(msg)
480 return errors
483def _get_sagemaker_home_efs_id(region: str, domain_id: str) -> str:
484 """Get the HomeEfsFileSystemId from the SageMaker domain."""
485 sm = boto3.client("sagemaker", region_name=region)
486 try:
487 resp = sm.describe_domain(DomainId=domain_id)
488 return str(resp.get("HomeEfsFileSystemId", ""))
489 except ClientError as e:
490 logger.warning("Failed to get HomeEfsFileSystemId for %s: %s", domain_id, e)
491 return ""
494def _delete_efs_resource_policy(region: str, efs_id: str) -> None:
495 """Delete the resource policy on the CDK-managed EFS.
497 The EFS resource policy triggers the intersection authorization model
498 which blocks DescribeFileSystems/DescribeAccessPoints calls even when
499 the caller has IAM Resource:* permissions. Removing the policy before
500 other EFS operations ensures they succeed.
501 """
502 efs_client = boto3.client("efs", region_name=region)
503 try:
504 efs_client.delete_file_system_policy(FileSystemId=efs_id)
505 logger.info("Deleted EFS resource policy on %s", efs_id)
506 except ClientError as e:
507 # PolicyNotFound is fine — means there's no policy to delete.
508 if "PolicyNotFound" not in str(e):
509 logger.warning("Failed to delete EFS resource policy on %s: %s", efs_id, e)
512def _delete_sagemaker_managed_efs(region: str, domain_id: str) -> list[str]:
513 """Delete the SageMaker-managed EFS created internally by the domain.
515 Uses sagemaker:DescribeDomain to get the HomeEfsFileSystemId directly,
516 avoiding DescribeFileSystems which is blocked by the EFS resource
517 policy intersection model. The domain still exists when this Lambda
518 runs (it's deleted after the custom resource completes).
519 """
520 errors: list[str] = []
521 sm = boto3.client("sagemaker", region_name=region)
522 efs = boto3.client("efs", region_name=region)
524 try:
525 # Get the EFS ID from the domain itself — no DescribeFileSystems needed.
526 domain_info = sm.describe_domain(DomainId=domain_id)
527 target_fs = domain_info.get("HomeEfsFileSystemId")
528 if not target_fs:
529 logger.info("No HomeEfsFileSystemId found for domain %s", domain_id)
530 return errors
532 logger.info("Found SageMaker-managed EFS: %s", target_fs)
534 # Delete all mount targets first
535 mt_response = efs.describe_mount_targets(FileSystemId=target_fs)
536 for mt in mt_response.get("MountTargets", []):
537 mt_id = mt["MountTargetId"]
538 try:
539 efs.delete_mount_target(MountTargetId=mt_id)
540 logger.info("Deleted mount target: %s", mt_id)
541 except ClientError as e:
542 msg = f"Failed to delete mount target {mt_id}: {e}"
543 logger.error(msg)
544 errors.append(msg)
546 # Wait for mount targets to be deleted.
547 for _ in range(_poll_iterations(MOUNT_TARGET_DELETE_WAIT_SECONDS)):
548 time.sleep(DRAIN_POLL_INTERVAL_SECONDS)
549 remaining = efs.describe_mount_targets(FileSystemId=target_fs)
550 if not remaining.get("MountTargets"):
551 break
553 # Delete the file system
554 try:
555 efs.delete_file_system(FileSystemId=target_fs)
556 logger.info("Deleted SageMaker-managed EFS: %s", target_fs)
557 except ClientError as e:
558 msg = f"Failed to delete EFS {target_fs}: {e}"
559 logger.error(msg)
560 errors.append(msg)
562 except ClientError as e:
563 msg = f"Failed to find/delete SageMaker-managed EFS: {e}"
564 logger.error(msg)
565 errors.append(msg)
567 return errors