Merge the individuals archives without holding them in memory

The merge accumulated every individual's merged content before writing any of
it. On a laptop that is merely wasteful; under the 1GiB a container is allowed
on the cluster it is fatal, and chr6's fifteen inputs -- 109MB compressed,
several times that expanded -- killed the task outright with OOMKilled.

Every input carries the same members in the same order, so they are now walked
in lockstep and each individual's merged content is written as it is formed.
Reading each input once is what makes it affordable: a .tar.gz is not seekable,
so looking members up by name would decompress the whole archive per lookup.

Building the archive from a fully-joined mapping had the same shape of problem
in miniature, doubling peak memory at the moment of writing, so write_archive
now joins one member at a time and accepts either text or a sequence of strings.

Merging chr7's three chunks: 36.6MB peak and 3.31s before, 16.5MB and 1.27s
after, with the output archive identical member for member.

Verified against the previous implementation on chr11 split three ways: all
three individuals archives and the merged one match on member names, types,
modes and content hashes.

Worker image 1.7.

Author: balis

.github/workflows/equivalence.yml

@@ -76,7 +76,7 @@ jobs:
           make -C worker-image image
           docker build --platform linux/amd64 \
             -f engines/nextflow/worker-nf.Dockerfile \
-            -t 1000genome-worker-nf:1.6 engines/nextflow/
+            -t 1000genome-worker-nf:1.7 engines/nextflow/
 
       - name: Install Nextflow
         run: |

engines/hyperflow/harness/docker-compose.yml

@@ -32,7 +32,7 @@ services:
       - REDIS_URL=redis://redis:6379
 
       # HyperFlow configuration
-      - HF_VAR_WORKER_CONTAINER=${HF_VAR_WORKER_CONTAINER:-hyperflowwms/1000genome-worker:1.6-je1.4.2}
+      - HF_VAR_WORKER_CONTAINER=${HF_VAR_WORKER_CONTAINER:-hyperflowwms/1000genome-worker:1.7-je1.4.2}
       - HF_VAR_WORK_DIR=${WORKFLOW_DIR}
       - HF_VAR_HFLOW_IN_CONTAINER=true
       - HF_VAR_function=redisCommand

engines/nextflow/worker-nf.Dockerfile

@@ -6,7 +6,7 @@
 # Nextflow does not use. Both engines therefore run the same scripts from the
 # same layer, so "the science is identical" is a property of the image graph
 # rather than a claim about two copies staying in sync.
-FROM hyperflowwms/1000genome-worker-base:1.6
+FROM hyperflowwms/1000genome-worker-base:1.7
 
 # Nextflow requires bash in the container; the base is Alpine, which ships ash.
 RUN apk add --no-cache bash

worker-base-image/Makefile

@@ -1,6 +1,6 @@
 REPO_NAME = 1000genome-worker-base
 PREFIX = hyperflowwms
-VERSION = 1.6
+VERSION = 1.7
 TAG = $(VERSION)
 
 all: push

worker-base-image/scripts/archive.py

@@ -41,7 +41,13 @@ def _member(name, mode, mtime, uid, gid):
 
 
 def write_archive(archive, contents):
-    """Write `contents`, a mapping of member name to member text, as a .tar.gz."""
+    """Write `contents` as a .tar.gz.
+
+    Values may be a member's text or a sequence of strings to concatenate. The
+    joining happens one member at a time: asking the caller to build a second,
+    fully-joined mapping doubles peak memory, which is what pushed the merge of
+    a large region past a 1GiB container limit.
+    """
     mtime = int(time.time())
     uid, gid = os.getuid(), os.getgid()
 
@@ -51,12 +57,63 @@ def write_archive(archive, contents):
         handle.addfile(root)
 
         for name in sorted(contents):
-            payload = contents[name].encode()
+            payload = contents[name]
+            if not isinstance(payload, str):
+                payload = ''.join(payload)
+            payload = payload.encode()
             member = _member(name, FILE_MODE, mtime, uid, gid)
             member.size = len(payload)
             handle.addfile(member, io.BytesIO(payload))
 
 
+def merge_archives(archive, inputs):
+    """Concatenate the same-named members of `inputs` into one archive.
+
+    Holding the merged result in memory does not scale: the fifteen individuals
+    archives of a chr6 region are 109MB compressed and several times that
+    expanded, well past the 1GiB a container is allowed here.
+
+    Every input carries the same members in the same sorted order, because each
+    was written from the same column list, so the inputs are walked in lockstep
+    and only one individual's merged content exists at a time. Walking them once
+    each is what makes this affordable -- a .tar.gz is not seekable, so looking
+    members up by name would decompress the whole archive on every lookup.
+    """
+    mtime = int(time.time())
+    uid, gid = os.getuid(), os.getgid()
+    handles = [tarfile.open(path, "r:*") for path in inputs]
+
+    def members(handle):
+        # A named function, not a generator expression: an expression closes over
+        # the loop variable and every stream would end up reading the last
+        # archive opened.
+        for member in handle:
+            if member.isfile():
+                yield member.name, handle.extractfile(member)
+
+    try:
+        # Each input yields its members in the same order; zip walks them together.
+        streams = [members(h) for h in handles]
+        with tarfile.open(archive, "w:gz") as out:
+            root = _member("", DIR_MODE, mtime, uid, gid)
+            root.type = tarfile.DIRTYPE
+            out.addfile(root)
+
+            for group in zip(*streams):
+                names = {name for name, _ in group}
+                if len(names) != 1:
+                    raise ValueError(
+                        "archives disagree on member order: %s" % sorted(names))
+                name = names.pop()
+                payload = b''.join(fh.read() for _, fh in group if fh is not None)
+                member = _member(name, FILE_MODE, mtime, uid, gid)
+                member.size = len(payload)
+                out.addfile(member, io.BytesIO(payload))
+    finally:
+        for h in handles:
+            h.close()
+
+
 def read_archive(archive):
     """Yield `(name, lines)` for every regular member, never touching the disk.
 

worker-base-image/scripts/individuals_merge.py

@@ -10,29 +10,18 @@ def merging(c, tar_files):
     print('= Merging chromosome {}...'.format(c))
     tic = time.perf_counter()
 
-    # Each input is read through memory and the result written straight back
-    # out. Extracting an input into a temporary directory and reading the files
-    # back cost three filesystem operations per individual per input, which
-    # dominates the stage on a shared network volume; it also needed a
-    # chromosome-specific staging directory to keep parallel merges apart, and
-    # with no directory there is nothing left to collide. See archive.py.
-    data = {}
-
-    for tar in tar_files:
-        tic_iter = time.perf_counter()
-        for name, lines in archive.read_archive(tar):
-            if name in data:
-                data[name] += lines
-            else:
-                data[name] = lines
-
-        print("Merged {} in {:0.2f} sec".format(tar, time.perf_counter()-tic_iter))
-
+    # The inputs are streamed in lockstep and the result written as it goes, so
+    # only one individual's merged content is ever held. Accumulating the whole
+    # merge first does not fit: chr6's fifteen inputs are 109MB compressed and
+    # several times that expanded, against the 1GiB a container gets here.
+    #
+    # This also needs no staging directory, so the chromosome-specific one that
+    # used to keep parallel merges apart is gone, and with it anything to
+    # collide over. See archive.py.
     outputfile = "chr{}n.tar.gz".format(c)
-    print("== Done. Zipping {} files into {}.".format(len(data), outputfile))
+    print("== Merging {} archives into {}.".format(len(tar_files), outputfile))
 
-    archive.write_archive(
-        outputfile, {name: ''.join(lines) for name, lines in data.items()})
+    archive.merge_archives(outputfile, tar_files)
 
     print("= Chromosome {} merged in {:0.2f} seconds.".format(
         c, time.perf_counter() - tic))

worker-image/Dockerfile

@@ -1,4 +1,4 @@
-FROM hyperflowwms/1000genome-worker-base:1.6
+FROM hyperflowwms/1000genome-worker-base:1.7
 
 # Version of the job executor should be passed via docker build, e.g.:
 # docker build --build-arg hf_job_executor_version="1.3.4"

worker-image/Makefile

@@ -1,6 +1,6 @@
 REPO_NAME = 1000genome-worker
 PREFIX = hyperflowwms
-VERSION = 1.6
+VERSION = 1.7
 # ?= so CI and Renovate can override the version from the environment
 HF_JOB_EXECUTOR_VERSION ?= 1.4.2
 TAG = $(VERSION)-je$(HF_JOB_EXECUTOR_VERSION)