Batch - TaskDefinition improvements

This commit is contained in:
Bert Blommers 2022-02-20 19:53:25 -01:00
parent ca7bc9273a
commit 24afea36c0
6 changed files with 281 additions and 64 deletions

View File

@ -199,7 +199,6 @@ class JobQueue(CloudFormationModel):
priority=properties["Priority"],
state=properties.get("State", "ENABLED"),
compute_env_order=compute_envs,
backend=backend,
)
arn = queue[1]
@ -218,29 +217,39 @@ class JobDefinition(CloudFormationModel):
revision=0,
retry_strategy=0,
timeout=None,
backend=None,
platform_capabilities=None,
propagate_tags=None,
):
self.name = name
self.retries = retry_strategy
self.retry_strategy = retry_strategy
self.type = _type
self.revision = revision
self._region = region_name
self.container_properties = container_properties
self.arn = None
self.status = "ACTIVE"
self.tagger = TaggingService()
self.parameters = parameters or {}
self.timeout = timeout
self.backend = backend
self.platform_capabilities = platform_capabilities
self.propagate_tags = propagate_tags
if "resourceRequirements" not in self.container_properties:
self.container_properties["resourceRequirements"] = []
if "secrets" not in self.container_properties:
self.container_properties["secrets"] = []
self._validate()
self._update_arn()
tags = self._format_tags(tags or {})
# Validate the tags before proceeding.
errmsg = self.tagger.validate_tags(tags)
errmsg = self.backend.tagger.validate_tags(tags)
if errmsg:
raise ValidationError(errmsg)
self.tagger.tag_resource(self.arn, tags)
self.backend.tagger.tag_resource(self.arn, tags)
def _format_tags(self, tags):
return [{"Key": k, "Value": v} for k, v in tags.items()]
@ -314,20 +323,24 @@ class JobDefinition(CloudFormationModel):
if vcpus <= 0:
raise ClientException("container vcpus limit must be greater than 0")
def deregister(self):
self.status = "INACTIVE"
def update(
self, parameters, _type, container_properties, retry_strategy, tags, timeout
):
if parameters is None:
parameters = self.parameters
if self.status != "INACTIVE":
if parameters is None:
parameters = self.parameters
if _type is None:
_type = self.type
if _type is None:
_type = self.type
if container_properties is None:
container_properties = self.container_properties
if container_properties is None:
container_properties = self.container_properties
if retry_strategy is None:
retry_strategy = self.retries
if retry_strategy is None:
retry_strategy = self.retry_strategy
return JobDefinition(
self.name,
@ -339,6 +352,9 @@ class JobDefinition(CloudFormationModel):
retry_strategy=retry_strategy,
tags=tags,
timeout=timeout,
backend=self.backend,
platform_capabilities=self.platform_capabilities,
propagate_tags=self.propagate_tags,
)
def describe(self):
@ -349,12 +365,13 @@ class JobDefinition(CloudFormationModel):
"revision": self.revision,
"status": self.status,
"type": self.type,
"tags": self.tagger.get_tag_dict_for_resource(self.arn),
"tags": self.backend.tagger.get_tag_dict_for_resource(self.arn),
"platformCapabilities": self.platform_capabilities,
"retryStrategy": self.retry_strategy,
"propagateTags": self.propagate_tags,
}
if self.container_properties is not None:
result["containerProperties"] = self.container_properties
if self.retries is not None and self.retries > 0:
result["retryStrategy"] = {"attempts": self.retries}
if self.timeout:
result["timeout"] = self.timeout
@ -387,6 +404,8 @@ class JobDefinition(CloudFormationModel):
retry_strategy=lowercase_first_key(properties["RetryStrategy"]),
container_properties=lowercase_first_key(properties["ContainerProperties"]),
timeout=lowercase_first_key(properties.get("timeout", {})),
platform_capabilities=None,
propagate_tags=None,
)
arn = res[1]
@ -1373,16 +1392,17 @@ class BatchBackend(BaseBackend):
retry_strategy,
container_properties,
timeout,
platform_capabilities,
propagate_tags,
):
if def_name is None:
raise ClientException("jobDefinitionName must be provided")
job_def = self.get_job_definition_by_name(def_name)
if retry_strategy is not None:
try:
retry_strategy = retry_strategy["attempts"]
except Exception:
raise ClientException("retryStrategy is malformed")
if retry_strategy is not None and "evaluateOnExit" in retry_strategy:
for strat in retry_strategy["evaluateOnExit"]:
if "action" in strat:
strat["action"] = strat["action"].lower()
if not tags:
tags = {}
if job_def is None:
@ -1395,6 +1415,9 @@ class BatchBackend(BaseBackend):
region_name=self.region_name,
retry_strategy=retry_strategy,
timeout=timeout,
backend=self,
platform_capabilities=platform_capabilities,
propagate_tags=propagate_tags,
)
else:
# Make new jobdef
@ -1413,7 +1436,7 @@ class BatchBackend(BaseBackend):
job_def = self.get_job_definition_by_name_revision(name, revision)
if job_def is not None:
del self._job_definitions[job_def.arn]
self._job_definitions[job_def.arn].deregister()
def describe_job_definitions(
self,

View File

@ -182,6 +182,8 @@ class BatchResponse(BaseResponse):
retry_strategy = self._get_param("retryStrategy")
_type = self._get_param("type")
timeout = self._get_param("timeout")
platform_capabilities = self._get_param("platformCapabilities")
propagate_tags = self._get_param("propagateTags")
try:
name, arn, revision = self.batch_backend.register_job_definition(
def_name=def_name,
@ -191,6 +193,8 @@ class BatchResponse(BaseResponse):
retry_strategy=retry_strategy,
container_properties=container_properties,
timeout=timeout,
platform_capabilities=platform_capabilities,
propagate_tags=propagate_tags,
)
except AWSError as err:
return err.response()

View File

@ -9,6 +9,7 @@ TestAccAWSAPIGatewayV2VpcLink
TestAccAWSAppsyncApiKey
TestAccAWSAppsyncGraphqlApi
TestAccAWSAvailabilityZones
TestAccAWSBatchJobDefinition
TestAccAWSBatchJobQueue
TestAccAWSBillingServiceAccount
TestAccAWSCallerIdentity

View File

@ -0,0 +1,67 @@
from . import _get_clients
import sure # noqa # pylint: disable=unused-import
from moto import mock_batch
from uuid import uuid4
container_properties = {
"image": "busybox",
"command": ["sleep", "1"],
"memory": 128,
"vcpus": 1,
}
@mock_batch
def test_list_tags_with_job_definition():
_, _, _, _, batch_client = _get_clients()
definition_name = str(uuid4())[0:6]
job_def_arn = batch_client.register_job_definition(
jobDefinitionName=definition_name,
type="container",
containerProperties=container_properties,
tags={"foo": "123", "bar": "456"},
)["jobDefinitionArn"]
my_queue = batch_client.list_tags_for_resource(resourceArn=job_def_arn)
my_queue.should.have.key("tags").equals({"foo": "123", "bar": "456"})
@mock_batch
def test_tag_job_definition():
_, _, _, _, batch_client = _get_clients()
definition_name = str(uuid4())[0:6]
job_def_arn = batch_client.register_job_definition(
jobDefinitionName=definition_name,
type="container",
containerProperties=container_properties,
)["jobDefinitionArn"]
batch_client.tag_resource(resourceArn=job_def_arn, tags={"k1": "v1", "k2": "v2"})
my_queue = batch_client.list_tags_for_resource(resourceArn=job_def_arn)
my_queue.should.have.key("tags").equals({"k1": "v1", "k2": "v2"})
@mock_batch
def test_untag_job_queue():
_, _, _, _, batch_client = _get_clients()
definition_name = str(uuid4())[0:6]
job_def_arn = batch_client.register_job_definition(
jobDefinitionName=definition_name,
type="container",
containerProperties=container_properties,
tags={"k1": "v1", "k2": "v2"},
)["jobDefinitionArn"]
batch_client.tag_resource(resourceArn=job_def_arn, tags={"k3": "v3"})
batch_client.untag_resource(resourceArn=job_def_arn, tagKeys=["k2"])
my_queue = batch_client.list_tags_for_resource(resourceArn=job_def_arn)
my_queue.should.have.key("tags").equals({"k1": "v1", "k3": "v3"})

View File

@ -1,19 +1,15 @@
from . import _get_clients, _setup
from . import _get_clients
import random
import pytest
import sure # noqa # pylint: disable=unused-import
from moto import mock_batch, mock_iam, mock_ec2, mock_ecs
from moto import mock_batch
from uuid import uuid4
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
@pytest.mark.parametrize("use_resource_reqs", [True, False])
def test_register_task_definition(use_resource_reqs):
ec2_client, iam_client, _, _, batch_client = _get_clients()
_setup(ec2_client, iam_client)
_, _, _, _, batch_client = _get_clients()
resp = register_job_def(batch_client, use_resource_reqs=use_resource_reqs)
@ -26,34 +22,87 @@ def test_register_task_definition(use_resource_reqs):
)
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
def test_register_task_definition_with_tags():
ec2_client, iam_client, _, _, batch_client = _get_clients()
_setup(ec2_client, iam_client)
@pytest.mark.parametrize("propagate_tags", [None, True, False])
def test_register_task_definition_with_tags(propagate_tags):
_, _, _, _, batch_client = _get_clients()
resp = register_job_def_with_tags(batch_client)
resp = register_job_def_with_tags(batch_client, propagate_tags=propagate_tags)
resp.should.contain("jobDefinitionArn")
resp.should.contain("jobDefinitionName")
resp.should.contain("revision")
resp = batch_client.describe_job_definitions(
jobDefinitionName=resp["jobDefinitionName"]
)
job_def = resp["jobDefinitions"][0]
if propagate_tags is None:
job_def.shouldnt.have.key("propagateTags")
else:
job_def.should.have.key("propagateTags").equals(propagate_tags)
assert resp["jobDefinitionArn"].endswith(
"{0}:{1}".format(resp["jobDefinitionName"], resp["revision"])
@mock_batch
@pytest.mark.parametrize("platform_capability", ["EC2", "FARGATE"])
def test_register_task_definition_with_platform_capability(platform_capability):
_, _, _, _, batch_client = _get_clients()
def_name = str(uuid4())[0:6]
batch_client.register_job_definition(
jobDefinitionName=def_name,
type="container",
containerProperties={
"image": "busybox",
"vcpus": 1,
"memory": 4,
"command": ["exit", "0"],
},
platformCapabilities=[platform_capability],
)
resp = batch_client.describe_job_definitions(jobDefinitionName=def_name)
resp["jobDefinitions"][0].should.have.key("platformCapabilities").equals(
[platform_capability]
)
@mock_batch
def test_register_task_definition_with_retry_strategies():
_, _, _, _, batch_client = _get_clients()
def_name = str(uuid4())[0:6]
batch_client.register_job_definition(
jobDefinitionName=def_name,
type="container",
containerProperties={
"image": "busybox",
"vcpus": 1,
"memory": 4,
"command": ["exit", "0"],
},
retryStrategy={
"attempts": 4,
"evaluateOnExit": [
{"onStatusReason": "osr", "action": "RETRY"},
{"onStatusReason": "osr2", "action": "Exit"},
],
},
)
resp = batch_client.describe_job_definitions(jobDefinitionName=def_name)
resp["jobDefinitions"][0].should.have.key("retryStrategy").equals(
{
"attempts": 4,
"evaluateOnExit": [
{"onStatusReason": "osr", "action": "retry"},
{"onStatusReason": "osr2", "action": "exit"},
],
}
)
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
@pytest.mark.parametrize("use_resource_reqs", [True, False])
def test_reregister_task_definition(use_resource_reqs):
# Reregistering task with the same name bumps the revision number
ec2_client, iam_client, _, _, batch_client = _get_clients()
_setup(ec2_client, iam_client)
_, _, _, _, batch_client = _get_clients()
job_def_name = str(uuid4())[0:6]
resp1 = register_job_def(
@ -94,14 +143,63 @@ def test_reregister_task_definition(use_resource_reqs):
resp4["jobDefinitionArn"].should_not.equal(resp3["jobDefinitionArn"])
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
def test_reregister_task_definition_should_not_reuse_parameters_from_inactive_definition():
# Reregistering task with the same name bumps the revision number
_, _, _, _, batch_client = _get_clients()
job_def_name = str(uuid4())[0:6]
# Register job definition with parameters
resp = batch_client.register_job_definition(
jobDefinitionName=job_def_name,
type="container",
containerProperties={
"image": "busybox",
"vcpus": 1,
"memory": 48,
"command": ["sleep", "0"],
},
parameters={"param1": "val1"},
)
job_def_arn = resp["jobDefinitionArn"]
definitions = batch_client.describe_job_definitions(jobDefinitionName=job_def_name)[
"jobDefinitions"
]
definitions.should.have.length_of(1)
definitions[0].should.have.key("parameters").equals({"param1": "val1"})
# Deactivate the definition
batch_client.deregister_job_definition(jobDefinition=job_def_arn)
# Second job definition does not provide any parameters
batch_client.register_job_definition(
jobDefinitionName=job_def_name,
type="container",
containerProperties={
"image": "busybox",
"vcpus": 1,
"memory": 96,
"command": ["sleep", "0"],
},
)
definitions = batch_client.describe_job_definitions(jobDefinitionName=job_def_name)[
"jobDefinitions"
]
definitions.should.have.length_of(2)
# Only the inactive definition should have the parameters
actual = [(d["revision"], d["status"], d.get("parameters")) for d in definitions]
actual.should.contain((1, "INACTIVE", {"param1": "val1"}))
actual.should.contain((2, "ACTIVE", {}))
@mock_batch
@pytest.mark.parametrize("use_resource_reqs", [True, False])
def test_delete_task_definition(use_resource_reqs):
ec2_client, iam_client, _, _, batch_client = _get_clients()
_setup(ec2_client, iam_client)
_, _, _, _, batch_client = _get_clients()
resp = register_job_def(
batch_client, definition_name=str(uuid4()), use_resource_reqs=use_resource_reqs
@ -111,17 +209,21 @@ def test_delete_task_definition(use_resource_reqs):
batch_client.deregister_job_definition(jobDefinition=resp["jobDefinitionArn"])
all_defs = batch_client.describe_job_definitions()["jobDefinitions"]
[jobdef["jobDefinitionName"] for jobdef in all_defs].shouldnt.contain(name)
[jobdef["jobDefinitionName"] for jobdef in all_defs].should.contain(name)
definitions = batch_client.describe_job_definitions(jobDefinitionName=name)[
"jobDefinitions"
]
definitions.should.have.length_of(1)
definitions[0].should.have.key("revision").equals(1)
definitions[0].should.have.key("status").equals("INACTIVE")
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
@pytest.mark.parametrize("use_resource_reqs", [True, False])
def test_delete_task_definition_by_name(use_resource_reqs):
ec2_client, iam_client, _, _, batch_client = _get_clients()
_setup(ec2_client, iam_client)
_, _, _, _, batch_client = _get_clients()
resp = register_job_def(
batch_client, definition_name=str(uuid4()), use_resource_reqs=use_resource_reqs
@ -131,17 +233,31 @@ def test_delete_task_definition_by_name(use_resource_reqs):
batch_client.deregister_job_definition(jobDefinition=f"{name}:{resp['revision']}")
all_defs = batch_client.describe_job_definitions()["jobDefinitions"]
[jobdef["jobDefinitionName"] for jobdef in all_defs].shouldnt.contain(name)
# We should still see our job definition as INACTIVE, as it is kept for 180 days
[jobdef["jobDefinitionName"] for jobdef in all_defs].should.contain(name)
# Registering the job definition again should up the revision number
register_job_def(
batch_client, definition_name=name, use_resource_reqs=use_resource_reqs
)
definitions = batch_client.describe_job_definitions(jobDefinitionName=name)[
"jobDefinitions"
]
definitions.should.have.length_of(2)
revision_status = [
{"revision": d["revision"], "status": d["status"]} for d in definitions
]
revision_status.should.contain({"revision": 1, "status": "INACTIVE"})
revision_status.should.contain({"revision": 2, "status": "ACTIVE"})
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
@pytest.mark.parametrize("use_resource_reqs", [True, False])
def test_describe_task_definition(use_resource_reqs):
ec2_client, iam_client, _, _, batch_client = _get_clients()
_setup(ec2_client, iam_client)
_, _, _, _, batch_client = _get_clients()
sleep_def_name = f"sleep10_{str(uuid4())[0:6]}"
other_name = str(uuid4())[0:6]
@ -183,6 +299,8 @@ def test_describe_task_definition(use_resource_reqs):
for job_definition in resp["jobDefinitions"]:
job_definition["status"].should.equal("ACTIVE")
job_definition.shouldnt.have.key("platformCapabilities")
job_definition.shouldnt.have.key("retryStrategy")
def register_job_def(batch_client, definition_name="sleep10", use_resource_reqs=True):
@ -212,7 +330,10 @@ def register_job_def(batch_client, definition_name="sleep10", use_resource_reqs=
)
def register_job_def_with_tags(batch_client, definition_name="sleep10"):
def register_job_def_with_tags(
batch_client, definition_name="sleep10", propagate_tags=False
):
kwargs = {} if propagate_tags is None else {"propagateTags": propagate_tags}
return batch_client.register_job_definition(
jobDefinitionName=definition_name,
type="container",
@ -223,4 +344,5 @@ def register_job_def_with_tags(batch_client, definition_name="sleep10"):
"command": ["sleep", "10"],
},
tags={"foo": "123", "bar": "456",},
**kwargs,
)