From f95d72c37c69bf1a4b042d36bd623973d7970218 Mon Sep 17 00:00:00 2001 From: Terry Cain Date: Fri, 29 Sep 2017 23:29:36 +0100 Subject: [PATCH] Finialised create compute environment + describe environments --- moto/batch/models.py | 160 +++++++++++++++++++++++++++++---- moto/batch/responses.py | 19 +++- moto/batch/urls.py | 1 + tests/test_batch/test_batch.py | 84 +++++++++++++++-- 4 files changed, 237 insertions(+), 27 deletions(-) diff --git a/moto/batch/models.py b/moto/batch/models.py index c7def48d1..7ed75e749 100644 --- a/moto/batch/models.py +++ b/moto/batch/models.py @@ -1,13 +1,18 @@ from __future__ import unicode_literals import boto3 import re +from itertools import cycle +import six +import uuid from moto.core import BaseBackend, BaseModel from moto.iam import iam_backends from moto.ec2 import ec2_backends +from moto.ecs import ecs_backends from .exceptions import InvalidParameterValueException, InternalFailure from .utils import make_arn_for_compute_env from moto.ec2.exceptions import InvalidSubnetIdError +from moto.ec2.models import INSTANCE_TYPES as EC2_INSTANCE_TYPES from moto.iam.exceptions import IAMNotFoundException @@ -17,13 +22,22 @@ COMPUTE_ENVIRONMENT_NAME_REGEX = re.compile(r'^[A-Za-z0-9_]{1,128}$') class ComputeEnvironment(BaseModel): def __init__(self, compute_environment_name, _type, state, compute_resources, service_role, region_name): - self.compute_environment_name = compute_environment_name + self.name = compute_environment_name self.type = _type self.state = state self.compute_resources = compute_resources self.service_role = service_role self.arn = make_arn_for_compute_env(DEFAULT_ACCOUNT_ID, compute_environment_name, region_name) + self.instances = [] + self.ecs_arn = None + + def add_instance(self, instance): + self.instances.append(instance) + + def set_ecs_arn(self, arn): + self.ecs_arn = arn + class BatchBackend(BaseBackend): def __init__(self, region_name=None): @@ -48,6 +62,14 @@ class BatchBackend(BaseBackend): """ return ec2_backends[self.region_name] + @property + def ecs_backend(self): + """ + :return: ECS Backend + :rtype: moto.ecs.models.EC2ContainerServiceBackend + """ + return ecs_backends[self.region_name] + def reset(self): region_name = self.region_name self.__dict__ = {} @@ -62,6 +84,33 @@ class BatchBackend(BaseBackend): return comp_env return None + def describe_compute_environments(self, environments=None, max_results=None, next_token=None): + envs = set() + if environments is not None: + envs = set(environments) + + result = [] + for arn, environment in self._compute_environments.items(): + # Filter shortcut + if len(envs) > 0 and arn not in envs and environment.name not in envs: + continue + + json_part = { + 'computeEnvironmentArn': arn, + 'computeEnvironmentName': environment.name, + 'ecsClusterArn': environment.ecs_arn, + 'serviceRole': environment.service_role, + 'state': environment.state, + 'type': environment.type, + 'status': 'VALID' + } + if environment.type == 'MANAGED': + json_part['computeResources'] = environment.compute_resources + + result.append(json_part) + + return result + def create_compute_environment(self, compute_environment_name, _type, state, compute_resources, service_role): # Validate if COMPUTE_ENVIRONMENT_NAME_REGEX.match(compute_environment_name) is None: @@ -95,21 +144,53 @@ class BatchBackend(BaseBackend): ) self._compute_environments[new_comp_env.arn] = new_comp_env - # TODO scale out if MANAGED and we have compute instance types + # Ok by this point, everything is legit, so if its Managed then start some instances + if _type == 'MANAGED': + cpus = int(compute_resources.get('desiredvCpus', compute_resources['minvCpus'])) + instance_types = compute_resources['instanceTypes'] + needed_instance_types = self.find_min_instances_to_meet_vcpus(instance_types, cpus) + # Create instances + + # Will loop over and over so we get decent subnet coverage + subnet_cycle = cycle(compute_resources['subnets']) + + for instance_type in needed_instance_types: + reservation = self.ec2_backend.add_instances( + image_id='ami-ecs-optimised', # Todo import AMIs + count=1, + user_data=None, + security_group_names=[], + instance_type=instance_type, + region_name=self.region_name, + subnet_id=six.next(subnet_cycle), + key_name=compute_resources.get('ec2KeyPair', 'AWS_OWNED'), + security_group_ids=compute_resources['securityGroupIds'] + ) + + new_comp_env.add_instance(reservation.instances[0]) + + # Create ECS cluster + # Should be of format P2OnDemand_Batch_UUID + cluster_name = 'OnDemand_Batch_' + str(uuid.uuid4()) + ecs_cluster = self.ecs_backend.create_cluster(cluster_name) + new_comp_env.set_ecs_arn(ecs_cluster.arn) return compute_environment_name, new_comp_env.arn def _validate_compute_resources(self, cr): - if 'instanceRole' not in cr: - raise InvalidParameterValueException('computeResources must contain instanceRole') - elif self.iam_backend.get_role_by_arn(cr['instanceRole']) is None: + """ + Checks contents of sub dictionary for managed clusters + + :param cr: computeResources + :type cr: dict + """ + for param in ('instanceRole', 'maxvCpus', 'minvCpus', 'instanceTypes', 'securityGroupIds', 'subnets', 'type'): + if param not in cr: + raise InvalidParameterValueException('computeResources must contain {0}'.format(param)) + + if self.iam_backend.get_role_by_arn(cr['instanceRole']) is None: raise InvalidParameterValueException('could not find instanceRole {0}'.format(cr['instanceRole'])) - # TODO move the not in checks to a loop, or create a json schema validator class - if 'maxvCpus' not in cr: - raise InvalidParameterValueException('computeResources must contain maxVCpus') - if 'minvCpus' not in cr: - raise InvalidParameterValueException('computeResources must contain minVCpus') if cr['maxvCpus'] < 0: raise InvalidParameterValueException('maxVCpus must be positive') if cr['minvCpus'] < 0: @@ -117,22 +198,18 @@ class BatchBackend(BaseBackend): if cr['maxvCpus'] < cr['minvCpus']: raise InvalidParameterValueException('maxVCpus must be greater than minvCpus') - # TODO check instance types when that logic exists - if 'instanceTypes' not in cr: - raise InvalidParameterValueException('computeResources must contain instanceTypes') if len(cr['instanceTypes']) == 0: raise InvalidParameterValueException('At least 1 instance type must be provided') + for instance_type in cr['instanceTypes']: + if instance_type not in EC2_INSTANCE_TYPES: + raise InvalidParameterValueException('Instance type {0} does not exist'.format(instance_type)) - if 'securityGroupIds' not in cr: - raise InvalidParameterValueException('computeResources must contain securityGroupIds') for sec_id in cr['securityGroupIds']: if self.ec2_backend.get_security_group_from_id(sec_id) is None: raise InvalidParameterValueException('security group {0} does not exist'.format(sec_id)) if len(cr['securityGroupIds']) == 0: raise InvalidParameterValueException('At least 1 security group must be provided') - if 'subnets' not in cr: - raise InvalidParameterValueException('computeResources must contain subnets') for subnet_id in cr['subnets']: try: self.ec2_backend.get_subnet(subnet_id) @@ -141,14 +218,59 @@ class BatchBackend(BaseBackend): if len(cr['subnets']) == 0: raise InvalidParameterValueException('At least 1 subnet must be provided') - if 'type' not in cr: - raise InvalidParameterValueException('computeResources must contain type') if cr['type'] not in ('EC2', 'SPOT'): raise InvalidParameterValueException('computeResources.type must be either EC2 | SPOT') if cr['type'] == 'SPOT': raise InternalFailure('SPOT NOT SUPPORTED YET') + @staticmethod + def find_min_instances_to_meet_vcpus(instance_types, target): + """ + Finds the minimum needed instances to meed a vcpu target + + :param instance_types: Instance types, like ['t2.medium', 't2.small'] + :type instance_types: list of str + :param target: VCPU target + :type target: float + :return: List of instance types + :rtype: list of str + """ + # vcpus = [ (vcpus, instance_type), (vcpus, instance_type), ... ] + instance_vcpus = [] + instances = [] + + for instance_type in instance_types: + instance_vcpus.append( + (EC2_INSTANCE_TYPES[instance_type]['vcpus'], instance_type) + ) + + instance_vcpus = sorted(instance_vcpus, key=lambda item: item[0], reverse=True) + # Loop through, + # if biggest instance type smaller than target, and len(instance_types)> 1, then use biggest type + # if biggest instance type bigger than target, and len(instance_types)> 1, then remove it and move on + + # if biggest instance type bigger than target and len(instan_types) == 1 then add instance and finish + # if biggest instance type smaller than target and len(instan_types) == 1 then loop adding instances until target == 0 + # ^^ boils down to keep adding last till target vcpus is negative + # #Algorithm ;-) ... Could probably be done better with some quality lambdas + while target > 0: + current_vcpu, current_instance = instance_vcpus[0] + + if len(instance_vcpus) > 1: + if current_vcpu <= target: + target -= current_vcpu + instances.append(current_instance) + else: + # try next biggest instance + instance_vcpus.pop(0) + else: + # Were on the last instance + target -= current_vcpu + instances.append(current_instance) + + return instances + available_regions = boto3.session.Session().get_available_regions("batch") batch_backends = {region: BatchBackend(region_name=region) for region in available_regions} diff --git a/moto/batch/responses.py b/moto/batch/responses.py index 0368906f0..80aedcf70 100644 --- a/moto/batch/responses.py +++ b/moto/batch/responses.py @@ -14,11 +14,17 @@ class BatchResponse(BaseResponse): @property def batch_backend(self): + """ + :return: Batch Backend + :rtype: moto.batch.models.BatchBackend + """ return batch_backends[self.region] @property def json(self): - if not hasattr(self, '_json'): + if self.body is None: + self._json = {} + elif not hasattr(self, '_json'): self._json = json.loads(self.body) return self._json @@ -56,3 +62,14 @@ class BatchResponse(BaseResponse): } return json.dumps(result) + + # DescribeComputeEnvironments + def describecomputeenvironments(self): + compute_environments = self._get_param('computeEnvironments') + max_results = self._get_param('maxResults') # Ignored, should be int + next_token = self._get_param('nextToken') # Ignored + + envs = self.batch_backend.describe_compute_environments(compute_environments, max_results=max_results, next_token=next_token) + + result = {'computeEnvironments': envs} + return json.dumps(result) diff --git a/moto/batch/urls.py b/moto/batch/urls.py index 93f8a2f23..9ad3db06f 100644 --- a/moto/batch/urls.py +++ b/moto/batch/urls.py @@ -7,4 +7,5 @@ url_bases = [ url_paths = { '{0}/v1/createcomputeenvironment': BatchResponse.dispatch, + '{0}/v1/describecomputeenvironments': BatchResponse.dispatch, } diff --git a/tests/test_batch/test_batch.py b/tests/test_batch/test_batch.py index 3aae48e1e..aceb95804 100644 --- a/tests/test_batch/test_batch.py +++ b/tests/test_batch/test_batch.py @@ -2,7 +2,7 @@ from __future__ import unicode_literals import boto3 import sure # noqa -from moto import mock_batch, mock_iam, mock_ec2 +from moto import mock_batch, mock_iam, mock_ec2, mock_ecs DEFAULT_REGION = 'eu-central-1' @@ -11,6 +11,7 @@ DEFAULT_REGION = 'eu-central-1' def _get_clients(): return boto3.client('ec2', region_name=DEFAULT_REGION), \ boto3.client('iam', region_name=DEFAULT_REGION), \ + boto3.client('ecs', region_name=DEFAULT_REGION), \ boto3.client('batch', region_name=DEFAULT_REGION) @@ -46,10 +47,11 @@ def _setup(ec2_client, iam_client): # Yes, yes it talks to all the things @mock_ec2 +@mock_ecs @mock_iam @mock_batch -def test_create_compute_environment(): - ec2_client, iam_client, batch_client = _get_clients() +def test_create_managed_compute_environment(): + ec2_client, iam_client, ecs_client, batch_client = _get_clients() vpc_id, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) compute_name = 'test_compute_env' @@ -59,11 +61,12 @@ def test_create_compute_environment(): state='ENABLED', computeResources={ 'type': 'EC2', - 'minvCpus': 123, - 'maxvCpus': 123, - 'desiredvCpus': 123, + 'minvCpus': 5, + 'maxvCpus': 10, + 'desiredvCpus': 5, 'instanceTypes': [ - 'some_instance_type', + 't2.small', + 't2.medium' ], 'imageId': 'some_image_id', 'subnets': [ @@ -85,4 +88,71 @@ def test_create_compute_environment(): resp.should.contain('computeEnvironmentArn') resp['computeEnvironmentName'].should.equal(compute_name) + # Given a t2.medium is 2 vcpu and t2.small is 1, therefore 2 mediums and 1 small should be created + resp = ec2_client.describe_instances() + resp.should.contain('Reservations') + len(resp['Reservations']).should.equal(3) + + # Should have created 1 ECS cluster + resp = ecs_client.list_clusters() + resp.should.contain('clusterArns') + len(resp['clusterArns']).should.equal(1) + + +@mock_ec2 +@mock_ecs +@mock_iam +@mock_batch +def test_create_unmanaged_compute_environment(): + ec2_client, iam_client, ecs_client, batch_client = _get_clients() + vpc_id, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) + + compute_name = 'test_compute_env' + resp = batch_client.create_compute_environment( + computeEnvironmentName=compute_name, + type='UNMANAGED', + state='ENABLED', + serviceRole=iam_arn + ) + resp.should.contain('computeEnvironmentArn') + resp['computeEnvironmentName'].should.equal(compute_name) + + # Its unmanaged so no instances should be created + resp = ec2_client.describe_instances() + resp.should.contain('Reservations') + len(resp['Reservations']).should.equal(0) + + # Should have created 1 ECS cluster + resp = ecs_client.list_clusters() + resp.should.contain('clusterArns') + len(resp['clusterArns']).should.equal(1) + # TODO create 1000s of tests to test complex option combinations of create environment + + +@mock_ec2 +@mock_ecs +@mock_iam +@mock_batch +def test_describe_compute_environment(): + ec2_client, iam_client, ecs_client, batch_client = _get_clients() + vpc_id, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) + + compute_name = 'test_compute_env' + batch_client.create_compute_environment( + computeEnvironmentName=compute_name, + type='UNMANAGED', + state='ENABLED', + serviceRole=iam_arn + ) + + resp = batch_client.describe_compute_environments() + len(resp['computeEnvironments']).should.equal(1) + resp['computeEnvironments'][0]['computeEnvironmentName'].should.equal(compute_name) + + # Test filtering + resp = batch_client.describe_compute_environments( + computeEnvironments=['test1'] + ) + len(resp['computeEnvironments']).should.equal(0) +