Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions docs/further.md
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,13 @@ inherit the run-specific tags. AWS Batch allows at most 50 tags per job;
malformed pairs (missing `=` or an empty key) raise an error at submission
time.

When tags are present the plugin also sets `propagateTags=True` on
`submit_job` so that the tags reach the underlying ECS task. Without this,
tags are visible on the Batch job object but absent from the ECS task that
incurs the actual EC2/ECS spend, making them invisible in Cost Explorer.
Depending on your account's ECS tag-authorization settings, `ecs:TagResource`
may be required on the executor role for propagation to succeed.

# Pre-existing Job Definitions

By default the plugin registers a fresh AWS Batch job definition for every job
Expand Down
7 changes: 5 additions & 2 deletions snakemake_executor_plugin_aws_batch/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,11 @@ class ExecutorSettings(ExecutorSettingsBase):
default=None,
metadata={
"help": (
"The tags that should be applied to all of the batch tasks,"
"of the form KEY=VALUE"
"The tags that should be applied to all of the batch tasks, "
"of the form KEY=VALUE. Tags are propagated to the underlying ECS "
"tasks so that cost-allocation tags reach the actual compute spend "
"in Cost Explorer. Note: ecs:TagResource may be required on the "
"executor role depending on account ECS tagging-authorization settings."
),
"env_var": False,
"required": False,
Expand Down
16 changes: 14 additions & 2 deletions snakemake_executor_plugin_aws_batch/batch_job_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -535,14 +535,19 @@ def _submit_with_preexisting_definition(self, job_definition: str) -> dict:
tags = self._build_job_tags()
if tags:
job_params["tags"] = tags
# Propagate tags from the Batch job to the underlying ECS task so that
# cost-allocation tags reach the actual compute layer. Without this,
# tags are visible on the Batch job object but silently absent from the
# ECS task that incurs the billed EC2/ECS spend.
# Note: ecs:TagResource may be required on the executor role depending
# on the account's ECS tag-authorization settings.
job_params["propagateTags"] = True

# Mirror the optional kwargs from the dynamic path. The dynamic path bakes
# the timeout into the registered definition; here we have no definition to
# register, so the timeout travels as SubmitJob's top-level `timeout`
# field, which overrides any timeout on the pre-existing definition.
# Scheduling priority applies equally to pre-existing definitions.
# NOTE: any future kwarg added to submit() (e.g. propagateTags) must also
# be mirrored here.
task_timeout = self._resolve_task_timeout()
if task_timeout is not None:
job_params["timeout"] = {"attemptDurationSeconds": task_timeout}
Expand Down Expand Up @@ -576,6 +581,13 @@ def submit(self):
tags = self._build_job_tags()
if tags:
job_params["tags"] = tags
# Propagate tags from the Batch job to the underlying ECS task so that
# cost-allocation tags reach the actual compute layer. Without this,
# tags are visible on the Batch job object but silently absent from the
# ECS task that incurs the billed EC2/ECS spend.
# Note: ecs:TagResource may be required on the executor role depending
# on the account's ECS tag-authorization settings.
job_params["propagateTags"] = True

priority = self._resolve_scheduling_priority()
if priority is not None:
Expand Down
58 changes: 58 additions & 0 deletions tests/test_batch_job_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -255,6 +255,35 @@ def test_per_rule_queue_override_passed_to_submit_job(self):
_, call_args = self._run_submit(builder)
assert call_args.kwargs["jobQueue"] == "override-queue"

def test_propagate_tags_set_when_tags_present(self):
"""propagateTags=True must be included in submit_job kwargs when tags exist."""
builder = _make_builder(tags={"Env": "prod"})
_, call_args = self._run_submit(builder)
assert _extract_propagate_tags(call_args) is True

def test_propagate_tags_set_when_only_env_var_tags(self):
"""propagateTags=True must appear even when tags come only from the env var."""
builder = _make_builder(tags=None)
with patch.dict(
os.environ, {SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR: "Team=data"}
):
_, call_args = self._run_submit(builder)
assert _extract_propagate_tags(call_args) is True

def test_propagate_tags_absent_when_no_tags(self):
"""propagateTags must not appear in submit_job kwargs when there are no tags."""
builder = _make_builder(tags=None)
env = {
k: v
for k, v in os.environ.items()
if k != SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR
}
with patch.dict(os.environ, env, clear=True):
_, call_args = self._run_submit(builder)
# Assert key absence directly: an explicit propagateTags=None would also
# satisfy `.get(...) is None` but is not what we want to send to the SDK.
assert "propagateTags" not in call_args.kwargs


def _extract_tags(call_args) -> dict | None:
"""Extract the 'tags' value from a mock call_args, or None if not present."""
Expand All @@ -265,6 +294,14 @@ def _extract_tags(call_args) -> dict | None:
return kwargs.get("tags")


def _extract_propagate_tags(call_args) -> bool | None:
"""Extract the 'propagateTags' value from a mock call_args, or None if absent."""
if call_args is None:
return None
kwargs = call_args.kwargs if hasattr(call_args, "kwargs") else call_args[1]
return kwargs.get("propagateTags")


# ---------------------------------------------------------------------------
# Tests for _get_platform_from_queue — exception handling
# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -1181,6 +1218,27 @@ def test_tags_forwarded_in_preexisting_mode(self):
builder.submit()
call_kwargs = builder.batch_client.submit_job.call_args.kwargs
assert call_kwargs["tags"] == {"Env": "prod", "Project": "fgumi"}
# propagateTags must be mirrored here just like the dynamic path so the
# tags reach the underlying ECS task (see submit()).
assert call_kwargs["propagateTags"] is True

def test_propagate_tags_absent_in_preexisting_mode_when_no_tags(self):
"""propagateTags must not appear when there are no tags in pre-existing mode."""
builder = _make_builder_with_preexisting(job_definition="my-job-def")
builder.settings.tags = None
builder.batch_client.submit_job.return_value = {
"jobName": "snakejob-test",
"jobId": "abc-123",
}
env = {
k: v
for k, v in os.environ.items()
if k != SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR
}
with patch.dict(os.environ, env, clear=True):
builder.submit()
call_kwargs = builder.batch_client.submit_job.call_args.kwargs
assert "propagateTags" not in call_kwargs

def test_container_overrides_no_gpu_when_zero(self):
"""GPU must NOT appear in resourceRequirements when gpu == 0."""
Expand Down
Loading