Finialised create compute environment + describe environments
This commit is contained in:
Terry Cain
parent
ea10c4dfb6
commit
f95d72c37c
@ -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}
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -7,4 +7,5 @@ url_bases = [
|
||||
|
||||
url_paths = {
|
||||
'{0}/v1/createcomputeenvironment': BatchResponse.dispatch,
|
||||
'{0}/v1/describecomputeenvironments': BatchResponse.dispatch,
|
||||
}
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user