diff --git a/docs/further.md b/docs/further.md index ea11953..9103fd2 100644 --- a/docs/further.md +++ b/docs/further.md @@ -210,6 +210,11 @@ infrastructure tooling (Terraform, CloudFormation) can opt out of this with with the supplied definition, pushing per-job specifics (command, environment variables, vcpu/mem/gpu) through `containerOverrides`. +Because the pre-existing definition supplies its own role, `--aws-batch-job-role` +is **not required** in this mode (it *is* required in the default register-per-job +mode). Providing it alongside `--aws-batch-job-definition` is rejected — see +*Incompatible combinations* below. + ```bash snakemake --executor aws-batch \ --aws-batch-job-definition my-snakemake-def:3 \ @@ -250,3 +255,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..7c59563 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 --- @@ -47,9 +57,19 @@ class ExecutorSettings(ExecutorSettingsBase): job_role: Optional[str] = field( default=None, metadata={ - "help": "The AWS job role ARN that is used for running the tasks", + "help": ( + "The AWS job role ARN used to run the tasks. Required in the " + "default mode (it is baked into each per-job definition the " + "plugin registers), but optional when --aws-batch-job-definition " + "is set: a pre-existing job definition carries its own role, and " + "combining the two is rejected." + ), "env_var": True, - "required": True, + # Not `required: True`: a pre-existing job definition + # (--aws-batch-job-definition) supplies its own role, so job_role + # must be omittable there. Presence is enforced conditionally in + # _preflight_validate for the default (register-per-job) mode. + "required": False, }, ) tags: Optional[dict] = field( @@ -158,6 +178,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 +353,165 @@ 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 raises, before any job is submitted: + combining ``--aws-batch-job-role`` with ``--aws-batch-job-definition``, + or omitting ``--aws-batch-job-role`` without a pre-existing job + definition (both checked first), a disabled/invalid queue or compute + environment, ``maxvCpus=0``, or a non-existent job role. + """ + job_definition = getattr(self.settings, "job_definition", None) + job_role = getattr(self.settings, "job_role", None) + + # Reject --aws-batch-job-role + --aws-batch-job-definition early, before + # _validate_job_role wastes an iam:GetRole call on a role that can never + # be used: a pre-existing job definition bakes in its own role. The + # per-job check in BatchJobBuilder must remain — a per-rule + # aws_batch_job_definition resource is not visible here at preflight time. + if job_definition and job_role: + raise WorkflowError( + "Cannot combine --aws-batch-job-role with " + "--aws-batch-job-definition: the job role is baked into a " + "pre-existing job definition and cannot be overridden per job. " + "Remove --aws-batch-job-role or omit --aws-batch-job-definition." + ) + # job_role is not `required` at the settings layer so it can be omitted + # in pre-existing-definition mode; enforce it here for the default + # register-per-job path, where it is baked into each definition. + if not job_definition and not job_role: + raise WorkflowError( + "AWS Batch requires a job role in the default mode: pass " + "--aws-batch-job-role (or set SNAKEMAKE_AWS_BATCH_JOB_ROLE) so it " + "can be registered into each per-job definition. It is only " + "optional when you use a pre-existing job definition via " + "--aws-batch-job-definition, which carries its own role." + ) + + 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..6df08da --- /dev/null +++ b/tests/test_preflight.py @@ -0,0 +1,393 @@ +"""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. + + Defaults to a valid ``job_role`` — the realistic default (register-per-job) + configuration. Tests exercising pre-existing-definition mode pass + ``job_role=None`` explicitly. + """ + base = { + "region": "us-east-1", + "job_queue": "arn:q", + "job_role": "arn:aws:iam::1:role/default", + } + 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 TestPreflightJobDefinitionJobRoleCombo: + """`job_definition` + `job_role` must fail fast, before any queue/role call. + + A pre-existing job definition bakes in its own role, so combining it with + `--aws-batch-job-role` can never work. The per-rule resource is caught later + in BatchJobBuilder; this covers the setting-level combo at preflight. + """ + + def test_both_set_raises_before_queue_or_role_checks(self): + ex = _executor( + job_definition="my-job-def", + job_role="arn:aws:iam::1:role/my-role", + ) + ex._queue_problems = MagicMock() + ex._validate_job_role = MagicMock() + with pytest.raises(WorkflowError) as excinfo: + ex._preflight_validate() + msg = str(excinfo.value) + assert "--aws-batch-job-role" in msg + assert "--aws-batch-job-definition" in msg + # Fail fast: neither the queue describe nor the iam:GetRole check runs. + ex._queue_problems.assert_not_called() + ex._validate_job_role.assert_not_called() + + def test_job_definition_only_does_not_raise(self): + # Pre-existing-definition mode: job_role omitted (it carries its own). + ex = _with_queue( + _executor(job_definition="my-job-def", job_role=None), + {"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_no_job_role_without_job_definition_raises(self): + # Default (register-per-job) mode requires a job role; the setting is no + # longer `required` at the framework layer, so preflight enforces it. + ex = _executor(job_role=None) + ex._queue_problems = MagicMock() + with pytest.raises(WorkflowError, match="requires a job role"): + ex._preflight_validate() + # Fail fast: the config error surfaces before any queue describe. + ex._queue_problems.assert_not_called() + + def test_job_role_only_preserves_existing_behaviour(self): + ex = _with_queue( + _executor(job_role="arn:aws:iam::1:role/my-role"), + {"state": "ENABLED", "status": "VALID", "computeEnvironmentOrder": []}, + ) + 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"},