mirror of
https://github.com/cactus-compute/needle.git
synced 2026-10-21 06:09:38 +08:00
280 lines
11 KiB
Python
280 lines
11 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
def test_lib_path_env_override(tmp_path, monkeypatch):
|
|
import needle
|
|
|
|
fake = tmp_path / "libneedle.dylib"
|
|
fake.write_bytes(b"x")
|
|
monkeypatch.setenv("NEEDLE_LIB_PATH", str(fake))
|
|
assert needle._library_path() == str(fake)
|
|
|
|
|
|
def test_legacy_lib_override_cannot_capture_v3(tmp_path, monkeypatch):
|
|
import needle
|
|
from needle.agent import fetch
|
|
|
|
v2 = tmp_path / "libneedle-v2.dylib"
|
|
v3 = tmp_path / "libneedle-v3.dylib"
|
|
v2.write_bytes(b"v2")
|
|
v3.write_bytes(b"v3")
|
|
monkeypatch.setenv("NEEDLE_LIB_PATH", str(v2))
|
|
monkeypatch.setenv("NEEDLE3_LIB_PATH", str(v3))
|
|
assert needle._library_path(2) == str(v2)
|
|
assert needle._library_path(3) == str(v3)
|
|
|
|
monkeypatch.delenv("NEEDLE3_LIB_PATH")
|
|
monkeypatch.setenv("HOME", str(tmp_path))
|
|
monkeypatch.setenv("USERPROFILE", str(tmp_path))
|
|
monkeypatch.setattr(fetch, "fetch_library", lambda *args, **kwargs: str(v3))
|
|
assert needle._library_path(3) == str(v3)
|
|
|
|
|
|
def test_download_target_kinds():
|
|
import pytest
|
|
from needle.cli import _download_target
|
|
|
|
assert _download_target("macos-arm64") == ("platform", "macos-arm64")
|
|
assert _download_target("needle3") == ("base", 3)
|
|
assert _download_target("needle2.cact") == ("base", 2)
|
|
assert _download_target("needle3.safetensors") == ("checkpoint", "needle3.safetensors")
|
|
assert _download_target("acme/tuned/model.cact") == ("hub", "acme/tuned/model.cact")
|
|
with pytest.raises(SystemExit, match="unknown download"):
|
|
_download_target("needle4")
|
|
|
|
|
|
def test_fetch_weights_copies_the_base_archive_and_reuses_it(tmp_path, monkeypatch):
|
|
from needle.agent import fetch
|
|
|
|
archive = tmp_path / "needle3.cact"
|
|
archive.write_bytes(b"weights")
|
|
calls = []
|
|
|
|
def fake_download(**kwargs):
|
|
calls.append(kwargs["filename"])
|
|
return str(archive)
|
|
|
|
monkeypatch.setattr(fetch, "_register_download", lambda generation: None)
|
|
monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download)
|
|
dest = tmp_path / "cache"
|
|
out = fetch.fetch_weights(3, str(dest))
|
|
assert out == str(dest / "needle3.cact") and open(out, "rb").read() == b"weights"
|
|
assert calls == ["needle3.cact"]
|
|
assert fetch.fetch_weights(3, str(dest)) == out and calls == ["needle3.cact"]
|
|
|
|
|
|
def test_fetch_checkpoint_prefers_the_checkpoints_folder(tmp_path, monkeypatch):
|
|
from huggingface_hub.errors import EntryNotFoundError
|
|
from needle.agent import fetch
|
|
|
|
source = tmp_path / "needle3.safetensors"
|
|
source.write_bytes(b"ckpt")
|
|
|
|
def fake_download(**kwargs):
|
|
if kwargs["filename"] != "checkpoints/needle3.safetensors":
|
|
raise EntryNotFoundError("missing")
|
|
return str(source)
|
|
|
|
monkeypatch.setattr(fetch, "_register_download", lambda generation: None)
|
|
monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download)
|
|
out = fetch.fetch_checkpoint("needle3.safetensors", str(tmp_path / "dl" / "checkpoints"))
|
|
assert out.endswith("checkpoints/needle3.safetensors") and open(out, "rb").read() == b"ckpt"
|
|
|
|
|
|
def test_weights_spec_parsing():
|
|
from needle.cli import _weights_spec
|
|
|
|
assert _weights_spec("acme/tuned/model.cact") == ("acme/tuned", "model.cact")
|
|
assert _weights_spec("acme/tuned") == ("acme/tuned", None)
|
|
assert _weights_spec("acme/tuned/sub/dir/m.cact") == ("acme/tuned", "sub/dir/m.cact")
|
|
|
|
|
|
def test_lib_name_for_tags():
|
|
from needle.agent.fetch import _lib_name_for
|
|
|
|
assert _lib_name_for("macosx_11_0_arm64") == "libneedle.dylib"
|
|
assert _lib_name_for("win_amd64") == "libneedle.dll"
|
|
assert _lib_name_for("manylinux2014_aarch64") == "libneedle.so"
|
|
assert _lib_name_for("musllinux_1_2_x86_64") == "libneedle.so"
|
|
|
|
|
|
def test_component_platform_is_downloadable():
|
|
from needle.agent.fetch import PLATFORMS
|
|
|
|
assert "wasm-component" in PLATFORMS
|
|
|
|
|
|
def test_unpublished_engine_wheels_lists_every_missing_tag(monkeypatch):
|
|
from needle.agent import fetch
|
|
|
|
published = {}
|
|
for generation, version in fetch.ENGINE_VERSIONS.items():
|
|
# needle3 hosts both its own engine and needle3_whistle, so the sets merge.
|
|
published.setdefault(fetch.ENGINE_REPOS[generation], set()).update(
|
|
fetch.engine_wheel(version, tag, generation) for tag in fetch.WHEEL_TAGS)
|
|
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]
|
|
|
|
audio = fetch.engine_wheel(fetch.ENGINE_VERSIONS[fetch.NEEDLE3_WHISTLE], "win_arm64", fetch.NEEDLE3_WHISTLE)
|
|
published[repo3].remove(audio)
|
|
assert fetch.unpublished_engine_wheels() == [repo3 + "/" + dropped, repo3 + "/" + audio]
|
|
published[repo3].add(audio)
|
|
|
|
whistle = fetch.ENGINE_REPOS[fetch.WHISTLE]
|
|
published[whistle].clear()
|
|
missing = fetch.unpublished_engine_wheels()
|
|
assert len(missing) == 1 + len(fetch.WHEEL_TAGS)
|
|
assert missing[-1] == f"{whistle}/python/cactus_whistle-{fetch.ENGINE_VERSIONS[fetch.WHISTLE]}-py3-none-win_arm64.whl"
|
|
|
|
|
|
def test_a_repo_that_cannot_be_listed_holds_none_of_its_wheels(monkeypatch):
|
|
"""The whistle repo is private until release; the gate must report, not raise."""
|
|
from needle.agent import fetch
|
|
|
|
def refuse(repo):
|
|
raise OSError("401 private repo")
|
|
|
|
monkeypatch.setattr("huggingface_hub.list_repo_files", refuse)
|
|
missing = fetch.unpublished_engine_wheels()
|
|
assert len(missing) == len(fetch.ENGINE_VERSIONS) * len(fetch.WHEEL_TAGS)
|
|
|
|
|
|
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
|
|
|
|
wheel = tmp_path / "engine.whl"
|
|
with zipfile.ZipFile(wheel, "w") as archive:
|
|
archive.writestr("needle/libneedle.so", b"engine")
|
|
monkeypatch.setattr(fetch, "_register_download", lambda generation: None)
|
|
monkeypatch.setattr("huggingface_hub.hf_hub_download",
|
|
lambda **kwargs: str(wheel))
|
|
out = fetch.fetch_library("2.0.4", tmp_path / "new", tag="manylinux2014_x86_64")
|
|
|
|
assert (tmp_path / "new" / "libneedle.so").read_bytes() == b"engine"
|
|
assert out == str(tmp_path / "new" / "libneedle.so")
|
|
|
|
|
|
def test_engine_gate_finds_the_cache_the_runtime_loads_from(tmp_path, monkeypatch):
|
|
import needle
|
|
from needle.agent import fetch
|
|
from conftest import _engine_available
|
|
|
|
package = tmp_path / "pkg"
|
|
package.mkdir()
|
|
monkeypatch.setattr(needle, "__file__", str(package / "__init__.py"))
|
|
monkeypatch.setenv("HOME", str(tmp_path))
|
|
monkeypatch.setenv("USERPROFILE", str(tmp_path))
|
|
monkeypatch.delenv("NEEDLE_LIB_PATH", raising=False)
|
|
monkeypatch.delenv("NEEDLE3_LIB_PATH", raising=False)
|
|
|
|
assert not _engine_available()
|
|
|
|
cache = tmp_path / ".cache" / "cactus-needle" / "v3" / fetch.engine_version(3)
|
|
cache.mkdir(parents=True)
|
|
(cache / fetch._lib_name()).write_bytes(b"")
|
|
|
|
assert _engine_available()
|
|
|
|
|
|
def test_engine_gate_honours_the_library_override(tmp_path, monkeypatch):
|
|
import needle
|
|
from needle.agent import fetch
|
|
from conftest import _engine_available
|
|
|
|
package = tmp_path / "pkg"
|
|
package.mkdir()
|
|
monkeypatch.setattr(needle, "__file__", str(package / "__init__.py"))
|
|
monkeypatch.setenv("HOME", str(tmp_path))
|
|
monkeypatch.setenv("USERPROFILE", str(tmp_path))
|
|
|
|
engine = tmp_path / fetch._lib_name()
|
|
engine.write_bytes(b"")
|
|
monkeypatch.setenv("NEEDLE3_LIB_PATH", str(engine))
|
|
assert _engine_available()
|
|
|
|
monkeypatch.setenv("NEEDLE3_LIB_PATH", str(tmp_path / "gone"))
|
|
assert not _engine_available()
|
|
|
|
|
|
def test_musl_detected_when_libc_ver_claims_glibc(monkeypatch, tmp_path):
|
|
import platform
|
|
from needle.agent import fetch
|
|
|
|
monkeypatch.setattr(platform, "libc_ver", lambda *a, **k: ("glibc", "2.9"))
|
|
monkeypatch.setattr(sys, "platform", "linux")
|
|
maps = tmp_path / "maps"
|
|
maps.write_bytes(b"7f0000000000-7f0000001000 r-xp /lib/ld-musl-x86_64.so.1\n")
|
|
real_open = open
|
|
monkeypatch.setattr("builtins.open",
|
|
lambda p, *a, **k: real_open(maps, *a, **k) if p == "/proc/self/maps"
|
|
else real_open(p, *a, **k))
|
|
assert fetch._is_musl() is True
|
|
|
|
|
|
def test_glibc_detected_from_maps(monkeypatch, tmp_path):
|
|
import platform
|
|
from needle.agent import fetch
|
|
|
|
monkeypatch.setattr(platform, "libc_ver", lambda *a, **k: ("", ""))
|
|
monkeypatch.setattr(sys, "platform", "linux")
|
|
maps = tmp_path / "maps"
|
|
maps.write_bytes(b"7f0000000000-7f0000001000 r-xp /usr/lib/x86_64-linux-gnu/libc.so.6\n")
|
|
real_open = open
|
|
monkeypatch.setattr("builtins.open",
|
|
lambda p, *a, **k: real_open(maps, *a, **k) if p == "/proc/self/maps"
|
|
else real_open(p, *a, **k))
|
|
assert fetch._is_musl() is False
|
|
|
|
|
|
def test_other_libc_tag_swaps_families(monkeypatch):
|
|
from needle.agent import fetch
|
|
|
|
monkeypatch.setattr(fetch, "_platform_tag", lambda: "manylinux2014_aarch64")
|
|
assert fetch.other_libc_tag() == "musllinux_1_2_aarch64"
|
|
monkeypatch.setattr(fetch, "_platform_tag", lambda: "musllinux_1_2_x86_64")
|
|
assert fetch.other_libc_tag() == "manylinux2014_x86_64"
|
|
monkeypatch.setattr(fetch, "_platform_tag", lambda: "macosx_11_0_arm64")
|
|
assert fetch.other_libc_tag() is None
|
|
|
|
|
|
def test_engine_load_falls_back_to_the_other_libc(monkeypatch, tmp_path):
|
|
import ctypes
|
|
import needle
|
|
from needle.agent import fetch
|
|
|
|
good = tmp_path / "musl" / "libneedle.so"
|
|
good.parent.mkdir()
|
|
good.write_bytes(b"x")
|
|
monkeypatch.setattr(needle, "_library_path", lambda generation=2: str(tmp_path / "libneedle.so"))
|
|
monkeypatch.setattr(fetch, "other_libc_tag", lambda: "musllinux_1_2_x86_64")
|
|
monkeypatch.setattr(fetch, "fetch_library",
|
|
lambda version=None, dest_dir=None, tag=None, generation=2: str(good))
|
|
loaded = {}
|
|
|
|
def fake_cdll(path):
|
|
if path.endswith("musl/libneedle.so"):
|
|
loaded["path"] = path
|
|
return "handle"
|
|
raise OSError("Error relocating libneedle.so: strtoll_l: symbol not found")
|
|
|
|
monkeypatch.setattr(ctypes, "CDLL", fake_cdll)
|
|
with pytest.warns(UserWarning, match="using the musllinux"):
|
|
assert needle._load_cdll(3) == "handle"
|
|
assert loaded["path"].endswith("musl/libneedle.so")
|