fix(novita): resolve whole ARG names and stop leaking the sandbox

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:")

src/openenv/core/containers/runtime/novita_provider.py

@@ -62,6 +62,7 @@
 _ARG_DEFAULT_RE = re.compile(
     r"^\s*ARG\s+(?P<name>\w+)=(?P<value>\S+)\s*$", re.IGNORECASE
 )
+_ARG_REFERENCE_RE = re.compile(r"\$(?:\{(?P<braced>\w+)\}|(?P<bare>\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 ``<base value>_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)
 

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"
+    )

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