diff --git a/src/groundhog_hpc/templates/shell_command.sh.jinja b/src/groundhog_hpc/templates/shell_command.sh.jinja index d27f2a7..ec9ea54 100644 --- a/src/groundhog_hpc/templates/shell_command.sh.jinja +++ b/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) diff --git a/tests/test_templating.py b/tests/test_templating.py index 7ba9003..1bdb9c6 100644 --- a/tests/test_templating.py +++ b/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,13 +1420,17 @@ 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 @@ -1427,10 +1438,12 @@ def test_uv_is_validated_by_running_it(self, tmp_path): 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)