[develop] Add support for Capacity Blocks for ML

CHANGELOG.md

@@ -7,6 +7,7 @@ This file is used to list changes made in each version of the aws-parallelcluste
 ------
 
 **ENHANCEMENTS**
+- Add support for EC2 Capacity Blocks for ML.
 
 **CHANGES**
 - Perform job-level scaling by default for all jobs, using information in the `SLURM_RESUME_FILE`. Job-level scaling

src/aws/__init__.py

@@ -0,0 +1,10 @@
+# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
+# limitations under the License.

src/aws/common.py

@@ -0,0 +1,157 @@
+# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
+# limitations under the License.
+
+import functools
+import logging
+import time
+from enum import Enum
+
+import boto3
+from botocore.config import Config
+from botocore.exceptions import BotoCoreError, ClientError, ParamValidationError
+
+LOGGER = logging.getLogger(__name__)
+
+
+class AWSClientError(Exception):
+    """Error during execution of some AWS calls."""
+
+    class ErrorCode(Enum):
+        """Error codes for AWS ClientError."""
+
+        VALIDATION_ERROR = "ValidationError"
+        REQUEST_LIMIT_EXCEEDED = "RequestLimitExceeded"
+        THROTTLING_EXCEPTION = "ThrottlingException"
+        CONDITIONAL_CHECK_FAILED_EXCEPTION = "ConditionalCheckFailedException"
+
+        @classmethod
+        def throttling_error_codes(cls):
+            """Return a set of error codes returned when service rate limits are exceeded."""
+            return {cls.REQUEST_LIMIT_EXCEEDED.value, cls.THROTTLING_EXCEPTION.value}
+
+    def __init__(self, function_name: str, message: str, error_code: str = None):
+        super().__init__(message)
+        self.message = message
+        self.error_code = error_code
+        self.function_name = function_name
+
+
+class LimitExceededError(AWSClientError):
+    """Error caused by exceeding the limits of a downstream AWS service."""
+
+    def __init__(self, function_name: str, message: str, error_code: str = None):
+        super().__init__(function_name=function_name, message=message, error_code=error_code)
+
+
+class BadRequestError(AWSClientError):
+    """Error caused by a problem in the request."""
+
+    def __init__(self, function_name: str, message: str, error_code: str = None):
+        super().__init__(function_name=function_name, message=message, error_code=error_code)
+
+
+class AWSExceptionHandler:
+    """AWS Exception handler."""
+
+    @staticmethod
+    def handle_client_exception(func):
+        """Handle Boto3 errors, can be used as a decorator."""
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except ParamValidationError as validation_error:
+                error = BadRequestError(
+                    func.__name__,
+                    "Error validating parameter. Failed with exception: {0}".format(str(validation_error)),
+                )
+            except BotoCoreError as e:
+                error = AWSClientError(func.__name__, str(e))
+            except ClientError as e:
+                # add request id
+                message = e.response["Error"]["Message"]
+                error_code = e.response["Error"]["Code"]
+
+                if error_code in AWSClientError.ErrorCode.throttling_error_codes():
+                    error = LimitExceededError(func.__name__, message, error_code)
+                elif error_code == AWSClientError.ErrorCode.VALIDATION_ERROR:
+                    error = BadRequestError(func.__name__, message, error_code)
+                else:
+                    error = AWSClientError(func.__name__, message, error_code)
+            LOGGER.error("Encountered error when performing boto3 call in %s: %s", error.function_name, error.message)
+            raise error
+
+        return wrapper
+
+    @staticmethod
+    def retry_on_boto3_throttling(func):
+        """Retry boto3 calls on throttling, can be used as a decorator."""
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            while True:
+                try:
+                    return func(*args, **kwargs)
+                except ClientError as e:
+                    if e.response["Error"]["Code"] != "Throttling":
+                        raise
+                    LOGGER.debug("Throttling when calling %s function. Will retry in %d seconds.", func.__name__, 5)
+                    time.sleep(5)
+
+        return wrapper
+
+
+def _log_boto3_calls(params, **kwargs):
+    service = kwargs["event_name"].split(".")[-2]
+    operation = kwargs["event_name"].split(".")[-1]
+    region = kwargs["context"].get("client_region", boto3.session.Session().region_name)
+    LOGGER.info(
+        "Executing boto3 call: region=%s, service=%s, operation=%s, params=%s", region, service, operation, params
+    )
+
+
+class Boto3Client:
+    """Boto3 client Class."""
+
+    def __init__(self, client_name: str, config: Config, region: str = None):
+        region = region if region else get_region()
+        self._client = boto3.client(client_name, region_name=region, config=config if config else None)
+        self._client.meta.events.register("provide-client-params.*.*", _log_boto3_calls)
+
+    def _paginate_results(self, method, **kwargs):
+        """
+        Return a generator for a boto3 call, this allows pagination over an arbitrary number of responses.
+
+        :param method: boto3 method
+        :param kwargs: arguments to method
+        :return: generator with boto3 results
+        """
+        paginator = self._client.get_paginator(method.__name__)
+        for page in paginator.paginate(**kwargs).result_key_iters():
+            for result in page:
+                yield result
+
+
+class Boto3Resource:
+    """Boto3 resource Class."""
+
+    def __init__(self, resource_name: str):
+        self._resource = boto3.resource(resource_name)
+        self._resource.meta.client.meta.events.register("provide-client-params.*.*", _log_boto3_calls)
+
+
+def get_region():
+    """Get region used internally for all the AWS calls."""
+    region = boto3.session.Session().region_name
+    if region is None:
+        raise AWSClientError("get_region", "AWS region not configured")
+    return region

src/aws/ec2.py

@@ -0,0 +1,89 @@
+# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import List
+
+from aws.common import AWSExceptionHandler, Boto3Client
+
+
+class CapacityReservationInfo:
+    """
+    Data object wrapping the result of a describe-capacity-reservations call.
+
+    {
+        "CapacityReservationId": "cr-123456",
+        "OwnerId": "123",
+        "CapacityReservationArn": "arn:aws:ec2:us-east-2:123:capacity-reservation/cr-123456",
+        "AvailabilityZoneId": "use2-az1",
+        "InstanceType": "t3.large",
+        "InstancePlatform": "Linux/UNIX",
+        "AvailabilityZone": "eu-west-1a",
+        "Tenancy": "default",
+        "TotalInstanceCount": 1,
+        "AvailableInstanceCount": 1,
+        "EbsOptimized": false,
+        "EphemeralStorage": false,
+        "State": "active",
+        "StartDate": "2023-11-15T11:30:00+00:00",
+        "EndDate": "2023-11-16T11:30:00+00:00",  # capacity-block only
+        "EndDateType": "limited",
+        "InstanceMatchCriteria": "targeted",
+        "CreateDate": "2023-10-25T20:40:13+00:00",
+        "Tags": [
+            {
+                "Key": "aws:ec2capacityreservation:incrementalRequestedQuantity",
+                "Value": "1"
+            },
+            {
+                "Key": "aws:ec2capacityreservation:capacityReservationType",
+                "Value": "capacity-block"
+            }
+        ],
+        "CapacityAllocations": [],
+        "ReservationType": "capacity-block"  # capacity-block only
+    }
+    """
+
+    def __init__(self, capacity_reservation_data):
+        self.capacity_reservation_data = capacity_reservation_data
+
+    def capacity_reservation_id(self):
+        """Return the id of the Capacity Reservation."""
+        return self.capacity_reservation_data.get("CapacityReservationId")
+
+    def state(self):
+        """Return the state of the Capacity Reservation."""
+        return self.capacity_reservation_data.get("State")
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+
+class Ec2Client(Boto3Client):
+    """Implement EC2 Boto3 client."""
+
+    def __init__(self, config=None, region=None):
+        super().__init__("ec2", region=region, config=config)
+
+    @AWSExceptionHandler.handle_client_exception
+    def describe_capacity_reservations(self, capacity_reservation_ids: List[str]) -> List[CapacityReservationInfo]:
+        """Accept a space separated list of reservation ids. Return a list of CapacityReservationInfo."""
+        result = []
+        response = list(
+            self._paginate_results(
+                self._client.describe_capacity_reservations,
+                CapacityReservationIds=capacity_reservation_ids,
+                # ReservationType=reservation_type,  # not yet available
+            )
+        )
+        for capacity_reservation in response:
+            result.append(CapacityReservationInfo(capacity_reservation))
+
+        return result

src/common/schedulers/slurm_commands.py

@@ -63,11 +63,12 @@
 SCONTROL_OUTPUT_AWK_PARSER = (
     'awk \'BEGIN{{RS="\\n\\n" ; ORS="######\\n";}} {{print}}\' | '
     + "grep -oP '^(NodeName=\\S+)|(NodeAddr=\\S+)|(NodeHostName=\\S+)|(?<!Next)(State=\\S+)|"
-    + "(Partitions=\\S+)|(SlurmdStartTime=\\S+)|(LastBusyTime=\\S+)|(Reason=.*)|(######)'"
+    + "(Partitions=\\S+)|(SlurmdStartTime=\\S+)|(LastBusyTime=\\S+)|(ReservationName=\\S+)|(Reason=.*)|(######)'"
 )
 
 # Set default timeouts for running different slurm commands.
 # These timeouts might be needed when running on large scale
+DEFAULT_SCONTROL_COMMAND_TIMEOUT = 30
 DEFAULT_GET_INFO_COMMAND_TIMEOUT = 30
 DEFAULT_UPDATE_COMMAND_TIMEOUT = 60
 
@@ -389,15 +390,16 @@ def _parse_nodes_info(slurm_node_info: str) -> List[SlurmNode]:
     """Parse slurm node info into SlurmNode objects."""
     # [ec2-user@ip-10-0-0-58 ~]$ /opt/slurm/bin/scontrol show nodes compute-dy-c5xlarge-[1-3],compute-dy-c5xlarge-50001\
     # | awk 'BEGIN{{RS="\n\n" ; ORS="######\n";}} {{print}}' | grep -oP "^(NodeName=\S+)|(NodeAddr=\S+)
-    # |(NodeHostName=\S+)|(?<!Next)(State=\S+)|(Partitions=\S+)|(SlurmdStartTime=\S+)|(LastBusyTime=\\S+)|(Reason=.*)\
-    # |(######)"
+    # |(NodeHostName=\S+)|(?<!Next)(State=\S+)|(Partitions=\S+)|(SlurmdStartTime=\S+)|(LastBusyTime=\\S+)
+    # |(ReservationName=\S+)|(Reason=.*)|(######)"
     # NodeName=compute-dy-c5xlarge-1
     # NodeAddr=1.2.3.4
     # NodeHostName=compute-dy-c5xlarge-1
     # State=IDLE+CLOUD+POWER
     # Partitions=compute,compute2
     # SlurmdStartTime=2023-01-26T09:57:15
     # Reason=some reason
+    # ReservationName=root_1
     # ######
     # NodeName=compute-dy-c5xlarge-2
     # NodeAddr=1.2.3.4
@@ -430,6 +432,7 @@ def _parse_nodes_info(slurm_node_info: str) -> List[SlurmNode]:
         "Reason": "reason",
         "SlurmdStartTime": "slurmdstarttime",
         "LastBusyTime": "lastbusytime",
+        "ReservationName": "reservation_name",
     }
 
     date_fields = ["SlurmdStartTime", "LastBusyTime"]