Skip to content

Commit 3b26b40

Browse files
authored
feat: support pre-existing job definitions via containerOverrides (#44)
## Summary The plugin always registered (and later deregistered) a fresh job definition per job — flexible, but it requires `batch:RegisterJobDefinition`/`DeregisterJobDefinition` IAM that security-conscious accounts refuse, churns thousands of single-use revisions, and keeps resource config out of infra review. A new optional `--aws-batch-job-definition <name|name:rev|ARN>` setting (with a per-rule `aws_batch_job_definition` resource override, mirroring `batch_queue`) switches to static mode: - No `RegisterJobDefinition` call; per-job specifics travel via `containerOverrides` (`command`, `environment`, vcpu/mem `resourceRequirements`, GPU only when >0 — same resource extraction as the dynamic path) - Submitted jobs are marked `_preexisting_job_definition` (carried through `run_job`'s `aux=dict(job_info)`), and every deregister/cleanup path skips them — we never deregister a definition we didn't register - Fail-fast `WorkflowError` when combined with `job_role` or `shared_memory_size_mb` (build-only knobs); `container_image` is documented-ignored since its default is indistinguishable from explicit - Tags flow through the same `_build_job_tags()` as the dynamic path - The optional submit kwargs the dynamic path applies are mirrored via shared resolvers: task timeout travels as `SubmitJob`'s top-level `timeout` field (the dynamic path bakes it into the definition), and scheduling priority forwards `schedulingPriorityOverride` identically - Platform (EC2/FARGATE) discovery is resolved lazily, so pre-existing mode never queries the queue — keeping its IAM surface free of `DescribeJobQueues`/`DescribeComputeEnvironments` - The default (dynamic) path is byte-identical to before ## Test Plan - [x] New tests (`TestPreExistingJobDefinition`, `TestDeregisterSkipsPreexisting`), including an end-to-end marker-survival test through the real `submit → SubmittedJobInfo(aux=…) → _interpret_job_status aux write → _deregister_job` chain, plus timeout/scheduling-priority mirroring and lazy-platform coverage - [x] Full builder suite green (98 tests); black clean; flake8 clean
1 parent 05d87f6 commit 3b26b40

4 files changed

Lines changed: 783 additions & 66 deletions

File tree

docs/further.md

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,3 +199,54 @@ coordinator job can set the variable so that all child jobs it submits
199199
inherit the run-specific tags. AWS Batch allows at most 50 tags per job;
200200
malformed pairs (missing `=` or an empty key) raise an error at submission
201201
time.
202+
203+
# Pre-existing Job Definitions
204+
205+
By default the plugin registers a fresh AWS Batch job definition for every job
206+
and deregisters it afterward. Accounts where job definitions are managed by
207+
infrastructure tooling (Terraform, CloudFormation) can opt out of this with
208+
`--aws-batch-job-definition`. When set, the plugin skips
209+
`RegisterJobDefinition`/`DeregisterJobDefinition` entirely and instead submits
210+
with the supplied definition, pushing per-job specifics (command, environment
211+
variables, vcpu/mem/gpu) through `containerOverrides`.
212+
213+
```bash
214+
snakemake --executor aws-batch \
215+
--aws-batch-job-definition my-snakemake-def:3 \
216+
...
217+
```
218+
219+
The value can be a bare name (`my-def`), a name:revision pair (`my-def:3`),
220+
or a full ARN
221+
(`arn:aws:batch:us-east-1:123456789012:job-definition/my-def:3`).
222+
223+
A rule can override the setting for a specific job with the
224+
`aws_batch_job_definition` resource (mirrors the `batch_queue` pattern):
225+
226+
```python
227+
rule align:
228+
resources:
229+
aws_batch_job_definition="gpu-enabled-def:2"
230+
...
231+
```
232+
233+
**Incompatible combinations** — the following raise a `WorkflowError` at
234+
submission time because they are only meaningful when the plugin builds the
235+
definition:
236+
237+
- `--aws-batch-job-role` (`job_role`): the job role is baked into the
238+
definition at registration time and cannot be overridden via
239+
`containerOverrides`.
240+
- The per-rule `shared_memory_size_mb` resource: `linuxParameters.sharedMemorySize`
241+
is a definition-level field.
242+
243+
The `--aws-batch-container-image` (`container_image`) setting and the per-rule
244+
`aws_batch_container_image` resource are both silently ignored in this mode — the
245+
container image is taken from the pre-existing definition.
246+
247+
Task timeout and scheduling priority **are** still honored: `--aws-batch-task-timeout`
248+
(and the per-rule `aws_batch_task_timeout` resource) and the scheduling priority
249+
(`--aws-batch-scheduling-priority` / the per-rule `aws_batch_scheduling_priority`
250+
resource) travel as `SubmitJob`'s top-level `timeout` and `schedulingPriorityOverride`
251+
fields, so they apply to pre-existing definitions just as they do to dynamically
252+
registered ones.

snakemake_executor_plugin_aws_batch/__init__.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,23 @@ class ExecutorSettings(ExecutorSettingsBase):
8989
"required": False,
9090
},
9191
)
92+
job_definition: Optional[str] = field(
93+
default=None,
94+
metadata={
95+
"help": (
96+
"Use a pre-existing AWS Batch job definition instead of registering "
97+
"one per job. Accepts a definition name (e.g. my-def), a name:revision "
98+
"pair (my-def:3), or a full ARN. When set, resource configuration "
99+
"(container image, job role, shared memory) is managed externally — "
100+
"the container_image setting is ignored. Per-job specifics (command, "
101+
"environment variables, vcpu/mem/gpu) are passed via "
102+
"containerOverrides. Cannot be combined with --aws-batch-job-role "
103+
"or the per-rule shared_memory_size_mb resource."
104+
),
105+
"env_var": False,
106+
"required": False,
107+
},
108+
)
92109

93110

94111
# Required:
@@ -277,6 +294,15 @@ def _terminate_job(self, job: SubmittedJobInfo):
277294
def _deregister_job(self, job: SubmittedJobInfo):
278295
"""deregister batch job definition"""
279296
try:
297+
# Pre-existing job definitions are not owned by this executor;
298+
# deregistering them would break other workflows using the same definition.
299+
if job.aux.get("_preexisting_job_definition"):
300+
self.logger.debug(
301+
"skipping deregistration of pre-existing Batch job definition "
302+
f"{job.aux.get('job_definition_arn')} "
303+
"(externally managed, not owned by this executor)"
304+
)
305+
return
280306
job_def_arn = job.aux.get("job_definition_arn")
281307
if job_def_arn is not None:
282308
self.logger.debug(f"de-registering Batch job definition {job_def_arn}")

0 commit comments

Comments
 (0)