Extract shared atomic-publish helper in shell command template

src/groundhog_hpc/templates/shell_command.sh.jinja

@@ -9,6 +9,51 @@ TASK_DIR=$(mktemp -d)
 trap 'rm -rf "$TASK_DIR" ${{ENV_TMP:+"$ENV_TMP" "$ENV_TMP.uv.toml"}} ${{UV_BOOT_TMP:+"$UV_BOOT_TMP"}}' EXIT
 {% endraw %}
 
+# groundhog_publish TMP FINAL [VALIDATE_CMD...]: publish a finished build
+# atomically by renaming a private temp dir into its final shared path,
+# tolerating concurrent tasks racing to publish the same path. A validation
+# command, when given, is run with a candidate dir appended and decides
+# whether that copy is usable; without one, an existing dir counts as
+# usable. Subshell function bodies keep state local and avoid literal
+# braces (Globus Compute runs .format() on this whole command).
+groundhog_publish() (
+    tmp="$1"
+    final="$2"
+    shift 2
+
+    usable() (
+        dir="$1"
+        shift
+        if [ "$#" -eq 0 ]; then [ -d "$dir" ]; else "$@" "$dir"; fi
+    )
+
+    if ! usable "$tmp" "$@"; then
+        # build failed or produced an unusable result; discard
+        rm -rf "$tmp"
+    elif usable "$final" "$@"; then
+        # another task already published a usable copy; use theirs
+        rm -rf "$tmp"
+    else
+        if [ "$#" -gt 0 ]; then
+            # with a validation command, an existing final dir can be
+            # stale/broken (e.g. scratch purges can delete its contents but
+            # leave the dir, which would otherwise wedge publishing
+            # forever); clear leftovers before the rename. A racer's fresh
+            # publish could be cleared in this tiny window, but the racer
+            # that follows republishes a working copy either way.
+            rm -rf "$final"
+        fi
+        if ! mv "$tmp" "$final" 2>/dev/null; then
+            # lost the publish race; use the winner's copy
+            rm -rf "$tmp"
+        elif [ -d "$final/$(basename "$tmp")" ]; then
+            # mv raced with another publisher and nested our copy inside
+            # the winner's dir; remove the stray copy.
+            rm -rf "$final/$(basename "$tmp")"
+        fi
+    fi
+)
+
 # try SCRATCH -> TMPDIR -> /tmp for uv cache to avoid filesystem locking issues
 # GROUNDHOG_CACHE_DIR allows users to override if needed, UV_* vars also respected
 {% raw %}
@@ -50,36 +95,24 @@ else
     {% raw %}
     UV_BOOT="${{GROUNDHOG_CACHE_BASE}}/${{USER:-$(id -un)}}/uv-bootstrap-$(uname -m)-$GROUNDHOG_VERSION"
     {% endraw %}
-    if ! "$UV_BOOT/bin/uv" --version &> /dev/null; then
+
+    # uv_works DIR: the install at DIR is usable. Run the binary rather
+    # than trust mode bits: -x checks pass for wrong-libc, noexec-mounted,
+    # or truncated binaries.
+    uv_works() ( "$1/bin/uv" --version &> /dev/null )
+
+    if ! uv_works "$UV_BOOT"; then
         UV_BOOT_TMP="$UV_BOOT.tmp.$(hostname).$$.$RANDOM"
         mkdir -p "$(dirname "$UV_BOOT")"
         $PYTHON -m pip install --target "$UV_BOOT_TMP" uv &> /dev/null || true
         # Publish only a working install, atomically; if another task
-        # published a working install first, discard ours and use theirs.
-        if ! "$UV_BOOT_TMP/bin/uv" --version &> /dev/null; then
-            # install failed or binary unusable; discard
-            rm -rf "$UV_BOOT_TMP"
-        elif "$UV_BOOT/bin/uv" --version &> /dev/null; then
-            # another task already published a working install; use theirs
-            rm -rf "$UV_BOOT_TMP"
-        else
-            # clear stale/broken leftovers (scratch purges can delete bin/uv
-            # but leave the dir, which would otherwise wedge bootstrap
-            # forever), then publish. A racer's fresh publish could be
-            # cleared in the tiny window below, but the racer that follows
-            # republishes a working copy either way.
-            rm -rf "$UV_BOOT"
-            if ! mv "$UV_BOOT_TMP" "$UV_BOOT" 2>/dev/null; then
-                rm -rf "$UV_BOOT_TMP"
-            elif [ -d "$UV_BOOT/$(basename "$UV_BOOT_TMP")" ]; then
-                # mv raced with another publisher and nested our install
-                # inside the winner's dir; remove the stray copy.
-                rm -rf "$UV_BOOT/$(basename "$UV_BOOT_TMP")"
-            fi
-        fi
+        # published a working install first, ours is discarded and theirs
+        # used. uv_works also lets the helper clear a stale/broken $UV_BOOT
+        # before the rename.
+        groundhog_publish "$UV_BOOT_TMP" "$UV_BOOT" uv_works
     fi
 
-    if "$UV_BOOT/bin/uv" --version &> /dev/null; then
+    if uv_works "$UV_BOOT"; then
         UV_BIN="$UV_BOOT/bin/uv"
     else
         # fall back e.g. to a uv PyPI package installed some other way
@@ -184,15 +217,10 @@ UV_CONFIG_EOF
 {{ '}}' }}
 META_EOF
 
-    # Publish atomically; if another task published the same env first,
-    # discard our build and use theirs.
-    if [ -d "$ENV_DIR" ] || ! mv "$ENV_TMP" "$ENV_DIR" 2>/dev/null; then
-        rm -rf "$ENV_TMP"
-    elif [ -d "$ENV_DIR/$(basename "$ENV_TMP")" ]; then
-        # mv raced with another publisher and nested our build inside the
-        # winner's env dir; remove the stray copy.
-        rm -rf "$ENV_DIR/$(basename "$ENV_TMP")"
-    fi
+    # Publish atomically; if another task published the same env first, our
+    # build is discarded and theirs used. No validation command: an existing
+    # $ENV_DIR counts as usable, and is never cleared.
+    groundhog_publish "$ENV_TMP" "$ENV_DIR"
 fi
 
 # Run using the cached environment's Python directly (bypasses uv resolution)

tests/test_templating.py

@@ -725,15 +725,22 @@ def func():
         # dependencies are installed into the temp env, not the final path
         assert '--python "$ENV_TMP/bin/python"' in shell_command
         assert '--python "$ENV_DIR/bin/python"' not in shell_command
-        # published by rename; the loser of a race discards its build
-        assert 'mv "$ENV_TMP" "$ENV_DIR"' in shell_command
-        assert 'rm -rf "$ENV_TMP"' in shell_command
+        # published via the shared atomic-rename helper, with no validation
+        # command (an existing $ENV_DIR counts as usable and is never
+        # cleared); the helper renames into place and the loser of a race
+        # discards its build
+        assert 'groundhog_publish "$ENV_TMP" "$ENV_DIR"\n' in shell_command
+        assert 'mv "$tmp" "$final"' in shell_command
+        assert 'rm -rf "$tmp"' in shell_command
+        # the helper is defined once and shared by both publish sites
+        assert shell_command.count("groundhog_publish() (") == 1
+        assert shell_command.count('groundhog_publish "') == 2
         # everything written into the env goes via the temp path; nothing
         # between venv creation and publish should touch $ENV_DIR/ directly.
         # Assert each slice marker is unique so future template edits fail
         # loudly here instead of silently shifting the inspected window.
         venv_marker = '"$UV_BIN" venv $UV_VENV_RELOCATABLE "$ENV_TMP"'
-        publish_marker = 'mv "$ENV_TMP" "$ENV_DIR"'
+        publish_marker = 'groundhog_publish "$ENV_TMP" "$ENV_DIR"'
         assert shell_command.count(venv_marker) == 1
         assert shell_command.count(publish_marker) == 1
         create_branch = shell_command.split(venv_marker)[1].split(publish_marker)[0]
@@ -1413,24 +1420,30 @@ def test_uv_is_validated_by_running_it(self, tmp_path):
         # a uv on PATH may be half-written by a concurrent pip install; the
         # fast path must run it rather than trust `command -v` alone
         assert '"$UV_BIN" --version &> /dev/null' in shell_command
-        # the reuse gate runs the published binary (mode-bit -x checks pass
-        # for wrong-libc, noexec-mounted, or truncated binaries)
-        assert 'if ! "$UV_BOOT/bin/uv" --version &> /dev/null; then' in shell_command
-        # the publish gate requires the freshly installed binary to run
-        assert (
-            'if ! "$UV_BOOT_TMP/bin/uv" --version &> /dev/null; then' in shell_command
-        )
+        # bootstrap installs are validated by running the binary (mode-bit
+        # -x checks pass for wrong-libc, noexec-mounted, or truncated
+        # binaries), via the uv_works helper
+        assert 'uv_works() ( "$1/bin/uv" --version &> /dev/null )' in shell_command
+        # the reuse gate runs the published binary
+        assert 'if ! uv_works "$UV_BOOT"; then' in shell_command
+        # publish requires the freshly installed binary to run: uv_works is
+        # passed to the shared helper as its validation command
+        assert 'groundhog_publish "$UV_BOOT_TMP" "$UV_BOOT" uv_works' in shell_command
+        # the selection gate re-runs the published binary before using it
+        assert 'if uv_works "$UV_BOOT"; then' in shell_command
         # no bare mode-bit trust anywhere in uv resolution
         assert '[ ! -x "$UV_BOOT_TMP/bin/uv" ]' not in shell_command
         assert '[ ! -x "$UV_BOOT/bin/uv" ]' not in shell_command
 
     def test_bootstrap_publishes_by_rename_and_discards_on_lost_race(self, tmp_path):
         shell_command = self._shell_command(tmp_path)
 
-        # only a working install is published, atomically
-        assert 'mv "$UV_BOOT_TMP" "$UV_BOOT"' in shell_command
+        # only a working install is published, atomically, via the shared
+        # helper (uv_works validates the fresh install before the rename)
+        assert 'groundhog_publish "$UV_BOOT_TMP" "$UV_BOOT" uv_works' in shell_command
+        assert 'mv "$tmp" "$final"' in shell_command
         # the loser of a publish race discards its install
-        assert 'rm -rf "$UV_BOOT_TMP"' in shell_command
+        assert 'rm -rf "$tmp"' in shell_command
         # the published binary is used, with the old discovery as fallback
         assert 'UV_BIN="$UV_BOOT/bin/uv"' in shell_command
         assert "uv.find_uv_bin()" in shell_command
@@ -1439,13 +1452,16 @@ def test_bootstrap_repairs_stale_destination_before_publish(self, tmp_path):
         shell_command = self._shell_command(tmp_path)
 
         # scratch purges can delete bin/uv but leave $UV_BOOT itself; the
-        # publisher must clear the stale dir before the rename or bootstrap
-        # wedges forever
-        assert 'rm -rf "$UV_BOOT"\n' in shell_command
-        stale_clear_pos = shell_command.find('rm -rf "$UV_BOOT"\n')
-        mv_pos = shell_command.find('mv "$UV_BOOT_TMP" "$UV_BOOT"')
+        # shared helper must clear the stale destination before the rename
+        # or bootstrap wedges forever
+        assert 'rm -rf "$final"\n' in shell_command
+        stale_clear_pos = shell_command.find('rm -rf "$final"\n')
+        mv_pos = shell_command.find('mv "$tmp" "$final"')
         assert stale_clear_pos != -1 and mv_pos != -1
         assert stale_clear_pos < mv_pos
+        # ...but only when a validation command can distinguish broken from
+        # usable; the env publish passes none, so $ENV_DIR is never cleared
+        assert 'if [ "$#" -gt 0 ]; then' in shell_command
 
     def test_exit_trap_cleans_up_bootstrap_tmp_dir(self, tmp_path):
         shell_command = self._shell_command(tmp_path)