diff --git a/examples/novita_tbench2_simple.py b/examples/novita_tbench2_simple.py index f3544591d1..563be53551 100644 --- a/examples/novita_tbench2_simple.py +++ b/examples/novita_tbench2_simple.py @@ -29,9 +29,12 @@ async def main() -> int: # First start builds a Novita template from the Dockerfile, which can take a # few minutes; later starts reuse the cached template. base_url = provider.start_container(image=image) - provider.wait_for_ready(base_url, timeout_s=300) try: + # Inside the try: a readiness timeout must still reach stop_container(), + # or the sandbox keeps billing after the example gives up. + provider.wait_for_ready(base_url, timeout_s=300) + async with Tbench2Env(base_url=base_url, provider=provider) as env: result = await env.reset(task_id=task_id) print("Instruction head:") diff --git a/src/openenv/core/containers/runtime/novita_provider.py b/src/openenv/core/containers/runtime/novita_provider.py index cdaa45a23a..e4fcd4d8dc 100644 --- a/src/openenv/core/containers/runtime/novita_provider.py +++ b/src/openenv/core/containers/runtime/novita_provider.py @@ -62,6 +62,7 @@ _ARG_DEFAULT_RE = re.compile( r"^\s*ARG\s+(?P\w+)=(?P\S+)\s*$", re.IGNORECASE ) +_ARG_REFERENCE_RE = re.compile(r"\$(?:\{(?P\w+)\}|(?P\w+))") def _strip_mount_flags(content: str) -> str: @@ -116,6 +117,23 @@ def _strip_mount_flags(content: str) -> str: return "\n".join(out) +def _substitute_args(reference: str, arg_defaults: Dict[str, str]) -> str: + """Expand ``$NAME`` and ``${NAME}`` in one pass against *arg_defaults*. + + Substituting name by name would let a shorter ARG consume the prefix of a + longer one: with ``ARG BASE`` declared before ``ARG BASE_IMAGE``, an + unbraced ``FROM $BASE_IMAGE`` became ``_IMAGE``. Matching the + whole name in a single pass makes the result independent of declaration + order. Names with no default are left verbatim, as before. + """ + + def replace(match: "re.Match[str]") -> str: + name = match.group("braced") or match.group("bare") + return arg_defaults.get(name, match.group(0)) + + return _ARG_REFERENCE_RE.sub(replace, reference) + + def _resolve_from_references(content: str) -> str: """Substitute ``ARG`` defaults into ``FROM`` lines and drop ``--platform``. @@ -147,9 +165,7 @@ def _resolve_from_references(content: str) -> str: reference = match.group("rest").strip() reference = re.sub(r"^(--platform=\S+\s*)+", "", reference).strip() - for name, value in arg_defaults.items(): - reference = reference.replace(f"${{{name}}}", value) - reference = reference.replace(f"${name}", value) + reference = _substitute_args(reference, arg_defaults) out.append(f"FROM {reference}") return "\n".join(out) diff --git a/tests/scripts/test_novita_examples_release_sandbox.py b/tests/scripts/test_novita_examples_release_sandbox.py new file mode 100644 index 0000000000..af704be34d --- /dev/null +++ b/tests/scripts/test_novita_examples_release_sandbox.py @@ -0,0 +1,46 @@ +"""Every Novita example must release its sandbox if readiness never arrives.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +EXAMPLES_DIR = Path(__file__).resolve().parents[2] / "examples" +NOVITA_EXAMPLES = sorted(EXAMPLES_DIR.glob("novita_*.py")) + + +def _calls(node: ast.AST, name: str) -> bool: + return any( + isinstance(child, ast.Attribute) and child.attr == name + for child in ast.walk(node) + ) + + +def _guarded_by_stop_container(tree: ast.AST) -> list[ast.Try]: + return [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Try) + and any(_calls(handler, "stop_container") for handler in node.finalbody) + ] + + +@pytest.mark.parametrize("path", NOVITA_EXAMPLES, ids=lambda path: path.name) +def test_wait_for_ready_runs_under_the_stop_container_guard(path: Path): + tree = ast.parse(path.read_text()) + + if not _calls(tree, "wait_for_ready"): + pytest.skip(f"{path.name} does not wait for readiness") + + guards = _guarded_by_stop_container(tree) + assert guards, f"{path.name} calls wait_for_ready without a stop_container finally" + assert any( + _calls(statement, "wait_for_ready") + for guard in guards + for statement in guard.body + ), ( + f"{path.name} calls wait_for_ready outside the try that stops the sandbox; " + "a readiness timeout would leak a paid sandbox" + ) diff --git a/tests/test_core/test_novita_provider.py b/tests/test_core/test_novita_provider.py index 05318beacb..0656b0f071 100644 --- a/tests/test_core/test_novita_provider.py +++ b/tests/test_core/test_novita_provider.py @@ -871,6 +871,61 @@ def test_arg_after_first_from_not_in_scope(self): out = _resolve_from_references("FROM python:3.11\nARG X=alpine\nRUN echo hi\n") assert "FROM python:3.11" in out + @pytest.mark.parametrize("reference", ["$BASE_IMAGE", "${BASE_IMAGE}"]) + def test_shorter_arg_does_not_consume_a_longer_name(self, reference): + from openenv.core.containers.runtime.novita_provider import ( + _resolve_from_references, + ) + + # ARG BASE is declared first, so name-by-name substitution used to turn + # `FROM $BASE_IMAGE` into `python:3.12_IMAGE`. + out = _resolve_from_references( + "ARG BASE=python:3.12\n" + "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest\n" + f"FROM {reference} AS builder\n" + ) + assert "FROM ghcr.io/huggingface/openenv-base:latest AS builder" in out + + def test_substitution_is_independent_of_declaration_order(self): + from openenv.core.containers.runtime.novita_provider import ( + _resolve_from_references, + ) + + declarations = [ + "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest", + "ARG BASE=python:3.12", + ] + outputs = { + _resolve_from_references("\n".join([*order, "FROM $BASE_IMAGE\n"])) + for order in (declarations, list(reversed(declarations))) + } + assert outputs == { + "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest\n" + "ARG BASE=python:3.12\n" + "FROM ghcr.io/huggingface/openenv-base:latest\n", + "ARG BASE=python:3.12\n" + "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest\n" + "FROM ghcr.io/huggingface/openenv-base:latest\n", + } + + def test_undeclared_arg_reference_is_left_verbatim(self): + from openenv.core.containers.runtime.novita_provider import ( + _resolve_from_references, + ) + + out = _resolve_from_references("ARG BASE=python:3.12\nFROM $UNSET_IMAGE\n") + assert "FROM $UNSET_IMAGE" in out + + def test_substituted_value_is_not_expanded_again(self): + from openenv.core.containers.runtime.novita_provider import ( + _resolve_from_references, + ) + + out = _resolve_from_references( + "ARG REGISTRY=example.com\nARG IMAGE=$REGISTRY/app\nFROM $IMAGE\n" + ) + assert "FROM $REGISTRY/app" in out + def test_flatten_drops_same_path_copy(self): from openenv.core.containers.runtime.novita_provider import _flatten_multistage