diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 6b77d2d..3206f65 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -91,6 +91,19 @@ jobs: next="${{ steps.version.outputs.next }}" sed -i "s/^version = \".*\"/version = \"$next\"/" pyproject.toml sed -i "s/^__version__ = \".*\"/__version__ = \"$next\"/" needle/__init__.py + # The engine wheels live on the Hub and are uploaded separately. Publishing a + # pin whose wheels are not there yet makes every fresh install fail at first use. + - if: steps.version.outputs.skip != 'true' + run: | + pip install huggingface_hub + python - <<'PY' + from needle.agent import fetch + missing = fetch.unpublished_engine_wheels() + if missing: + raise SystemExit("engine wheels missing from the Hub, refusing to publish:\n " + + "\n ".join(missing)) + print("every pinned engine is published") + PY - if: steps.version.outputs.skip != 'true' run: python -m build - if: steps.version.outputs.skip != 'true' diff --git a/needle/agent/fetch.py b/needle/agent/fetch.py index 52564c9..dd7055f 100644 --- a/needle/agent/fetch.py +++ b/needle/agent/fetch.py @@ -22,6 +22,11 @@ CHECKPOINT_PREFIX = "checkpoints" # Backwards-compatible aliases for callers that explicitly fetch Needle 2. HF_REPO = ENGINE_REPOS[2] +WHEEL_TAGS = ("macosx_11_0_arm64", "macosx_11_0_x86_64", + "manylinux2014_aarch64", "manylinux2014_x86_64", + "musllinux_1_2_aarch64", "musllinux_1_2_x86_64", + "win_amd64", "win_arm64") + PLATFORMS = ("macos-arm64", "linux-x86_64", "linux-arm64", "linux-armv7", "linux-riscv64", "linux-mipsel", "windows-x86_64", "windows-arm64", "android-arm64", "android-armv7", "android-riscv64", @@ -95,6 +100,30 @@ def other_libc_tag(): return None +def engine_wheel(version, tag): + return "python/cactus_needle-{}-py3-none-{}.whl".format(version, tag) + + +def unpublished_engine_wheels(): + """Pinned engine wheels that are missing from the Hub, as repo-qualified paths. + + ENGINE_VERSIONS and the wheels it names are published separately, so a pin + that lands before its upload leaves every fresh install with nothing to + fetch. The release train calls this before it publishes. + """ + from huggingface_hub import list_repo_files + + missing = [] + for generation in sorted(ENGINE_VERSIONS): + repo = ENGINE_REPOS[generation] + present = set(list_repo_files(repo)) + for tag in WHEEL_TAGS: + wheel = engine_wheel(ENGINE_VERSIONS[generation], tag) + if wheel not in present: + missing.append(repo + "/" + wheel) + return missing + + def engine_repo(generation=2): try: return ENGINE_REPOS[int(generation)] @@ -209,10 +238,9 @@ def fetch_library(version=None, dest_dir=None, tag=None, generation=2): raise TypeError("dest_dir is required") version = version or engine_version(generation) tag = tag or _platform_tag() - wheel = "cactus_needle-{}-py3-none-{}.whl".format(version, tag) repo = engine_repo(generation) _register_download(generation) - path = hf_hub_download(repo_id=repo, filename="python/" + wheel, repo_type="model") + path = hf_hub_download(repo_id=repo, filename=engine_wheel(version, tag), repo_type="model") lib = _lib_name_for(tag) stem, suffix = os.path.splitext(lib) member = f"{stem}{generation}{suffix}" if int(generation) >= 3 else lib diff --git a/tests/test_fetch.py b/tests/test_fetch.py index 159b5eb..5e59cca 100644 --- a/tests/test_fetch.py +++ b/tests/test_fetch.py @@ -105,6 +105,32 @@ def test_component_platform_is_downloadable(): assert "wasm-component" in PLATFORMS +def test_unpublished_engine_wheels_lists_every_missing_tag(monkeypatch): + from needle.agent import fetch + + published = { + repo: {fetch.engine_wheel(version, tag) for tag in fetch.WHEEL_TAGS} + for repo, version in ((fetch.ENGINE_REPOS[2], fetch.ENGINE_VERSIONS[2]), + (fetch.ENGINE_REPOS[3], fetch.ENGINE_VERSIONS[3])) + } + monkeypatch.setattr("huggingface_hub.list_repo_files", + lambda repo: sorted(published[repo]) + ["needle3.cact"]) + assert fetch.unpublished_engine_wheels() == [] + + repo3 = fetch.ENGINE_REPOS[3] + dropped = fetch.engine_wheel(fetch.ENGINE_VERSIONS[3], "win_arm64") + published[repo3].remove(dropped) + assert fetch.unpublished_engine_wheels() == [repo3 + "/" + dropped] + + +def test_release_gate_runs_before_the_publish_step(): + from pathlib import Path + + workflow = Path(__file__).parents[1] / ".github/workflows/release.yaml" + text = workflow.read_text() + assert text.index("unpublished_engine_wheels") < text.index("pypa/gh-action-pypi-publish") + + def test_fetch_library_creates_destination(tmp_path, monkeypatch): import zipfile from needle.agent import fetch