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

snakemake_executor_plugin_aws_batch_flex/orphans.py

@@ -0,0 +1,51 @@
+"""Jobs left running by an earlier run of the same workflow (Snakemake killed, laptop lost its
+connection, ...). Snakemake would submit them again, so they'd run twice."""
+
+import hashlib
+import os
+
+ACTIVE = ("SUBMITTED", "PENDING", "RUNNABLE", "STARTING", "RUNNING")
+
+
+def workflow_id(workdir: str, snakefile: str, storage_prefix: str) -> str:
+    key = "|".join([os.path.abspath(workdir), os.path.abspath(snakefile), storage_prefix or ""])
+    return hashlib.sha256(key.encode()).hexdigest()[:16]
+
+
+def list_jobs(batch, **kw) -> list:
+    """The job summaries list_jobs(**kw) gives, from every page."""
+    out, token = [], None
+    while True:
+        resp = batch.list_jobs(**kw, **({"nextToken": token} if token else {}))
+        out += resp.get("jobSummaryList", [])
+        token = resp.get("nextToken")
+        if not token:
+            return out
+
+
+def describe_jobs(batch, ids) -> list:
+    """The jobs described (with their tags and job definition), 100 at a time as Batch takes them."""
+    ids, jobs = list(ids), []
+    for i in range(0, len(ids), 100):
+        jobs += batch.describe_jobs(jobs=ids[i : i + 100])["jobs"]
+    return jobs
+
+
+def active_jobs(batch, queue: str) -> list:
+    """Every active job in `queue`, described. The summaries have no job definition or tags
+    (despite the API model), so every one is described."""
+    return describe_jobs(batch, [j["jobId"] for status in ACTIVE
+                                 for j in list_jobs(batch, jobQueue=queue, jobStatus=status, maxResults=1000)])
+
+
+def find_orphans(batch, queue: str, jobdef_prefix: str, wf_id: str, run_id: str) -> list:
+    """[(job id, job name, status)] of active jobs in `queue` from this workflow but another run."""
+    found = []
+    for d in active_jobs(batch, queue):
+        tags = d.get("tags") or {}
+        jobdef = d.get("jobDefinition", "").rsplit("/", 1)[-1].split(":")[0]
+        if (jobdef.startswith(jobdef_prefix + "-")
+                and tags.get("snakemake-workflow") == wf_id
+                and tags.get("snakemake-run-id") != run_id):
+            found.append((d["jobId"], d.get("jobName", ""), d.get("status", "")))
+    return found

snakemake_executor_plugin_aws_batch_flex/patches.py

@@ -0,0 +1,196 @@
+"""Workarounds for bugs in Snakemake and the S3 storage plugin, each limited to what this
+plugin needs. They can go once the bugs are fixed upstream.
+
+Each one first checks that the code it patches is still what it was written for (the
+parameters it takes); if a new version changed it, the workaround isn't applied and the
+executor says so at start (NOT_APPLIED), instead of failing quietly or calling it wrongly.
+tests/test_units.py fails as soon as one doesn't apply, and CI runs the newest Snakemake weekly."""
+
+import inspect
+import os
+import re
+
+from snakemake_interface_common.exceptions import WorkflowError
+
+from .settings import MAX_SOURCES_MB, PLUGIN_NAME
+from .sources import too_big, workflow_files
+
+
+SOURCES = ("sending the workflow folder to the jobs (Snakemake's DAG.get_sources): jobs get only "
+           "what git knows about")
+MISSING = "the hint for inputs that are only on this machine (Snakemake's MissingInputException)"
+S3_DIRS = ("directory() outputs on S3 (the S3 storage plugin's StorageObject.exists): they can be "
+           "reported missing after the job made them")
+# Workarounds this Snakemake (or S3 plugin) doesn't let us apply; the executor warns about them
+NOT_APPLIED: list = []
+
+
+def fits(func, params, what) -> bool:
+    """Does func take these parameters first, as the code a workaround was written for did?
+    If not, `what` goes into NOT_APPLIED."""
+    try:
+        ok = list(inspect.signature(func).parameters)[: len(params)] == list(params)
+    except (TypeError, ValueError):
+        ok = False
+    if not ok and what not in NOT_APPLIED:
+        NOT_APPLIED.append(what)
+    return ok
+
+
+def skip_remote_module_scripts():
+    """Two fixes to what Snakemake packs for the jobs. (1) Work around Snakemake (9.27): packing the workflow's sources for the jobs crashes with
+    "TypeError: ... not GithubFile" when a rule comes from a module fetched from GitHub (as
+    snakedeploy sets up), because it treats the rule's script as a local file. Those scripts
+    don't need packing: each job fetches them from GitHub itself. So they're hidden while the
+    sources are collected. (2) The whole workflow folder is packed, see below."""
+    try:
+        from snakemake.dag import DAG
+        from snakemake.sourcecache import LocalSourceFile, SourceFile
+    except ImportError:
+        NOT_APPLIED.append(SOURCES)
+        return
+    orig = getattr(DAG, "get_sources", None)
+    if getattr(orig, "_skips_remote", False) or not fits(orig, ["self"], SOURCES):
+        return
+
+    def get_sources(self):
+        # Snakemake imports every installed executor plugin, so this runs for all of them:
+        # leave other executors' runs as they are
+        plugin = getattr(getattr(self, "workflow", None), "executor_plugin", None)
+        if getattr(plugin, "name", None) != PLUGIN_NAME:
+            return get_sources.__wrapped__(self)
+        # Jobs only get files below the folder Snakemake works in; Snakemake leaves the rest out
+        # with a warning, and every job then fails for want of the Snakefile
+        main = str(getattr(self.workflow, "main_snakefile", "") or "")
+        if main and os.path.exists(main) and os.path.relpath(main).startswith(".."):
+            raise WorkflowError(
+                f"The Snakefile ({os.path.relpath(main)}) is outside the folder Snakemake works in "
+                f"({os.getcwd()}), set by a workdir: directive or --directory, so the jobs couldn't get it. "
+                "Remove workdir: and start Snakemake in the workflow's folder: outputs go under the "
+                "storage prefix anyway (a project of its own: smkflex init in its own folder)."
+            )
+        hidden = []
+        for rule in getattr(self.workflow, "rules", []):
+            base = getattr(rule, "basedir", None)
+            if isinstance(base, SourceFile) and not isinstance(base, LocalSourceFile):
+                for attr in ("script", "notebook"):
+                    value = getattr(rule, attr, None)
+                    if value:
+                        hidden.append((rule, attr, value))
+                        setattr(rule, attr, None)
+        try:
+            files = get_sources.__wrapped__(self)
+        finally:
+            for rule, attr, value in hidden:
+                setattr(rule, attr, value)
+        # Snakemake sends jobs only what git knows about, besides the Snakefile and the files
+        # named in script:/conda:. Scripts a shell: command runs, modules the Snakefile imports,
+        # sample sheets and the like are missing in a folder that isn't a git checkout (such as
+        # one made by snakedeploy), or when never `git add`ed. Send the folder as smkflex run
+        # does: what git would track, or everything .gitignore doesn't exclude, and config/.
+        if isinstance(files, set):
+            files.update(workflow_files("."))
+            settings = getattr(self.workflow, "executor_settings", None)
+            limit = getattr(settings, "max_sources_mb", None) or MAX_SOURCES_MB
+            problem = too_big(files, limit)
+            if problem:
+                raise WorkflowError(problem)
+        return files
+
+    get_sources._skips_remote = True
+    get_sources.__wrapped__ = orig
+    DAG.get_sources = get_sources
+
+
+# Snakemake packs the sources before it creates the executor, so this can't wait for it.
+skip_remote_module_scripts()
+
+
+def explain_local_only_inputs():
+    """Snakemake only says "Missing input files ... (in storage)" for an input that is on this
+    machine but was never uploaded (data/, or an absolute path): jobs read their inputs from
+    storage. Say so, with the smkflex push command that uploads it."""
+    try:
+        from snakemake.exceptions import MissingInputException
+    except ImportError:
+        NOT_APPLIED.append(MISSING)
+        return
+    orig = MissingInputException.__init__
+    if getattr(orig, "_smkflex", False) or not fits(orig, ["self", "job", "files"], MISSING):
+        return
+
+    def __init__(self, job, files, *args, **kwargs):
+        orig(self, job, files, *args, **kwargs)
+        try:
+            hint = local_only_hint(job.rule.workflow, files)
+        except Exception:  # only an extra line: never let it break Snakemake's own error
+            hint = None
+        if hint and self.args:
+            self.args = (self.args[0] + hint, *self.args[1:])
+
+    __init__._smkflex = True
+    MissingInputException.__init__ = __init__
+
+
+def local_only_hint(workflow, files):
+    if getattr(getattr(workflow, "executor_plugin", None), "name", None) != PLUGIN_NAME:
+        return None
+    prefix = str(getattr(workflow.storage_settings, "default_storage_prefix", "") or "").rstrip("/")
+    here, stored = [], False
+    for f in files:
+        query = getattr(getattr(f, "storage_object", None), "query", None) if getattr(f, "is_storage", False) else None
+        if not prefix or not query or not str(query).startswith(prefix + "/"):
+            continue
+        stored = True
+        rel = str(query)[len(prefix) + 1:]
+        # a relative path goes to prefix/path, an absolute one to prefix/path-without-its-/
+        for local in (rel, "/" + rel):
+            if os.path.exists(local):
+                here.append(local)
+                break
+    if not here:
+        return (f"\n    Jobs read their inputs from {prefix}/, not from this machine: upload them there, "
+                "at the paths your rules and sample sheets use, with: smkflex push <files or folders>"
+                if stored else None)
+    push = sorted({os.path.dirname(p) if os.path.isabs(p) else p.split("/")[0] for p in here})
+    files_text = ", ".join(here[:5]) + (f" and {len(here) - 5} more" if len(here) > 5 else "")
+    return (f"\n    On this machine but not in storage, where the jobs read their inputs: {files_text}. "
+            f"Upload with: smkflex push {' '.join(push)}")
+
+
+explain_local_only_inputs()
+
+
+def recheck_s3_directories():
+    """Work around snakemake-storage-plugin-s3 (0.3.6): a storage object remembers is_dir() =
+    False from before the job ran, so a directory() output the job then uploads is reported
+    "missing in storage". exists() looks a remembered False up again (it only asks is_dir()
+    when there's no plain object of that name). Only there: mtime() and size() of every file
+    ask is_dir() too, and keep the remembered answer, as without the patch."""
+    try:
+        from snakemake_storage_plugin_s3 import StorageObject
+    except ImportError:  # not installed: nothing to work around
+        return
+    orig = getattr(StorageObject, "exists", None)
+    if getattr(orig, "_rechecks", False) or not fits(orig, ["self"], S3_DIRS):
+        return
+    # what it works around: the remembered is_dir() answer, in _is_dir
+    if "_is_dir" not in inspect.getsource(StorageObject):
+        NOT_APPLIED.append(S3_DIRS)
+        return
+
+    def exists(self):
+        if getattr(self, "_is_dir", None) is False:
+            self._is_dir = None
+        return orig(self)
+
+    exists._rechecks = True
+    StorageObject.exists = exists
+
+
+def keep_unneeded_temp(cmd: str) -> str:
+    """Drop --unneeded-temp-files from a job's command. Snakemake (9.x) tells the job to delete
+    temp() outputs that no other job needs, but then waits for every output in storage after
+    the job, and fails with "Missing files". Without the option the job uploads them, and
+    Snakemake removes them from storage itself as with any other temp file."""
+    return re.sub(r" --unneeded-temp-files(?: (?!-)(?:'[^']*'|\S+))+", "", cmd)

snakemake_executor_plugin_aws_batch_flex/preflight.py

@@ -0,0 +1,165 @@
+"""Startup checks on the queue, its compute environments and the job role.
+
+A missing permission to look only gives a warning. Skip with --aws-batch-flex-skip-preflight."""
+
+import weakref
+from typing import Optional
+
+from botocore.exceptions import BotoCoreError, ClientError
+
+
+def access_denied(e: ClientError) -> bool:
+    code = e.response.get("Error", {}).get("Code", "")
+    return "AccessDenied" in code or "UnauthorizedOperation" in code or code == "AccessDeniedException"
+
+
+# What describe_queue found, per Batch client (one region and login), for as long as it lives
+_QUEUES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
+
+
+def describe_queue(batch, queue: str):
+    """(queue, its compute environments) as Batch describes them, looked up once per client:
+    the startup checks, Spot and priority all need them. (None, None) if there's no such
+    queue; the environments are None if they couldn't be described. Batch's errors for the
+    queue itself are raised."""
+    try:
+        cache = _QUEUES.setdefault(batch, {})
+    except TypeError:  # a stand-in client that can't be weakly referenced (tests): no cache
+        cache = {}
+    if queue in cache:
+        return cache[queue]
+    found = batch.describe_job_queues(jobQueues=[queue])["jobQueues"]
+    if not found:
+        return None, None
+    q = found[0]
+    ces = [o["computeEnvironment"] for o in q.get("computeEnvironmentOrder", [])]
+    envs = []
+    if ces:
+        try:
+            envs = batch.describe_compute_environments(computeEnvironments=ces)["computeEnvironments"]
+        except (ClientError, BotoCoreError):
+            envs = None
+    cache[queue] = (q, envs)
+    return q, envs
+
+
+def check_queue(batch, queue: str, platform: Optional[str] = None):
+    """Returns (errors, warnings, detected platform or None) for the job queue and the compute
+    environments behind it. A mismatch with `platform` is only an error when `platform` is given."""
+    errors, warnings, detected = [], [], None
+    try:
+        q, envs = describe_queue(batch, queue)
+    except ClientError as e:
+        if access_denied(e):
+            return [], [f"cannot verify queue '{queue}' (no batch:DescribeJobQueues permission); continuing"], None
+        return [f"AWS Batch refused the queue lookup: {e}"], [], None
+    except BotoCoreError as e:
+        return [], [f"cannot verify queue '{queue}': {e}; continuing"], None
+    if q is None:
+        hint = ""
+        try:
+            names = sorted(q["jobQueueName"] for q in batch.describe_job_queues()["jobQueues"])
+            hint = f" Queues in this region: {', '.join(names[:10]) or 'none'}."
+        except (ClientError, BotoCoreError):
+            pass
+        return [f"Job queue '{queue}' not found in this region (check the queue name and the region)." + hint], [], None
+    if q.get("state") != "ENABLED":
+        errors.append(f"Job queue '{queue}' is {q.get('state')}, not ENABLED.")
+    if q.get("status") != "VALID":
+        errors.append(f"Job queue '{queue}' has status {q.get('status')}: {q.get('statusReason', 'no reason given')}")
+    if not q.get("computeEnvironmentOrder"):
+        return errors + [f"Job queue '{queue}' has no compute environment attached."], warnings, None
+    if envs is None:
+        return errors, warnings + [f"cannot inspect the compute environments of '{queue}'; continuing"], None
+    for ce in envs:
+        name, kind = ce.get("computeEnvironmentName"), ce.get("computeResources", {}).get("type", "")
+        if ce.get("status") not in (None, "VALID"):
+            errors.append(f"Compute environment '{name}' has status {ce.get('status')}: {ce.get('statusReason', '')}")
+        if kind:
+            actual = "fargate" if kind.startswith("FARGATE") else "ec2"
+            detected = detected or actual
+            if platform and actual != platform:
+                errors.append(
+                    f"Queue '{queue}' runs on {actual.upper()} (compute environment '{name}') but "
+                    f"--aws-batch-flex-platform is '{platform}'. Set --aws-batch-flex-platform {actual} "
+                    "(or leave it unset: the platform is then detected per queue)."
+                )
+    return errors, warnings, detected
+
+
+def queue_uses_spot(batch, queue: str) -> bool:
+    """Best effort: is any compute environment behind the queue Spot (EC2 or Fargate)?"""
+    try:
+        _q, envs = describe_queue(batch, queue)
+    except (ClientError, BotoCoreError):
+        return False
+    return any(ce.get("computeResources", {}).get("type", "").endswith("SPOT") for ce in envs or [])
+
+
+def check_role(iam, role_arn: str, label: str):
+    """Best effort: does the role exist? Only a definite 'no such role' is an error."""
+    name = role_arn.rsplit("/", 1)[-1]
+    try:
+        iam.get_role(RoleName=name)
+    except ClientError as e:
+        if e.response.get("Error", {}).get("Code") == "NoSuchEntity":
+            return [f"{label} '{role_arn}' does not exist."], []
+        return [], []  # cannot look (no iam:GetRole): not our business
+    except BotoCoreError:
+        return [], []
+    return [], []
+
+
+SUBMIT_HINT = (
+    "The identity running snakemake needs batch:SubmitJob (on the queue and the snakemake-* job definitions), "
+    "batch:RegisterJobDefinition, batch:DescribeJobDefinitions, batch:DescribeJobs, batch:ListJobs, batch:TerminateJob, "
+    "batch:TagResource (jobs are tagged), logs:GetLogEvents, and iam:PassRole on the job role."
+)
+
+
+def ami_snakemake_version(batch, ec2, queue: str) -> Optional[str]:
+    """The SnakemakeVersion tag (set by bake-ami.sh) of the AMI the queue's instances start from.
+    None when it can't be found out: no tag (older AMIs), no permission, or several AMIs."""
+    try:
+        _q, envs = describe_queue(batch, queue)
+        images = set()
+        for ce in envs or []:
+            cr = ce.get("computeResources") or {}
+            # Batch uses an image override first, then imageId, then the launch template's image
+            found = {c.get("imageIdOverride") for c in cr.get("ec2Configuration") or []} - {None}
+            if not found and cr.get("imageId"):
+                found = {cr["imageId"]}
+            lt = cr.get("launchTemplate") or {}
+            if not found and (lt.get("launchTemplateId") or lt.get("launchTemplateName")):
+                which = ({"LaunchTemplateId": lt["launchTemplateId"]} if lt.get("launchTemplateId")
+                         else {"LaunchTemplateName": lt["launchTemplateName"]})
+                for v in ec2.describe_launch_template_versions(
+                        Versions=[lt.get("version") or "$Default"], **which)["LaunchTemplateVersions"]:
+                    found |= {v.get("LaunchTemplateData", {}).get("ImageId")} - {None}
+            images |= found
+        if not images:
+            return None
+        versions = {t["Value"] for img in ec2.describe_images(ImageIds=sorted(images))["Images"]
+                    for t in img.get("Tags", []) if t["Key"] == "SnakemakeVersion"}
+        return versions.pop() if len(versions) == 1 else None
+    except (ClientError, BotoCoreError, KeyError, IndexError):
+        return None
+
+
+def _version_tuple(v: str):
+    try:
+        return tuple(int(x) for x in v.split(".")[:3])
+    except ValueError:
+        return None
+
+
+def version_note(ami: Optional[str], local: str):
+    """(level, message) comparing the AMI's Snakemake with this machine's, or (None, None)."""
+    if not ami or ami == "unknown" or ami == local:
+        return None, None
+    a, l = _version_tuple(ami), _version_tuple(local)
+    msg = f"the AMI has Snakemake {ami}, this machine has {local}."
+    if a and l and a[0] == l[0] and min(a, l) >= (9, 5):
+        return "info", "Note: " + msg + " Versions from 9.5 on are known to work together."
+    return "warning", ("Snakemake version gap: " + msg + " Jobs may fail with 'unrecognized "
+                       "arguments'; install the AMI's version here or bake a new AMI.")

snakemake_executor_plugin_aws_batch_flex/profiles.py

@@ -0,0 +1,98 @@
+"""smkflex: the profile (as Snakemake finds it) and the AWS clients its commands use."""
+
+import os
+import sys
+
+import boto3
+import yaml
+from botocore.config import Config as BotoConfig
+
+
+class Fail(Exception):
+    pass
+
+class _Log:
+    def info(self, msg):
+        print(msg, file=sys.stderr)
+
+    warning = info
+
+def workflow_profile_file(create=False):
+    """The project's workflow profile, as Snakemake finds it when started here: the first of
+    default/, profiles/default/ and (next to workflow/Snakefile) workflow/profiles/default/ that
+    exists. With create, profiles/default/config.yaml when there's none (Snakemake reads that
+    for a Snakefile at the top or in workflow/)."""
+    candidates = ["default", os.path.join("profiles", "default")]
+    if os.path.isfile(os.path.join("workflow", "Snakefile")):
+        candidates.append(os.path.join("workflow", "profiles", "default"))
+    for d in candidates:
+        if os.path.isdir(d):
+            for name in ("profile.yaml", "config.yaml"):
+                if os.path.isfile(os.path.join(d, name)):
+                    return os.path.join(d, name)
+            return os.path.join(d, "config.yaml")
+    return os.path.join("profiles", "default", "config.yaml") if create else None
+
+def _only_installed_profile():
+    """The name of the one profile in ~/.config/snakemake that uses this plugin, if there is
+    exactly one (setup.sh installs it there)."""
+    base = os.path.join(os.environ.get("XDG_CONFIG_HOME") or os.path.expanduser("~/.config"), "snakemake")
+    names = []
+    try:
+        entries = sorted(os.listdir(base))
+    except OSError:
+        return None
+    for n in entries:
+        for f in ("profile.yaml", "config.yaml"):
+            cfg = os.path.join(base, n, f)
+            if os.path.isfile(cfg):
+                try:
+                    with open(cfg) as fh:
+                        if (yaml.safe_load(fh) or {}).get("executor") == "aws-batch-flex":
+                            names.append(n)
+                except (OSError, yaml.YAMLError):
+                    pass
+                break
+    if len(names) == 1:
+        print(f"(using profile {names[0]} from {base})", file=sys.stderr)
+        return names[0]
+    if names:
+        raise Fail(f"there are {len(names)} profiles for this plugin in {base} ({', '.join(names)}): "
+                   f"say which with --profile {names[0]} (or export SNAKEMAKE_PROFILE={names[0]})")
+    return None
+
+def _load_profile(path):
+    """(profile folder, settings) for --profile: a folder, or a name that Snakemake looks up in
+    ~/.config/snakemake/ (as `snakemake --profile NAME` does). Default: ./profile, then
+    $SNAKEMAKE_PROFILE."""
+    if path is None:
+        path = "profile" if os.path.isdir("profile") else os.environ.get("SNAKEMAKE_PROFILE")
+    if not path:
+        path = _only_installed_profile()
+    if not path:
+        raise Fail("no profile: pass --profile (a folder with config.yaml, or a profile name)")
+    found = None
+    try:
+        from snakemake.cli import get_profile_dir
+
+        found = get_profile_dir(path)
+    except ImportError:
+        cfg = os.path.join(path, "config.yaml")
+        found = (path, cfg) if os.path.isfile(cfg) else None
+    if not found:
+        raise Fail(f"profile {path} not found (looked for a folder with config.yaml or profile.yaml "
+                   "here and in ~/.config/snakemake/)")
+    folder, cfg = found
+    with open(cfg) as f:
+        settings = yaml.safe_load(f) or {}
+    # The project's own workflow profile, which Snakemake loads on top (e.g. its storage prefix)
+    wp = workflow_profile_file()
+    if wp and os.path.isfile(wp):
+        with open(wp) as f:
+            settings.update(yaml.safe_load(f) or {})
+    return str(folder), settings
+
+def _clients(profile):
+    cfg = BotoConfig(region_name=profile.get("aws-batch-flex-region"),
+                     retries={"max_attempts": 10, "mode": "adaptive"})
+    return boto3.client("batch", config=cfg), boto3.client("logs", config=cfg), cfg.region_name