snakemake-executor-plugin-aws-batch-flex 0.1.0__py3-none-any.whl

snakemake_executor_plugin_aws_batch_flex/imageinfo.py

@@ -0,0 +1,207 @@
+"""Images with an ENTRYPOINT. Batch can't set or clear an image's entrypoint (a job definition
+has no field for it), and a job's command is Docker's Cmd, which Docker appends to the
+entrypoint: an image with ENTRYPOINT ["python3"] would run `python3 /bin/sh -c <script>`.
+
+So the plugin reads each image's entrypoint from its registry, once per run, and gives its jobs
+a command that entrypoint runs: `-c <script>` for a shell, a `-c` that hands over to /bin/sh for
+Python, and `/bin/sh -c <script>` for everything else. That includes no entrypoint, wrappers such
+as tini, and entrypoint scripts, which end with `exec "$@"` (continuumio/miniconda3's,
+mambaorg/micromamba's). If the registry can't be read, jobs get `/bin/sh -c` too, as before."""
+
+import base64
+import json
+import os
+import re
+import urllib.error
+import urllib.parse
+import urllib.request
+from typing import Optional
+
+import boto3
+from snakemake_interface_common.exceptions import WorkflowError
+
+from .jobdefs import image_registry, is_ecr
+
+SHELLS = {"sh", "bash", "dash", "ash", "zsh", "ksh"}
+WRAPPERS = {"tini", "dumb-init", "catatonit", "env", "env-execute"}  # run what follows their options
+WITH_VALUE = {"-u", "--unset", "-C", "--chdir", "-e", "-p", "-r", "--rewrite"}  # env's, tini's, dumb-init's
+PYTHON = re.compile(r"(python|pypy)[0-9.]*")
+RUNS_NEXT = re.compile(r"-[A-Za-z]*c")  # -c, -lc, -ec: the shell or Python runs the next argument
+# Python's part: become /bin/sh -c <script> (the argument after it), so that Batch's stop signal
+# goes to the job and not to a Python waiting for it
+EXEC_SH = "import os,sys;os.execv('/bin/sh',['/bin/sh','-c']+sys.argv[1:])"
+MANIFESTS = ", ".join([
+    "application/vnd.docker.distribution.manifest.v2+json",
+    "application/vnd.docker.distribution.manifest.list.v2+json",
+    "application/vnd.oci.image.manifest.v1+json",
+    "application/vnd.oci.image.index.v1+json",
+])
+TIMEOUT = 10  # seconds per registry request
+
+
+def command(entrypoint, script):
+    """(Docker's Cmd that makes `entrypoint` run `script`, whether the plugin knows that
+    entrypoint). One it doesn't know gets /bin/sh -c, as an entrypoint script expects."""
+    sh = ["/bin/sh", "-c", script]
+    if not entrypoint:
+        return sh, True
+    exe, args = os.path.basename(entrypoint[0]), list(entrypoint[1:])
+    options = all(a.startswith("-") for a in args)
+    if exe in SHELLS and options:
+        return ([script] if args and RUNS_NEXT.fullmatch(args[-1]) else ["-c", script]), True
+    if PYTHON.fullmatch(exe) and options:  # e.g. ["python3"] or ["python3", "-u"]
+        return ([EXEC_SH, script] if args and RUNS_NEXT.fullmatch(args[-1]) else ["-c", EXEC_SH, script]), True
+    if exe in WRAPPERS:
+        rest = args
+        while rest and (rest[0].startswith("-") or (exe == "env" and "=" in rest[0])):
+            rest = rest[2:] if rest[0] in WITH_VALUE else rest[1:]
+        return command(rest, script) if rest else (sh, True)
+    return sh, False
+
+
+def parse_image(image: str):
+    """(registry API host, repository, tag or digest) of an image reference."""
+    name, ref = image.split("@", 1) if "@" in image else (image, "latest")
+    if "@" not in image and name.rfind(":") > name.rfind("/"):
+        name, ref = name.rsplit(":", 1)
+    first, _, rest = name.partition("/")
+    host, repo = (first, rest) if rest and ("." in first or ":" in first or first == "localhost") else ("docker.io", name)
+    if host in ("docker.io", "index.docker.io", "registry-1.docker.io"):
+        host = "registry-1.docker.io"
+        if "/" not in repo:  # Docker Hub's own images
+            repo = f"library/{repo}"
+    return host, repo, ref
+
+
+class _NoAuthElsewhere(urllib.request.HTTPRedirectHandler):
+    """Registries send blobs from other hosts (ECR from S3, with a presigned URL that refuses a
+    request that also has an Authorization header): it isn't sent there."""
+
+    def redirect_request(self, req, fp, code, msg, headers, newurl):
+        new = super().redirect_request(req, fp, code, msg, headers, newurl)
+        if new is not None and urllib.parse.urlsplit(newurl).netloc != urllib.parse.urlsplit(req.full_url).netloc:
+            new.remove_header("Authorization")
+        return new
+
+
+_OPENER = urllib.request.build_opener(_NoAuthElsewhere())
+
+
+def _get(url, headers):
+    """(status, headers, body)."""
+    try:
+        with _OPENER.open(urllib.request.Request(url, headers=headers), timeout=TIMEOUT) as r:
+            return r.status, r.headers, r.read()
+    except urllib.error.HTTPError as e:
+        return e.code, e.headers, e.read()
+
+
+class Entrypoints:
+    """Image entrypoints from their registries: anonymously, with ECR's token, or with the
+    registry-credentials secret. Each image is looked up once (a failure too)."""
+
+    def __init__(self, credentials: Optional[dict] = None, logger=None):
+        self.credentials = credentials or {}
+        self.logger = logger
+        self._found: dict = {}  # image -> entrypoint (list) or the reason it couldn't be read (str)
+        self._auth: dict = {}  # registry -> Authorization header or None
+        self._said: set = set()
+
+    def command(self, image: str, script: str, override=None, what: str = "") -> list:
+        """The job command for `image`. `override` is the batch_entrypoint resource: the
+        entrypoint as a JSON list, for an image whose registry can't be read."""
+        if override:
+            try:
+                entrypoint = json.loads(override)
+                assert isinstance(entrypoint, list) and all(isinstance(a, str) for a in entrypoint)
+            except (ValueError, AssertionError):
+                raise WorkflowError(f"{what}: batch_entrypoint must be the image's ENTRYPOINT as a JSON list, "
+                                    f"such as '[\"python3\"]' ('[]' for none), got {override!r}.") from None
+            return command(entrypoint, script)[0]
+        if image not in self._found:
+            try:
+                self._found[image] = self._lookup(image)
+            except Exception as e:  # network, permissions, an unexpected registry answer
+                self._found[image] = str(e) or type(e).__name__
+        entrypoint = self._found[image]
+        if isinstance(entrypoint, str):
+            self._say(image, "warning", f"Could not read the ENTRYPOINT of image {image} ({entrypoint}), so its "
+                      "jobs run `/bin/sh -c ...` as if it had none. If it has one that doesn't run that (such "
+                      "as [\"python3\"]), give its rules batch_entrypoint='[\"python3\"]'.")
+            return command(None, script)[0]
+        cmd, known = command(entrypoint, script)
+        if entrypoint:
+            how = "" if known else " (as an entrypoint script that runs the command it's given)"
+            self._say(image, "info", f"Image {image} has ENTRYPOINT {entrypoint}: its jobs start through it{how}.")
+        return cmd
+
+    def _say(self, image, level, text):
+        if self.logger and image not in self._said:
+            self._said.add(image)
+            getattr(self.logger, level)(text)
+
+    def _lookup(self, image: str) -> list:
+        host, repo, ref = parse_image(image)
+
+        def get(path):
+            return json.loads(self._fetch(host, repo, path, image_registry(image)))
+
+        manifest = get(f"manifests/{ref}")
+        if "manifests" in manifest:  # one image per platform: Batch runs linux/amd64
+            digest = next((m["digest"] for m in manifest["manifests"]
+                           if m.get("platform", {}).get("os") == "linux"
+                           and m.get("platform", {}).get("architecture") == "amd64"), None)
+            if not digest:
+                raise ValueError("no linux/amd64 image")
+            manifest = get(f"manifests/{digest}")
+        config = get(f"blobs/{manifest['config']['digest']}")
+        return list((config.get("config") or {}).get("Entrypoint") or [])
+
+    def _fetch(self, host, repo, path, registry) -> bytes:
+        url = f"https://{host}/v2/{repo}/{path}"
+        headers = {"Accept": MANIFESTS}
+        auth = self._basic(registry)
+        if auth:
+            headers["Authorization"] = auth
+        status, got, body = _get(url, headers)
+        if status == 401:  # Docker Hub, quay.io, ghcr.io, public.ecr.aws: a token for this repository
+            token = self._token(got.get("WWW-Authenticate", ""), repo, auth)
+            if token:
+                status, got, body = _get(url, {**headers, "Authorization": token})
+        if status != 200:
+            raise ValueError(f"{host} answered HTTP {status} for {repo}/{path}")
+        return body
+
+    def _token(self, challenge, repo, auth):
+        scheme, _, params = challenge.partition(" ")
+        fields = dict(re.findall(r'(\w+)="([^"]*)"', params))
+        if scheme.lower() != "bearer" or "realm" not in fields:
+            return None
+        query = {"scope": fields.get("scope") or f"repository:{repo}:pull"}
+        if "service" in fields:
+            query["service"] = fields["service"]
+        realm = fields["realm"]
+        status, _, body = _get(f"{realm}{'&' if '?' in realm else '?'}{urllib.parse.urlencode(query)}",
+                               {"Authorization": auth} if auth else {})
+        doc = json.loads(body) if status == 200 else {}
+        token = doc.get("token") or doc.get("access_token")
+        return f"Bearer {token}" if token else None
+
+    def _basic(self, registry):
+        """Authorization for a private registry: ECR's token, or the registry-credentials secret."""
+        if registry not in self._auth:
+            header = None
+            if is_ecr(registry):
+                account, region = registry.split(".")[0], registry.split(".")[3]
+                data = boto3.client("ecr", region_name=region).get_authorization_token(
+                    registryIds=[account])["authorizationData"][0]
+                header = "Basic " + data["authorizationToken"]
+            elif registry in self.credentials:
+                arn = self.credentials[registry]
+                raw = boto3.client("secretsmanager", region_name=arn.split(":")[3]).get_secret_value(
+                    SecretId=arn)["SecretString"]
+                secret = json.loads(raw)  # ECS's format: {"username": ..., "password": ...}
+                pair = f"{secret['username']}:{secret['password']}"
+                header = "Basic " + base64.b64encode(pair.encode()).decode()
+            self._auth[registry] = header
+        return self._auth[registry]

snakemake_executor_plugin_aws_batch_flex/jobdefs.py

@@ -0,0 +1,221 @@
+"""Job definitions, registered on demand: one per (image, platform), since Batch can't
+override the image at submit time. The name includes a hash of the definition, so runs reuse
+existing ones."""
+
+import hashlib
+import json
+import re
+from dataclasses import dataclass
+from typing import Optional
+
+from botocore.exceptions import ClientError
+from snakemake_interface_common.exceptions import WorkflowError
+
+from .tags import MANAGED_BY
+
+DEFAULT_LOG_GROUP = "/aws/batch/job"
+FARGATE_PLATFORM = "1.4.0"
+
+
+def log_group_of(container: dict, default: str = DEFAULT_LOG_GROUP) -> str:
+    """The CloudWatch log group of a described job's container: its job definition's, or default."""
+    return ((container.get("logConfiguration") or {}).get("options") or {}).get("awslogs-group") or default
+
+
+def normalize_image(url: str) -> str:
+    if url.startswith("docker://"):
+        return url[len("docker://") :]
+    if re.match(r"^[a-z]+://", url):
+        raise WorkflowError(
+            f"Unsupported container URL '{url}': AWS Batch only runs Docker/OCI images from a "
+            "registry such as Docker Hub, quay.io, ghcr.io or ECR (e.g. "
+            "'docker://quay.io/biocontainers/samtools:1.21--h50ea8bc_0'). Singularity/Apptainer "
+            "images (.sif, oras://, shub://) can't be used."
+        )
+    return url
+
+
+def mutable_tag(image: str) -> bool:
+    """True for an image without a version: no tag, or :latest, and no digest."""
+    if "@" in image:
+        return False
+    tag = image.rsplit("/", 1)[-1].partition(":")[2]
+    return tag in ("", "latest")
+
+
+def image_registry(image: str) -> str:
+    """Registry host of an image reference; Docker Hub images (no host part) give docker.io."""
+    first, _, rest = image.partition("/")
+    if rest and ("." in first or ":" in first or first == "localhost"):
+        return "docker.io" if first in ("index.docker.io", "registry-1.docker.io") else first
+    return "docker.io"
+
+
+def is_ecr(registry: str) -> bool:
+    return bool(re.match(r"^\d{12}\.dkr\.ecr(-fips)?\.[a-z0-9-]+\.amazonaws\.com(\.cn)?$", registry))
+
+
+class JobDefinitions:
+    def __init__(
+        self,
+        batch_client,
+        logger,
+        *,
+        prefix: str,
+        platform: str,
+        job_role_arn: Optional[str],
+        execution_role_arn: Optional[str],
+        log_group: str,
+        assign_public_ip: str,
+        ephemeral_storage_gib: Optional[int],
+        tags: dict,
+        mounts: Optional[list] = None,
+        registry_credentials: Optional[dict] = None,
+    ):
+        self.batch = batch_client
+        self.logger = logger
+        self.prefix = re.sub(r"[^A-Za-z0-9_-]", "-", prefix)[:80]
+        self.platform = platform
+        self.job_role_arn = job_role_arn
+        self.execution_role_arn = execution_role_arn
+        self.log_group = log_group
+        self.assign_public_ip = assign_public_ip
+        self.ephemeral_storage_gib = ephemeral_storage_gib
+        self.tags = tags
+        # [(host_path, container_path, read_only)] bind mounts (EC2 only)
+        self.mounts = list(mounts or [])
+        # registry host -> Secrets Manager secret ARN (username/password) for private images
+        self.registry_credentials = dict(registry_credentials or {})
+        self._cache: dict[tuple, str] = {}
+
+    def container_properties(
+        self,
+        image: str,
+        platform: Optional[str] = None,
+        shm_mb: Optional[int] = None,
+        disk_gib: Optional[int] = None,
+    ) -> dict:
+        platform = platform or self.platform
+        # vCPU/memory here are placeholders; they are overridden on every submit.
+        props: dict = {
+            "image": image,
+            "resourceRequirements": [
+                {"type": "VCPU", "value": "1"},
+                {"type": "MEMORY", "value": "2048"},
+            ],
+        }
+        if self.job_role_arn:
+            props["jobRoleArn"] = self.job_role_arn
+        secret = self.registry_credentials.get(image_registry(image))
+        if secret:
+            props["repositoryCredentials"] = {"credentialsParameter": secret}
+        # ECS reads the registry secret with the execution role, on EC2 too
+        if self.execution_role_arn and (platform == "fargate" or secret):
+            props["executionRoleArn"] = self.execution_role_arn
+        if self.mounts and platform == "ec2":  # Fargate has no host paths
+            props["volumes"] = [
+                {"name": f"vol{i}", "host": {"sourcePath": host}}
+                for i, (host, _c, _ro) in enumerate(self.mounts)
+            ]
+            props["mountPoints"] = [
+                {"sourceVolume": f"vol{i}", "containerPath": cont, "readOnly": ro}
+                for i, (_h, cont, ro) in enumerate(self.mounts)
+            ]
+        if self.log_group != DEFAULT_LOG_GROUP:
+            options = {"awslogs-group": self.log_group}
+            region = getattr(getattr(self.batch, "meta", None), "region_name", None)
+            if region:  # Fargate requires it
+                options["awslogs-region"] = region
+            props["logConfiguration"] = {"logDriver": "awslogs", "options": options}
+        if shm_mb and platform == "ec2":  # /dev/shm size; Fargate doesn't support it
+            props["linuxParameters"] = {"sharedMemorySize": int(shm_mb)}
+        if platform == "fargate":
+            props["networkConfiguration"] = {"assignPublicIp": self.assign_public_ip}
+            # a fixed version, not LATEST, so jobs don't change platform from one run to the next
+            props["fargatePlatformConfiguration"] = {"platformVersion": FARGATE_PLATFORM}
+            disk = max(disk_gib or 0, self.ephemeral_storage_gib or 0)
+            if disk:
+                props["ephemeralStorage"] = {"sizeInGiB": disk}
+        return props
+
+    def ensure(
+        self,
+        image: str,
+        platform: Optional[str] = None,
+        shm_mb: Optional[int] = None,
+        disk_gib: Optional[int] = None,
+    ) -> str:
+        """Return the job definition name for `image` on `platform` (and /dev/shm size,
+        Fargate disk size), registering it if needed."""
+        platform = platform or self.platform
+        key = (image, platform, shm_mb, disk_gib)
+        if key in self._cache:
+            return self._cache[key]
+        props = self.container_properties(image, platform, shm_mb, disk_gib)
+        digest = hashlib.sha256(
+            json.dumps([props, platform], sort_keys=True).encode()
+        ).hexdigest()[:12]
+        name = f"{self.prefix}-{digest}"
+        try:
+            existing = self.batch.describe_job_definitions(
+                jobDefinitionName=name, status="ACTIVE"
+            )["jobDefinitions"]
+            if not existing:
+                self.logger.info(f"Registering AWS Batch job definition {name} for {image} ({platform})")
+                kwargs = dict(
+                    jobDefinitionName=name,
+                    type="container",
+                    containerProperties=props,
+                    platformCapabilities=[platform.upper()],
+                )
+                # ManagedBy lets scripts/cleanup-jobdefs.sh tell the plugin's definitions apart
+                kwargs["tags"] = {**self.tags, **MANAGED_BY}
+                self.batch.register_job_definition(**kwargs)
+        except ClientError as e:
+            raise WorkflowError(f"Could not register AWS Batch job definition: {e}") from e
+        self._cache[key] = name
+        return name
+
+
+@dataclass
+class ExistingJobDefinition:
+    """What a job definition made outside the plugin fixes for its jobs."""
+    arn: str
+    image: str
+    platform: str
+    container_paths: set
+    job_role: Optional[str]
+
+
+def describe_existing(batch, ref: str) -> ExistingJobDefinition:
+    """`ref` is a name (latest ACTIVE revision), name:revision, or an ARN."""
+    try:
+        if ref.startswith("arn:") or ":" in ref:
+            found = batch.describe_job_definitions(jobDefinitions=[ref])["jobDefinitions"]
+        else:
+            found, token = [], None
+            while True:
+                kw = dict(jobDefinitionName=ref, status="ACTIVE")
+                if token:
+                    kw["nextToken"] = token
+                resp = batch.describe_job_definitions(**kw)
+                found += resp["jobDefinitions"]
+                token = resp.get("nextToken")
+                if not token:
+                    break
+    except ClientError as e:
+        raise WorkflowError(f"Could not look up AWS Batch job definition '{ref}': {e}") from e
+    found = [d for d in found if d.get("status", "ACTIVE") == "ACTIVE"]
+    if not found:
+        raise WorkflowError(f"AWS Batch job definition '{ref}' not found (or not ACTIVE).")
+    d = max(found, key=lambda d: d.get("revision", 0))
+    props = d.get("containerProperties")
+    if d.get("type", "container") != "container" or not props:
+        raise WorkflowError(f"Job definition '{ref}' is not a single-container job definition.")
+    return ExistingJobDefinition(
+        arn=d["jobDefinitionArn"],
+        image=props.get("image", ""),
+        platform="fargate" if "FARGATE" in (d.get("platformCapabilities") or ["EC2"]) else "ec2",
+        container_paths={m.get("containerPath") for m in props.get("mountPoints") or []},
+        job_role=props.get("jobRoleArn"),
+    )

snakemake_executor_plugin_aws_batch_flex/jobfixes.py

@@ -0,0 +1,65 @@
+"""Fixes the job's launcher applies before Snakemake starts in the job (it loads this file with
+condacache.py and jobmonitor.py; job_tools_source). Only the standard library."""
+
+import os
+import re
+import sys
+
+MAX_LIBRARIES = 10  # missing system libraries to load before giving up
+
+
+def _conda_root():
+    """The AMI's conda mounted into the job, found by its condabin on PATH (bootstrap.conda_snippet)."""
+    for d in os.environ.get("PATH", "").split(os.pathsep):
+        if d.endswith(os.sep + "condabin"):
+            return os.path.dirname(d)
+    return None
+
+
+def _load(root, lib):
+    """Load the AMI conda's copy of a library into this process (for every later lookup)."""
+    import ctypes
+
+    path = os.path.join(root, "lib", lib)
+    if not os.path.exists(path):
+        return False
+    try:
+        ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
+        return True
+    except OSError:
+        return False
+
+
+def make_numpy_loadable():
+    """numpy (and pandas) in Snakemake's env need libraries of the system, such as libstdc++.so.6
+    and libz.so.1, which many images don't have where Linux looks (bioconda's keep theirs inside
+    their conda env). A Snakefile that imports pandas then fails in every job, and Snakemake
+    itself tries pandas in every script: job.
+
+    The AMI's conda, mounted into the job, has conda-forge's copies (made to work on old and new
+    systems): load the missing ones into this process, so numpy finds them already loaded. Only
+    Snakemake's own process gets them, not the rule's commands. If numpy still can't load
+    (no conda mounted), numpy and pandas look not installed, which Snakemake shrugs off, and the
+    log says why."""
+    root = _conda_root()
+    if root:
+        _load(root, "libgcc_s.so.1")  # what libstdc++ itself needs
+        _load(root, "libstdc++.so.6")
+    for _ in range(MAX_LIBRARIES):
+        try:
+            import numpy  # noqa: F401
+            return
+        except ModuleNotFoundError:
+            return
+        except ImportError as e:
+            error = e
+            missing = re.search(r"(\S+\.so[.\d]*): cannot open shared object file", str(e))
+            if not (root and missing and _load(root, os.path.basename(missing.group(1)))):
+                break
+            # what's left of the failed import goes, so the next try starts afresh
+            for name in [m for m in sys.modules if m == "numpy" or m.startswith("numpy.")]:
+                del sys.modules[name]
+    sys.modules["numpy"] = sys.modules["pandas"] = None
+    why = (str(error).strip().splitlines() or [""])[-1][:200]
+    print(f"[snakemake-batch] numpy and pandas of Snakemake's env can't load in this image, so this job "
+          f"runs without them: {why}", file=sys.stderr, flush=True)

snakemake_executor_plugin_aws_batch_flex/jobmonitor.py

@@ -0,0 +1,146 @@
+"""Measure a job while it runs, for the table at the end of the run.
+
+Runs inside the job (in Snakemake's Python, started by the job's launcher). A thread samples
+the job's container:
+
+- memory, every 0.5 seconds: anonymous plus shared memory from the cgroup's memory.stat (v2
+  anon/shmem, v1 rss/shmem), i.e. without file cache, which the kernel frees before it kills
+  anything. A spike between two samples isn't seen, so it's a lower bound; the cgroup's count
+  of processes killed for lack of memory (memory.events oom_kill) says when it fell short;
+- CPUs, every 2 seconds: how many were busy at once, from the cgroup's CPU time (cpu.stat
+  usage_usec, or v1 cpuacct.usage) over each 2-second window;
+- disk: the size of the job's working folder (its inputs, outputs and conda downloads),
+  every 20 seconds.
+
+When Snakemake exits, one line goes to the job's log:
+[snakemake-batch] peak-memory-bytes=... peak-cpus=... avg-cpus=... peak-disk-bytes=... oom-kills=...
+
+At the start, the ID of the image the job runs in (what a tag pointed to at the time) goes there
+too, from ECS's container metadata: [snakemake-batch] image ... id=sha256:...
+
+Only the standard library is used."""
+
+import atexit
+import json
+import os
+import sys
+import threading
+import time
+
+CGROUP = "/sys/fs/cgroup"
+
+
+def _memory():
+    path = os.environ.get("SMK_FLEX_MEMSTAT_FILE", os.path.join(CGROUP, "memory.stat"))
+    total = 0
+    with open(path) as f:
+        for line in f:
+            key, _, value = line.partition(" ")
+            if key in ("anon", "rss", "shmem"):
+                total += int(value)
+    return total
+
+
+def _cpu_microseconds():
+    path = os.environ.get("SMK_FLEX_CPUSTAT_FILE", os.path.join(CGROUP, "cpu.stat"))
+    try:
+        with open(path) as f:
+            for line in f:
+                if line.startswith("usage_usec "):
+                    return int(line.split()[1])
+    except OSError:
+        pass
+    try:
+        with open(os.path.join(CGROUP, "cpuacct", "cpuacct.usage")) as f:  # cgroup v1, ns
+            return int(f.read()) // 1000
+    except OSError:
+        return None
+
+
+def _oom_kills():
+    """Processes the kernel killed in this container for lack of memory, or None."""
+    for path in (os.path.join(CGROUP, "memory.events"), os.path.join(CGROUP, "memory", "memory.oom_control")):
+        try:
+            with open(path) as f:
+                for line in f:
+                    if line.startswith("oom_kill "):
+                        return int(line.split()[1])
+        except (OSError, ValueError):
+            pass
+    return None
+
+
+def _folder_size(path):
+    total = 0
+    for root, _dirs, files in os.walk(path):
+        for name in files:
+            try:
+                st = os.lstat(os.path.join(root, name))
+                total += st.st_blocks * 512 if hasattr(st, "st_blocks") else st.st_size
+            except OSError:
+                pass
+    return total
+
+
+def log_image():
+    url = os.environ.get("ECS_CONTAINER_METADATA_URI_V4")
+    if not url:
+        return
+    try:
+        from urllib.request import urlopen
+
+        with urlopen(url, timeout=2) as r:
+            c = json.load(r)
+        print(f"[snakemake-batch] image {c.get('Image')} id={c.get('ImageID')}", file=sys.stderr, flush=True)
+    except Exception:
+        pass
+
+
+def monitor(interval=0.5, cpu_every=4, disk_every=40):
+    log_image()
+    work = os.environ.get("SMK_WORK") or os.getcwd()
+    peak = {"mem": 0, "cpus": 0.0, "disk": 0}
+    start_cpu, start_t = _cpu_microseconds(), time.monotonic()
+
+    def sample_memory():
+        try:
+            peak["mem"] = max(peak["mem"], _memory())
+        except (OSError, ValueError):
+            pass
+
+    def loop():
+        last_cpu, last_t, n = start_cpu, start_t, 0
+        while True:
+            time.sleep(interval)
+            n += 1
+            try:
+                sample_memory()
+                if n % cpu_every == 0:
+                    cpu, t = _cpu_microseconds(), time.monotonic()
+                    if cpu is not None and last_cpu is not None and t > last_t:
+                        peak["cpus"] = max(peak["cpus"], (cpu - last_cpu) / 1e6 / (t - last_t))
+                    last_cpu, last_t = cpu, t
+                if n % disk_every == 1:
+                    peak["disk"] = max(peak["disk"], _folder_size(work))
+            except Exception:
+                pass
+
+    threading.Thread(target=loop, daemon=True).start()
+
+    def report():
+        try:
+            sample_memory()
+            peak["disk"] = max(peak["disk"], _folder_size(work))
+            parts = [f"peak-memory-bytes={peak['mem']}", f"peak-cpus={peak['cpus']:.2f}"]
+            cpu, t = _cpu_microseconds(), time.monotonic()
+            if cpu is not None and start_cpu is not None and t > start_t:
+                parts.append(f"avg-cpus={(cpu - start_cpu) / 1e6 / (t - start_t):.2f}")
+            parts.append(f"peak-disk-bytes={peak['disk']}")
+            kills = _oom_kills()
+            if kills is not None:
+                parts.append(f"oom-kills={kills}")
+            print("[snakemake-batch] " + " ".join(parts), file=sys.stderr, flush=True)
+        except Exception:
+            pass
+
+    atexit.register(report)

snakemake_executor_plugin_aws_batch_flex/naming.py

@@ -0,0 +1,16 @@
+"""AWS Batch job names: up to 128 chars, letters/digits/hyphens/underscores, starting with a letter or digit."""
+import hashlib
+import re
+import string
+
+MAX_LEN = 128
+
+
+def batch_job_name(name: str) -> str:
+    """Make any Snakemake job name valid for Batch; never rejects. Long names keep a hash so they stay distinct."""
+    clean = re.sub(r"[^A-Za-z0-9_-]", "_", name)
+    if not clean or clean[0] not in string.ascii_letters + string.digits:
+        clean = "snk-" + clean
+    if len(clean) > MAX_LEN:
+        clean = clean[: MAX_LEN - 9] + "-" + hashlib.sha1(name.encode()).hexdigest()[:8]
+    return clean