Fix #4228: support Fargate batch compute environment (#4477)

This commit is contained in:
Vincent Barbaresi 2021-10-26 14:27:24 +02:00 committed by GitHub
parent 5771dcf73b
commit 7e3db1ecac
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 66 additions and 41 deletions

View File

@ -859,7 +859,7 @@ class BatchBackend(BaseBackend):
self._compute_environments[new_comp_env.arn] = new_comp_env
# Ok by this point, everything is legit, so if its Managed then start some instances
if _type == "MANAGED":
if _type == "MANAGED" and "FARGATE" not in compute_resources["type"]:
cpus = int(
compute_resources.get("desiredvCpus", compute_resources["minvCpus"])
)
@ -902,49 +902,38 @@ class BatchBackend(BaseBackend):
:param cr: computeResources
:type cr: dict
"""
for param in (
"instanceRole",
"maxvCpus",
"minvCpus",
"instanceTypes",
"securityGroupIds",
"subnets",
"type",
):
if param not in cr:
pass # commenting out invalid check below - values may be missing (tf-compat)
# raise InvalidParameterValueException(
# "computeResources must contain {0}".format(param)
# )
for profile in self.iam_backend.get_instance_profiles():
if profile.arn == cr["instanceRole"]:
break
else:
raise InvalidParameterValueException(
"could not find instanceRole {0}".format(cr["instanceRole"])
)
if int(cr["maxvCpus"]) < 0:
raise InvalidParameterValueException("maxVCpus must be positive")
if int(cr["minvCpus"]) < 0:
raise InvalidParameterValueException("minVCpus must be positive")
if int(cr["maxvCpus"]) < int(cr["minvCpus"]):
raise InvalidParameterValueException(
"maxVCpus must be greater than minvCpus"
)
if len(cr["instanceTypes"]) == 0:
raise InvalidParameterValueException(
"At least 1 instance type must be provided"
)
for instance_type in cr["instanceTypes"]:
if instance_type == "optimal":
pass # Optimal should pick from latest of current gen
elif instance_type not in EC2_INSTANCE_TYPES:
if "FARGATE" not in cr["type"]:
# Most parameters are not applicable to jobs that are running on Fargate resources:
# non exhaustive list: minvCpus, instanceTypes, imageId, ec2KeyPair, instanceRole, tags
for profile in self.iam_backend.get_instance_profiles():
if profile.arn == cr["instanceRole"]:
break
else:
raise InvalidParameterValueException(
"Instance type {0} does not exist".format(instance_type)
"could not find instanceRole {0}".format(cr["instanceRole"])
)
if int(cr["minvCpus"]) < 0:
raise InvalidParameterValueException("minvCpus must be positive")
if int(cr["maxvCpus"]) < int(cr["minvCpus"]):
raise InvalidParameterValueException(
"maxVCpus must be greater than minvCpus"
)
if len(cr["instanceTypes"]) == 0:
raise InvalidParameterValueException(
"At least 1 instance type must be provided"
)
for instance_type in cr["instanceTypes"]:
if instance_type == "optimal":
pass # Optimal should pick from latest of current gen
elif instance_type not in EC2_INSTANCE_TYPES:
raise InvalidParameterValueException(
"Instance type {0} does not exist".format(instance_type)
)
for sec_id in cr["securityGroupIds"]:
if self.ec2_backend.get_security_group_from_id(sec_id) is None:
raise InvalidParameterValueException(
@ -965,9 +954,9 @@ class BatchBackend(BaseBackend):
if len(cr["subnets"]) == 0:
raise InvalidParameterValueException("At least 1 subnet must be provided")
if cr["type"] not in ("EC2", "SPOT"):
if cr["type"] not in {"EC2", "SPOT", "FARGATE", "FARGATE_SPOT"}:
raise InvalidParameterValueException(
"computeResources.type must be either EC2 | SPOT"
"computeResources.type must be either EC2 | SPOT | FARGATE | FARGATE_SPOT"
)
@staticmethod

View File

@ -1,4 +1,5 @@
from . import _get_clients, _setup
import pytest
import sure # noqa # pylint: disable=unused-import
from moto import mock_batch, mock_iam, mock_ec2, mock_ecs, settings
from uuid import uuid4
@ -232,3 +233,38 @@ def test_update_unmanaged_compute_environment_state():
our_envs = [e for e in all_envs if e["computeEnvironmentName"] == compute_name]
our_envs.should.have.length_of(1)
our_envs[0]["state"].should.equal("DISABLED")
@pytest.mark.parametrize("compute_env_type", ["FARGATE", "FARGATE_SPOT"])
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
def test_create_fargate_managed_compute_environment(compute_env_type):
ec2_client, iam_client, ecs_client, _, batch_client = _get_clients()
_, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client)
compute_name = str(uuid4())
resp = batch_client.create_compute_environment(
computeEnvironmentName=compute_name,
type="MANAGED",
state="ENABLED",
computeResources={
"type": compute_env_type,
"maxvCpus": 10,
"subnets": [subnet_id],
"securityGroupIds": [sg_id],
},
serviceRole=iam_arn,
)
resp.should.contain("computeEnvironmentArn")
resp["computeEnvironmentName"].should.equal(compute_name)
our_env = batch_client.describe_compute_environments(
computeEnvironments=[compute_name]
)["computeEnvironments"][0]
our_env["computeResources"]["type"].should.equal(compute_env_type)
# Should have created 1 ECS cluster
all_clusters = ecs_client.list_clusters()["clusterArns"]
all_clusters.should.contain(our_env["ecsClusterArn"])