SCHED-1918: Improve multi-node startup with concurrent Pyxis cache reads

helm/slurm-cluster/slurm_scripts/pyxis_caching_importer.sh

@@ -6,13 +6,57 @@ set -euo pipefail
 readonly cmd="$1"
 
 readonly cache_dir="${ENROOT_CONTAINER_IMAGES_CACHE_DIR:-/var/cache/enroot-container-images}"
-readonly squashfs_temp_path="${cache_dir}/${SLURM_JOB_ID}.${SLURM_STEP_ID}.sqsh"
-readonly lock_path="${squashfs_temp_path}.lock"
-trap 'rm -f "${lock_path}"' EXIT
+readonly node_id="${SLURMD_NODENAME:-${HOSTNAME:-unknown}}"
+readonly squashfs_temp_path="${cache_dir}/${SLURM_JOB_ID}.${SLURM_STEP_ID}.${node_id}.sqsh"
 
 # Since it's not an ephemeral squashfs file, we can use compression.
 export ENROOT_SQUASH_OPTIONS="-comp zstd -Xcompression-level 3 -b 1M"
 
+publish_image()
+{
+    # The cache may have been populated after the lock was requested.
+    if [ -e "${squashfs_path}" ]; then
+        return
+    fi
+
+    # Each node owns a separate temporary path, so cleanup from one node cannot
+    # remove another node's in-progress import.
+    rm -f "${squashfs_temp_path}"
+    trap 'rm -f "${squashfs_temp_path}"' EXIT
+
+    # TODO: use `digest` approach once 406 Not Acceptable is tolerated in enroot
+    # https://github.com/NVIDIA/enroot/pull/263
+    # if [[ "${image_uri}" == *"@${digest}" ]]; then
+    #     # URI already has the digest in it.
+    #     enroot import --output "${squashfs_temp_path}" "${image_uri}" >&2
+    # else
+    #     # Add the digest to the URI.
+    #     enroot import --output "${squashfs_temp_path}" "${image_uri}@${digest}" >&2
+    # fi
+    enroot import --output "${squashfs_temp_path}" "${image_uri}" >&2
+
+    # Save the URI as an extended attribute.
+    if command -v "setfattr" >/dev/null; then
+        setfattr -n user.image_uri -v "${image_uri}" "${squashfs_temp_path}"
+    fi
+
+    chmod 777 "${squashfs_temp_path}"
+
+    # The temporary and final paths are in the same directory, so rename
+    # publishes a complete SquashFS image atomically.
+    mv -n "${squashfs_temp_path}" "${squashfs_path}"
+
+    if [ ! -e "${squashfs_path}" ]; then
+        echo "error: could not publish image cache: ${squashfs_path}" >&2
+        exit 1
+    fi
+
+    # mv -n can leave the source in place if an external writer published the
+    # destination without following this locking protocol.
+    rm -f "${squashfs_temp_path}"
+    trap - EXIT
+}
+
 case "${cmd}" in
     get)
         if [ $# -ne 2 ]; then
@@ -30,30 +74,39 @@ case "${cmd}" in
             exit 1
         fi
         readonly squashfs_path="${cache_dir}/${digest}.sqsh"
+        readonly lock_path="${squashfs_path}.lock"
 
-        # Serialize access to shared temp/output paths for this SLURM job step.
-        exec 9>"${lock_path}"
-        flock -x 9
-
-        if [ ! -e "${squashfs_path}" ]; then
-            # TODO: use `digest` approach once 406 Not Acceptable is tolerated in enroot
-            # https://github.com/NVIDIA/enroot/pull/263
-            # if [[ "${image_uri}" == *"@${digest}" ]]; then
-            #     # URI already has the digest in it.
-            #     enroot import --output "${squashfs_temp_path}" "${image_uri}" >&2
-            # else
-            #     # Add the digest to the URI.
-            #     enroot import --output "${squashfs_temp_path}" "${image_uri}@${digest}" >&2
-            # fi
-            enroot import --output "${squashfs_temp_path}" "${image_uri}" >&2
-
-            # Save the URI as an extended attribute.
-            if command -v "setfattr" >/dev/null; then
-                setfattr -n user.image_uri -v "${image_uri}" "${squashfs_temp_path}"
-            fi
+        # Warm-cache fast path: completed images are immutable and can be read
+        # concurrently, so readers do not need to join the exclusive lock queue.
+        if [ -e "${squashfs_path}" ]; then
+            echo "${squashfs_path}"
+            exit 0
+        fi
+
+        # Use a stable, digest-scoped lock. Do not unlink it: flock locks the
+        # opened inode, and unlinking it can let callers lock different inodes.
+        # The cache directory is shared between Slurm users, so lock files must
+        # be writable by callers other than their creator.
+        readonly previous_umask=$(umask)
+        umask 000
+        exec 9>>"${lock_path}"
+        umask "${previous_umask}"
 
-            chmod 777 "${squashfs_temp_path}"
-            mv -n "${squashfs_temp_path}" "${squashfs_path}"
+        if flock -n -x 9; then
+            # This caller is the cold-cache producer.
+            publish_image
+        else
+            # A producer is already active. Shared waiters are released
+            # together when the producer closes its exclusive lock.
+            flock -s 9
+
+            if [ ! -e "${squashfs_path}" ]; then
+                # The producer exited without publishing. Queue for exclusive
+                # ownership so one waiter retries while the others remain safe.
+                flock -u 9
+                flock -x 9
+                publish_image
+            fi
         fi
 
         # Output the squashfs path on stdout for pyxis to read
@@ -65,11 +118,8 @@ case "${cmd}" in
             exit 1
         fi
 
-        # Serialize cleanup with "get" so temp path isn't removed mid-import.
-        exec 9>"${lock_path}"
-        flock -x 9
-
-        # Remove temporary file if still present (e.g. "get" was interrupted)
+        # Temporary paths are node-specific, so releases do not need to be
+        # serialized and cannot remove another node's active import.
         rm -f "${squashfs_temp_path}"
         ;;
     *)

helm/slurm-cluster/slurm_scripts/pyxis_caching_importer_test.py

@@ -0,0 +1,195 @@
+import os
+import stat
+import subprocess
+import tempfile
+import textwrap
+import unittest
+from pathlib import Path
+
+
+IMPORTER_PATH = Path(__file__).with_name("pyxis_caching_importer.sh")
+DIGEST = "sha256-test"
+
+
+FAKE_ENROOT = """\
+#!/usr/bin/env python3
+import fcntl
+import os
+import sys
+import time
+from pathlib import Path
+
+command = sys.argv[1]
+if command == "digest":
+    print(os.environ["FAKE_ENROOT_DIGEST"])
+    sys.exit(0)
+
+if command != "import":
+    raise SystemExit(f"unsupported enroot command: {command}")
+
+output = Path(sys.argv[sys.argv.index("--output") + 1])
+count_path = Path(os.environ["FAKE_ENROOT_IMPORT_COUNT"])
+with count_path.open("a+") as count_file:
+    fcntl.flock(count_file.fileno(), fcntl.LOCK_EX)
+    count_file.seek(0)
+    count = int(count_file.read() or "0") + 1
+    count_file.seek(0)
+    count_file.truncate()
+    count_file.write(str(count))
+    count_file.flush()
+
+output.write_text("incomplete")
+if os.environ.get("FAKE_ENROOT_FAIL") == "1":
+    sys.exit(1)
+time.sleep(float(os.environ.get("FAKE_ENROOT_IMPORT_DELAY", "0")))
+output.write_text("complete")
+"""
+
+
+FAKE_FLOCK = """\
+#!/usr/bin/env python3
+import fcntl
+import os
+import sys
+
+arguments = sys.argv[1:]
+operation = fcntl.LOCK_UN
+if "-x" in arguments:
+    operation = fcntl.LOCK_EX
+elif "-s" in arguments:
+    operation = fcntl.LOCK_SH
+if "-n" in arguments:
+    operation |= fcntl.LOCK_NB
+
+fd = int(arguments[-1])
+try:
+    fcntl.flock(fd, operation)
+except BlockingIOError:
+    sys.exit(1)
+
+log_path = os.environ.get("FAKE_FLOCK_LOG")
+if log_path:
+    with open(log_path, "a") as log:
+        log.write(" ".join(arguments) + "\\n")
+"""
+
+
+class PyxisCachingImporterTest(unittest.TestCase):
+    def setUp(self):
+        self.temporary_directory = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temporary_directory.cleanup)
+        self.root = Path(self.temporary_directory.name)
+        self.cache = self.root / "cache"
+        self.cache.mkdir()
+        self.bin = self.root / "bin"
+        self.bin.mkdir()
+        self.import_count = self.root / "import-count"
+        self.flock_log = self.root / "flock.log"
+
+        self.write_executable("enroot", FAKE_ENROOT)
+        self.write_executable("flock", FAKE_FLOCK)
+        self.write_executable("setfattr", "#!/bin/sh\nexit 0\n")
+
+        self.environment = os.environ.copy()
+        self.environment.update(
+            {
+                "PATH": f"{self.bin}{os.pathsep}{self.environment['PATH']}",
+                "ENROOT_CONTAINER_IMAGES_CACHE_DIR": str(self.cache),
+                "FAKE_ENROOT_DIGEST": DIGEST,
+                "FAKE_ENROOT_IMPORT_COUNT": str(self.import_count),
+                "FAKE_FLOCK_LOG": str(self.flock_log),
+                "SLURM_JOB_ID": "100",
+                "SLURM_STEP_ID": "0",
+            }
+        )
+
+    def write_executable(self, name, contents):
+        path = self.bin / name
+        path.write_text(textwrap.dedent(contents))
+        path.chmod(0o755)
+
+    def run_get(self, node):
+        environment = self.environment.copy()
+        environment["SLURMD_NODENAME"] = node
+        return subprocess.Popen(
+            ["bash", str(IMPORTER_PATH), "get", "docker://example/image:tag"],
+            env=environment,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+
+    def collect(self, processes):
+        results = [process.communicate(timeout=10) for process in processes]
+        for process, (_, stderr) in zip(processes, results):
+            self.assertEqual(0, process.returncode, stderr)
+        return results
+
+    def test_warm_cache_readers_do_not_call_flock(self):
+        final_path = self.cache / f"{DIGEST}.sqsh"
+        final_path.write_text("complete")
+
+        results = self.collect([self.run_get(f"node-{index}") for index in range(10)])
+
+        self.assertEqual(
+            [str(final_path)] * 10,
+            [stdout.strip() for stdout, _ in results],
+        )
+        self.assertFalse(self.import_count.exists())
+        self.assertFalse(self.flock_log.exists())
+
+    def test_cold_cache_has_one_producer_and_concurrent_waiters(self):
+        self.environment["FAKE_ENROOT_IMPORT_DELAY"] = "0.3"
+
+        results = self.collect([self.run_get(f"node-{index}") for index in range(10)])
+
+        final_path = self.cache / f"{DIGEST}.sqsh"
+        self.assertEqual("1", self.import_count.read_text())
+        self.assertEqual("complete", final_path.read_text())
+        self.assertEqual(
+            0o666,
+            stat.S_IMODE((self.cache / f"{DIGEST}.sqsh.lock").stat().st_mode),
+        )
+        self.assertEqual(
+            [str(final_path)] * 10,
+            [stdout.strip() for stdout, _ in results],
+        )
+
+        flock_calls = self.flock_log.read_text().splitlines()
+        self.assertIn("-n -x 9", flock_calls)
+        self.assertIn("-s 9", flock_calls)
+
+    def test_release_removes_only_the_calling_nodes_temporary_file(self):
+        node_1_temp = self.cache / "100.0.node-1.sqsh"
+        node_2_temp = self.cache / "100.0.node-2.sqsh"
+        node_1_temp.write_text("node 1")
+        node_2_temp.write_text("node 2")
+        environment = self.environment.copy()
+        environment["SLURMD_NODENAME"] = "node-1"
+
+        result = subprocess.run(
+            ["bash", str(IMPORTER_PATH), "release"],
+            env=environment,
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+
+        self.assertEqual(0, result.returncode, result.stderr)
+        self.assertFalse(node_1_temp.exists())
+        self.assertTrue(node_2_temp.exists())
+        self.assertFalse(self.flock_log.exists())
+
+    def test_failed_import_does_not_publish_or_leave_a_temporary_file(self):
+        self.environment["FAKE_ENROOT_FAIL"] = "1"
+
+        process = self.run_get("node-1")
+        _, stderr = process.communicate(timeout=10)
+
+        self.assertNotEqual(0, process.returncode, stderr)
+        self.assertFalse((self.cache / f"{DIGEST}.sqsh").exists())
+        self.assertFalse((self.cache / "100.0.node-1.sqsh").exists())
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)