Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
File size: 6,804 Bytes
ee0b555 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | """The in-process query layer must mirror ``pwc ... --json`` exactly."""
from __future__ import annotations
import argparse
import io
import json
from contextlib import redirect_stdout
import pytest
from pwc_cli import queries
from pwc_cli.cli import UsageError, build_parser, main
from pwc_cli.transport import Response
NON_RESEARCH_COMMANDS = {
("skills", "add"),
("version",),
("auth", "login"),
("auth", "status"),
("auth", "logout"),
("paper", "edit", "export"),
("paper", "edit", "preview"),
("paper", "edit", "submit"),
}
EVALUATION_ROWS = [
{
"id": "1",
"model_name": "Small",
"metrics": {"mAP": 50.0},
"best_rank": 2,
"num_parameters": 100_000_000,
"paper_id": "1",
"task_id": "t",
"dataset_id": "9",
},
{
"id": "2",
"model_name": "Smaller",
"metrics": {"mAP": 60.0},
"best_rank": 1,
"num_parameters": 50_000_000,
"paper_id": "2",
"task_id": "t",
"dataset_id": "9",
},
]
BENCHMARK_ROUTES = {
"datasets/": {"results": [{"id": "9", "name": "COCO", "slug": "coco"}]},
"datasets/9/evaluations/": {"count": 2, "results": EVALUATION_ROWS},
"evaluations/": {
"count": 2,
"next_page": None,
"parameter_coverage_known": 2,
"parameter_coverage_total": 2,
"results": EVALUATION_ROWS,
},
}
class StubClient:
def __init__(self, routes):
self.routes = routes
self.calls = []
def get(self, path, params=None):
self.calls.append((path, dict(params or {})))
return Response(json.dumps(self.routes[path]).encode(), {})
def _leaf_commands(parser, prefix=()):
for action in parser._actions:
if not isinstance(action, argparse._SubParsersAction):
continue
for name, child in action.choices.items():
path = (*prefix, name)
if "handler" in child._defaults:
yield path
yield from _leaf_commands(child, path)
def test_read_only_commands_cover_every_research_command():
research = set(_leaf_commands(build_parser())) - NON_RESEARCH_COMMANDS
assert set(queries.READ_ONLY_COMMANDS) == research
for command in queries.READ_ONLY_COMMANDS:
assert "handler" in queries.command_parser(command)._defaults
def test_query_options_exclude_presentation_flags_and_selectors():
options = queries.query_options(("benchmark", "list"))
assert set(options) == {
"page",
"page_size",
"search",
"task",
"group_by_area",
"area",
"benchmarks_per_task",
"include_descendants",
"min_eval_count",
"is_open",
"order_by",
"order_dir",
}
assert queries.query_options(("search",))["query"].option_strings == []
def test_build_argv_translates_structured_options():
argv = queries.build_argv(
("benchmark",),
{
"name": "COCO",
"limit": 5,
"is_open": True,
"require_metrics": ["mAP", "FPS"],
"minimum_metrics": {"mAP": 40},
"maximum_metrics": {"FPS": 10.5},
"pareto": ["mAP:higher", "FPS:lower"],
"max_parameters": None,
},
)
assert argv == [
"benchmark",
"--name",
"COCO",
"--limit",
"5",
"--is-open",
"true",
"--require-metrics",
"mAP,FPS",
"--min",
"mAP=40",
"--max",
"FPS=10.5",
"--pareto",
"mAP:higher,FPS:lower",
]
assert queries.build_argv(
("paper", "list"),
{"author": ["Ada", "@ada"], "all_versions": True, "include_resources": False},
) == ["paper", "list", "--author", "Ada", "--author", "@ada", "--all-versions"]
argv = queries.build_argv(("search",), {"query": "-attention", "limit": 3})
assert argv == ["search", "--limit", "3", "--", "-attention"]
assert build_parser().parse_args(argv).query == "-attention"
def test_query_returns_exactly_the_cli_json_data(monkeypatch, capsys):
monkeypatch.setattr("pwc_cli.cli.Client", lambda: StubClient(BENCHMARK_ROUTES))
output = io.StringIO()
with redirect_stdout(output):
assert (
main(
[
"benchmark",
"--name",
"COCO",
"--max-parameters",
"500M",
"--sort",
"mAP",
"--limit",
"1",
"--json",
]
)
== 0
)
printed = json.loads(output.getvalue())
client = StubClient(BENCHMARK_ROUTES)
data = queries.query(
("benchmark",),
{"name": "COCO", "max_parameters": "500M", "sort_metric": "mAP", "limit": 1},
client,
)
assert {"schema_version": "v1", "data": data} == printed
assert data["matched_count"] == 2
assert [row["model_name"] for row in data["results"]] == ["Smaller"]
assert client.calls[1][1]["max_parameters_exclusive"] == 500_000_001
assert capsys.readouterr().out == ""
def test_query_raises_usage_errors_instead_of_exiting():
client = StubClient({})
with pytest.raises(UsageError, match="--limit must be 1-100"):
queries.query(("search",), {"query": "x", "limit": 500}, client)
with pytest.raises(UsageError, match="unknown option"):
queries.query(("search",), {"query": "x", "page_size": 5}, client)
with pytest.raises(UsageError, match="invalid"):
queries.query(("search",), {"query": "x", "start_date": "2026-13-01"}, client)
with pytest.raises(UsageError, match="cannot be combined"):
queries.query(("task", "list"), {"group_by_area": True, "page": 2}, client)
with pytest.raises(UsageError, match="unknown metric"):
queries.query(
("benchmark",),
{"name": "COCO", "sort_metric": "Latency"},
StubClient(BENCHMARK_ROUTES),
)
assert client.calls == []
def test_query_refuses_commands_that_are_not_read_only():
for command in (("skills", "add"), ("version",), ("paper", "edit", "export")):
with pytest.raises(UsageError, match="not a read-only"):
queries.query(command, {}, StubClient({}))
def test_query_resolves_metric_aliases_against_the_leaderboard():
# Agents ask for "AP" on COCO leaderboards that report "mAP".
data = queries.query(
("benchmark",),
{"name": "COCO", "sort_metric": "AP:asc"},
StubClient(BENCHMARK_ROUTES),
)
assert data["matched_count"] == 2
assert [row["model_name"] for row in data["results"]] == ["Small", "Smaller"]
|