diff --git a/src/groundhog_hpc/app/init.py b/src/groundhog_hpc/app/init.py index cbb7cb4..53d77f3 100644 --- a/src/groundhog_hpc/app/init.py +++ b/src/groundhog_hpc/app/init.py @@ -15,6 +15,7 @@ get_endpoint_schema_comments, parse_endpoint_spec, ) +from groundhog_hpc.configuration.models import _default_exclude_newer from groundhog_hpc.configuration.pep723 import ( Pep723Metadata, add_endpoint_to_script, @@ -81,8 +82,10 @@ def init( else: python = default_meta.requires_python - assert default_meta.tool and default_meta.tool.uv - exclude_newer = default_meta.tool.uv.exclude_newer + # exclude-newer is unset by default on the model (so it never churns env + # hashes or gets baked in by rewrites); new scripts still get an explicit + # timestamp for reproducibility. + exclude_newer = _default_exclude_newer() # Parse endpoint specs if provided endpoint_specs = [] diff --git a/src/groundhog_hpc/configuration/models.py b/src/groundhog_hpc/configuration/models.py index e12c5de..3c29206 100644 --- a/src/groundhog_hpc/configuration/models.py +++ b/src/groundhog_hpc/configuration/models.py @@ -76,7 +76,10 @@ class UvMetadata(BaseModel, extra="allow", serialize_by_alias=True): See uv documentation for full precedence hierarchy. Attributes: - exclude_newer: Limit packages to versions uploaded before cutoff (ISO 8601 timestamp) + exclude_newer: Limit packages to versions uploaded before cutoff (ISO 8601 + timestamp). Defaults to None (unset); when unset, an effective default + (the build time) is injected at command-templating time rather than here, + so the volatile value never leaks into env hashes or file rewrites. python_preference: Control system vs managed Python ("managed" | "only-managed" | "system" | "only-system") index_url: Primary package index URL (default: PyPI) extra_index_url: Additional package indexes (searched after index_url) @@ -84,9 +87,7 @@ class UvMetadata(BaseModel, extra="allow", serialize_by_alias=True): offline: Disable all network access (use only cache and local files) """ - exclude_newer: str | None = Field( - default_factory=_default_exclude_newer, alias="exclude-newer" - ) + exclude_newer: str | None = Field(default=None, alias="exclude-newer") python_preference: str | None = Field( default="only-managed", alias="python-preference" ) diff --git a/src/groundhog_hpc/templates/shell_command.sh.jinja b/src/groundhog_hpc/templates/shell_command.sh.jinja index c3e08e8..411a1ea 100644 --- a/src/groundhog_hpc/templates/shell_command.sh.jinja +++ b/src/groundhog_hpc/templates/shell_command.sh.jinja @@ -1,7 +1,13 @@ set -euo pipefail TASK_DIR=$(mktemp -d) -trap 'rm -rf "$TASK_DIR"' EXIT +# Clean up the task dir plus any in-progress temp build dirs if we exit +# before publishing (the :+ expansions are no-ops while unset, and after a +# successful publish the mv has renamed the temp path away, so the trap +# never touches published state). +{% raw %} +trap 'rm -rf "$TASK_DIR" ${{ENV_TMP:+"$ENV_TMP" "$ENV_TMP.uv.toml"}} ${{UV_BOOT_TMP:+"$UV_BOOT_TMP"}}' EXIT +{% endraw %} if command -v uv &> /dev/null; then UV_BIN=$(command -v uv) @@ -79,32 +85,50 @@ else fi {% endraw %} + # Build into a unique temp dir next to the final location, then rename + # into place at the end, so concurrent tasks that share an env hash + # never see (or reuse) a half-built environment. + # $RANDOM guards against pid collisions across PID namespaces (e.g. + # containerized tasks sharing a hostname and bind-mounted scratch). + ENV_TMP="$ENV_DIR.tmp.$(hostname).$$.$RANDOM" + mkdir -p "$(dirname "$ENV_DIR")" + {% if uv_config_toml %} # Both `uv venv` and `uv pip install` will receive --config-file, then - # the file is moved into $ENV_DIR as a human-readable audit trail. - {% raw %}UV_TOML="${{ENV_DIR}}.uv.toml"{% endraw %} - mkdir -p "$(dirname "$UV_TOML")" + # the file is moved into the environment as a human-readable audit trail. + UV_TOML="$ENV_TMP.uv.toml" cat > "$UV_TOML" << 'UV_CONFIG_EOF' {{ uv_config_toml | escape_braces }} UV_CONFIG_EOF {% endif %} - "$UV_BIN" venv "$ENV_DIR"{% if requires_python %} --python "{{ requires_python }}"{% endif %}{% if uv_config_toml %} --config-file "$UV_TOML"{% endif %} + # --relocatable keeps entry-point shebangs valid after the rename below, + # but only exists in uv >= 0.2.31; probe --help so older endpoint uv + # installs degrade gracefully instead of hard-failing on an unknown flag. + if "$UV_BIN" venv --help 2>/dev/null | grep -q -- --relocatable; then + UV_VENV_RELOCATABLE="--relocatable" + else + UV_VENV_RELOCATABLE="" + fi + + # unquoted $UV_VENV_RELOCATABLE is intentional: it must expand to zero + # arguments when empty + "$UV_BIN" venv $UV_VENV_RELOCATABLE "$ENV_TMP"{% if requires_python %} --python "{{ requires_python }}"{% endif %}{% if uv_config_toml %} --config-file "$UV_TOML"{% endif %} {% if uv_config_toml %} - mv "$UV_TOML" "$ENV_DIR/uv.toml" + mv "$UV_TOML" "$ENV_TMP/uv.toml" {% endif %} # Install dependencies - "$UV_BIN" pip install --python "$ENV_DIR/bin/python" \ - {% if uv_config_toml %}--config-file "$ENV_DIR/uv.toml" {% endif %}\ + "$UV_BIN" pip install --python "$ENV_TMP/bin/python" \ + {% if uv_config_toml %}--config-file "$ENV_TMP/uv.toml" {% endif %}\ --exclude-newer-package groundhog-hpc={{ groundhog_timestamp }} \ {% for dep in dependencies %}"{{ dep }}" {% endfor %}{{ version_spec }} # Write metadata for posterity - cat > "$ENV_DIR/groundhog-meta.json" << 'META_EOF' + cat > "$ENV_TMP/groundhog-meta.json" << 'META_EOF' {{ '{{' }} "created_at": "{{ groundhog_timestamp }}", "requires_python": "{{ requires_python }}", @@ -112,6 +136,16 @@ UV_CONFIG_EOF "groundhog_version": "{{ groundhog_version }}" {{ '}}' }} 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 fi # Run using the cached environment's Python directly (bypasses uv resolution) diff --git a/src/groundhog_hpc/templating.py b/src/groundhog_hpc/templating.py index 020483c..7c142c1 100644 --- a/src/groundhog_hpc/templating.py +++ b/src/groundhog_hpc/templating.py @@ -42,6 +42,12 @@ def compute_env_hash(metadata: Pep723Metadata) -> str: a script can have many endpoints and worker_init content is not always environment-affecting. + An unset exclude-newer never affects the hash: the model defaults it to + None (excluded from the dump below), and the effective default — the + build-time clock — is injected only when templating the shell command. + Only an exclude-newer the user actually pinned in the script header + affects the hash. + Args: metadata: PEP 723 metadata from the user script @@ -134,7 +140,7 @@ def template_shell_command(script_path: str, function_name: str) -> str: local_log_level = local_log_level.upper() logger.debug(f"Propagating log level to remote: {local_log_level}") - uv_config_toml = _serialize_uv_toml(metadata) + uv_config_toml = _serialize_uv_toml(metadata, groundhog_timestamp) shell_template = jinja_env.get_template("shell_command.sh.jinja") shell_command_string = shell_template.render( @@ -155,18 +161,25 @@ def template_shell_command(script_path: str, function_name: str) -> str: return shell_command_string -def _serialize_uv_toml(metadata: Pep723Metadata | None) -> str: +def _serialize_uv_toml( + metadata: Pep723Metadata | None, default_exclude_newer: str +) -> str: """Serialize [tool.uv] settings to uv.toml format for uv pip install. Returns a TOML string containing all non-None settings from the user's - [tool.uv] block, or an empty string if there are no settings. + [tool.uv] block, or an empty string if there are no settings (or no + metadata at all). + + If the user did not pin exclude-newer in their script, the provided + default (the build-time clock) is injected here so fresh environment + builds still resolve against a fixed cutoff — without the volatile + value ever entering the env hash or the user's script file. """ if not metadata or not metadata.tool or not metadata.tool.uv: return "" uv_dict = metadata.tool.uv.model_dump(by_alias=True, exclude_none=True) - if not uv_dict: - return "" + uv_dict.setdefault("exclude-newer", default_exclude_newer) return tomlkit.dumps(uv_dict).strip() diff --git a/tests/test_pep723.py b/tests/test_pep723.py index aa01df7..f202ba2 100644 --- a/tests/test_pep723.py +++ b/tests/test_pep723.py @@ -107,7 +107,9 @@ def test_create_with_defaults(self): assert metadata.dependencies == [] assert metadata.tool is not None assert metadata.tool.uv is not None - assert metadata.tool.uv.exclude_newer is not None + # exclude-newer stays unset by default; the effective default (build + # time) is injected at command-templating time, not on the model + assert metadata.tool.uv.exclude_newer is None def test_create_with_explicit_values(self): """Test creating metadata with explicit values.""" diff --git a/tests/test_templating.py b/tests/test_templating.py index aeccb66..e6a65c3 100644 --- a/tests/test_templating.py +++ b/tests/test_templating.py @@ -531,6 +531,100 @@ def test_hash_unchanged_by_tool_hog_config(self, tmp_path): assert hash1 == hash2 + def test_defaulted_exclude_newer_does_not_affect_hash(self): + """Scripts that don't pin exclude-newer hash identically across parses. + + An unset exclude-newer stays None on the model (the effective default + is injected at command-templating time instead), so it never appears + in the hash input and cannot churn the env hash. + """ + from groundhog_hpc.configuration.pep723 import read_pep723 + from groundhog_hpc.templating import compute_env_hash + + script = """# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# /// +""" + metadata1 = read_pep723(script) + metadata2 = read_pep723(script) + assert metadata1 is not None and metadata2 is not None + + # unpinned exclude-newer stays unset on the model + assert metadata1.tool.uv.exclude_newer is None + assert metadata2.tool.uv.exclude_newer is None + + assert compute_env_hash(metadata1) == compute_env_hash(metadata2) + + def test_underscore_typo_exclude_newer_does_not_churn_hash(self): + """An extra key spelled exclude_newer (underscore typo) is hash-stable. + + UvMetadata has extra="allow", so a [tool.uv] key literally named + `exclude_newer` is stored as an extra field rather than the real + exclude-newer setting. It must not reintroduce per-parse hash churn. + """ + from groundhog_hpc.configuration.pep723 import read_pep723 + from groundhog_hpc.templating import compute_env_hash + + script = """# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# +# [tool.uv] +# exclude_newer = "2025-01-01T00:00:00Z" +# /// +""" + metadata1 = read_pep723(script) + metadata2 = read_pep723(script) + assert metadata1 is not None and metadata2 is not None + + assert compute_env_hash(metadata1) == compute_env_hash(metadata2) + + def test_rewrite_roundtrip_does_not_inject_exclude_newer(self): + """A read/write round-trip must not bake exclude-newer into the script. + + CLI file rewrites (hog add, hog run metadata prompt) dump the parsed + metadata back to the script. If parsing defaulted exclude-newer to the + wall clock, the rewrite would permanently pin it (and change the env + hash); with default=None it must simply never appear. + """ + from groundhog_hpc.configuration.pep723 import read_pep723, write_pep723 + from groundhog_hpc.templating import compute_env_hash + + script = """# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# /// +""" + metadata = read_pep723(script) + assert metadata is not None + + rewritten = write_pep723(metadata) + assert "exclude-newer" not in rewritten + + reparsed = read_pep723(rewritten) + assert reparsed is not None + assert compute_env_hash(reparsed) == compute_env_hash(metadata) + + def test_user_pinned_exclude_newer_affects_hash(self): + """An exclude-newer pinned in the script header still affects the hash.""" + from groundhog_hpc.configuration.pep723 import read_pep723 + from groundhog_hpc.templating import compute_env_hash + + script_template = """# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# +# [tool.uv] +# exclude-newer = "{}" +# /// +""" + metadata1 = read_pep723(script_template.format("2025-01-01T00:00:00Z")) + metadata2 = read_pep723(script_template.format("2025-06-01T00:00:00Z")) + assert metadata1 is not None and metadata2 is not None + + assert compute_env_hash(metadata1) != compute_env_hash(metadata2) + class TestEnvReuseTemplating: """Test environment reuse in shell command templating.""" @@ -598,6 +692,80 @@ def func(): assert '"$UV_BIN" venv' in shell_command assert '"$UV_BIN" pip install' in shell_command + def test_env_built_in_temp_dir_and_published_by_rename(self, tmp_path): + """Environment creation is safe under concurrent same-hash tasks. + + The env is built into a unique temp dir and renamed into place, so a + task can never observe (or reuse) a half-built environment; a task + that loses the publish race discards its build and uses the winner's. + """ + script_path = tmp_path / "script.py" + script_content = """# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# /// + +import groundhog_hpc as hog + +@hog.function() +def func(): + return 1 +""" + script_path.write_text(script_content) + + shell_command = template_shell_command(str(script_path), "func") + + # unique per-task build dir (hostname + pid + $RANDOM, so tasks in + # separate PID namespaces sharing a hostname/scratch don't collide) + assert 'ENV_TMP="$ENV_DIR.tmp.$(hostname).$$.$RANDOM"' in shell_command + # venv is created at the temp path; --relocatable (probed, since old + # uv lacks it) keeps entry-point shebangs valid across the rename + assert '"$UV_BIN" venv $UV_VENV_RELOCATABLE "$ENV_TMP"' in shell_command + assert 'UV_VENV_RELOCATABLE="--relocatable"' in shell_command + # 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 + # 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"' + 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] + assert '"$ENV_DIR/' not in create_branch + + def test_exit_trap_cleans_up_env_tmp(self, tmp_path): + """The EXIT trap removes a partially-built env if the build fails. + + Under set -euo pipefail any failure between venv creation and publish + would otherwise leak $ENV_TMP (a full venv) and $ENV_TMP.uv.toml. + After a successful publish the mv has removed $ENV_TMP, so the trap + never touches the published env. + """ + script_path = tmp_path / "script.py" + script_path.write_text(MINIMAL_SCRIPT) + + shell_command = template_shell_command(str(script_path), "func") + + # pre-.format(): braces are doubled for Globus Compute's .format() call + # (the trap also names UV_BOOT_TMP so the line stays identical across + # branches that harden the uv bootstrap the same way) + assert ( + 'trap \'rm -rf "$TASK_DIR" ${{ENV_TMP:+"$ENV_TMP" "$ENV_TMP.uv.toml"}} ' + '${{UV_BOOT_TMP:+"$UV_BOOT_TMP"}}\' EXIT' in shell_command + ) + # post-.format(): the shell sees single braces + formatted = shell_command.format(payload="test") + assert ( + 'trap \'rm -rf "$TASK_DIR" ${ENV_TMP:+"$ENV_TMP" "$ENV_TMP.uv.toml"} ' + '${UV_BOOT_TMP:+"$UV_BOOT_TMP"}\' EXIT' in formatted + ) + def test_shell_command_runs_python_directly(self, tmp_path): """Shell command runs Python directly instead of uv run.""" script_path = tmp_path / "script.py" @@ -678,7 +846,7 @@ class TestSerializeUvToml: def test_returns_empty_string_for_none_metadata(self): from groundhog_hpc.templating import _serialize_uv_toml - result = _serialize_uv_toml(None) + result = _serialize_uv_toml(None, "2099-01-01T00:00:00Z") assert result == "" @@ -692,7 +860,7 @@ def test_returns_empty_string_when_tool_is_none(self): tool=None, ) - result = _serialize_uv_toml(metadata) + result = _serialize_uv_toml(metadata, "2099-01-01T00:00:00Z") assert result == "" @@ -718,7 +886,7 @@ def test_serializes_string_values(self): ), ) - result = _serialize_uv_toml(metadata) + result = _serialize_uv_toml(metadata, "2099-01-01T00:00:00Z") assert 'exclude-newer = "2025-01-01T00:00:00Z"' in result assert 'python-preference = "only-managed"' in result @@ -747,7 +915,7 @@ def test_serializes_list_values(self): ), ) - result = _serialize_uv_toml(metadata) + result = _serialize_uv_toml(metadata, "2099-01-01T00:00:00Z") assert "extra-index-url" in result assert '"https://download.pytorch.org/whl/cpu"' in result @@ -767,7 +935,7 @@ def test_serializes_bool_values(self): tool=ToolMetadata(uv=UvMetadata(**{"offline": True})), ) - result = _serialize_uv_toml(metadata) + result = _serialize_uv_toml(metadata, "2099-01-01T00:00:00Z") assert "offline = true" in result @@ -789,7 +957,7 @@ def test_fields_defaulting_to_none_are_excluded(self): ), ) - result = _serialize_uv_toml(metadata) + result = _serialize_uv_toml(metadata, "2099-01-01T00:00:00Z") assert "index-url" not in result assert "extra-index-url" not in result @@ -813,7 +981,7 @@ def test_extra_fields_are_included(self): ), ) - result = _serialize_uv_toml(metadata) + result = _serialize_uv_toml(metadata, "2099-01-01T00:00:00Z") assert 'find-links = "https://example.com/wheels"' in result @@ -842,7 +1010,7 @@ def func(): shell_command = template_shell_command(str(script_path), "func") - assert '"$ENV_DIR/uv.toml"' in shell_command + assert '"$ENV_TMP/uv.toml"' in shell_command assert 'exclude-newer = "2025-01-01T00:00:00Z"' in shell_command assert '"https://download.pytorch.org/whl/cpu"' in shell_command @@ -866,7 +1034,7 @@ def func(): shell_command = template_shell_command(str(script_path), "func") - assert '--config-file "$ENV_DIR/uv.toml"' in shell_command + assert '--config-file "$ENV_TMP/uv.toml"' in shell_command def test_exclude_newer_not_passed_as_cli_flag(self, tmp_path): """--exclude-newer is no longer a CLI flag; it lives in uv.toml.""" @@ -917,12 +1085,17 @@ def func(): shell_command = template_shell_command(str(script_path), "func") - # uv venv line should carry --config-file + # the venv *creation* line (not the --relocatable --help probe) + # should carry --config-file venv_line = next( - (line for line in shell_command.splitlines() if '"$UV_BIN" venv' in line), + ( + line + for line in shell_command.splitlines() + if '"$UV_BIN" venv $UV_VENV_RELOCATABLE "$ENV_TMP"' in line + ), None, ) - assert venv_line is not None, "No uv venv line found" + assert venv_line is not None, "No uv venv creation line found" assert "--config-file" in venv_line def test_uv_toml_written_before_venv_creation(self, tmp_path): @@ -953,6 +1126,75 @@ def func(): "uv.toml must be written before uv venv creates the directory" ) + def test_uv_toml_gets_build_time_exclude_newer_when_not_pinned(self, tmp_path): + """Unpinned scripts still get an exclude-newer cutoff in uv.toml. + + The env hash ignores an unset exclude-newer, but the uv.toml handed to + uv venv / uv pip install carries the build-time default so fresh + builds resolve against a fixed cutoff. + """ + import re + + script_path = tmp_path / "script.py" + script_path.write_text("""# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy"] +# /// + +import groundhog_hpc as hog + +@hog.function() +def func(): + return 1 +""") + + shell_command = template_shell_command(str(script_path), "func") + + # inspect only the uv.toml heredoc contents (opener + closing delimiter) + assert shell_command.count("UV_CONFIG_EOF") == 2 + uv_toml_body = shell_command.split("UV_CONFIG_EOF")[1] + assert re.search( + r'exclude-newer = "\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z"', uv_toml_body + ), "uv.toml should contain an injected build-time exclude-newer" + + def test_uv_toml_keeps_pinned_exclude_newer(self, tmp_path): + """A user-pinned exclude-newer appears verbatim in the uv.toml heredoc.""" + script_path = tmp_path / "script.py" + script_path.write_text("""# /// script +# requires-python = ">=3.11" +# dependencies = [] +# +# [tool.uv] +# exclude-newer = "2025-01-01T00:00:00Z" +# /// + +import groundhog_hpc as hog + +@hog.function() +def func(): + return 1 +""") + + shell_command = template_shell_command(str(script_path), "func") + + uv_toml_body = shell_command.split("UV_CONFIG_EOF")[1] + assert 'exclude-newer = "2025-01-01T00:00:00Z"' in uv_toml_body + + def test_no_pep723_metadata_no_injected_exclude_newer(self, tmp_path): + """Scripts without metadata get no uv.toml and no injected exclude-newer.""" + script_path = tmp_path / "script.py" + script_path.write_text("""import groundhog_hpc as hog + +@hog.function() +def func(): + return 1 +""") + + shell_command = template_shell_command(str(script_path), "func") + + assert "UV_CONFIG_EOF" not in shell_command + assert "exclude-newer =" not in shell_command + def test_no_uv_toml_written_for_script_without_pep723_metadata(self, tmp_path): """Scripts without PEP 723 metadata don't write a uv.toml.""" script_path = tmp_path / "script.py"