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

snakemake_executor_plugin_aws_batch_flex/tags.py

@@ -0,0 +1,76 @@
+"""Job tags: global (--aws-batch-flex-tags), automatic, and per-rule (`batch_tags`).
+
+Precedence, lowest to highest: automatic < global < per-rule.
+
+AWS Batch tag limits: max 50 tags, key 1-128 chars, value up to 256, characters limited
+to letters, digits, spaces and _ . : / = + - @ (anything else is replaced with `_`).
+Keys starting with `aws:` are reserved.
+"""
+
+import re
+from typing import Iterable, Mapping, Optional
+
+from snakemake_interface_common.exceptions import WorkflowError
+
+_BAD = re.compile(r"[^A-Za-z0-9 _.:/=+\-@]")
+MAX_TAGS = 50
+# On every job, even without the automatic tags: the IAM policy in terraform/ lets users stop
+# only jobs that have it
+MANAGED_BY = {"ManagedBy": "snakemake-batch-flex"}
+
+
+def _clean(text: str, limit: int) -> str:
+    return _BAD.sub("_", str(text))[:limit]
+
+
+def parse_pairs(items: Iterable[str], source: str) -> dict:
+    """['k=v', ...] -> dict, validating each entry."""
+    out = {}
+    for item in items:
+        item = item.strip()
+        if not item:
+            continue
+        if "=" not in item:
+            raise WorkflowError(f"Invalid tag '{item}' in {source}, expected KEY=VALUE.")
+        k, v = item.split("=", 1)
+        k = k.strip()
+        if not k:
+            raise WorkflowError(f"Empty tag key in {source}: '{item}'.")
+        if k.lower().startswith("aws:"):
+            raise WorkflowError(f"Tag key '{k}' in {source} uses the reserved 'aws:' prefix.")
+        out[_clean(k, 128)] = _clean(v.strip(), 256)
+    return out
+
+
+def parse_resource_string(value: Optional[str]) -> dict:
+    """The `batch_tags` resource: 'k1=v1,k2=v2'."""
+    if not value:
+        return {}
+    return parse_pairs(str(value).split(","), "resource batch_tags")
+
+
+def job_tags(
+    *,
+    run_id: str,
+    workflow_id: str = "",
+    rule_names: Iterable[str],
+    jobid: int,
+    global_tags: Mapping[str, str],
+    rule_tags: Mapping[str, str],
+    auto: bool,
+) -> dict:
+    tags: dict = {}
+    if auto:
+        tags["snakemake-run-id"] = _clean(run_id, 256)
+        if workflow_id:
+            tags["snakemake-workflow"] = _clean(workflow_id, 256)
+        tags["snakemake-rule"] = _clean("+".join(sorted(set(rule_names))), 256)
+        tags["snakemake-job-id"] = _clean(str(jobid), 256)
+    tags.update(global_tags)
+    tags.update(rule_tags)
+    tags.update(MANAGED_BY)
+    if len(tags) > MAX_TAGS:
+        raise WorkflowError(
+            f"{len(tags)} tags exceed the AWS Batch limit of {MAX_TAGS}: {sorted(tags)}"
+        )
+    return tags

snakemake_executor_plugin_aws_batch_flex/transfer.py

@@ -0,0 +1,193 @@
+"""smkflex push and pull: files to and from the workflow's storage, at the paths Snakemake maps
+them to, 16 at a time; unchanged ones are skipped, as aws s3 sync does."""
+
+import os
+import shutil
+import sys
+import time
+
+import boto3
+from botocore.config import Config as BotoConfig
+
+from .profiles import Fail, _load_profile
+from .sources import SKIP_DIRS, workflow_files
+
+
+def _remote_files(prefix, path, s3):
+    """(relative path, size, modification time) of the files under prefix/path, in storage."""
+    path = path.strip("/")
+    if prefix.startswith("s3://"):
+        bucket, _, base = prefix[len("s3://"):].partition("/")
+        base = base.strip("/")
+        key = f"{base}/{path}" if base else path
+        out = []
+        for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=key):
+            for o in page.get("Contents", []):
+                rel = o["Key"][len(base) + 1:] if base else o["Key"]
+                # "results" must not also match "results2/..."
+                if rel == path or rel.startswith(path + "/") or not path:
+                    out.append((rel, o["Size"], o["LastModified"].timestamp()))
+        return out
+    root = os.path.join(prefix, path)
+    if os.path.isfile(root):
+        return [(path, os.path.getsize(root), os.path.getmtime(root))]
+    out = []
+    for d, _dirs, names in os.walk(root):
+        for n in names:
+            full = os.path.join(d, n)
+            out.append((os.path.relpath(full, prefix), os.path.getsize(full), os.path.getmtime(full)))
+    return sorted(out)
+
+def _same(size, mtime, other_size, other_mtime):
+    """A copy is up to date if it has the same size and was made after the original last
+    changed, as `aws s3 sync` decides (a copy's time is when it was copied)."""
+    return other_size == size and other_mtime >= mtime
+
+def _bucket_key(prefix, rel):
+    """(bucket, key) of the file at `rel` under an s3:// prefix."""
+    bucket, _, base = prefix[len("s3://"):].partition("/")
+    base = base.strip("/")
+    return bucket, f"{base}/{rel}" if base else rel
+
+TRANSFER_THREADS = 16  # files at once; big files are also split into parallel parts by boto3
+
+def _storage_path(path):
+    """Where Snakemake puts a workflow path in storage, relative to the prefix: as it is, or
+    for an absolute path without its leading /."""
+    norm = os.path.normpath(path)
+    if norm == ".." or norm.startswith(".." + os.sep):
+        raise Fail(f"{path}: Snakemake can't map paths outside the workflow folder to storage; "
+                   "use the path the workflow names (relative, or absolute like /shared/refs)")
+    return norm.lstrip("/") if os.path.isabs(norm) else norm
+
+def _transfer(jobs, work, verb):
+    """Run work(item) for every item, TRANSFER_THREADS at a time; prints each as it finishes."""
+    from concurrent.futures import ThreadPoolExecutor, as_completed
+
+    done = total = 0
+    start = time.time()
+    with ThreadPoolExecutor(TRANSFER_THREADS) as pool:
+        futures = {pool.submit(work, item): item for item in jobs}
+        for f in as_completed(futures):
+            name, size = futures[f][0], futures[f][-1]
+            f.result()  # re-raises the first error
+            done += 1
+            total += size
+            print(f"{verb} {name}")
+    secs = max(time.time() - start, 0.001)
+    return done, total, secs
+
+def _rate(total, secs):
+    mb = total / 1e6
+    return f"{mb:.1f} MB in {secs:.0f} s ({mb / secs:.1f} MB/s)" if total else ""
+
+def cmd_pull(a):
+    _dir, profile = _load_profile(a.profile)
+    prefix = str(a.prefix or profile.get("default-storage-prefix") or "").rstrip("/")
+    if not prefix:
+        raise Fail("no default-storage-prefix in the profile")
+    region = profile.get("aws-batch-flex-region")
+    s3 = boto3.client("s3", config=BotoConfig(region_name=region, max_pool_connections=TRANSFER_THREADS * 2)) \
+        if prefix.startswith("s3://") else None
+    if not a.paths:
+        tops = sorted({rel.split("/")[0] for rel, *_ in _remote_files(prefix, "", s3)
+                       if not rel.startswith((".smkflex/", "snakemake-workflow-sources."))})
+        raise Fail("say what to download, e.g. `smkflex pull results`. In storage: " + (", ".join(tops) or "nothing"))
+    todo, skipped = [], 0
+    for path in a.paths:
+        files = _remote_files(prefix, _storage_path(path), s3)
+        if not files:
+            print(f"nothing in storage under {path}", file=sys.stderr)
+        for rel, size, mtime in files:
+            dest = os.path.join(a.to, rel)
+            if not a.force and os.path.isfile(dest) and _same(size, mtime, os.path.getsize(dest), os.path.getmtime(dest)):
+                skipped += 1
+                continue
+            if a.dry_run:
+                print(f"would download {rel} ({size} bytes)")
+                continue
+            todo.append((rel, dest, size))
+
+    def get(item):
+        rel, dest, _size = item
+        os.makedirs(os.path.dirname(dest) or ".", exist_ok=True)
+        if s3:
+            s3.download_file(*_bucket_key(prefix, rel), dest)
+        else:
+            shutil.copyfile(os.path.join(prefix, rel), dest)
+
+    if not a.dry_run:
+        got, total, secs = _transfer(todo, get, "downloaded")
+        rate = _rate(total, secs)
+        print(f"{got} downloaded" + (f" ({rate})" if rate else "")
+              + (f", {skipped} already here (unchanged; --force downloads them again)" if skipped else ""))
+    return 0
+
+def cmd_push(a):
+    """Upload local input files to where the workflow's rules look for them in storage: a
+    relative path to the same path under the prefix, an absolute one (/shared/refs) to
+    prefix/shared/refs, as Snakemake maps them."""
+    _dir, profile = _load_profile(a.profile)
+    prefix = str(profile.get("default-storage-prefix") or "").rstrip("/")
+    if a.prefix:
+        prefix = a.prefix.rstrip("/")
+    if not prefix:
+        raise Fail("no default-storage-prefix in the profile")
+    files = {}
+    for path in a.paths:
+        target = _storage_path(path)
+        if not os.path.exists(path):
+            raise Fail(f"{path}: no such file or folder")
+        if os.path.isfile(path):
+            files[target] = path
+            continue
+        for d, dirs, names in os.walk(path):
+            dirs[:] = [x for x in dirs if x not in SKIP_DIRS]
+            for n in names:
+                local = os.path.join(d, n)
+                files[os.path.normpath(os.path.join(target, os.path.relpath(local, path)))] = local
+    s3 = boto3.client("s3", config=BotoConfig(region_name=profile.get("aws-batch-flex-region"),
+                                              max_pool_connections=TRANSFER_THREADS * 2)) \
+        if prefix.startswith("s3://") else None
+    there = {}
+    for top in {f.split("/")[0] for f in files}:
+        there.update({rel: (size, mtime) for rel, size, mtime in _remote_files(prefix, top, s3)})
+    todo, skipped = [], 0
+    for rel, local in sorted(files.items()):
+        size = os.path.getsize(local)
+        if not a.force and rel in there and _same(size, os.path.getmtime(local), *there[rel]):
+            skipped += 1
+            continue
+        if a.dry_run:
+            print(f"would upload {rel} ({size} bytes)")
+            continue
+        todo.append((rel, local, size))
+
+    def put(item):
+        rel, local, _size = item
+        if s3:
+            s3.upload_file(local, *_bucket_key(prefix, rel))
+        else:
+            os.makedirs(os.path.dirname(os.path.join(prefix, rel)), exist_ok=True)
+            shutil.copyfile(local, os.path.join(prefix, rel))
+
+    if not a.dry_run:
+        sent, total, secs = _transfer(todo, put, "uploaded")
+        rate = _rate(total, secs)
+        print(f"{sent} uploaded to {prefix}" + (f" ({rate})" if rate else "")
+              + (f", {skipped} already there (unchanged; --force sends them again)" if skipped else ""))
+        _hint_gitignore([p for p in a.paths if not os.path.isabs(p)])
+    return 0
+
+def _hint_gitignore(paths):
+    """Data uploaded with push reaches the jobs through storage; if smkflex run would also
+    pack it with the workflow, say how to leave it out."""
+    if not paths:
+        return
+    sent = set(workflow_files("."))
+    tops = sorted({os.path.normpath(p).split(os.sep)[0] for p in paths
+                   if any(f == os.path.normpath(p) or f.startswith(os.path.normpath(p) + "/") for f in sent)})
+    if tops:
+        lines = "".join(f"{t}/\n" if os.path.isdir(t) else f"{t}\n" for t in tops)
+        print(f"\nNote: smkflex run would also send {', '.join(tops)} with the workflow (it doesn't need to: the "
+              f"jobs read it from storage). To leave it out, add to .gitignore:\n{lines}", end="")

snakemake_executor_plugin_aws_batch_flex/usage.py

@@ -0,0 +1,208 @@
+"""What a run used on Batch, per rule, from Batch's own start and stop times of every attempt
+(Spot retries included) and the figures each job logs (jobmonitor.py): vCPU- and GiB-hours,
+memory, CPUs and disk, with the resources to paste into the profile. Printed when the run ends.
+
+What it cost is AWS's to say: every job carries the tags snakemake-run-id and snakemake-rule,
+which AWS Billing (with split cost allocation data) breaks costs down by."""
+
+import math
+import re
+import statistics
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Optional
+
+from botocore.exceptions import BotoCoreError, ClientError
+
+# A Batch job costs time on top of its command: measured on real runs, about 15 s from submitted to
+# running on an instance that's already up (a minute or more when one has to start), and about
+# 10 s for Snakemake in the job and its files. Many short jobs of one rule are better grouped.
+SHORT_JOB_SECONDS = 60
+MANY_JOBS = 20
+
+# Usage figures: each job logs them when it ends (jobmonitor.py). Reading the logs costs an
+# API call per job, so big runs read a sample of each rule's jobs.
+USAGE_RE = re.compile(r"\[snakemake-batch\] (peak-memory-bytes=\S+.*)")
+PEAK_SAMPLE = 25  # jobs per rule, for runs with more than ALL_PEAKS jobs
+ALL_PEAKS = 300
+HEADROOM = 1.3  # suggested memory = peak * HEADROOM
+THREAD_HEADROOM = 1.25  # suggested threads = most in use at once * THREAD_HEADROOM
+FEW_THREADS = 0.6  # suggest fewer threads when a rule never had more than this share busy
+BIG_DISK_GIB = 20  # say so when a rule's job used more disk than this
+
+
+def _seconds(d: dict) -> float:
+    attempts = d.get("attempts") or []
+    spans = [(a.get("startedAt"), a.get("stoppedAt")) for a in attempts] or [(d.get("startedAt"), d.get("stoppedAt"))]
+    return sum(max(0, stop - start) / 1000 for start, stop in spans if start is not None and stop is not None)
+
+
+@dataclass
+class Row:
+    """One rule's jobs."""
+    jobs: int = 0
+    vcpu_h: float = 0.0
+    gib_h: float = 0.0
+
+
+class Usage:
+    def __init__(self):
+        self.rows = defaultdict(Row)  # rule -> Row
+        self.seconds = defaultdict(list)  # rule -> run time of each job that ran
+        self.asked = defaultdict(int)  # rule -> most memory a job asked for (MiB)
+        self.logs = defaultdict(list)  # rule -> [(log group, stream)] of jobs that succeeded
+        self.peaks = defaultdict(list)  # rule -> peak memory of jobs (MiB) read from their logs
+        self.busy = defaultdict(list)  # rule -> most vCPUs in use at once, per job
+        self.disk = defaultdict(list)  # rule -> most disk a job's working folder took (GiB)
+        self.threads = defaultdict(float)  # rule -> most vCPUs a job asked for
+        self.oom_kills = defaultdict(int)  # rule -> its jobs' processes killed for lack of memory
+        self.platform = {}  # rule -> ec2 | fargate
+        self.sampled = False
+
+    def add(self, rule: str, d: dict, vcpu: float, mem_mib: int, platform: str):
+        secs = _seconds(d)
+        if secs:
+            self.seconds[rule].append(secs)
+        hours = secs / 3600
+        self.asked[rule] = max(self.asked[rule], mem_mib)
+        self.threads[rule] = max(self.threads[rule], vcpu)
+        self.platform[rule] = platform
+        row = self.rows[rule]
+        row.jobs += 1
+        row.vcpu_h += vcpu * hours
+        row.gib_h += mem_mib / 1024 * hours
+
+    def add_log(self, rule: str, group: str, stream: str):
+        self.logs[rule].append((group, stream))
+
+    def fetch_peaks(self, logs_client):
+        """Read the peak-memory line from the end of each job's log (a sample per rule on big
+        runs). Best effort: jobs whose line can't be read are left out."""
+        total = sum(len(v) for v in self.logs.values())
+        todo = []
+        for rule, streams in self.logs.items():
+            if total > ALL_PEAKS and len(streams) > PEAK_SAMPLE:
+                self.sampled = True
+                streams = streams[-PEAK_SAMPLE:]
+            todo += [(rule, g, st) for g, st in streams]
+
+        def one(item):
+            rule, group, stream = item
+            try:
+                events = logs_client.get_log_events(
+                    logGroupName=group, logStreamName=stream, limit=20, startFromHead=False
+                )["events"]
+            except (ClientError, BotoCoreError):
+                return rule, None
+            for e in reversed(events):
+                m = USAGE_RE.search(e.get("message", ""))
+                if m:
+                    return rule, dict(kv.split("=", 1) for kv in m.group(1).split() if "=" in kv)
+            return rule, None
+
+        with ThreadPoolExecutor(8) as pool:
+            for rule, fields in pool.map(one, todo):
+                if not fields:
+                    continue
+                try:
+                    if int(fields.get("peak-memory-bytes", 0)):
+                        self.peaks[rule].append(int(fields["peak-memory-bytes"]) / 1024 / 1024)
+                    if "peak-cpus" in fields:
+                        self.busy[rule].append(float(fields["peak-cpus"]))
+                    if "peak-disk-bytes" in fields:
+                        self.disk[rule].append(int(fields["peak-disk-bytes"]) / 1024 ** 3)
+                    self.oom_kills[rule] += int(fields.get("oom-kills", 0))
+                except ValueError:
+                    pass
+
+    def suggestions(self, overhead_mib: int) -> dict:
+        """rule -> mem_mb that gives its jobs their peak plus headroom (at least 512 MiB in
+        all), for every rule whose peak is known. Not for rules whose jobs ran out: what was
+        seen before that is less than they need."""
+        out = {}
+        for rule, peaks in self.peaks.items():
+            if self.oom_kills.get(rule):
+                continue
+            want = max(512, math.ceil(max(peaks) * HEADROOM / 100) * 100)  # container memory, MiB
+            # the plugin adds overhead_mib for Snakemake itself; mem_mb is in MB
+            out[rule] = max(100, math.ceil((want - overhead_mib) * 1.048576 / 100) * 100)
+        return out
+
+    def thread_suggestions(self) -> dict:
+        """rule -> threads, for rules that never had most of their vCPUs busy at once."""
+        out = {}
+        for rule, busy in self.busy.items():
+            asked, most = self.threads.get(rule, 0), max(busy)
+            if asked >= 2 and most <= asked * FEW_THREADS:
+                out[rule] = max(1, min(int(asked) - 1, math.ceil(most * THREAD_HEADROOM)))
+        return out
+
+    def table(self, overhead_mib: int = 256) -> Optional[str]:
+        if not self.rows:
+            return None
+        width = max(10, *(len(r) for r in self.rows))
+        lines = [
+            "AWS Batch usage of this run (costs: AWS Billing, by the jobs' tags snakemake-run-id and "
+            "snakemake-rule):",
+            f"  {'rule':<{width}}  {'jobs':>5}  {'vCPU-h':>8}  {'GiB-h':>8}"
+            f"  {'vCPUs':>5}  {'busy':>5}  {'mem MiB':>8}  {'peak':>6}  {'disk GiB':>8}",
+        ]
+        total = Row()
+        # the rules that took the most first
+        for rule in sorted(self.rows, key=lambda r: -self.rows[r].vcpu_h):
+            row = self.rows[rule]
+            n, v, g = row.jobs, row.vcpu_h, row.gib_h
+            peak = f"{max(self.peaks[rule]):.0f}" if self.peaks.get(rule) else "-"
+            busy = f"{max(self.busy[rule]):.1f}" if self.busy.get(rule) else "-"
+            disk = f"{max(self.disk[rule]):.1f}" if self.disk.get(rule) else "-"
+            vcpus = f"{self.threads.get(rule, 0):g}"
+            lines.append(f"  {rule:<{width}}  {n:>5}  {v:>8.2f}  {g:>8.2f}"
+                         f"  {vcpus:>5}  {busy:>5}  {self.asked.get(rule, 0):>8}  {peak:>6}  {disk:>8}")
+            total = Row(total.jobs + n, total.vcpu_h + v, total.gib_h + g)
+        if len(self.rows) > 1:
+            lines.append(f"  {'total':<{width}}  {total.jobs:>5}  {total.vcpu_h:>8.2f}  {total.gib_h:>8.2f}")
+        if self.peaks:
+            note = f" (from up to {PEAK_SAMPLE} jobs per rule)" if self.sampled else ""
+            lines.append(f"vCPUs, mem MiB: what a job had; busy: the most vCPUs in use at once (on EC2 a job can use "
+                         "idle vCPUs of its instance, so this can be more than it had); peak: the most "
+                         f"memory its processes used (without file cache); disk GiB: the most its working "
+                         f"folder took. Each the highest of the rule's jobs{note}; memory sampled every "
+                         f"0.5 s (a shorter spike isn't seen), CPUs every 2 s.")
+        tips = self.suggestions(overhead_mib)
+        if tips:
+            lines.append("Memory for each rule (its peak + 30%; mem_mb leaves out the plugin's own "
+                         f"{overhead_mib} MiB, added on top), to paste into the project's "
+                         "profiles/default/config.yaml (smkflex init makes it):")
+            lines.append("set-resources:")
+            for rule, mb in sorted(tips.items()):
+                had = self.asked.get(rule)
+                gets = int(mb / 1.048576) + overhead_mib  # what the job is given for that, MiB
+                lines += [f"  {rule}:",
+                          f"    mem_mb: {mb:<8}# job gets {gets} MiB; had {had}, peak {max(self.peaks[rule]):.0f}"]
+        threads = self.thread_suggestions()
+        if threads:
+            lines.append("Threads for rules that never had most of theirs busy (the most in use at once "
+                         "+ 25%; a burst shorter than 2 s isn't seen, so check the tool first):")
+            lines.append("set-threads:")
+            for rule, n in sorted(threads.items()):
+                lines.append(f"  {rule}: {n:<6}# had {self.threads[rule]:g}, at most {max(self.busy[rule]):.1f} busy")
+        for rule, n in sorted(self.oom_kills.items()):
+            if n:
+                lines.append(f"Hint: in {rule}'s jobs, Linux killed {n} process(es) for lack of memory (a job "
+                             "can still succeed if the tool retried). Its peak is from before that, so it "
+                             "needs more: raise its mem_mb.")
+        for rule, gib in sorted((r, max(d)) for r, d in self.disk.items() if d and max(d) > BIG_DISK_GIB):
+            where = ("Fargate gives a job 20 GiB unless disk_mb asks for more"
+                     if self.platform.get(rule) == "fargate" else
+                     "on EC2 the jobs on an instance share its disk (root_volume_gib in Terraform)")
+            lines.append(f"Hint: {rule}'s jobs used up to {gib:.0f} GiB of disk each; {where}.")
+        for rule, secs in sorted(self.seconds.items()):
+            if len(secs) >= MANY_JOBS and statistics.median(secs) < SHORT_JOB_SECONDS:
+                lines.append(
+                    f"Hint: {rule} ran {len(secs)} jobs of {statistics.median(secs):.0f} s each (median). Every "
+                    "Batch job adds about 25 s (a minute or more when a new instance has to start), so "
+                    "grouping them makes the run faster: "
+                    "see \"Many small jobs\" in the plugin's README."
+                )
+        return "\n".join(lines)