From a07a8aa2b05128a22602256ead9ac11d9127ec6b Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Wed, 15 Jul 2026 19:23:43 -0400 Subject: [PATCH 1/2] feat: preflight-validate the job queue, compute environment, and job role Run a best-effort validation in __post_init__, before any job is submitted, so a definitively misconfigured setup fails fast with a clear error instead of leaving jobs stuck forever: - _queue_problems(): inspect the configured job queue and its compute environment(s); report a disabled/invalid queue or compute environment, a fatal queue status (INVALID/DELETING/DELETED), or maxvCpus=0. - _preflight_validate(): raise WorkflowError on any confirmed problem, then check the job role. - _validate_job_role(): iam:GetRole existence check; NoSuchEntity fails fast, anything else (most importantly a missing iam:GetRole permission) degrades. The check is conservative about uncertainty: a transient API error, a queue mid-update, or a missing describe/iam permission degrades to a debug message and the workflow proceeds. Only a confirmed-bad configuration blocks. Documented in docs/further.md; the workflow-level mocked test short-circuits preflight the same way it short-circuits platform detection. --- docs/further.md | 17 + .../__init__.py | 147 ++++++++ tests/test_preflight.py | 330 ++++++++++++++++++ tests/tests_mocked_api.py | 7 + 4 files changed, 501 insertions(+) create mode 100644 tests/test_preflight.py diff --git a/docs/further.md b/docs/further.md index ea11953..5f4b697 100644 --- a/docs/further.md +++ b/docs/further.md @@ -250,3 +250,20 @@ Task timeout and scheduling priority **are** still honored: `--aws-batch-task-ti resource) travel as `SubmitJob`'s top-level `timeout` and `schedulingPriorityOverride` fields, so they apply to pre-existing definitions just as they do to dynamically registered ones. + +# Preflight Validation + +At startup (before submitting any job) the executor verifies that the +configured job queue is `ENABLED` and `VALID`, that its compute environment(s) +are `ENABLED`/`VALID` with `maxvCpus > 0`, and — when `--aws-batch-job-role` is +set and `iam:GetRole` is available — that the job role exists. A confirmed +misconfiguration (a disabled/invalid queue or compute environment, `maxvCpus=0`, +or a non-existent job role) fails fast with a clear error, so you don't wait for +jobs that could never start. + +The check is deliberately conservative about *uncertainty*: a transient API +error, a queue mid-update (`status` `CREATING`/`UPDATING`), or a missing +`batch:Describe*` / `iam:GetRole` permission is treated as a warning and the +workflow proceeds. It uses the `batch:DescribeJobQueues` / +`batch:DescribeComputeEnvironments` permissions the executor already needs, plus +the optional `iam:GetRole` for the job-role check. diff --git a/snakemake_executor_plugin_aws_batch/__init__.py b/snakemake_executor_plugin_aws_batch/__init__.py index 8db8e2c..6ceacc4 100644 --- a/snakemake_executor_plugin_aws_batch/__init__.py +++ b/snakemake_executor_plugin_aws_batch/__init__.py @@ -6,6 +6,10 @@ from dataclasses import dataclass, field from pprint import pformat from typing import List, AsyncGenerator, Optional + +import boto3 +from botocore.exceptions import ClientError + from snakemake_executor_plugin_aws_batch.batch_client import BatchClient from snakemake_executor_plugin_aws_batch.batch_job_builder import BatchJobBuilder from snakemake_interface_executor_plugins.executors.base import SubmittedJobInfo @@ -20,6 +24,12 @@ from snakemake_interface_common.exceptions import WorkflowError +# Batch job-queue / compute-environment `status` values that definitively +# prevent jobs from running. CREATING/UPDATING are transient and must NOT fail +# the preflight check (a queue mid-update is recoverable), so only these abort. +FATAL_BATCH_STATUSES = frozenset({"INVALID", "DELETING", "DELETED"}) + + # Optional: # Define additional settings for your executor. # They will occur in the Snakemake CLI as --- @@ -158,6 +168,10 @@ def __post_init__(self): except Exception as e: raise WorkflowError(f"Failed to initialize AWS Batch client: {e}") from e + # Fail fast on a definitively misconfigured queue/compute environment/role + # before submitting any jobs (degrades to a no-op if state is uncertain). + self._preflight_validate() + def run_job(self, job: JobExecutorInterface): # Implement here how to run a job. # You can access the job's resources, etc. @@ -329,3 +343,136 @@ def cancel_jobs(self, active_jobs: List[SubmittedJobInfo]): # cleanup jobs for j in active_jobs: self.cleanup_job_resources(j) + + def _preflight_validate(self) -> None: + """Fail fast on a definitively misconfigured queue / compute env / role. + + Best-effort about *uncertainty*: a transient API error or a missing + describe/iam permission degrades to a warning and the workflow proceeds. + Only a confirmed-bad configuration (disabled/invalid queue or compute + environment, ``maxvCpus=0``, or a non-existent job role) raises, before + any job is submitted. + """ + problems = self._queue_problems() + if problems: + raise WorkflowError( + "AWS Batch preflight check failed — jobs would never start: " + + "; ".join(problems) + + ". Check the configured --aws-batch-job-queue and its compute " + "environment(s)." + ) + self._validate_job_role() + + def _queue_problems(self) -> Optional[List[str]]: + """Return definitive job-queue / compute-environment misconfigurations. + + Returns an empty list when everything looks healthy, a non-empty list of + problem descriptions when the queue or a compute environment is in a + state that would prevent jobs from ever starting (disabled/invalid, + ``maxvCpus=0``), or None when the state can't be determined (no queue + configured, or the queue describe itself failed — never block the + workflow on a transient failure or a missing describe permission). + + The queue and compute-environment lookups fail independently: a failure + describing the compute environment(s) (e.g. a missing + ``batch:DescribeComputeEnvironments`` permission) still returns any + confirmed queue-level problems collected so far rather than discarding + them, so a definitively disabled/invalid queue is never masked by an + unrelated failure downstream. + + Compute environments are judged collectively: AWS Batch falls back + across a queue's ``computeEnvironmentOrder``, so a compute-environment + problem is reported only when *every* attached environment is unusable — + one healthy environment is enough for jobs to run. + """ + queue_arn = getattr(self.settings, "job_queue", None) + if not queue_arn: + return None + # If we can't even describe the queue, the state is unknown — degrade to + # a no-op rather than blocking the workflow on a transient failure or a + # missing batch:DescribeJobQueues permission. + try: + queues = self.batch_client.describe_job_queues(jobQueues=[queue_arn]).get( + "jobQueues", [] + ) + except Exception as e: + self.logger.debug(f"could not check job queue: {e}") + return None + if not queues: + return ["job queue not found"] + jq = queues[0] + problems: List[str] = [] + if jq.get("state") != "ENABLED": + problems.append(f"job queue is {jq.get('state')} (not ENABLED)") + if jq.get("status") in FATAL_BATCH_STATUSES: + problems.append(f"job queue status is {jq.get('status')}") + ce_arns = [ + o.get("computeEnvironment") for o in jq.get("computeEnvironmentOrder", []) + ] + ce_arns = [c for c in ce_arns if c] + if ce_arns: + # A failure here must NOT discard the confirmed queue-level problems + # already collected above — return them instead of None. + try: + ces = self.batch_client.describe_compute_environments( + computeEnvironments=ce_arns + ).get("computeEnvironments", []) + except Exception as e: + self.logger.debug(f"could not check compute environment(s): {e}") + return problems + # AWS Batch tries a queue's compute environments in order and falls + # back to the next, so jobs are only definitively blocked when EVERY + # attached compute environment is unusable. Collect per-CE reasons + # but surface them only when none is usable — one healthy compute + # environment is enough for jobs to run. + ce_reasons: List[str] = [] + any_usable = False + for ce in ces: + name = ce.get("computeEnvironmentName", "?") + reasons: List[str] = [] + if ce.get("state") != "ENABLED": + reasons.append(f"compute environment {name} is {ce.get('state')}") + if ce.get("status") in FATAL_BATCH_STATUSES: + reasons.append( + f"compute environment {name} status is {ce.get('status')}" + ) + if (ce.get("computeResources") or {}).get("maxvCpus") == 0: + reasons.append(f"compute environment {name} has maxvCpus=0") + if reasons: + ce_reasons.extend(reasons) + else: + any_usable = True + if ces and not any_usable: + problems.append( + "no usable compute environment on the job queue: " + + "; ".join(ce_reasons) + ) + return problems + + def _validate_job_role(self) -> None: + """Verify the configured job role exists (best-effort; needs iam:GetRole). + + A confirmed-missing role (``NoSuchEntity``) fails fast; anything else + (most importantly a missing ``iam:GetRole`` permission) degrades silently + so the check never blocks a workflow on uncertainty. + """ + role_arn = getattr(self.settings, "job_role", None) + if not role_arn or "/" not in role_arn: + return + # GetRole takes the bare role name, not the IAM path: for + # arn:aws:iam:::role// the name is the final segment. + role_name = role_arn.rsplit("/", 1)[-1] + try: + boto3.client("iam", region_name=self.settings.region).get_role( + RoleName=role_name + ) + except ClientError as e: + if e.response.get("Error", {}).get("Code") == "NoSuchEntity": + raise WorkflowError( + f"Configured AWS Batch job role does not exist: {role_arn}" + ) from e + self.logger.debug( + f"could not verify job role (likely missing iam:GetRole): {e}" + ) + except Exception as e: + self.logger.debug(f"could not verify job role: {e}") diff --git a/tests/test_preflight.py b/tests/test_preflight.py new file mode 100644 index 0000000..9f6ddad --- /dev/null +++ b/tests/test_preflight.py @@ -0,0 +1,330 @@ +"""Unit tests for the executor's startup preflight validation. + +Covers ``_preflight_validate``, ``_queue_problems``, and ``_validate_job_role``: +the best-effort checks run in ``__post_init__`` that fail fast on a definitively +misconfigured job queue / compute environment / job role, but degrade to a +no-op on uncertain state (a transient API error or a missing describe/iam +permission). +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from botocore.exceptions import ClientError +from snakemake_interface_common.exceptions import WorkflowError + +from snakemake_executor_plugin_aws_batch import Executor + + +def _executor(**settings) -> Executor: + """Build a bare Executor (bypassing __post_init__) with mocks in place.""" + base = {"region": "us-east-1", "job_queue": "arn:q", "job_role": None} + base.update(settings) + ex = Executor.__new__(Executor) + ex.logger = MagicMock() + ex.settings = SimpleNamespace(**base) + ex.batch_client = MagicMock() + return ex + + +def _with_queue(ex: Executor, queue, compute_envs=None) -> Executor: + """Stub describe_job_queues / describe_compute_environments on the client.""" + ex.batch_client.describe_job_queues.return_value = { + "jobQueues": [queue] if queue else [] + } + ex.batch_client.describe_compute_environments.return_value = { + "computeEnvironments": compute_envs or [] + } + return ex + + +class TestQueueProblems: + def test_no_queue_configured_returns_none(self): + ex = _executor(job_queue=None) + assert ex._queue_problems() is None + + def test_queue_not_found_reported(self): + ex = _with_queue(_executor(), None) + assert ex._queue_problems() == ["job queue not found"] + + def test_healthy_queue_has_no_problems(self): + ex = _with_queue( + _executor(), + {"state": "ENABLED", "status": "VALID", "computeEnvironmentOrder": []}, + ) + assert ex._queue_problems() == [] + + def test_disabled_queue_reported(self): + ex = _with_queue( + _executor(), {"state": "DISABLED", "computeEnvironmentOrder": []} + ) + assert any("DISABLED" in p for p in ex._queue_problems()) + + def test_fatal_queue_status_reported(self): + ex = _with_queue( + _executor(), + {"state": "ENABLED", "status": "INVALID", "computeEnvironmentOrder": []}, + ) + assert any("INVALID" in p for p in ex._queue_problems()) + + def test_transient_queue_status_not_reported(self): + # A queue mid-update (UPDATING) is recoverable and must NOT be flagged. + ex = _with_queue( + _executor(), + {"state": "ENABLED", "status": "UPDATING", "computeEnvironmentOrder": []}, + ) + assert ex._queue_problems() == [] + + def test_compute_env_maxvcpus_zero_reported(self): + ex = _with_queue( + _executor(), + { + "state": "ENABLED", + "status": "VALID", + "computeEnvironmentOrder": [{"computeEnvironment": "ce1"}], + }, + compute_envs=[ + { + "computeEnvironmentName": "ce1", + "state": "ENABLED", + "status": "VALID", + "computeResources": {"maxvCpus": 0}, + } + ], + ) + assert any("maxvCpus=0" in p for p in ex._queue_problems()) + + def test_compute_env_disabled_state_reported(self): + ex = _with_queue( + _executor(), + { + "state": "ENABLED", + "status": "VALID", + "computeEnvironmentOrder": [{"computeEnvironment": "ce1"}], + }, + compute_envs=[ + { + "computeEnvironmentName": "ce1", + "state": "DISABLED", + "status": "VALID", + } + ], + ) + assert any("ce1 is DISABLED" in p for p in ex._queue_problems()) + + def test_compute_env_fatal_status_reported(self): + ex = _with_queue( + _executor(), + { + "state": "ENABLED", + "status": "VALID", + "computeEnvironmentOrder": [{"computeEnvironment": "ce1"}], + }, + compute_envs=[ + { + "computeEnvironmentName": "ce1", + "state": "ENABLED", + "status": "INVALID", + } + ], + ) + assert any("ce1 status is INVALID" in p for p in ex._queue_problems()) + + def test_mixed_healthy_and_unhealthy_compute_envs_not_reported(self): + # AWS Batch falls back across compute environments, so one bad CE next + # to a healthy one does not block jobs — preflight must not flag it. + ex = _with_queue( + _executor(), + { + "state": "ENABLED", + "status": "VALID", + "computeEnvironmentOrder": [ + {"computeEnvironment": "bad"}, + {"computeEnvironment": "good"}, + ], + }, + compute_envs=[ + { + "computeEnvironmentName": "bad", + "state": "DISABLED", + "status": "VALID", + }, + { + "computeEnvironmentName": "good", + "state": "ENABLED", + "status": "VALID", + "computeResources": {"maxvCpus": 16}, + }, + ], + ) + assert ex._queue_problems() == [] + + def test_all_unhealthy_compute_envs_reported(self): + # Only when EVERY compute environment is unusable is the queue blocked. + ex = _with_queue( + _executor(), + { + "state": "ENABLED", + "status": "VALID", + "computeEnvironmentOrder": [ + {"computeEnvironment": "a"}, + {"computeEnvironment": "b"}, + ], + }, + compute_envs=[ + {"computeEnvironmentName": "a", "state": "DISABLED", "status": "VALID"}, + { + "computeEnvironmentName": "b", + "state": "ENABLED", + "status": "VALID", + "computeResources": {"maxvCpus": 0}, + }, + ], + ) + problems = ex._queue_problems() + assert problems + assert any("a is DISABLED" in p for p in problems) + assert any("maxvCpus=0" in p for p in problems) + + def test_compute_env_describe_failure_preserves_queue_problems(self): + # A failure describing the compute environment(s) must NOT discard a + # confirmed queue-level problem (e.g. a DISABLED queue) collected first. + ex = _with_queue( + _executor(), + { + "state": "DISABLED", + "status": "VALID", + "computeEnvironmentOrder": [{"computeEnvironment": "ce1"}], + }, + ) + ex.batch_client.describe_compute_environments.side_effect = Exception( + "AccessDenied" + ) + assert any("DISABLED" in p for p in ex._queue_problems()) + + def test_compute_env_describe_failure_healthy_queue_returns_empty(self): + # If the queue is healthy and only the CE describe fails, degrade to an + # empty problem list (not None) — nothing confirmed-bad was found. + ex = _with_queue( + _executor(), + { + "state": "ENABLED", + "status": "VALID", + "computeEnvironmentOrder": [{"computeEnvironment": "ce1"}], + }, + ) + ex.batch_client.describe_compute_environments.side_effect = Exception( + "AccessDenied" + ) + assert ex._queue_problems() == [] + + def test_api_error_returns_none(self): + ex = _executor() + ex.batch_client.describe_job_queues.side_effect = Exception("throttled") + assert ex._queue_problems() is None + + +class TestPreflightValidate: + def test_healthy_passes_and_checks_role(self): + ex = _with_queue( + _executor(), + {"state": "ENABLED", "status": "VALID", "computeEnvironmentOrder": []}, + ) + ex._validate_job_role = MagicMock() + ex._preflight_validate() # must not raise + ex._validate_job_role.assert_called_once() + + def test_disabled_queue_raises(self): + ex = _with_queue( + _executor(), {"state": "DISABLED", "computeEnvironmentOrder": []} + ) + with pytest.raises(WorkflowError, match="DISABLED"): + ex._preflight_validate() + + def test_maxvcpus_zero_raises(self): + ex = _with_queue( + _executor(), + { + "state": "ENABLED", + "status": "VALID", + "computeEnvironmentOrder": [{"computeEnvironment": "ce1"}], + }, + compute_envs=[ + { + "computeEnvironmentName": "ce1", + "state": "ENABLED", + "status": "VALID", + "computeResources": {"maxvCpus": 0}, + } + ], + ) + with pytest.raises(WorkflowError, match="maxvCpus=0"): + ex._preflight_validate() + + def test_disabled_queue_raises_even_when_compute_env_describe_fails(self): + # Regression: a missing batch:DescribeComputeEnvironments permission must + # not mask a definitively DISABLED queue. + ex = _with_queue( + _executor(), + { + "state": "DISABLED", + "status": "VALID", + "computeEnvironmentOrder": [{"computeEnvironment": "ce1"}], + }, + ) + ex.batch_client.describe_compute_environments.side_effect = Exception( + "AccessDenied" + ) + ex._validate_job_role = MagicMock() + with pytest.raises(WorkflowError, match="DISABLED"): + ex._preflight_validate() + + def test_uncertain_state_does_not_raise(self): + # An API error -> _queue_problems returns None -> preflight proceeds. + ex = _executor() + ex.batch_client.describe_job_queues.side_effect = Exception("throttled") + ex._validate_job_role = MagicMock() + ex._preflight_validate() # must not raise + ex._validate_job_role.assert_called_once() + + def test_no_queue_configured_does_not_raise(self): + ex = _executor(job_queue=None) + ex._validate_job_role = MagicMock() + ex._preflight_validate() # must not raise + ex._validate_job_role.assert_called_once() + + +class TestValidateJobRole: + def _iam(self, **get_role_kwargs) -> MagicMock: + return MagicMock(get_role=MagicMock(**get_role_kwargs)) + + def _client_error(self, code: str) -> ClientError: + return ClientError({"Error": {"Code": code}}, "GetRole") + + def test_missing_role_raises(self): + ex = _executor(job_role="arn:aws:iam::1:role/missing") + iam = self._iam(side_effect=self._client_error("NoSuchEntity")) + with patch("boto3.client", return_value=iam): + with pytest.raises(WorkflowError, match="does not exist"): + ex._validate_job_role() + + def test_access_denied_degrades(self): + ex = _executor(job_role="arn:aws:iam::1:role/maybe") + iam = self._iam(side_effect=self._client_error("AccessDenied")) + with patch("boto3.client", return_value=iam): + ex._validate_job_role() # must not raise + + def test_existing_role_passes_with_bare_name(self): + # A pathed ARN must resolve to the final segment for GetRole. + ex = _executor(job_role="arn:aws:iam::1:role/team/path/good") + iam = self._iam(return_value={"Role": {}}) + with patch("boto3.client", return_value=iam): + ex._validate_job_role() + iam.get_role.assert_called_once_with(RoleName="good") + + def test_no_role_configured_is_noop(self): + ex = _executor(job_role=None) + with patch("boto3.client") as mocked_client: + ex._validate_job_role() # returns before creating any client + mocked_client.assert_not_called() diff --git a/tests/tests_mocked_api.py b/tests/tests_mocked_api.py index ff30856..6948525 100644 --- a/tests/tests_mocked_api.py +++ b/tests/tests_mocked_api.py @@ -14,6 +14,13 @@ class TestWorkflowsMocked(TestWorkflowsBase): "BatchJobBuilder._get_platform_from_queue", return_value="EC2", ) + @patch( + # __post_init__ runs _preflight_validate(), which issues live + # describe_job_queues / iam:GetRole calls. Without AWS credentials (as in + # CI) these raise before any job runs, so short-circuit preflight here. + "snakemake_executor_plugin_aws_batch.Executor._preflight_validate", + return_value=None, + ) @patch( "snakemake_executor_plugin_aws_batch.batch_job_builder.BatchJobBuilder.submit", return_value={"jobName": "job_id", "jobId": "job_id", "jobQueue": "job_queue"}, From bfb28d83378aa7e53c429bbfaeda262c601c4b8f Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Thu, 16 Jul 2026 12:58:29 -0400 Subject: [PATCH 2/2] feat: preflight-check batch:TagResource when tags are configured When tags are set (via --aws-batch-tags or SNAKEMAKE_AWS_BATCH_JOB_TAGS), every job is tagged at submit time, so a missing batch:TagResource surfaces as an opaque AccessDenied an hour into the run. Extend _preflight_validate with a non-destructive tag/untag round-trip on the job queue ARN so the permission gap fails fast at startup instead. Only a permissions error raises; any other error (or no tags / no queue ARN) degrades to a no-op. Adds tag_resource/untag_resource wrappers to BatchClient and documents the batch:TagResource/UntagResource needs. --- docs/further.md | 13 +++ .../__init__.py | 84 ++++++++++++++++- .../batch_client.py | 18 ++++ tests/test_preflight.py | 92 ++++++++++++++++++- 4 files changed, 205 insertions(+), 2 deletions(-) diff --git a/docs/further.md b/docs/further.md index 5f4b697..ccbffd9 100644 --- a/docs/further.md +++ b/docs/further.md @@ -261,6 +261,19 @@ misconfiguration (a disabled/invalid queue or compute environment, `maxvCpus=0`, or a non-existent job role) fails fast with a clear error, so you don't wait for jobs that could never start. +When tags are configured (via `--aws-batch-tags` or the +`SNAKEMAKE_AWS_BATCH_JOB_TAGS` environment variable), the executor also runs a +tag/untag round-trip on the job queue as a best-effort *proxy* for the +`batch:TagResource` permission — otherwise a missing tag permission would only +surface as an opaque `AccessDenied` at submission time. This probes the queue, +whereas jobs are tagged on the job and job-definition resources, so an IAM +policy that scopes `batch:TagResource` per resource cannot be fully verified +this way — treat a pass or failure as a strong hint, not a guarantee. The check +needs `batch:TagResource` and `batch:UntagResource` (and possibly +`ecs:TagResource`, depending on your account's tag-authorization settings); if +the cleanup untag is denied, the throwaway `snakemake-preflight` tag is left on +the queue. + The check is deliberately conservative about *uncertainty*: a transient API error, a queue mid-update (`status` `CREATING`/`UPDATING`), or a missing `batch:Describe*` / `iam:GetRole` permission is treated as a warning and the diff --git a/snakemake_executor_plugin_aws_batch/__init__.py b/snakemake_executor_plugin_aws_batch/__init__.py index 6ceacc4..26b66f6 100644 --- a/snakemake_executor_plugin_aws_batch/__init__.py +++ b/snakemake_executor_plugin_aws_batch/__init__.py @@ -3,6 +3,7 @@ __email__ = "jake.vancampen7@gmail.com" __license__ = "MIT" +import os from dataclasses import dataclass, field from pprint import pformat from typing import List, AsyncGenerator, Optional @@ -11,7 +12,10 @@ from botocore.exceptions import ClientError from snakemake_executor_plugin_aws_batch.batch_client import BatchClient -from snakemake_executor_plugin_aws_batch.batch_job_builder import BatchJobBuilder +from snakemake_executor_plugin_aws_batch.batch_job_builder import ( + BatchJobBuilder, + SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR, +) from snakemake_interface_executor_plugins.executors.base import SubmittedJobInfo from snakemake_interface_executor_plugins.executors.remote import RemoteExecutor from snakemake_interface_executor_plugins.settings import ( @@ -30,6 +34,18 @@ FATAL_BATCH_STATUSES = frozenset({"INVALID", "DELETING", "DELETED"}) +def _is_access_denied(error: ClientError) -> bool: + """True if a botocore ClientError is a permissions failure (vs. transient).""" + response = getattr(error, "response", {}) or {} + code = response.get("Error", {}).get("Code", "") + status = response.get("ResponseMetadata", {}).get("HTTPStatusCode") + return code in { + "AccessDenied", + "AccessDeniedException", + "UnauthorizedOperation", + } or status in (401, 403) + + # Optional: # Define additional settings for your executor. # They will occur in the Snakemake CLI as --- @@ -362,6 +378,72 @@ def _preflight_validate(self) -> None: "environment(s)." ) self._validate_job_role() + self._preflight_check_tags() + + def _preflight_check_tags(self) -> None: + """Best-effort probe of ``batch:TagResource`` when tags are configured. + + When tags are set (via ``--aws-batch-tags`` or the + ``SNAKEMAKE_AWS_BATCH_JOB_TAGS`` env var) every job is tagged at submit + time, so a missing ``batch:TagResource`` would otherwise surface as an + opaque ``AccessDenied`` an hour into the run. This runs a tag/untag + round-trip on the job queue ARN as a *proxy* for that permission and + fails fast if it is denied. + + The probe is a heuristic, not a proof: jobs are tagged on the *job* and + *job-definition* resources (via ``submit_job`` / ``register_job_definition``), + not the queue, so an IAM policy that scopes ``batch:TagResource`` per + resource-type/ARN could grant it on the queue while denying it on the + job/job-definition resources (or the reverse). Treat a pass as "the + permission is very likely present" and a denial as "very likely + missing", not a guarantee either way. + + Best-effort: only a permissions denial raises; any other error (or a + missing queue ARN / no tags configured) is a no-op. If the cleanup + untag fails, the throwaway ``snakemake-preflight`` tag is left on the + job queue (logged at debug) — so the round-trip is not guaranteed to be + fully non-destructive. + """ + if not self._tags_configured(): + return + queue_arn = getattr(self.settings, "job_queue", None) + if not queue_arn: + return + + try: + self.batch_client.tag_resource( + resourceArn=queue_arn, tags={"snakemake-preflight": "1"} + ) + # Clean up the throwaway tag; failure to untag is non-fatal. + try: + self.batch_client.untag_resource( + resourceArn=queue_arn, tagKeys=["snakemake-preflight"] + ) + except Exception as e: + self.logger.debug( + "could not remove throwaway 'snakemake-preflight' tag from " + f"the job queue (left in place): {e}" + ) + except ClientError as e: + if _is_access_denied(e): + raise WorkflowError( + "AWS Batch preflight check failed — the executor role was " + "denied batch:TagResource on the job queue, but tags are " + "configured. Job tagging at submit time would very likely " + "fail with AccessDenied. Grant batch:TagResource (and " + "ecs:TagResource if your account requires it) on the job / " + "job-definition resources, or unset the tags." + ) from e + self.logger.debug(f"could not verify batch:TagResource: {e}") + except Exception as e: + self.logger.debug(f"could not verify batch:TagResource: {e}") + + def _tags_configured(self) -> bool: + """True if any tags are set via the setting or the env var.""" + settings_tags = getattr(self.settings, "tags", None) + if isinstance(settings_tags, dict) and settings_tags: + return True + return bool(os.environ.get(SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR, "").strip()) def _queue_problems(self) -> Optional[List[str]]: """Return definitive job-queue / compute-environment misconfigurations. diff --git a/snakemake_executor_plugin_aws_batch/batch_client.py b/snakemake_executor_plugin_aws_batch/batch_client.py index 98e3fec..50bf907 100644 --- a/snakemake_executor_plugin_aws_batch/batch_client.py +++ b/snakemake_executor_plugin_aws_batch/batch_client.py @@ -91,3 +91,21 @@ def describe_compute_environments(self, **kwargs): :return: The response from the describe_compute_environments method. """ return self.client.describe_compute_environments(**kwargs) + + def tag_resource(self, **kwargs): + """ + Add tags to an AWS Batch resource. + + :param kwargs: The keyword arguments to pass to the tag_resource method. + :return: The response from the tag_resource method. + """ + return self.client.tag_resource(**kwargs) + + def untag_resource(self, **kwargs): + """ + Remove tags from an AWS Batch resource. + + :param kwargs: The keyword arguments to pass to the untag_resource method. + :return: The response from the untag_resource method. + """ + return self.client.untag_resource(**kwargs) diff --git a/tests/test_preflight.py b/tests/test_preflight.py index 9f6ddad..7c37d08 100644 --- a/tests/test_preflight.py +++ b/tests/test_preflight.py @@ -7,6 +7,7 @@ permission). """ +import os from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -14,7 +15,16 @@ from botocore.exceptions import ClientError from snakemake_interface_common.exceptions import WorkflowError -from snakemake_executor_plugin_aws_batch import Executor +from snakemake_executor_plugin_aws_batch import Executor, _is_access_denied +from snakemake_executor_plugin_aws_batch.batch_job_builder import ( + SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR, +) + + +def _env_without_tags() -> dict: + return { + k: v for k, v in os.environ.items() if k != SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR + } def _executor(**settings) -> Executor: @@ -295,6 +305,86 @@ def test_no_queue_configured_does_not_raise(self): ex._validate_job_role.assert_called_once() +class TestPreflightCheckTags: + """The tag/untag round-trip that verifies batch:TagResource when tags are set.""" + + def _access_denied(self) -> ClientError: + return ClientError({"Error": {"Code": "AccessDenied"}}, "TagResource") + + def test_no_tags_configured_skips(self): + ex = _executor() + with patch.dict(os.environ, _env_without_tags(), clear=True): + ex._preflight_check_tags() + ex.batch_client.tag_resource.assert_not_called() + + def test_no_queue_skips(self): + ex = _executor(job_queue=None, tags={"Env": "prod"}) + ex._preflight_check_tags() + ex.batch_client.tag_resource.assert_not_called() + + def test_success_tags_then_untags(self): + ex = _executor(tags={"Env": "prod"}) + ex._preflight_check_tags() # must not raise + ex.batch_client.tag_resource.assert_called_once() + ex.batch_client.untag_resource.assert_called_once() + + def test_env_var_tags_trigger_check(self): + ex = _executor() + with patch.dict( + os.environ, {SNAKEMAKE_AWS_BATCH_JOB_TAGS_ENV_VAR: "Team=data"} + ): + ex._preflight_check_tags() + ex.batch_client.tag_resource.assert_called_once() + + def test_missing_permission_raises(self): + ex = _executor(tags={"Env": "prod"}) + ex.batch_client.tag_resource.side_effect = self._access_denied() + with pytest.raises(WorkflowError, match="batch:TagResource"): + ex._preflight_check_tags() + + def test_non_permission_error_degrades(self): + ex = _executor(tags={"Env": "prod"}) + ex.batch_client.tag_resource.side_effect = ClientError( + {"Error": {"Code": "ThrottlingException"}}, "TagResource" + ) + ex._preflight_check_tags() # must not raise + + @pytest.mark.parametrize("code", ["AccessDeniedException", "UnauthorizedOperation"]) + def test_alternate_access_denied_codes_raise(self, code): + ex = _executor(tags={"Env": "prod"}) + ex.batch_client.tag_resource.side_effect = ClientError( + {"Error": {"Code": code}}, "TagResource" + ) + with pytest.raises(WorkflowError, match="batch:TagResource"): + ex._preflight_check_tags() + + def test_http_forbidden_status_is_access_denied(self): + # A permissions failure signalled only by HTTP 403 (unrecognized code). + err = ClientError( + { + "Error": {"Code": "Forbidden"}, + "ResponseMetadata": {"HTTPStatusCode": 403}, + }, + "TagResource", + ) + assert _is_access_denied(err) is True + + def test_untag_failure_is_non_fatal(self): + # Tag succeeds but the cleanup untag fails: must not raise (the stray + # snakemake-preflight tag is left behind and only logged at debug). + ex = _executor(tags={"Env": "prod"}) + ex.batch_client.untag_resource.side_effect = Exception("no UntagResource") + ex._preflight_check_tags() # must not raise + ex.batch_client.tag_resource.assert_called_once() + + def test_untag_not_called_when_tagging_denied(self): + ex = _executor(tags={"Env": "prod"}) + ex.batch_client.tag_resource.side_effect = self._access_denied() + with pytest.raises(WorkflowError, match="batch:TagResource"): + ex._preflight_check_tags() + ex.batch_client.untag_resource.assert_not_called() + + class TestValidateJobRole: def _iam(self, **get_role_kwargs) -> MagicMock: return MagicMock(get_role=MagicMock(**get_role_kwargs))