Coverage for src/tagmania/iac_tools/clusterset.py: 85%
524 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-02 01:51 +0000
1"""ClusterSet - Core AWS EC2 Cluster Management.
3This module provides the ClusterSet class, which is the central component for managing
4collections of EC2 instances based on cluster tags. It provides comprehensive
5functionality for cluster operations including instance management, volume operations,
6and snapshot management with built-in safety features.
8The ClusterSet class serves as the foundation for all cluster operations in Tagmania,
9offering both simple operations (start/stop instances) and complex operations
10(snapshot creation/restoration with targeted instance filtering).
12Key Features:
13 - Tag-based resource identification and management
14 - Instance lifecycle management (start, stop, targeted operations)
15 - EBS volume and snapshot management
16 - Targeted operations using regex pattern matching
17 - Safety limits and automation tracking
18 - Comprehensive logging and error handling
20Safety Features:
21 - Maximum item limits to prevent performance issues
22 - Automation key tracking to avoid modifying unmanaged resources
23 - Built-in confirmation and validation mechanisms
24 - Comprehensive error handling and logging
26Example:
27 Basic cluster management:
29 ```python
30 # Create a cluster set
31 cluster = ClusterSet('production-web')
33 # Start all instances
34 cluster.start_instances()
36 # Create snapshots
37 cluster.create_snapshots('daily-backup')
39 # Targeted operations
40 cluster.stop_instances_targeted('.*-worker-.*')
41 ```
42"""
44from __future__ import annotations
46import datetime
47import logging
48import re
49import time
50from typing import Any
52import boto3
54from .filterset import FilterSet
55from .tagset import TagSet
56from .timing import log_duration
59class ClusterSet:
60 """Manages collections of EC2 instances based on cluster tags.
62 The ClusterSet class provides comprehensive cluster management functionality
63 including instance lifecycle management, volume operations, and snapshot
64 management. It uses tag-based resource identification to ensure operations
65 only affect intended resources.
67 Attributes:
68 cluster_names: The name(s) of the cluster(s) being managed
69 _MAX_ITEMS: Safety limit for collection operations (150)
70 AUTOMATION_KEY: Key used to track managed resources ('SNAPSHOT_MANAGER')
72 Note:
73 All operations rely on the "Cluster" tag to identify cluster membership.
74 The class includes built-in safety features to prevent accidental
75 modification of unmanaged resources.
76 """
78 def __init__(self, cluster_names: str | list[str], profile: str | None = None) -> None:
79 """Initialize ClusterSet for managing one or more clusters.
81 Creates a new ClusterSet instance for managing EC2 instances and related
82 resources (volumes, snapshots) for the specified cluster(s). Sets up
83 AWS connectivity and initializes safety limits and automation tracking.
85 Args:
86 cluster_names: The name of a cluster (str) or list of cluster names (list).
87 Must match the "Cluster" tag value on EC2 instances.
88 profile: AWS profile name to use for authentication (optional).
89 If None, uses default AWS credentials chain.
91 Example:
92 ```python
93 # Single cluster
94 cluster = ClusterSet('production-web')
96 # Multiple clusters
97 clusters = ClusterSet(['prod-web', 'prod-api'])
99 # With specific AWS profile
100 cluster = ClusterSet('staging', profile='dev-account')
101 ```
103 Note:
104 The cluster names must exactly match the "Cluster" tag values
105 on your EC2 instances for operations to work correctly.
106 """
107 # Collections operate lazily. Turning a collection into a list can cause
108 # performance issues if the collection is very large. Although this is
109 # unlikely in our environment, we protect against this by setting an
110 # upper bound on the number of items from the collection that can be
111 # placed in the list. If the number of items in the collection exceed
112 # this upper bound, then those items will not be processed. If this
113 # happens, then it can be addressed by increasing this value.
114 self._MAX_ITEMS = 150
116 # Used to ensure we don't clobber anything we don't make
117 self.AUTOMATION_KEY = "SNAPSHOT_MANAGER"
119 # cluster_names can be a string 'aws-dev5' or a list ['aws-dev1', 'aws-dev2', etc...]
120 # Normalize to always be a string for single-cluster usage
121 self.cluster_names = cluster_names
123 cluster_list = cluster_names if isinstance(cluster_names, list) else [cluster_names]
125 # Use this as a starting point to build filters that only return
126 # resources associated with this cluster. Use the 'get_cluster_filter'
127 # method to get a copy of it rather than doing direct assignment.
128 self._cluster_filter: list[dict[str, Any]] = [
129 {"Name": "tag:Cluster", "Values": cluster_list},
130 ]
132 # Set up logging
133 self._logger = logging.getLogger("tagmania")
134 self._logger.setLevel(logging.INFO)
135 self._logger.info("Logging initialized.")
137 # Create a boto3 session using the specified profile if provided.
138 if profile:
139 aws_session = boto3.Session(profile_name=profile)
140 self._logger.info(f"Using AWS profile: {profile}")
141 else:
142 aws_session = boto3.Session()
144 # Create EC2 resource and client from the session.
145 self._ec2 = aws_session.resource("ec2")
146 self._ec2_client = aws_session.client("ec2")
148 @property
149 def _cluster_name_str(self) -> str:
150 """Get cluster name as a string (uses first name if multiple)."""
151 if isinstance(self.cluster_names, list):
152 return self.cluster_names[0]
153 return self.cluster_names
155 def get_cluster_filter(self) -> list[dict[str, Any]]:
156 # Return a copy to defend against modifications
157 return self._cluster_filter.copy()
159 def _wait_instances_running(self, instance_ids: list[str]) -> None:
160 """Wait for instances to reach running state using 5s polling."""
161 waiter = self._ec2_client.get_waiter("instance_running")
162 waiter.wait(
163 InstanceIds=instance_ids,
164 WaiterConfig={"Delay": 5, "MaxAttempts": 120},
165 )
167 def _wait_instances_stopped(self, instance_ids: list[str]) -> None:
168 """Wait for instances to reach stopped state using 5s polling."""
169 waiter = self._ec2_client.get_waiter("instance_stopped")
170 waiter.wait(
171 InstanceIds=instance_ids,
172 WaiterConfig={"Delay": 5, "MaxAttempts": 120},
173 )
175 def get_instances(self) -> list[Any]:
176 """Get all EC2 instances belonging to this cluster set.
178 Retrieves all EC2 instances that have a "Cluster" tag matching any of the
179 cluster names specified during initialization. The instances are filtered
180 using the cluster filter and limited by the safety maximum (_MAX_ITEMS).
182 Returns:
183 list: List of EC2 instance objects from boto3. Each instance object
184 contains all AWS instance metadata including ID, state, tags, etc.
186 Example:
187 ```python
188 cluster = ClusterSet('production-web')
189 instances = cluster.get_instances()
190 for instance in instances:
191 print(f"Instance {instance.id} is {instance.state['Name']}")
192 ```
194 Note:
195 Returns a maximum of _MAX_ITEMS (150) instances for performance protection.
196 If your cluster has more instances, increase this limit or use filtering.
197 """
198 self._logger.debug("method_call: get_instances")
199 fs = FilterSet(self.get_cluster_filter())
200 # Exclude terminated instances
201 fs.add(
202 "instance-state-name",
203 ["pending", "running", "shutting-down", "stopped", "stopping"],
204 )
205 filters = fs.to_list()
206 instances = self._ec2.instances.filter(Filters=filters)
207 return list(instances.limit(self._MAX_ITEMS))
209 def get_deployed_clusters(self) -> dict[Any, list[Any]]:
210 """
211 Get a dictionary with all clusters already deployed.
213 Args:
214 none
215 Returns:
216 dictionary of clusters
217 """
218 self._logger.debug("method_call: get_deployed_clusters")
219 fs = FilterSet(self.get_cluster_filter())
220 # Exclude terminated instances
221 fs.add(
222 "instance-state-name",
223 ["pending", "running", "shutting-down", "stopped", "stopping"],
224 )
225 filters = fs.to_list()
226 instances = self._ec2.instances.filter(Filters=filters)
227 instance_list = list(instances)
228 cluster_names = (
229 self.cluster_names if isinstance(self.cluster_names, list) else [self.cluster_names]
230 )
231 cluster_dict: dict[str, list[Any]] = {k: [] for k in cluster_names}
233 for i in instance_list:
234 cluster_tag = TagSet(i.tags).get("Cluster")
235 if cluster_tag is not None:
236 cluster_dict[cluster_tag].append(i)
238 return cluster_dict
240 def get_deployed_cluster_names(self) -> set[str | None]:
241 """
242 Get a set of all supported environment cluster names already deployed.
244 Args:
245 none
246 Returns:
247 set of deployed cluster_names
248 """
249 self._logger.debug("method_call: get_deployed_cluster_names")
250 fs = FilterSet(self.get_cluster_filter())
251 # Exclude terminated instances
252 fs.add(
253 "instance-state-name",
254 ["pending", "running", "shutting-down", "stopped", "stopping"],
255 )
256 filters = fs.to_list()
257 instances = self._ec2.instances.filter(Filters=filters)
258 instance_list = list(instances)
259 instance_cluster_names = []
261 for i in instance_list:
262 cluster_name = TagSet(i.tags).get("Cluster")
263 if cluster_name not in instance_cluster_names:
264 instance_cluster_names.append(cluster_name)
266 return set(instance_cluster_names)
268 def get_running_instances(self) -> list[Any]:
269 """
270 Get list of instances that are powered on. This includes instances in
271 a pending, running, or stopping state.
273 Args:
274 none
275 Returns:
276 list of instances
277 """
278 self._logger.debug("method_call: get_instances")
279 fs = FilterSet(self.get_cluster_filter())
280 # Only want instances that are pending, running, or stopping
281 fs.add("instance-state-name", ["pending", "running", "stopping"])
282 filters = fs.to_list()
283 instances = self._ec2.instances.filter(Filters=filters)
284 return list(instances.limit(self._MAX_ITEMS))
286 def get_running_clusters(self) -> dict[Any, list[Any]]:
287 """
288 Get a dictionary with all clusters that are powered on. This includes instances in
289 a pending, running, or stopping state.
291 Args:
292 none
293 Returns:
294 dictionary of running clusters
295 """
296 self._logger.debug("method_call: get_running_clusters")
297 fs = FilterSet(self.get_cluster_filter())
298 # Only want instances that are pending, running, or stopping
299 fs.add("instance-state-name", ["pending", "running", "stopping"])
300 filters = fs.to_list()
301 instances = self._ec2.instances.filter(Filters=filters)
302 instance_list = list(instances)
303 cluster_names = (
304 self.cluster_names if isinstance(self.cluster_names, list) else [self.cluster_names]
305 )
306 cluster_dict: dict[str, list[Any]] = {k: [] for k in cluster_names}
308 for i in instance_list:
309 cluster_tag = TagSet(i.tags).get("Cluster")
310 if cluster_tag is not None:
311 cluster_dict[cluster_tag].append(i)
313 return cluster_dict
315 def get_stopped_instances(self) -> list[Any]:
316 """
317 Get list of instances that are powered off.
319 Args:
320 none
321 Returns:
322 list of instances
323 """
324 self._logger.debug("method_call: get_instances")
325 fs = FilterSet(self.get_cluster_filter())
326 # Only want instances that are completely stopped
327 fs.add("instance-state-name", "stopped")
328 filters = fs.to_list()
329 instances = self._ec2.instances.filter(Filters=filters)
330 return list(instances.limit(self._MAX_ITEMS))
332 def get_stopped_clusters(self) -> dict[Any, list[Any]]:
333 """
334 Get a dictionary with all clusters that are powered off.
336 Args:
337 none
338 Returns:
339 dictionary of stopped clusters
340 """
341 self._logger.debug("method_call: get_stopped_clusters")
342 fs = FilterSet(self.get_cluster_filter())
343 # Only want instances that are completely stopped
344 fs.add("instance-state-name", "stopped")
345 filters = fs.to_list()
346 instances = self._ec2.instances.filter(Filters=filters)
347 instance_list = list(instances)
348 cluster_names = (
349 self.cluster_names if isinstance(self.cluster_names, list) else [self.cluster_names]
350 )
351 cluster_dict: dict[str, list[Any]] = {k: [] for k in cluster_names}
353 for i in instance_list:
354 cluster_tag = TagSet(i.tags).get("Cluster")
355 if cluster_tag is not None:
356 cluster_dict[cluster_tag].append(i)
358 return cluster_dict
360 def start_instances(self) -> None:
361 """Start all stopped EC2 instances in this cluster.
363 Identifies all instances in the cluster that are currently in 'stopped' state
364 and starts them. Running instances are not affected. The operation processes
365 instances in batches and provides feedback on progress.
367 Example:
368 ```python
369 cluster = ClusterSet('production-web')
370 cluster.start_instances() # Starts all stopped instances
371 ```
373 Note:
374 Only instances in 'stopped' state will be started. Instances in other
375 states (running, stopping, etc.) are ignored. The operation may take
376 several minutes for large clusters.
377 """
378 self._logger.debug("method_call: start_instances")
379 instances = self.get_stopped_instances()
380 if len(instances) == 0:
381 print("No instances to start.")
382 else:
383 # Start instances
384 for i in instances:
385 name = TagSet(i.tags).get("Name")
386 print(f"Starting {name} ({i.id})")
387 i.start()
388 # Wait for all instances in one batch call with 5s polling
389 print(f"Waiting for {len(instances)} instances to start...")
390 self._wait_instances_running([i.id for i in instances])
392 def stop_instances(self) -> None:
393 """
394 Stop instances associated with this cluster.
396 Args:
397 none
398 Returns:
399 none
400 """
401 self._logger.debug("method_call: stop_instances")
402 instances = self.get_running_instances()
403 if len(instances) == 0:
404 print("No instances to stop.")
405 else:
406 # Stop instances
407 for i in instances:
408 name = TagSet(i.tags).get("Name")
409 print(f"Stopping {name} ({i.id})")
410 i.stop()
411 # Wait for all instances in one batch call with 5s polling
412 print(f"Waiting for {len(instances)} instances to stop...")
413 self._wait_instances_stopped([i.id for i in instances])
415 def tag_instances(self, tags: list[dict[str, str]]) -> None:
416 # The resource API is somewhat less efficient than the low-level client
417 # API because it does one request per tag operation. For reasonably
418 # sized collections this should be ok.
419 # tags = {'key': 'values'}
420 instances = self.get_instances()
421 for i in instances:
422 print(f"Tagging instance {i.id}")
423 i.create_tags(Tags=tags)
425 def untag_instances(self, tags: list[dict[str, str]]) -> None:
426 # Might as well un-tag on one big batch since instance objects don't
427 # have direct support for un-tagging.
428 instances = self.get_instances()
429 instance_ids = []
430 for i in instances:
431 print(f"Un-tagging instance {i.id}")
432 instance_ids.append(i.id)
433 self._ec2_client.delete_tags(Resources=instance_ids, Tags=tags) # type: ignore[arg-type]
435 def get_volumes(self) -> list[Any]:
436 """
437 Get list of volumes associated with this cluster.
439 Args:
440 none
441 Returns:
442 list of volumes
443 """
444 self._logger.debug("method_call: get_volumes")
445 fs = FilterSet(self.get_cluster_filter())
446 # The tool deliberately only works with volumes that have the tag
447 # 'automation_key' set. It will ignore all other volumes so that it
448 # doesn't clobber volumes that it is not responsible for. It may be that
449 # we only need to worry about volumes created by snapshot manager
450 # because volumes created by the provisioner won't ever have a label,
451 # so when looking for restored volumes, we can just look for managed
452 # volumes that have a label.
453 fs.add("tag:automation_key", ["PROVISIONER", self.AUTOMATION_KEY])
454 filters = fs.to_list()
455 volumes = self._ec2.volumes.filter(Filters=filters)
456 return list(volumes.limit(self._MAX_ITEMS))
458 def get_kubernetes_volumes(self) -> list[Any]:
459 """
460 Get list of kubernetes volumes associated with this cluster.
462 Args:
463 none
464 Returns:
465 list of kubernetes volumes
466 """
467 self._logger.debug("method_call: get_kubernetes_volumes")
468 fs = FilterSet(
469 [
470 {
471 "Name": "tag:KubernetesCluster",
472 "Values": self.cluster_names
473 if isinstance(self.cluster_names, list)
474 else [self.cluster_names],
475 },
476 {
477 "Name": "tag:kubernetes.io/created-for/pvc/namespace",
478 "Values": ["openshift-logging"],
479 },
480 ]
481 )
482 filters = fs.to_list()
483 volumes = self._ec2.volumes.filter(Filters=filters)
484 return list(volumes)
486 def get_restored_volumes(self, label: str | None = None) -> list[Any]:
487 """
488 Get list of volumes that were previously created from snapshots.
490 Args:
491 - label: label of volumes being sought (optional)
492 Returns:
493 list of snapshots
494 """
495 self._logger.debug("method_call: get_restored_volumes")
496 fs = FilterSet(self.get_cluster_filter())
497 # These volumes can be in available state (if just created) or in-use state (if attached)
498 # We need to check for both states to properly detect restored volumes
499 # | Note: The 'status' filter corresponds to the 'state' attribute in
500 # | the AWS management console. Not sure about the reason behind that.
501 fs.add("tag:automation_key", self.AUTOMATION_KEY)
502 # Optionally, get volumes with the given label
503 # This might not really be needed because the way the snapshot manager
504 # works is that it deletes all managed volumes before creating new ones
505 # from snapshots. Therefore at any given time, there should only be one
506 # set of available volumes.
507 if label is not None:
508 fs.add("tag:Label", label)
509 filters = fs.to_list()
510 volumes = self._ec2.volumes.filter(Filters=filters)
511 return list(volumes.limit(self._MAX_ITEMS))
513 def attach_volumes(self, label: str) -> None:
514 """
515 Attach volumes to associated instances.
517 Args:
518 - label: label of volume to attach
519 Returns:
520 none
521 """
522 self._logger.debug("method_call: attach_volumes")
523 instances = self.get_instances()
524 # We should be able to work with get_volumes, but this just protects
525 # against the case when users are manually creating volumes. It filters
526 # out any volumes that were not created by the snapshot manager.
527 volumes = self.get_restored_volumes(label)
528 # This is for debugging purposes. Sometimes the algorithm below doesn't
529 # find all volumes.
530 # volume_list = list(volumes)
531 # print(f"Found {len(volume_list)} volumes to attach.")
533 # For each instance, find all associated volumes and attach them. The
534 # association is performed by matching the volume 'Instance' tag to
535 # the instance 'Name' tag.
536 volume_ids = []
537 for i in instances:
538 instance_name = TagSet(i.tags).get("Name") or ""
539 for volume in volumes:
540 ts = TagSet(volume.tags)
541 instance = ts.get("Instance")
542 device = ts.get("Device")
543 if instance == instance_name:
544 # Attach volume
545 shortname = instance_name.split(".")[0]
546 print(f"Attaching {device} ({volume.id}) to {shortname} ({i.id})")
547 volume.attach_to_instance(Device=device, InstanceId=i.id)
548 volume_ids.append(volume.id)
549 if len(volume_ids) == 0:
550 # This is probably an error. The expectation is that we have a set
551 # of newly created volumes from snapshots.
552 print("Error: No volumes to attach.")
553 else:
554 # Wait for the volumes to be attached
555 print(f"Waiting for {len(volume_ids)} volumes to be attached...")
556 self.wait_for_volumes(volume_ids, "volume_in_use")
558 def create_volumes(self, label: str) -> None:
559 """
560 Create new volumes from managed snapshots.
562 Args:
563 - label: label of snapshots to restore
564 Returns:
565 none
566 """
567 self._logger.debug("method_call: create_managed_volumes")
568 with log_duration(self._logger, "create_volumes"):
569 snapshots = self.get_snapshots(label)
570 # Check if snapshot list is empty (e.g. due to an invalid label)
571 if len(snapshots) == 0:
572 print(f"Error: No snapshots found with label '{label}'.")
573 # Create volumes
574 volume_ids = []
575 for snapshot in snapshots:
576 # Determine snapshot's associated instance and device. This is
577 # needed later on so that we know where to attach it.
578 ts = TagSet(snapshot.tags)
579 device = ts.get("Device")
580 instance = ts.get("Instance")
581 if not device:
582 raise Exception(
583 f"Error: create_volume: Can't find device tag for snapshot {snapshot.id}."
584 )
585 if not instance:
586 raise Exception(
587 f"Error: create_volume: Can't find instance tag for snapshot {snapshot.id}."
588 )
589 # Determine availability zone from one of the cluster instances
590 avail_zone = self.get_instances()[0].placement["AvailabilityZone"]
591 # Make tags
592 ts = TagSet()
593 ts.add("Cluster", self._cluster_name_str)
594 ts.add("Device", device)
595 ts.add("Instance", instance)
596 ts.add("Label", label)
597 ts.add("Name", f"{instance} - {device}")
598 ts.add("automation_key", self.AUTOMATION_KEY)
599 tags = ts.to_list()
600 # Create volume
601 print(f"Creating volume from snapshot {snapshot.id}")
602 volume = self._ec2.create_volume(
603 SnapshotId=snapshot.id,
604 AvailabilityZone=avail_zone,
605 VolumeInitializationRate=300,
606 TagSpecifications=[{"ResourceType": "volume", "Tags": tags}],
607 )
608 volume_ids.append(volume.id)
609 # Wait for the volumes to be created
610 if len(volume_ids) > 0:
611 print(f"Waiting for {len(volume_ids)} volumes to be created...")
612 self.wait_for_volumes(volume_ids, "volume_available")
613 self._wait_for_volume_tags(volume_ids)
615 def delete_volumes(self) -> None:
616 """
617 Delete all volumes associated with this cluster.
619 Args:
620 none
621 Returns:
622 none
623 """
624 self._logger.debug("method_call: delete_volumes")
625 with log_duration(self._logger, "delete_volumes"):
626 # We really only have to delete the previously detached volumes, but
627 # this will delete all managed volumes in the cluster. Its probably a
628 # good idea to do so as a matter of good housekeeping. Also, if there
629 # are any stale volumes hanging around with the same label that we are
630 # about to restore from that would cause problems because we wouldn't
631 # know which volumes to attach.
632 volumes = self.get_volumes()
633 if len(volumes) == 0:
634 print("No volumes to delete.")
635 else:
636 volume_ids = []
637 for volume in volumes:
638 print(f"Deleting volume {volume.id}")
639 volume.delete()
640 volume_ids.append(volume.id)
641 # Wait for the volumes to be deleted
642 print(f"Waiting for {len(volume_ids)} volumes to be deleted...")
643 self.wait_for_volumes(volume_ids, "volume_deleted")
645 def delete_kubernetes_volumes(self) -> None:
646 """
647 Delete all kubernetes volumes associated with this cluster.
649 Args:
650 none
651 Returns:
652 none
653 """
654 self._logger.debug("method_call: delete_kubernetes_volumes")
655 # We really only have to delete the previously detached volumes, but
656 # this will delete all managed volumes in the cluster. Its probably a
657 # good idea to do so as a matter of good housekeeping. Also, if there
658 # are any stale volumes hanging around with the same label that we are
659 # about to restore from that would cause problems because we wouldn't
660 # know which volumes to attach.
661 volumes = self.get_kubernetes_volumes()
662 if len(volumes) == 0:
663 print("No kubernetes volumes to delete.")
664 else:
665 volume_ids = []
666 for volume in volumes:
667 print(f"Deleting volume {volume.id}")
668 volume.delete()
669 volume_ids.append(volume.id)
670 # Wait for the volumes to be deleted
671 print(f"Waiting for {len(volume_ids)} volumes to be deleted...")
672 self.wait_for_volumes(volume_ids, "volume_deleted")
674 def detach_volumes(self) -> None:
675 """
676 Detach all currently attached volumes.
678 Args:
679 none
680 Returns:
681 none
682 """
683 self._logger.debug("method_call: detach_volumes")
684 # Build list of volume_ids so to pass to waiter in one big batch
685 volume_ids = []
686 # For each instance, detach all volumes
687 instances = self.get_instances()
688 for i in instances:
689 volumes = i.volumes.all()
690 for volume in volumes:
691 device = volume.attachments[0]["Device"]
692 instance_name = TagSet(i.tags).get("Name") or ""
693 shortname = instance_name.split(".")[0]
694 print(f"Detaching {device} ({volume.id}) from {shortname} ({i.id})")
695 volume.detach_from_instance(Device=device, InstanceId=i.id)
696 volume_ids.append(volume.id)
697 # Wait for volumes to detach
698 if len(volume_ids) > 0:
699 print(f"Waiting for {len(volume_ids)} volumes to be detached...")
700 self.wait_for_volumes(volume_ids, "volume_available")
702 def tag_volumes(self, tags: list[dict[str, str]]) -> None:
703 volumes = self.get_volumes()
704 for volume in volumes:
705 print(f"Tagging volume {volume.id}")
706 volume.create_tags(Tags=tags)
708 def untag_volumes(self, tags: list[dict[str, str]]) -> None:
709 volumes = self.get_volumes()
710 volume_ids = []
711 for volume in volumes:
712 print(f"Un-tagging volume {volume.id}")
713 volume_ids.append(volume.id)
714 # Might as well un-tag on one big batch since volume objects don't
715 # have direct support for un-tagging.
716 self._ec2_client.delete_tags(Resources=volume_ids, Tags=tags) # type: ignore[arg-type]
718 def wait_for_volumes(self, volume_ids: list[str], status: str) -> None:
719 """
720 Wait for an action to complete on volumes.
722 Args:
723 volume_ids - volume_ids: list of volume IDs
724 status - status to check for completion of action
725 Returns:
726 none
727 """
728 # This is a helper method - perhaps it should be static
729 self._logger.debug("method_call: wait_for_volumes")
730 waiter = self._ec2_client.get_waiter(status) # type: ignore[call-overload]
731 waiter.wait(
732 VolumeIds=volume_ids,
733 WaiterConfig={
734 "Delay": 5,
735 "MaxAttempts": 240,
736 },
737 )
739 def _wait_for_volume_tags(
740 self, volume_ids: list[str], expected_tag_key: str = "Cluster"
741 ) -> None:
742 """Poll until volumes have their tags propagated (up to 10s)."""
743 for _attempt in range(5):
744 response = self._ec2_client.describe_volumes(VolumeIds=volume_ids)
745 all_tagged = all(
746 any(t["Key"] == expected_tag_key for t in vol.get("Tags", []))
747 for vol in response["Volumes"]
748 )
749 if all_tagged:
750 return
751 time.sleep(2)
753 def get_snapshots(self, label: str | None = None) -> list[Any]:
754 """
755 Get a list of cluster snapshots.
757 Args:
758 label - label of snapshots being sought (optional)
759 Returns:
760 list of snapshots
761 """
762 self._logger.debug("method_call: get_all_snapshots")
763 fs = FilterSet(self.get_cluster_filter())
764 # Only get snapshots in a completed state. It is expected that users
765 # will only call this method after snapshots have been completed.
766 fs.add("status", "completed")
767 # Only get snapshots that were created by the snapshot manager.
768 fs.add("tag:automation_key", self.AUTOMATION_KEY)
769 # Optionally, get snapshots with a given label
770 if label is not None:
771 fs.add("tag:Label", label)
772 filters = fs.to_list()
773 snapshots = self._ec2.snapshots.filter(Filters=filters)
774 return list(snapshots.limit(self._MAX_ITEMS))
776 def create_snapshots(self, label: str) -> None:
777 """
778 Create snapshots of volumes.
780 Args:
781 - label: label to apply to each snapshot
782 Returns:
783 none
784 """
785 self._logger.debug("method_call: create_snapshots")
786 with log_duration(self._logger, "create_snapshots"):
787 # Check if any snapshots with the same label already exists. If so,
788 # delete them. Only one set of snapshots with a given label may
789 # exist at a time.
790 old_snapshots = self.get_snapshots(label)
791 if len(old_snapshots) > 0:
792 self.delete_snapshots(label)
793 snapshot_ids = []
794 # Get list of instances that need snapshots taken
795 instances = self.get_instances()
796 for i in instances:
797 instance_name = TagSet(i.tags).get("Name") or ""
798 # Get collection of volumes for the current instance
799 volumes = i.volumes.all()
800 for volume in volumes:
801 device = volume.attachments[0]["Device"]
802 # Make description
803 timestamp = datetime.datetime.now(tz=datetime.UTC)
804 date = timestamp.strftime("%Y-%m-%d")
805 timestr = timestamp.strftime("%H:%M:%S")
806 description = f"Managed snapshot taken on {date} at {timestr}"
807 # Make tags
808 ts = TagSet()
809 ts.add("Cluster", self._cluster_name_str)
810 ts.add("Device", device)
811 ts.add("Instance", instance_name)
812 ts.add("Label", label)
813 ts.add("Name", f"{instance_name} - {device}")
814 ts.add("automation_key", self.AUTOMATION_KEY)
815 tags = ts.to_list()
816 # Create shapshot
817 shortname = instance_name.split(".")[0]
818 print(f"Creating snapshot of {device} ({volume.id}) on {shortname} ({i.id})")
819 snapshot = volume.create_snapshot(
820 Description=description,
821 TagSpecifications=[{"ResourceType": "snapshot", "Tags": tags}],
822 )
823 snapshot_ids.append(snapshot.id)
824 # Wait for snapshots to complete
825 print(f"Waiting for {len(snapshot_ids)} snapshots to complete...")
826 waiter = self._ec2_client.get_waiter("snapshot_completed")
827 waiter.wait(
828 SnapshotIds=snapshot_ids,
829 WaiterConfig={
830 "Delay": 5,
831 "MaxAttempts": 720,
832 },
833 )
835 def delete_snapshots(self, label: str) -> None:
836 """
837 Delete cluster snapshots that have the given label.
839 Args:
840 label - label of snapshots to be deleted
841 Returns:
842 none
843 """
844 self._logger.debug("method_call: delete_snapshots")
845 with log_duration(self._logger, "delete_snapshots"):
846 snapshots = self.get_snapshots(label)
847 # Delete each snapshot
848 print(f"Deleting {len(snapshots)} snapshots...")
849 for snapshot in snapshots:
850 print(f"Deleting snapshot {snapshot.id}")
851 snapshot.delete()
852 # There is no waiter for snapshot deletion. Add a small delay to guard
853 # against any possible timing issues.
854 time.sleep(2)
856 def tag_snapshots(self, tags: list[dict[str, str]]) -> None:
857 snapshots = self.get_snapshots()
858 for snapshot in snapshots:
859 print(f"Tagging snapshot {snapshot.id}")
860 snapshot.create_tags(Tags=tags)
862 def untag_snapshots(self, tags: list[dict[str, str]]) -> None:
863 snapshots = self.get_snapshots()
864 snapshot_ids = []
865 for snapshot in snapshots:
866 print(f"Un-tagging snapshot {snapshot.id}")
867 snapshot_ids.append(snapshot.id)
868 # Might as well un-tag on one big batch since snapshot objects don't
869 # have direct support for un-tagging.
870 self._ec2_client.delete_tags(Resources=snapshot_ids, Tags=tags) # type: ignore[arg-type]
872 def get_subnet(self) -> Any:
873 """
874 Get subnet belonging to this cluster.
876 Args:
877 none
878 Returns:
879 subnet
880 """
881 self._logger.debug("method_call: get_subnet")
882 fs = FilterSet(self.get_cluster_filter())
883 filters = fs.to_list()
884 subnets = self._ec2.subnets.filter(Filters=filters)
885 subnet_list = list(subnets.limit(self._MAX_ITEMS))
886 # There should only be one subnet per cluster
887 if len(subnet_list) > 1:
888 raise Exception("Error (get_subnet): more than one subnet returned.")
889 return subnet_list[0]
891 def tag_subnet(self, tags: list[dict[str, str]]) -> None:
892 subnet = self.get_subnet()
893 print(f"Tagging subnet {subnet.id}")
894 subnet.create_tags(Tags=tags)
896 def untag_subnet(self, tags: list[dict[str, str]]) -> None:
897 subnet = self.get_subnet()
898 print(f"Un-tagging subnet {subnet.id}")
899 self._ec2_client.delete_tags(Resources=[subnet.id], Tags=tags) # type: ignore[arg-type]
901 def _filter_instances_by_name_regex(self, instances: list[Any], name_pattern: str) -> list[Any]:
902 """
903 Filter instances by matching their Name tag against a regex pattern.
905 Args:
906 instances: list of EC2 instance objects
907 name_pattern: regex pattern to match against Name tags
908 Returns:
909 list of filtered instances
910 """
911 if not name_pattern:
912 return instances
914 try:
915 pattern = re.compile(name_pattern)
916 except re.error as e:
917 raise ValueError(f"Invalid regex pattern '{name_pattern}': {e}") from e
919 filtered_instances = []
920 for instance in instances:
921 name_tag = None
922 if instance.tags:
923 for tag in instance.tags:
924 if tag["Key"] == "Name":
925 name_tag = tag["Value"]
926 break
928 if name_tag and pattern.search(name_tag):
929 filtered_instances.append(instance)
931 return filtered_instances
933 def stop_instances_targeted(self, name_pattern: str) -> None:
934 """
935 Stop instances that match the given Name tag regex pattern.
937 Args:
938 name_pattern: regex pattern to match instance Name tags
939 Returns:
940 none
941 """
942 self._logger.debug("method_call: stop_instances_targeted")
943 all_instances = self.get_running_instances()
944 instances = self._filter_instances_by_name_regex(all_instances, name_pattern)
946 if len(instances) == 0:
947 print(f"No running instances found matching pattern '{name_pattern}'.")
948 else:
949 # Stop instances
950 for i in instances:
951 name = TagSet(i.tags).get("Name")
952 print(f"Stopping {name} ({i.id})")
953 i.stop()
954 # Wait for all instances in one batch call with 5s polling
955 print(f"Waiting for {len(instances)} instances to stop...")
956 self._wait_instances_stopped([i.id for i in instances])
958 def start_instances_targeted(self, name_pattern: str) -> None:
959 """
960 Start instances that match the given Name tag regex pattern.
962 Args:
963 name_pattern: regex pattern to match instance Name tags
964 Returns:
965 none
966 """
967 self._logger.debug("method_call: start_instances_targeted")
968 all_instances = self.get_stopped_instances()
969 instances = self._filter_instances_by_name_regex(all_instances, name_pattern)
971 if len(instances) == 0:
972 print(f"No stopped instances found matching pattern '{name_pattern}'.")
973 else:
974 # Start instances
975 for i in instances:
976 name = TagSet(i.tags).get("Name")
977 print(f"Starting {name} ({i.id})")
978 i.start()
979 # Wait for all instances in one batch call with 5s polling
980 print(f"Waiting for {len(instances)} instances to start...")
981 self._wait_instances_running([i.id for i in instances])
983 def detach_volumes_targeted(self, name_pattern: str) -> None:
984 """
985 Detach volumes from instances that match the given Name tag regex pattern.
987 Args:
988 name_pattern: regex pattern to match instance Name tags
989 Returns:
990 none
991 """
992 self._logger.debug("method_call: detach_volumes_targeted")
993 all_instances = self.get_instances()
994 instances = self._filter_instances_by_name_regex(all_instances, name_pattern)
996 if len(instances) == 0:
997 print(f"No instances found matching pattern '{name_pattern}'.")
998 return
1000 # Build list of volume_ids so to pass to waiter in one big batch
1001 volume_ids = []
1002 # For each matching instance, detach all volumes
1003 for i in instances:
1004 volumes = i.volumes.all()
1005 for volume in volumes:
1006 device = volume.attachments[0]["Device"]
1007 instance_name = TagSet(i.tags).get("Name") or ""
1008 shortname = instance_name.split(".")[0]
1009 print(f"Detaching {device} ({volume.id}) from {shortname} ({i.id})")
1010 volume.detach_from_instance(Device=device, InstanceId=i.id)
1011 volume_ids.append(volume.id)
1012 # Wait for volumes to detach
1013 if len(volume_ids) > 0:
1014 print(f"Waiting for {len(volume_ids)} volumes to be detached...")
1015 self.wait_for_volumes(volume_ids, "volume_available")
1017 def delete_volumes_targeted(self, name_pattern: str) -> None:
1018 """
1019 Delete volumes associated with instances that match the given Name tag regex pattern.
1021 Args:
1022 name_pattern: regex pattern to match instance Name tags
1023 Returns:
1024 none
1025 """
1026 self._logger.debug("method_call: delete_volumes_targeted")
1027 # Get all volumes and filter by instance name pattern
1028 all_volumes = self.get_volumes()
1029 targeted_volumes = []
1031 for volume in all_volumes:
1032 volume_instance_tag = None
1033 if volume.tags:
1034 for tag in volume.tags:
1035 if tag["Key"] == "Instance":
1036 volume_instance_tag = tag["Value"]
1037 break
1039 if volume_instance_tag:
1040 try:
1041 pattern = re.compile(name_pattern)
1042 if pattern.search(volume_instance_tag):
1043 targeted_volumes.append(volume)
1044 except re.error as e:
1045 raise ValueError(f"Invalid regex pattern '{name_pattern}': {e}") from e
1047 if len(targeted_volumes) == 0:
1048 print(f"No volumes found for instances matching pattern '{name_pattern}'.")
1049 else:
1050 volume_ids = []
1051 for volume in targeted_volumes:
1052 print(f"Deleting volume {volume.id}")
1053 volume.delete()
1054 volume_ids.append(volume.id)
1055 # Wait for the volumes to be deleted
1056 print(f"Waiting for {len(volume_ids)} volumes to be deleted...")
1057 self.wait_for_volumes(volume_ids, "volume_deleted")
1059 def create_volumes_targeted(self, label: str, name_pattern: str) -> None:
1060 """
1061 Create new volumes from snapshots for instances matching the given Name tag regex pattern.
1063 Args:
1064 label: label of snapshots to restore
1065 name_pattern: regex pattern to match instance Name tags
1066 Returns:
1067 none
1068 """
1069 self._logger.debug("method_call: create_volumes_targeted")
1070 with log_duration(self._logger, "create_volumes_targeted"):
1071 # Validate regex pattern first
1072 try:
1073 pattern = re.compile(name_pattern)
1074 except re.error as e:
1075 raise ValueError(f"Invalid regex pattern '{name_pattern}': {e}") from e
1077 snapshots = self.get_snapshots(label)
1079 # Check if snapshot list is empty
1080 if len(snapshots) == 0:
1081 print(f"Error: No snapshots found with label '{label}'.")
1082 return
1084 # Filter snapshots by instance name pattern
1085 targeted_snapshots = []
1087 for snapshot in snapshots:
1088 ts = TagSet(snapshot.tags)
1089 instance = ts.get("Instance")
1090 if instance and pattern.search(instance):
1091 targeted_snapshots.append(snapshot)
1093 if len(targeted_snapshots) == 0:
1094 print(f"No snapshots found for instances matching pattern '{name_pattern}'.")
1095 return
1097 # Create volumes
1098 volume_ids = []
1099 for snapshot in targeted_snapshots:
1100 ts = TagSet(snapshot.tags)
1101 device = ts.get("Device")
1102 instance = ts.get("Instance")
1103 if not device:
1104 raise Exception(
1105 f"Error: create_volume: Can't find device tag for snapshot {snapshot.id}."
1106 )
1107 if not instance:
1108 raise Exception(
1109 f"Error: create_volume: Can't find instance tag for snapshot {snapshot.id}."
1110 )
1112 # Determine availability zone from one of the cluster instances
1113 avail_zone = self.get_instances()[0].placement["AvailabilityZone"]
1114 # Make tags
1115 ts = TagSet()
1116 ts.add("Cluster", self._cluster_name_str)
1117 ts.add("Device", device)
1118 ts.add("Instance", instance)
1119 ts.add("Label", label)
1120 ts.add("Name", f"{instance} - {device}")
1121 ts.add("automation_key", self.AUTOMATION_KEY)
1122 tags = ts.to_list()
1123 # Create volume
1124 print(f"Creating volume from snapshot {snapshot.id} for {instance}")
1125 volume = self._ec2.create_volume(
1126 SnapshotId=snapshot.id,
1127 AvailabilityZone=avail_zone,
1128 VolumeInitializationRate=300,
1129 TagSpecifications=[{"ResourceType": "volume", "Tags": tags}],
1130 )
1131 volume_ids.append(volume.id)
1133 # Wait for the volumes to be created
1134 if len(volume_ids) > 0:
1135 print(f"Waiting for {len(volume_ids)} volumes to be created...")
1136 self.wait_for_volumes(volume_ids, "volume_available")
1137 self._wait_for_volume_tags(volume_ids)
1139 def attach_volumes_targeted(self, label: str, name_pattern: str) -> None:
1140 """
1141 Attach volumes to instances that match the given Name tag regex pattern.
1143 Args:
1144 label: label of volumes to attach
1145 name_pattern: regex pattern to match instance Name tags
1146 Returns:
1147 none
1148 """
1149 self._logger.debug("method_call: attach_volumes_targeted")
1150 all_instances = self.get_instances()
1151 instances = self._filter_instances_by_name_regex(all_instances, name_pattern)
1153 if len(instances) == 0:
1154 print(f"No instances found matching pattern '{name_pattern}'.")
1155 return
1157 volumes = self.get_restored_volumes(label)
1159 # Note: regex pattern already validated in _filter_instances_by_name_regex
1161 volume_ids = []
1162 for i in instances:
1163 instance_name = TagSet(i.tags).get("Name") or ""
1164 for volume in volumes:
1165 ts = TagSet(volume.tags)
1166 volume_instance = ts.get("Instance")
1167 device = ts.get("Device")
1168 if volume_instance == instance_name:
1169 # Attach volume
1170 shortname = instance_name.split(".")[0]
1171 print(f"Attaching {device} ({volume.id}) to {shortname} ({i.id})")
1172 volume.attach_to_instance(Device=device, InstanceId=i.id)
1173 volume_ids.append(volume.id)
1175 if len(volume_ids) == 0:
1176 print(f"Error: No volumes to attach for instances matching pattern '{name_pattern}'.")
1177 else:
1178 # Wait for the volumes to be attached
1179 print(f"Waiting for {len(volume_ids)} volumes to be attached...")
1180 self.wait_for_volumes(volume_ids, "volume_in_use")