fix(scripts): push_predictions 的分页与令牌过期
两个在真实环境暴露的 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 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
9ead8c0a6e
commit
744e5fd18a
+56
-40
@@ -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()
|
||||
self.request("DELETE", f"/api/predictions/{item['id']}/")
|
||||
removed += 1
|
||||
if isinstance(payload, dict) and not payload.get("next"):
|
||||
break
|
||||
page += 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
|
||||
|
||||
@@ -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/"
|
||||
|
||||
Reference in New Issue
Block a user