feat: preflight-check batch:TagResource when tags are configured

docs/further.md

@@ -250,3 +250,33 @@ 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.
+
+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
+workflow proceeds. It uses the `batch:DescribeJobQueues` /
+`batch:DescribeComputeEnvironments` permissions the executor already needs, plus
+the optional `iam:GetRole` for the job-role check.

snakemake_executor_plugin_aws_batch/__init__.py

@@ -3,11 +3,19 @@
 __email__ = "jake.vancampen7@gmail.com"
 __license__ = "MIT"
 
+import os
 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_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 (
@@ -20,6 +28,24 @@
 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"})
+
+
+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 --<executor-name>-<param-name>
@@ -158,6 +184,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 +359,202 @@ 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()
+        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.
+
+        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::<acct>:role/<path>/<name> 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}")

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)