Finialised create compute environment + describe environments

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')
+        Checks contents of sub dictionary for managed clusters
-        elif self.iam_backend.get_role_by_arn(cr['instanceRole']) is None:
+
+        :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}

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)

moto/batch/urls.py

@@ -7,4 +7,5 @@ url_bases = [
 
 url_paths = {
     '{0}/v1/createcomputeenvironment': BatchResponse.dispatch,
+    '{0}/v1/describecomputeenvironments': BatchResponse.dispatch,
 }

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():
+def test_create_managed_compute_environment():
-    ec2_client, iam_client, batch_client = _get_clients()
+    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,
+            'minvCpus': 5,
-            'maxvCpus': 123,
+            'maxvCpus': 10,
-            'desiredvCpus': 123,
+            '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)
+