From 744e5fd18a3b169b784c251c0eb309d7262a2b76 Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Tue, 22 Sep 2026 13:17:53 +0800 Subject: [PATCH] =?UTF-8?q?fix(scripts):=20push=5Fpredictions=20=E7=9A=84?= =?UTF-8?q?=E5=88=86=E9=A1=B5=E4=B8=8E=E4=BB=A4=E7=89=8C=E8=BF=87=E6=9C=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 两个在真实环境暴露的 bug： 1. /api/predictions/ 根本不分页——Label Studio 的 REST_FRAMEWORK 没设 DEFAULT_PAGINATION_CLASS，PAGE_SIZE 形同虚设，它一次返回全部预测并忽略 page。 原先按页翻永远不终止，实测翻到第 1798 页才因别的原因报错。改成只请求一次。 2. access token 默认只有 5 分钟（LS 未覆盖 SimpleJWT 默认值），而推上万条预测远超 5 分钟。所有请求统一走 request()，碰到 401 自动用 refresh token 续期后重试一次。 另加 --replace-all：用 Data Manager 的批量操作一次删光项目里所有预测， 比逐条删同版本快得多（适合版本号变了、旧预测留着没用的情况）。 Co-Authored-By: Claude Opus 5 --- scripts/push_predictions.py | 102 +++++++++++++++++++-------------- tests/test_push_predictions.py | 78 +++++++++++++++++++------ 2 files changed, 118 insertions(+), 62 deletions(-) diff --git a/scripts/push_predictions.py b/scripts/push_predictions.py index 9ac9712..e714493 100644 --- a/scripts/push_predictions.py +++ b/scripts/push_predictions.py @@ -23,7 +23,8 @@ Label Studio 的 `LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT` 一致。 **预测是追加的，不会覆盖。** 同一个任务可以有多版预测，Label Studio 在项目设置里按 -model_version 选用哪一版。重复跑同一个版本会堆重复预测，用 `--replace` 先删掉同名版本。 +model_version 选用哪一版。重复跑同一个版本会堆重复预测：`--replace` 先删掉同名版本 +（换了权重时一般删不到东西，很快），`--replace-all` 则一次删光项目里的所有预测。 """ from __future__ import annotations @@ -53,37 +54,51 @@ class LabelStudio: self.url = url.rstrip("/") self.session = requests.Session() - self.session.headers["Authorization"] = self._authorization(token) + self.token = token + self._authorize() - def _authorization(self, token: str) -> str: - """区分两种令牌。 + def _authorize(self) -> None: + """设置 Authorization 头；JWT 令牌会换成 access token。 新版默认关闭了旧式 `Token xxx`（组织设置 legacy_api_tokens_enabled=false）， 个人访问令牌是一个 JWT refresh token，要先换成 access token 再以 Bearer 发送。 """ - if token.count(".") != 2: # 不是 JWT，按旧式 Token 处理 - return f"Token {token}" + if self.token.count(".") != 2: # 不是 JWT，按旧式 Token 处理 + self.session.headers["Authorization"] = f"Token {self.token}" + return response = self.session.post( - f"{self.url}/api/token/refresh/", json={"refresh": token}, timeout=60 + f"{self.url}/api/token/refresh/", json={"refresh": self.token}, timeout=60 ) if response.status_code >= 400: raise RuntimeError( f"用个人访问令牌换取 access token 失败（HTTP {response.status_code}）：" f"{response.text[:200]}" ) - return f"Bearer {response.json()['access']}" + self.session.headers["Authorization"] = f"Bearer {response.json()['access']}" + + def request(self, method: str, path: str, **kwargs): + """带自动续期的请求。 + + access token 默认只有 5 分钟有效期（Label Studio 没有覆盖 SimpleJWT 的默认值）， + 而推完上万条预测远不止 5 分钟，所以碰到 401 就用 refresh token 再换一个再试一次。 + """ + kwargs.setdefault("timeout", 120) + response = self.session.request(method, f"{self.url}{path}", **kwargs) + if response.status_code == 401: + self._authorize() + response = self.session.request(method, f"{self.url}{path}", **kwargs) + response.raise_for_status() + return response def tasks(self, project: int, page_size: int = 500) -> list[dict]: """列出项目里的全部任务（只取 id 与 data，分页拉完）。""" collected, page = [], 1 while True: - response = self.session.get( - f"{self.url}/api/tasks/", + response = self.request( + "GET", "/api/tasks/", params={"project": project, "page": page, "page_size": page_size, "fields": "task_only"}, - timeout=120, ) - response.raise_for_status() payload = response.json() # 不同版本的分页字段不一样：data manager 用 tasks/total，DRF 默认用 results/count batch = payload.get("tasks", payload.get("results", payload)) if isinstance( @@ -99,41 +114,35 @@ class LabelStudio: return collected def delete_predictions(self, project: int, model_version: str) -> int: - """删掉该项目下指定版本的预测，返回删除条数。""" + """删掉该项目下指定版本的预测，返回删除条数。 + + `/api/predictions/` **不分页**：Label Studio 的 REST_FRAMEWORK 没设 + DEFAULT_PAGINATION_CLASS，PAGE_SIZE 形同虚设，它一次返回全部预测并忽略 page 参数。 + 早先按页翻会翻不到头（实测翻到第 1798 页才因令牌过期报错），所以这里只请求一次。 + """ + payload = self.request("GET", "/api/predictions/", params={"project": project}).json() + items = payload.get("results", []) if isinstance(payload, dict) else payload removed = 0 - page = 1 - while True: - response = self.session.get( - f"{self.url}/api/predictions/", - params={"project": project, "page": page, "page_size": 500}, - timeout=120, - ) - if response.status_code == 404: # 翻到尾页，某些版本直接返回 404 - break - response.raise_for_status() - payload = response.json() - items = payload.get("results", payload) if isinstance(payload, dict) else payload - if not items: - break - for item in items: - if item.get("model_version") == model_version: - self.session.delete( - f"{self.url}/api/predictions/{item['id']}/", timeout=60 - ).raise_for_status() - removed += 1 - if isinstance(payload, dict) and not payload.get("next"): - break - page += 1 + for item in items: + if item.get("model_version") == model_version: + self.request("DELETE", f"/api/predictions/{item['id']}/") + removed += 1 return removed + def delete_all_predictions(self, project: int) -> dict: + """用 Data Manager 的批量操作一次删光项目里的所有预测（不分版本）。""" + return self.request( + "POST", "/api/dm/actions/", + params={"id": "delete_tasks_predictions", "project": project}, + json={"selectedItems": {"all": True, "excluded": []}}, + ).json() + def import_predictions(self, project: int, predictions: list[dict]) -> None: """批量导入，服务端自己再分批入库。""" - response = self.session.post( - f"{self.url}/api/projects/{project}/import/predictions", - json=predictions, - timeout=600, + self.request( + "POST", f"/api/projects/{project}/import/predictions", + json=predictions, timeout=600, ) - response.raise_for_status() def local_path(image_url: str, document_root: Path) -> Path: @@ -181,6 +190,10 @@ def main() -> int: "--replace", action="store_true", help="先删掉同一 model_version 的旧预测；否则会追加，同版本堆两份", ) + parser.add_argument( + "--replace-all", action="store_true", + help="先删光项目里的**所有**预测（不分版本），用 Data Manager 的批量操作，一次请求", + ) parser.add_argument("--dry-run", action="store_true", help="只推理不上传") args = parser.parse_args() @@ -231,9 +244,12 @@ def main() -> int: ) return 1 - if args.replace and not args.dry_run: + if args.replace_all and not args.dry_run: + detail = client.delete_all_predictions(args.project) + logger.info("删除项目里的全部旧预测：%s", detail.get("detail", detail)) + elif args.replace and not args.dry_run: removed = client.delete_predictions(args.project, model_version) - logger.info("删除同版本旧预测 %d 条", removed) + logger.info("删除同版本（%s）旧预测 %d 条", model_version, removed) from PIL import Image from ultralytics import YOLO diff --git a/tests/test_push_predictions.py b/tests/test_push_predictions.py index ca24ebe..dc37bdf 100644 --- a/tests/test_push_predictions.py +++ b/tests/test_push_predictions.py @@ -53,9 +53,9 @@ class _FakeSession: self.headers = {} self.calls = [] - def get(self, url, params=None, timeout=None): + def request(self, method, url, params=None, json=None, timeout=None): self.calls.append(params) - page = params["page"] + page = (params or {}).get("page", 1) return _FakeResponse(self.pages[page - 1] if page <= len(self.pages) else {"tasks": []}) @@ -96,36 +96,76 @@ def test_tasks_stops_on_empty_page(client): class _TokenSession: - def __init__(self, access="ACCESS", status=200): + """会 post 换取 access token 的假会话；request 按预设的状态码序列返回。""" + + def __init__(self, access="ACCESS", status=200, statuses=None, payload=None): self.headers = {} self.access, self.status, self.posted = access, status, [] + self.statuses = list(statuses or []) + self.payload = payload if payload is not None else [] + self.requests = [] + self.exchanges = 0 def post(self, url, json=None, timeout=None): self.posted.append((url, json)) - return _FakeResponse({"access": self.access}, self.status) + self.exchanges += 1 + return _FakeResponse({"access": f"{self.access}{self.exchanges}"}, self.status) + + def request(self, method, url, params=None, json=None, timeout=None): + self.requests.append((method, url, self.headers.get("Authorization"))) + status = self.statuses.pop(0) if self.statuses else 200 + return _FakeResponse(self.payload, status) -def test_legacy_token_is_sent_as_token(monkeypatch): +def _client(session, token="aaa.bbb.ccc"): client = pp.LabelStudio.__new__(pp.LabelStudio) client.url = "http://ls" - client.session = _TokenSession() - assert client._authorization("c657261234567890") == "Token c657261234567890" - assert client.session.posted == [] # 旧式令牌不需要换取 + client.session = session + client.token = token + client._authorize() + return client -def test_jwt_personal_token_is_exchanged_for_a_bearer(monkeypatch): +def test_legacy_token_is_sent_as_token(): + session = _TokenSession() + _client(session, token="c657261234567890") + assert session.headers["Authorization"] == "Token c657261234567890" + assert session.posted == [] # 旧式令牌不需要换取 + + +def test_jwt_personal_token_is_exchanged_for_a_bearer(): """新版默认禁用旧式 token，个人访问令牌要先换成 access token。""" - client = pp.LabelStudio.__new__(pp.LabelStudio) - client.url = "http://ls" - client.session = _TokenSession(access="AAA") - assert client._authorization("aaa.bbb.ccc") == "Bearer AAA" - assert client.session.posted[0][0] == "http://ls/api/token/refresh/" - assert client.session.posted[0][1] == {"refresh": "aaa.bbb.ccc"} + session = _TokenSession(access="AAA") + _client(session) + assert session.headers["Authorization"] == "Bearer AAA1" + assert session.posted[0] == ("http://ls/api/token/refresh/", {"refresh": "aaa.bbb.ccc"}) def test_failed_exchange_reports_clearly(): - client = pp.LabelStudio.__new__(pp.LabelStudio) - client.url = "http://ls" - client.session = _TokenSession(status=401) with pytest.raises(RuntimeError, match="access token 失败"): - client._authorization("aaa.bbb.ccc") + _client(_TokenSession(status=401)) + + +def test_expired_access_token_is_refreshed_and_retried(): + """access token 默认只有 5 分钟，推上万条预测必然跨过期，401 要能自愈。""" + session = _TokenSession(statuses=[401, 200]) + client = _client(session) + client.request("GET", "/api/tasks/") + + assert session.exchanges == 2 # 初次换取 + 过期后再换一次 + assert [auth for _, _, auth in session.requests] == ["Bearer ACCESS1", "Bearer ACCESS2"] + + +def test_predictions_endpoint_is_not_paginated(): + """回归：/api/predictions/ 一次返回全部并忽略 page，按页翻会翻不到头。""" + session = _TokenSession(payload=[ + {"id": 1, "model_version": "v1"}, + {"id": 2, "model_version": "v2"}, + {"id": 3, "model_version": "v1"}, + ]) + client = _client(session) + assert client.delete_predictions(39, "v1") == 2 + + methods = [method for method, _, _ in session.requests] + assert methods == ["GET", "DELETE", "DELETE"] # 只列一次，删两条 + assert session.requests[1][1] == "http://ls/api/predictions/1/"