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

snakemake_executor_plugin_aws_batch_flex/__init__.py

@@ -0,0 +1,8 @@
+"""AWS Batch executor for Snakemake that runs each rule in the image from its
+`container:` directive. Images don't need Snakemake installed."""
+
+from . import patches  # noqa: F401  (Snakemake packs the sources before it creates the executor)
+from .executor import Executor
+from .settings import ExecutorSettings, common_settings
+
+__all__ = ["Executor", "ExecutorSettings", "common_settings"]

snakemake_executor_plugin_aws_batch_flex/bootstrap.py

@@ -0,0 +1,376 @@
+"""Shell snippets that run before the Snakemake job command inside the container.
+
+"host" mode uses the env mounted from the instance; the other modes install Snakemake into a
+venv under $TMPDIR with uv. Either way Snakemake is only reachable through a shell function
+called `python`, so the rule's own commands still see the image's python and PATH."""
+
+import base64
+import importlib.metadata as md
+import re
+import shlex
+import textwrap
+import zlib
+from typing import Iterable, List, Optional
+
+MODES = ("auto", "always", "never", "host")
+# uv for jobs that install Snakemake at start, and for the AMI's env; its installer checks the
+# download's sha256
+UV_VERSION = "0.13.0"
+
+_TEMPLATE = r"""
+set -eu
+@@WORKDIR@@
+SMK_BOOT="${TMPDIR:-/tmp}/.snakemake-batch-bootstrap"
+SMK_NEED=1
+@@CHECK@@
+# the image's own matching Snakemake: still run it through the launcher (cache, usage figures)
+if [ "$SMK_NEED" = 0 ]; then SMK_PY=$(command -v python); fi
+if [ "$SMK_NEED" = 1 ]; then
+  echo "[snakemake-batch] bootstrapping Snakemake into $SMK_BOOT" >&2
+  mkdir -p "$SMK_BOOT"
+  export UV_CACHE_DIR="$SMK_BOOT/cache" UV_PYTHON_INSTALL_DIR="$SMK_BOOT/python"
+@@PREREQS@@
+@@UV@@
+  if [ -z "$UV" ]; then
+    echo "[snakemake-batch] cannot bootstrap: image needs curl, wget, or python3+pip (and internet access)" >&2
+    exit 127
+  fi
+  $UV venv --clear --quiet --python @@PYVER@@ --system-site-packages "$SMK_BOOT/venv" >&2
+  $UV pip install --quiet --python "$SMK_BOOT/venv/bin/python" @@NEWER@@@@REQS@@ >&2
+  SMK_PY="$SMK_BOOT/venv/bin/python"
+  python() { "$SMK_PY" "$@"; }
+fi
+"""
+
+
+def uv_snippet(bin_dir: str, lib_dir: str) -> str:
+    """Shell that sets UV to a uv to run: the image's own, else version UV_VERSION from its
+    installer (curl or wget; it checks the download's sha256) into bin_dir, else from pip into
+    lib_dir. UV stays empty if none of them worked. For jobs that install Snakemake at start
+    and for the AMI's env (hostenv.py)."""
+    return f"""UV=""
+if command -v uv >/dev/null 2>&1; then
+  UV=uv
+else
+  if command -v curl >/dev/null 2>&1; then
+    curl -fsSL https://astral.sh/uv/{UV_VERSION}/install.sh | UV_UNMANAGED_INSTALL="{bin_dir}" sh >&2 || true
+  elif command -v wget >/dev/null 2>&1; then
+    wget -qO- https://astral.sh/uv/{UV_VERSION}/install.sh | UV_UNMANAGED_INSTALL="{bin_dir}" sh >&2 || true
+  fi
+  if [ -x "{bin_dir}/uv" ]; then
+    UV="{bin_dir}/uv"
+  else
+    for p in python3 python; do
+      if command -v "$p" >/dev/null 2>&1 && "$p" -m pip install -q --target "{lib_dir}" uv=={UV_VERSION} >&2; then
+        UV="env PYTHONPATH={lib_dir} $p -m uv"
+        break
+      fi
+    done
+  fi
+fi"""
+
+
+# Installs curl (for uv) in an image without curl, wget or pip: only with install_prereqs
+_PREREQS = r'''
+  if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1 \
+     && ! { command -v python3 >/dev/null 2>&1 && python3 -m pip --version >/dev/null 2>&1; }; then
+    echo "[snakemake-batch] no curl/wget/pip in the image: installing curl, ca-certificates and bash" >&2
+    if [ "$(id -u)" != 0 ]; then
+      echo "[snakemake-batch] cannot install prerequisites: the image runs as a non-root user" >&2
+    elif command -v apt-get >/dev/null 2>&1; then
+      apt-get update -qq >&2 && DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends curl ca-certificates bash >&2 || true
+    elif command -v apk >/dev/null 2>&1; then apk add -q --no-cache curl ca-certificates bash >&2 || true
+    elif command -v dnf >/dev/null 2>&1; then dnf install -y -q curl ca-certificates bash >&2 || true
+    elif command -v microdnf >/dev/null 2>&1; then microdnf install -y curl ca-certificates bash >&2 || true
+    elif command -v yum >/dev/null 2>&1; then yum install -y -q curl ca-certificates bash >&2 || true
+    else echo "[snakemake-batch] no known package manager in the image" >&2
+    fi
+  fi
+'''
+
+
+WORKDIR = (
+    'SMK_WORK="${TMPDIR:-/tmp}/snakemake-job-$$"; mkdir -p "$SMK_WORK"; cd "$SMK_WORK"'
+)
+
+_HOST_TEMPLATE = r"""
+set -eu
+@@WORKDIR@@
+SMK_ENV=@@ENV@@
+SMK_WAITED=0
+while [ ! -f "$SMK_ENV/.ready" ]; do
+  if [ "$SMK_WAITED" -ge @@WAIT@@ ]; then
+    echo "[snakemake-batch] host env $SMK_ENV is missing or not ready. Install it on the" >&2
+    echo "  EC2 host/AMI (snakemake-batch-flex-host-env) and check the job definition mount." >&2
+    exit 127
+  fi
+  sleep 2; SMK_WAITED=$((SMK_WAITED + 2))
+done
+SMK_PY="$SMK_ENV/venv/bin/python"
+python() { "$SMK_PY" "$@"; }
+if ! python -c "import snakemake" >/dev/null 2>&1; then
+  echo "[snakemake-batch] the host env cannot run in this image (musl/Alpine? exotic libc?)." >&2
+  exit 127
+fi
+"""
+
+# Python `script:` rules normally run with Snakemake's own interpreter. This launcher makes them
+# use the image's python (>= 3.10, which Snakemake's script header needs) when it has one, the
+# same way Snakemake treats conda envs. Anything else falls through to the original code.
+_LAUNCHER = r"""
+import os, runpy, shutil, subprocess, sys
+a = sys.argv[1:]
+if a[:2] != ["-m", "snakemake"]:
+    os.execv(sys.executable, [sys.executable] + a)
+sys.argv = ["snakemake"] + a[2:]
+if "monitor" in globals() and "--deploy-sources" not in sys.argv:
+    # the job tools came with this launcher, in the same file (job_tools_source)
+    try:
+        make_numpy_loadable()
+        monitor()
+        _cc = os.environ.get("SMK_FLEX_CONDA_CACHE", "none")
+        if _cc != "none":
+            sys.argv = ["snakemake"] + install(sys.argv[1:], _cc)
+    except Exception as e:
+        print(f"[snakemake-batch] job tools off for this job (conda env cache, usage figures): {e!r}",
+              file=sys.stderr)
+if os.environ.get("SMK_FLEX_SCRIPT_PYTHON") == "host":
+    runpy.run_module("snakemake", run_name="__main__", alter_sys=True)
+    sys.exit()
+found = []
+def image_python():
+    if not found:
+        found.append(None)
+        p = shutil.which("python") or shutil.which("python3")
+        if p and os.path.realpath(p) != os.path.realpath(sys.executable):
+            try:
+                if subprocess.run([p, "-c", "import sys; sys.exit(sys.version_info < (3, 10))"]).returncode == 0:
+                    found[0] = p
+            except OSError:
+                pass
+    return found[0]
+def own_packages():
+    # Links to just this Snakemake's packages, put first on the script's path so that a
+    # Snakemake installed in the image can't shadow them. Everything else: image first.
+    import snakemake
+    site = os.path.dirname(os.path.dirname(snakemake.__file__))
+    d = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".smk-pkgs")
+    if not os.path.isdir(d):
+        os.makedirs(d)
+        for n in os.listdir(site):
+            if n.startswith("snakemake") and os.path.isdir(os.path.join(site, n)) and "." not in n:
+                os.symlink(os.path.join(site, n), os.path.join(d, n))
+    return d
+try:
+    import inspect
+    from snakemake import script as sm_script
+    orig = sm_script.PythonScript.execute_script
+    if list(inspect.signature(orig).parameters) != ["self", "fname", "edit"]:
+        raise RuntimeError("this Snakemake's PythonScript.execute_script changed")
+    def execute_script(self, fname, edit=False):
+        if self.container_img is None and self.conda_env is None and self.env_modules is None and image_python():
+            return self._execute_cmd(
+                "PYTHONPATH={pkgs:q}${{PYTHONPATH:+:$PYTHONPATH}} {py_exec:q} {fname:q}",
+                pkgs=own_packages(), py_exec=image_python(), fname=fname, is_python_script=True)
+        return orig(self, fname, edit)
+    sm_script.PythonScript.execute_script = execute_script
+except Exception as e:
+    print(f"[snakemake-batch] script: rules will use Snakemake's python ({e!r})", file=sys.stderr)
+runpy.run_module("snakemake", run_name="__main__", alter_sys=True)
+"""
+
+
+def job_tools_source() -> str:
+    """The file every job downloads and runs (launcher_snippet): condacache.py, jobmonitor.py and
+    jobfixes.py, then the launcher. Too long to send with each job; uploaded once per run."""
+    from pathlib import Path
+
+    here = Path(__file__).parent
+    tools = [(here / name).read_text() for name in ("condacache.py", "jobmonitor.py", "jobfixes.py")]
+    # the launcher runs when the job's loader runs the file, not when the tools are only loaded
+    launcher = 'if __name__ == "smkflex_launcher":\n' + textwrap.indent(_LAUNCHER, "    ")
+    return "\n\n".join(tools + [launcher])
+
+
+# Fetches the job tools (with the launcher) named in SMK_FLEX_JOB_CODE, checks them against the
+# sha256 in their name, and runs them; else plain Snakemake. In the job's command, so short.
+_LOADER = (
+    "import hashlib,os,runpy,sys\n"
+    "u=os.environ.get(\"SMK_FLEX_JOB_CODE\");a=sys.argv[1:];d=None\n"
+    "if a[:2]!=[\"-m\",\"snakemake\"]:os.execv(sys.executable,[sys.executable]+a)\n"
+    "if u and \"--deploy-sources\" not in a:\n"
+    " try:\n"
+    "  if u.startswith(\"s3://\"):\n"
+    "   import boto3;b,_,k=u[5:].partition(\"/\");d=boto3.client(\"s3\").get_object(Bucket=b,Key=k)[\"Body\"].read()\n"
+    "  else:d=open(u,\"rb\").read()\n"
+    "  if not u.endswith(\"jobtools-%s.py\"%hashlib.sha256(d).hexdigest()):raise ValueError(u+\" does not match its checksum\")\n"
+    " except Exception as e:\n"
+    "  d=None;print(\"[snakemake-batch] job tools off for this job (conda env cache, usage figures, image python for script:): %r\"%e,file=sys.stderr)\n"
+    "if d:exec(compile(d,\"smkflex-jobtools.py\",\"exec\"),{\"__name__\":\"smkflex_launcher\",\"__file__\":os.path.join(os.getcwd(),\".smk-launch.py\")})\n"
+    "else:sys.argv=[\"snakemake\"]+a[2:];runpy.run_module(\"snakemake\",run_name=\"__main__\",alter_sys=True)\n"
+)
+
+
+def launcher_snippet(job_code: bool) -> str:
+    """Point the `python` function (if the prelude defined one) at the launcher: the short loader
+    above when the job tools were uploaded (job_code), else the launcher itself, compressed (ECS
+    allows only 8192 characters of overrides per job)."""
+    if job_code:
+        return 'if [ -n "${SMK_PY:-}" ]; then\n' f'  python() {{ "$SMK_PY" -c {shlex.quote(_LOADER)} "$@"; }}\n' "fi\n"
+    return image_python_snippet()
+
+
+def image_python_snippet() -> str:
+    """Point the `python` function at the launcher above (only if the prelude defined one).
+    The launcher is sent compressed: ECS allows only 8192 characters of overrides per job."""
+    packed = base64.b64encode(zlib.compress(_LAUNCHER.encode(), 9)).decode()
+    unpack = "import base64,sys,zlib;open(sys.argv[1],'wb').write(zlib.decompress(base64.b64decode(sys.argv[2])))"
+    return (
+        'if [ -n "${SMK_PY:-}" ]; then\n'
+        f'  "$SMK_PY" -c "{unpack}" "$SMK_WORK/.smk-launch.py" {packed}\n'
+        '  python() { "$SMK_PY" "$SMK_WORK/.smk-launch.py" "$@"; }\n'
+        "fi\n"
+    )
+
+
+def _py_check(snakemake_version: str, package_names: Iterable[str]) -> str:
+    code = (
+        "import importlib.metadata as m\n"
+        f"assert m.version('snakemake') == {snakemake_version!r}\n"
+        + "".join(f"m.version({n!r})\n" for n in package_names)
+    )
+    return shlex.quote(code)
+
+
+def _compact(script: str) -> str:
+    """The script without indentation and blank lines: the job's command has to fit in
+    ECS's 8192 characters, and sh doesn't need them."""
+    return "\n".join(line.strip() for line in script.splitlines() if line.strip()) + "\n"
+
+
+def build_prelude(
+    mode: str,
+    snakemake_version: str,
+    requirements: list[str],
+    python_version: str,
+    install_prereqs: bool = False,
+    resolved_at: Optional[str] = None,
+) -> str:
+    """Return a POSIX-sh snippet; after it runs, `python -m snakemake` works. With resolved_at
+    (an RFC 3339 time), packages are installed as they were then: the same versions in every
+    job of a run, however long it takes."""
+    if mode not in MODES:
+        raise ValueError(f"bootstrap mode must be one of {MODES}, got {mode!r}")
+    if mode == "never":
+        return f"set -eu\n{WORKDIR}"
+
+    names = [r.split("==")[0] for r in requirements]
+    check = ""
+    if mode == "auto":  # use the image's own Snakemake if it's this version
+        check = (
+            "if command -v python >/dev/null 2>&1 && "
+            f"python -c {_py_check(snakemake_version, names)} >/dev/null 2>&1; "
+            "then SMK_NEED=0; fi"
+        )
+    reqs = " ".join(shlex.quote(r) for r in requirements)
+    return _compact(
+        _TEMPLATE.replace("@@WORKDIR@@", WORKDIR)
+        .replace("@@CHECK@@", check)
+        .replace("@@PYVER@@", shlex.quote(python_version))
+        .replace("@@PREREQS@@", _PREREQS if install_prereqs else "")
+        .replace("@@REQS@@", reqs)
+        .replace("@@NEWER@@", f"--exclude-newer {shlex.quote(resolved_at)} " if resolved_at else "")
+        .replace("@@UV@@", uv_snippet("$SMK_BOOT/bin", "$SMK_BOOT/uvlib"))
+    )
+
+
+def build_host_prelude(env_path: str, wait_seconds: int = 120) -> str:
+    """Prelude for the host-mounted env. Waits for `.ready` because an instance that is
+    still installing the env (user-data) may already have been given a job."""
+    return _compact(
+        _HOST_TEMPLATE.replace("@@WORKDIR@@", WORKDIR)
+        .replace("@@ENV@@", shlex.quote(env_path))
+        .replace("@@WAIT@@", str(int(wait_seconds)))
+    )
+
+
+def conda_snippet(conda_path: str) -> str:
+    """Make a host-mounted conda usable in any image: only `condabin` (conda itself, no python)
+    goes at the END of PATH, so an image's own conda wins; packages cache to a writable dir.
+    The AMI's git (smkflex-bin, see hostenv) goes there too, for images without one."""
+    return (
+        f'export PATH="$PATH:{conda_path}/condabin:{conda_path}/smkflex-bin"\n'
+        'export CONDA_PKGS_DIRS="$SMK_WORK/conda-pkgs"\n'
+        '[ -w "${HOME:-/nonexistent}" ] || export HOME="$SMK_WORK"\n'
+    )
+
+
+def pinned(package: str) -> str:
+    """A pip requirement for the version installed here (unpinned if it isn't installed)."""
+    try:
+        return f"{package}=={md.version(package)}"
+    except md.PackageNotFoundError:
+        return package
+
+
+def locked(packages: Iterable[str]):
+    """(pins, skipped): name==version of everything installed here that `packages` need, all
+    the way down, so the AMI's env has exactly this machine's versions; for a constraints file,
+    so pins the env's Python doesn't need are left alone. A version that doesn't fit what the
+    packages ask for (an environment mixed by hand) isn't pinned: it's in `skipped`."""
+    from packaging.requirements import Requirement
+    from packaging.specifiers import SpecifierSet
+
+    def key(name):
+        return re.sub(r"[-_.]+", "-", name).lower()
+
+    found, wanted, todo = {}, {}, list(packages)
+    while todo:
+        name = todo.pop()
+        if key(name) in found:
+            continue
+        try:
+            dist = md.distribution(name)
+        except md.PackageNotFoundError:
+            continue
+        found[key(name)] = dist.version
+        for line in dist.requires or []:
+            req = Requirement(line)
+            if req.marker is None or req.marker.evaluate({"extra": ""}):
+                wanted[key(req.name)] = wanted.get(key(req.name), SpecifierSet()) & req.specifier
+                todo.append(req.name)
+    pins, skipped = [], []
+    for k, v in sorted(found.items()):
+        # a local build ("+...") isn't on PyPI
+        fits = "+" not in v and wanted.get(k, SpecifierSet()).contains(v, prereleases=True)
+        (pins if fits else skipped).append(f"{k}=={v}")
+    return pins, skipped
+
+
+def storage_requirements(default_provider: Optional[str] = None) -> List[str]:
+    """pip requirements for the storage plugins installed here, pinned to these versions."""
+    packages = set()
+    try:
+        from snakemake_interface_storage_plugins.registry import StoragePluginRegistry
+
+        reg = StoragePluginRegistry()
+        packages = {reg.get_plugin_package_name(n) for n in reg.get_registered_plugins()}
+    except Exception:  # registry API drift: fall back to the default provider
+        if default_provider:
+            packages = {f"snakemake-storage-plugin-{default_provider}"}
+    return [pinned(p) for p in sorted(packages)]
+
+
+def snakemake_prelude(s, mode: str, default_storage_provider: Optional[str] = None,
+                      resolved_at: Optional[str] = None) -> str:
+    """The prelude that gives a job Snakemake, for settings `s` (settings.ExecutorSettings):
+    from the host env (mode "host") or installed at start, the same version as here, with the
+    other packages as they were at resolved_at."""
+    if mode == "host":
+        return build_host_prelude(s.host_env_path, s.host_env_wait_seconds)
+    import sys
+
+    version = md.version("snakemake")
+    reqs = [f"snakemake=={version}"] + storage_requirements(default_storage_provider) + list(s.extra_packages or [])
+    pyver = str(s.python_version or f"{sys.version_info[0]}.{sys.version_info[1]}")
+    return build_prelude(mode, version, reqs, pyver, s.install_prereqs, resolved_at)

snakemake_executor_plugin_aws_batch_flex/bundle.py

@@ -0,0 +1,142 @@
+"""smkflex run: the workflow folder packed for the head job (with --add files, or a --git clone),
+the commit it records, and the head job's command."""
+
+import io
+import os
+import re
+import shlex
+import shutil
+import subprocess
+import tarfile
+from datetime import datetime, timezone
+
+from . import headjob
+from .bootstrap import conda_snippet, snakemake_prelude
+from .profiles import Fail
+from .settings import from_profile
+from .sources import SKIP_DIRS, workflow_files
+
+# Runs in the head job before anything else: fetch and unpack the bundle. Kept small, since
+# everything sent with a job must fit in 8192 characters.
+_FETCH = (
+    "import io,sys,tarfile\n"
+    "u,d=sys.argv[1],sys.argv[2]\n"
+    "if u.startswith('s3://'):\n"
+    " import boto3;b,_,k=u[5:].partition('/')\n"
+    " data=boto3.client('s3').get_object(Bucket=b,Key=k)['Body'].read()\n"
+    "else:\n"
+    " data=open(u,'rb').read()\n"
+    "t=tarfile.open(fileobj=io.BytesIO(data))\n"
+    "t.extractall(d,**({'filter':'data'} if hasattr(tarfile,'data_filter') else {}))\n"
+)
+
+def added_files(paths, base="."):
+    """{path in the bundle: local file} for --add: files, or every file in a folder, given
+    relative to `base` (the current folder) and sent to the same place."""
+    out = {}
+    for p in paths:
+        norm = os.path.normpath(p)
+        if os.path.isabs(p) or norm == ".." or norm.startswith(".." + os.sep):
+            raise Fail(f"--add {p}: give a path inside this folder (it's sent to the same place)")
+        full = os.path.join(base, norm)
+        if os.path.isfile(full):
+            out[norm] = full
+        elif os.path.isdir(full):
+            for d, dirs, names in os.walk(full):
+                dirs[:] = [x for x in dirs if x not in SKIP_DIRS]
+                for n in names:
+                    f = os.path.join(d, n)
+                    out[os.path.relpath(f, base)] = f
+        else:
+            raise Fail(f"--add {p}: no such file or folder here")
+    return out
+
+def make_bundle(root, profile_dir, max_mb, extra=None):
+    """The workflow folder `root` as a tar.gz, with `extra` ({path in the bundle: local file})
+    laid over it, plus the profile, this plugin and the head job's script."""
+    buf = io.BytesIO()
+    total = 0
+    sources = {f: os.path.join(root, f) for f in workflow_files(root)}
+    sources.update(extra or {})
+    files = sorted(sources)
+    with tarfile.open(fileobj=buf, mode="w:gz") as t:
+        for f in files:
+            total += os.path.getsize(sources[f])
+            if total > max_mb * 1024 * 1024:
+                raise Fail(f"the workflow folder has more than {max_mb} MB of files (at {f}); data "
+                           "belongs in storage. Add big files to .gitignore, or raise --max-bundle-mb.")
+            t.add(sources[f], arcname=f)
+        # what `git ls-files` would say, for the head job's stand-in git
+        listing = "".join(f + "\n" for f in files).encode()
+        info = tarfile.TarInfo(".smkflex/files.txt")
+        info.size = len(listing)
+        t.addfile(info, io.BytesIO(listing))
+        for d, dirs, names in os.walk(profile_dir):
+            dirs[:] = [x for x in dirs if x != "__pycache__"]
+            for n in names:
+                p = os.path.join(d, n)
+                t.add(p, arcname=os.path.join(".smkflex", "profile", os.path.relpath(p, profile_dir)))
+        # this plugin as installed here, so the head job runs the same version
+        pkg = os.path.dirname(os.path.abspath(__file__))
+        for n in sorted(os.listdir(pkg)):
+            if n.endswith(".py"):
+                t.add(os.path.join(pkg, n), arcname=f".smkflex/site/{os.path.basename(pkg)}/{n}")
+        t.add(headjob.__file__, arcname=".smkflex/head.py")
+    return buf.getvalue()
+
+def head_command(profile, platform, bundle_url, state_url, snakemake_args):
+    s = from_profile(profile)
+    mode = "auto" if s.bootstrap == "host" and platform == "fargate" else s.bootstrap
+    resolved_at = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+    prelude = snakemake_prelude(s, mode, profile.get("default-storage-provider"), resolved_at)
+    if s.conda_path and platform == "ec2":
+        prelude += "\n" + conda_snippet(s.conda_path)
+    args = " ".join(shlex.quote(a) for a in snakemake_args)
+    return (
+        f"{prelude}\n"
+        'SMKFLEX_DIR="${TMPDIR:-/tmp}/smkflex-workflow"; rm -rf "$SMKFLEX_DIR"; mkdir -p "$SMKFLEX_DIR"\n'
+        'SMKFLEX_PY="${SMK_PY:-python}"\n'
+        f'"$SMKFLEX_PY" -c {shlex.quote(_FETCH)} {shlex.quote(bundle_url)} "$SMKFLEX_DIR"\n'
+        'cd "$SMKFLEX_DIR"\n'
+        f'exec "$SMKFLEX_PY" .smkflex/head.py {shlex.quote(state_url)} {args}\n'
+    )
+
+def _git(*args, cwd=None):
+    try:
+        return subprocess.run(["git", *args], cwd=cwd, capture_output=True, text=True, check=True).stdout.strip()
+    except FileNotFoundError:
+        raise Fail("git isn't installed here; it's needed for --git") from None
+    except subprocess.CalledProcessError as e:
+        raise Fail(f"git {' '.join(args[:2])} failed: {(e.stderr or e.stdout).strip()}") from None
+
+def clone(url, rev, dest):
+    """Shallow clone of `rev` (branch, tag; default: the default branch), or a full clone
+    when it's a commit. Returns the commit."""
+    opts = ["--quiet", "--recurse-submodules", "--shallow-submodules"]
+    try:
+        _git("clone", *opts, "--depth", "1", *(["--branch", rev] if rev else []), url, dest)
+    except Fail:
+        if not rev:
+            raise
+        shutil.rmtree(dest, ignore_errors=True)
+        _git("clone", "--quiet", "--recurse-submodules", url, dest)
+        try:
+            _git("checkout", "--quiet", rev, cwd=dest)
+        except Fail:
+            raise Fail(f"{url} has no branch, tag or commit '{rev}'") from None
+        _git("submodule", "update", "--quiet", "--init", "--recursive", cwd=dest)
+    return _git("rev-parse", "HEAD", cwd=dest)
+
+def local_commit(root):
+    """The commit of a local git checkout, with -modified if it has uncommitted changes,
+    including new files that aren't git-ignored (the bundle sends those too)."""
+    try:
+        sha = subprocess.run(["git", "rev-parse", "HEAD"], cwd=root, capture_output=True, text=True, check=True).stdout.strip()
+        dirty = subprocess.run(["git", "status", "--porcelain", "."], cwd=root,
+                               capture_output=True, text=True, check=True).stdout.strip()
+    except (OSError, subprocess.CalledProcessError):
+        return None
+    return sha[:12] + ("-modified" if dirty else "")
+
+def _tag(value):
+    return re.sub(r"[^A-Za-z0-9 _.:/=+\-@]", "_", str(value))[:256]

snakemake_executor_plugin_aws_batch_flex/check.py

@@ -0,0 +1,129 @@
+"""smkflex check: the setup, checked from this machine, one line per check."""
+
+import os
+
+import boto3
+from botocore.config import Config as BotoConfig
+from botocore.exceptions import BotoCoreError, ClientError
+
+from . import headjob
+from .preflight import ami_snakemake_version, check_queue, check_role, queue_uses_spot, version_note
+from .profiles import Fail, _load_profile
+
+
+def cmd_check(a):
+    """Check the setup from this machine, one line per check, with what to do about failures."""
+    results = []
+
+    def report(level, what, fix=""):
+        results.append(level)
+        mark = {"ok": "  ok  ", "warn": " warn ", "fail": " FAIL "}[level]
+        print(f"[{mark}] {what}" + (f"\n         -> {fix}" if fix else ""))
+
+    try:
+        profile_dir, profile = _load_profile(a.profile)
+    except Fail as e:
+        report("fail", f"profile: {e}", "run this in your workflow folder, or pass --profile DIR")
+        return 1
+    report("ok", f"profile {profile_dir}")
+    if profile.get("executor") != "aws-batch-flex":
+        report("warn", f"the profile's executor is {profile.get('executor')!r}, not aws-batch-flex",
+               "add `executor: aws-batch-flex`, or pass --executor aws-batch-flex to snakemake")
+
+    import importlib.metadata as md
+    try:
+        local = md.version("snakemake")
+        report("ok", f"Snakemake {local} installed here")
+    except md.PackageNotFoundError:
+        local = None
+        report("fail", "Snakemake isn't installed here", "pip install snakemake")
+    provider = profile.get("default-storage-provider")
+    prefix = str(profile.get("default-storage-prefix") or "").rstrip("/")
+    if not provider or provider == "local" or not prefix:
+        report("fail", "no remote storage in the profile",
+               "set default-storage-provider: s3 and default-storage-prefix: s3://BUCKET/PREFIX")
+    else:
+        try:
+            md.version(f"snakemake-storage-plugin-{provider}")
+            report("ok", f"storage plugin {provider} installed here")
+        except md.PackageNotFoundError:
+            report("fail", f"snakemake-storage-plugin-{provider} isn't installed here",
+                   f"pip install snakemake-storage-plugin-{provider}")
+
+    region = profile.get("aws-batch-flex-region")
+    cfg = BotoConfig(region_name=region, retries={"max_attempts": 3})
+    try:
+        who = boto3.client("sts", config=cfg).get_caller_identity()["Arn"]
+        report("ok", f"AWS login: {who}")
+    except (ClientError, BotoCoreError) as e:
+        report("fail", f"AWS login: {e}", "run `aws configure` (or `aws sso login`), or set AWS_PROFILE")
+        return 1
+    batch = boto3.client("batch", config=cfg)
+    region = batch.meta.region_name if hasattr(batch, "meta") else region
+    if not region:
+        report("fail", "no AWS region", "set aws-batch-flex-region in the profile")
+        return 1
+
+    queue = profile.get("aws-batch-flex-queue")
+    platform, queue_ok = None, False
+    if not queue:
+        report("fail", "no queue in the profile", "set aws-batch-flex-queue (Terraform prints it as queue_name)")
+    else:
+        errors, warnings, detected = check_queue(batch, queue, profile.get("aws-batch-flex-platform"))
+        platform = profile.get("aws-batch-flex-platform") or detected or "ec2"
+        for e in errors:
+            report("fail", f"queue: {e}")
+        for w in warnings:
+            report("warn", f"queue: {w}")
+        queue_ok = not errors
+        if not errors:
+            spot = " (Spot)" if queue_uses_spot(batch, queue) else ""
+            report("ok", f"queue {queue} in {region}: enabled, {platform.upper()}{spot}")
+
+    iam = boto3.client("iam", config=cfg)
+    for key, label in (("aws-batch-flex-job-role-arn", "job role"), ("aws-batch-flex-head-role-arn", "head role"),
+                       ("aws-batch-flex-execution-role-arn", "execution role")):
+        arn = profile.get(key)
+        if arn:
+            errors, _w = check_role(iam, arn, label.capitalize())
+            if errors:
+                report("fail", errors[0], f"check {key} in the profile (Terraform prints the right one)")
+            else:
+                report("ok", f"{label} {arn.rsplit('/', 1)[-1]}")
+    if not profile.get("aws-batch-flex-job-role-arn") and platform == "fargate":
+        report("fail", "Fargate jobs need a job role (there's no instance role to reach storage with)",
+               "set aws-batch-flex-job-role-arn")
+    if platform == "fargate" and not profile.get("aws-batch-flex-execution-role-arn"):
+        report("fail", "Fargate needs an execution role", "set aws-batch-flex-execution-role-arn")
+    if not profile.get("aws-batch-flex-head-role-arn"):
+        report("warn", "no aws-batch-flex-head-role-arn: fine for snakemake, but smkflex run needs it",
+               "add the head_role_arn that Terraform prints")
+
+    if prefix and provider and provider != "local":
+        probe = f"{prefix}/.smkflex/check-{os.getpid()}"
+        try:
+            headjob.put(probe, b"smkflex check")
+            ok = headjob.get(probe) == b"smkflex check"
+            if probe.startswith("s3://"):
+                s3, bucket, key = headjob._s3(probe)
+                s3.delete_object(Bucket=bucket, Key=key)
+            else:
+                os.remove(probe)
+            report("ok" if ok else "fail", f"storage {prefix}: write, read, delete")
+        except (ClientError, BotoCoreError, OSError) as e:
+            report("fail", f"storage {prefix}: {e}",
+                   "check the bucket name and that your AWS login may read and write under the prefix")
+
+    if queue_ok and platform == "ec2" and profile.get("aws-batch-flex-bootstrap") == "host" and local:
+        ami = ami_snakemake_version(batch, boto3.client("ec2", config=cfg), queue)
+        level, note = version_note(ami, local)
+        if not ami:
+            report("warn", "couldn't find the AMI's Snakemake version (an older AMI, or no ec2:Describe* permission)")
+        elif level == "warning":
+            report("warn", note, f"pip install snakemake=={ami}, or bake a new AMI")
+        else:
+            report("ok", f"the AMI has Snakemake {ami}" + (f" (here: {local}; known to work together)" if level else ""))
+
+    fails, warns = results.count("fail"), results.count("warn")
+    print(f"\n{fails} problem(s), {warns} warning(s)." if fails or warns else "\nAll good.")
+    return 1 if fails else 0