414 lines
11 KiB
Python
414 lines
11 KiB
Python
# -*- coding: utf-8 -*-
|
|
# pylint: disable=protected-access
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from pineagents.local_models.download_manager import (
|
|
DownloadTaskResult,
|
|
DownloadTaskStatus,
|
|
)
|
|
from pineagents.local_models.model_manager import ModelManager, DownloadSource
|
|
|
|
|
|
class _FakeController:
|
|
def __init__(self) -> None:
|
|
self.started_spec = None
|
|
self.cancel_called = False
|
|
self.active = False
|
|
self.snapshot_value = {
|
|
"status": "idle",
|
|
"model_name": None,
|
|
"downloaded_bytes": 0,
|
|
"total_bytes": None,
|
|
"speed_bytes_per_sec": 0.0,
|
|
"source": None,
|
|
"error": None,
|
|
"local_path": None,
|
|
}
|
|
|
|
def start(self, spec) -> None:
|
|
self.started_spec = spec
|
|
self.active = True
|
|
|
|
def cancel(self) -> None:
|
|
self.cancel_called = True
|
|
self.active = False
|
|
|
|
def snapshot(self) -> dict:
|
|
return self.snapshot_value
|
|
|
|
def is_active(self) -> bool:
|
|
return self.active
|
|
|
|
|
|
def test_start_download_uses_reachable_source(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
downloader = ModelManager()
|
|
controller = _FakeController()
|
|
downloader.__dict__["_download_controller"] = controller
|
|
downloader.__dict__["_download_tmp_dir"] = tmp_path / "tmp"
|
|
target_dir = tmp_path / "resolved-model-dir"
|
|
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"get_model_dir",
|
|
lambda repo_id: target_dir,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_resolve_download_source",
|
|
lambda: DownloadSource.MODELSCOPE,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_estimate_download_size",
|
|
lambda **kwargs: 100,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_check_gguf_exists",
|
|
lambda **kwargs: (True, ""),
|
|
)
|
|
|
|
downloader.start_download("Qwen/Qwen2-0.5B-Instruct-GGUF")
|
|
|
|
assert controller.started_spec is not None
|
|
assert controller.started_spec.command == [
|
|
"qwenpaw-model-download",
|
|
"Qwen/Qwen2-0.5B-Instruct-GGUF",
|
|
"modelscope",
|
|
]
|
|
assert (
|
|
controller.started_spec.model_name == "Qwen/Qwen2-0.5B-Instruct-GGUF"
|
|
)
|
|
assert controller.started_spec.source == "modelscope"
|
|
assert controller.started_spec.total_bytes == 100
|
|
assert controller.started_spec.task.payload == {
|
|
"repo_id": "Qwen/Qwen2-0.5B-Instruct-GGUF",
|
|
"source": "modelscope",
|
|
"staging_dir": str(
|
|
(tmp_path / "tmp").joinpath(
|
|
Path(
|
|
controller.started_spec.task.payload["staging_dir"],
|
|
).name,
|
|
),
|
|
),
|
|
}
|
|
progress = controller.started_spec.task.probe_progress()
|
|
assert progress is not None
|
|
assert (
|
|
Path(controller.started_spec.task.payload["staging_dir"]).parent
|
|
== tmp_path / "tmp"
|
|
)
|
|
assert progress.total_bytes == 100
|
|
assert progress.model_name == "Qwen/Qwen2-0.5B-Instruct-GGUF"
|
|
assert progress.source == "modelscope"
|
|
|
|
|
|
def test_download_model_is_wrapper(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
downloader = ModelManager()
|
|
calls: list[tuple[str, DownloadSource | None]] = []
|
|
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"start_download",
|
|
lambda model_id, source=None: calls.append((model_id, source)),
|
|
)
|
|
|
|
downloader.download_model(
|
|
"Qwen/Qwen2-0.5B-Instruct-GGUF",
|
|
source=DownloadSource.HUGGINGFACE,
|
|
)
|
|
|
|
assert calls == [
|
|
(
|
|
"Qwen/Qwen2-0.5B-Instruct-GGUF",
|
|
DownloadSource.HUGGINGFACE,
|
|
),
|
|
]
|
|
|
|
|
|
def test_get_download_progress_returns_idle_by_default() -> None:
|
|
downloader = ModelManager()
|
|
|
|
assert downloader.get_download_progress() == {
|
|
"status": "idle",
|
|
"model_name": None,
|
|
"downloaded_bytes": 0,
|
|
"total_bytes": None,
|
|
"speed_bytes_per_sec": 0.0,
|
|
"source": None,
|
|
"error": None,
|
|
"local_path": None,
|
|
}
|
|
|
|
|
|
def test_download_model_rejects_repo_without_gguf(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
downloader = ModelManager()
|
|
controller = _FakeController()
|
|
downloader.__dict__["_download_controller"] = controller
|
|
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"get_model_dir",
|
|
lambda repo_id: tmp_path / repo_id,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_resolve_download_source",
|
|
lambda: DownloadSource.MODELSCOPE,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_estimate_download_size",
|
|
lambda **kwargs: 100,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_check_gguf_exists",
|
|
lambda **kwargs: (
|
|
False,
|
|
(
|
|
"Repository demo/no-gguf does not contain any .gguf "
|
|
"files on ModelScope."
|
|
),
|
|
),
|
|
)
|
|
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="does not contain any .gguf files",
|
|
):
|
|
downloader.start_download("demo/no-gguf")
|
|
|
|
assert controller.started_spec is None
|
|
|
|
|
|
def test_cancel_download_delegates_to_controller() -> None:
|
|
downloader = ModelManager()
|
|
controller = _FakeController()
|
|
downloader.__dict__["_download_controller"] = controller
|
|
|
|
downloader.cancel_download()
|
|
|
|
assert controller.cancel_called is True
|
|
|
|
|
|
def test_start_download_uses_explicit_source_without_probe(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
downloader = ModelManager()
|
|
controller = _FakeController()
|
|
downloader.__dict__["_download_controller"] = controller
|
|
downloader.__dict__["_download_tmp_dir"] = tmp_path / "tmp"
|
|
target_dir = tmp_path / "resolved-model-dir"
|
|
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"get_model_dir",
|
|
lambda repo_id: target_dir,
|
|
)
|
|
|
|
def _unexpected_probe() -> DownloadSource:
|
|
raise AssertionError("source probing should be skipped")
|
|
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_resolve_download_source",
|
|
_unexpected_probe,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_estimate_download_size",
|
|
lambda **kwargs: 100,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloader,
|
|
"_check_gguf_exists",
|
|
lambda **kwargs: (True, ""),
|
|
)
|
|
|
|
downloader.start_download(
|
|
"Qwen/Qwen2-0.5B-Instruct-GGUF",
|
|
source=DownloadSource.HUGGINGFACE,
|
|
)
|
|
|
|
assert controller.started_spec is not None
|
|
assert controller.started_spec.command == [
|
|
"qwenpaw-model-download",
|
|
"Qwen/Qwen2-0.5B-Instruct-GGUF",
|
|
"huggingface",
|
|
]
|
|
assert controller.started_spec.source == "huggingface"
|
|
|
|
|
|
def test_finalize_download_result_promotes_staging_dir(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
downloader = ModelManager()
|
|
staging_dir = tmp_path / "staging"
|
|
final_dir = tmp_path / "final"
|
|
staging_dir.mkdir()
|
|
(staging_dir / "model.gguf").write_bytes(b"123")
|
|
|
|
result, downloaded_bytes = downloader._finalize_download_result(
|
|
DownloadTaskResult(
|
|
status=DownloadTaskStatus.COMPLETED,
|
|
local_path=str(staging_dir),
|
|
),
|
|
staging_dir=staging_dir,
|
|
final_dir=final_dir,
|
|
)
|
|
|
|
assert result.status == DownloadTaskStatus.COMPLETED
|
|
assert result.local_path == str(final_dir)
|
|
assert downloaded_bytes == 3
|
|
assert not staging_dir.exists()
|
|
assert (final_dir / "model.gguf").exists()
|
|
|
|
|
|
def test_get_model_dir_preserves_repo_id_path() -> None:
|
|
downloader = ModelManager()
|
|
|
|
model_dir = downloader.get_model_dir("Qwen/Qwen3-0.6B-GGUF")
|
|
|
|
assert model_dir.parts[-2:] == ("Qwen", "Qwen3-0.6B-GGUF")
|
|
|
|
|
|
def test_list_and_remove_downloaded_models_with_repo_id_layout(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
downloader = ModelManager()
|
|
downloader.__dict__["_model_dir"] = tmp_path / "models"
|
|
|
|
repo_dir = downloader.get_model_dir("Qwen/Qwen3-0.6B-GGUF")
|
|
repo_dir.mkdir(parents=True)
|
|
(repo_dir / "model.gguf").write_bytes(b"123")
|
|
(repo_dir / "README.md").write_text("demo", encoding="utf-8")
|
|
|
|
models = downloader.list_downloaded_models()
|
|
|
|
assert len(models) == 1
|
|
assert models[0].id == "Qwen/Qwen3-0.6B-GGUF"
|
|
assert models[0].name == "Qwen/Qwen3-0.6B-GGUF"
|
|
|
|
downloader.remove_downloaded_model("Qwen/Qwen3-0.6B-GGUF")
|
|
|
|
assert not repo_dir.exists()
|
|
assert not (tmp_path / "models" / "Qwen").exists()
|
|
|
|
|
|
def test_list_downloaded_models_ignores_temporary_download_dirs(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
downloader = ModelManager()
|
|
downloader.__dict__["_model_dir"] = tmp_path / "models"
|
|
|
|
completed_dir = downloader.get_model_dir("Qwen/Qwen3-0.6B-GGUF")
|
|
completed_dir.mkdir(parents=True)
|
|
(completed_dir / "model.gguf").write_bytes(b"123")
|
|
|
|
staging_dir = (
|
|
completed_dir.parent / ".Qwen3-0.6B-GGUF.1234abcd.downloading"
|
|
)
|
|
staging_dir.mkdir(parents=True)
|
|
(staging_dir / "partial.gguf").write_bytes(b"12")
|
|
|
|
models = downloader.list_downloaded_models()
|
|
|
|
assert [model.id for model in models] == ["Qwen/Qwen3-0.6B-GGUF"]
|
|
|
|
|
|
def test_download_worker_sanitizes_standard_streams(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
queue_messages: list[dict[str, object | dict[str, object]]] = []
|
|
calls: list[str] = []
|
|
|
|
class _Queue:
|
|
def put(self, item):
|
|
queue_messages.append(item)
|
|
|
|
monkeypatch.setattr(
|
|
"pineagents.local_models.model_manager.ensure_standard_streams",
|
|
lambda: calls.append("sanitized"),
|
|
)
|
|
monkeypatch.setattr(
|
|
ModelManager,
|
|
"_download_to_directory",
|
|
staticmethod(lambda **kwargs: str(tmp_path / "downloaded")),
|
|
)
|
|
|
|
getattr(ModelManager, "_download_worker")(
|
|
{
|
|
"repo_id": "AgentScope/demo",
|
|
"source": "modelscope",
|
|
"staging_dir": str(tmp_path / "staging"),
|
|
},
|
|
_Queue(),
|
|
)
|
|
|
|
assert calls == ["sanitized"]
|
|
assert queue_messages[0]["type"] == "result"
|
|
payload = queue_messages[0]["payload"]
|
|
assert isinstance(payload, dict)
|
|
assert payload["status"] == "completed"
|
|
|
|
|
|
def test_download_from_modelscope_redirects_broken_standard_streams(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
class _BrokenStream:
|
|
encoding = "utf-8"
|
|
|
|
def write(self, _text: str) -> int:
|
|
raise OSError(22, "Invalid argument")
|
|
|
|
def flush(self) -> None:
|
|
raise OSError(22, "Invalid argument")
|
|
|
|
captured: dict[str, object] = {}
|
|
|
|
def _fake_snapshot_download(*, model_id: str, local_dir: str) -> str:
|
|
captured["model_id"] = model_id
|
|
captured["local_dir"] = local_dir
|
|
captured["stdout"] = sys.stdout
|
|
captured["stderr"] = sys.stderr
|
|
sys.stderr.write("")
|
|
sys.stderr.flush()
|
|
return local_dir
|
|
|
|
monkeypatch.setattr(sys, "stdout", _BrokenStream())
|
|
monkeypatch.setattr(sys, "stderr", _BrokenStream())
|
|
monkeypatch.setattr(
|
|
ModelManager,
|
|
"_get_modelscope_snapshot_download",
|
|
staticmethod(lambda: _fake_snapshot_download),
|
|
)
|
|
|
|
result = ModelManager._download_from_modelscope(
|
|
repo_id="AgentScope/demo",
|
|
local_dir=tmp_path / "download",
|
|
)
|
|
|
|
assert result == str(tmp_path / "download")
|
|
assert captured["model_id"] == "AgentScope/demo"
|
|
assert captured["local_dir"] == str(tmp_path / "download")
|
|
assert captured["stdout"] is not sys.stdout
|
|
assert captured["stderr"] is not sys.stderr
|