mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-10-25 09:35:17 +08:00
536 lines
26 KiB
Python
536 lines
26 KiB
Python
#!/usr/bin/env python3
|
|
"""仅使用 Python 标准库的 tasks HTTP 示例：JSON 清单、逐文件上传、轮询与独立结果下载。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import http.client
|
|
import json
|
|
import os
|
|
import socket
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import zipfile
|
|
from collections.abc import Callable, Generator, Iterator, Mapping
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.parse import quote, urljoin, urlsplit
|
|
|
|
TERMINAL = frozenset({"completed", "partial", "failed", "canceled"})
|
|
STATUSES = TERMINAL | {"waiting_upload", "queued", "running"}
|
|
RETRYABLE_HTTP = frozenset({408, 429, 500, 502, 503, 504})
|
|
REDIRECTS = frozenset({301, 302, 303, 307, 308})
|
|
CHUNK_BYTES = 1024 * 1024
|
|
MAX_JSON_BYTES = 4 * 1024 * 1024
|
|
BodyFactory = Callable[[], Generator[bytes, None, None]]
|
|
|
|
|
|
class ClientError(RuntimeError):
|
|
"""保存可重试属性，避免将错误 hash、鉴权或非法响应反复提交。"""
|
|
|
|
def __init__(self, message: str, *, retryable: bool = False, status: int | None = None) -> None:
|
|
"""区分网络临时错误、HTTP 状态和确定性的协议错误。"""
|
|
super().__init__(message)
|
|
self.retryable = retryable
|
|
self.status = status
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class InputFile:
|
|
"""固定文件顺序和客户端标识，恢复任务时仍按同一份清单对应文件。"""
|
|
|
|
path: Path
|
|
client_file_id: str
|
|
size: int
|
|
sha256sum: str | None
|
|
|
|
def task_file(self) -> dict[str, Any]:
|
|
"""生成 JSON 清单条目；关闭 hash 时完全省略该可选字段。"""
|
|
item: dict[str, Any] = {"client_file_id": self.client_file_id, "file_name": self.path.name, "bytes": self.size}
|
|
if self.sha256sum is not None:
|
|
item["sha256sum"] = self.sha256sum
|
|
return item
|
|
|
|
|
|
def _file_body(path: Path) -> Generator[bytes, None, None]:
|
|
"""文件上传按固定块读取，中断或重试均关闭本次文件句柄。"""
|
|
try:
|
|
with path.open("rb") as source:
|
|
while chunk := source.read(CHUNK_BYTES):
|
|
yield chunk
|
|
except OSError as exc:
|
|
raise ClientError(f"Cannot read input file: {path.name}: {exc}") from exc
|
|
|
|
|
|
def _bytes_body(data: bytes) -> Generator[bytes, None, None]:
|
|
"""为小型 JSON 主体提供可重新创建的发送迭代器。"""
|
|
yield data
|
|
|
|
|
|
def _origin(url: str) -> tuple[str, str, int]:
|
|
"""严格比较协议、主机和有效端口，阻止跨源转发 API Key。"""
|
|
parts = urlsplit(url)
|
|
if parts.scheme not in {"http", "https"} or not parts.hostname or parts.username or parts.password:
|
|
raise ClientError("Expected an http(s) URL without embedded credentials")
|
|
return parts.scheme.lower(), parts.hostname.lower(), parts.port or (443 if parts.scheme == "https" else 80)
|
|
|
|
|
|
def _integer(name: str, default: int, minimum: int = 1) -> int:
|
|
"""把环境变量明确校验为有界客户端预算，错误配置立即退出。"""
|
|
try:
|
|
value = int(os.environ.get(name, str(default)))
|
|
except ValueError as exc:
|
|
raise ClientError(f"{name} must be an integer") from exc
|
|
if value < minimum:
|
|
raise ClientError(f"{name} must be at least {minimum}")
|
|
return value
|
|
|
|
|
|
class TaskClient:
|
|
"""公共客户端只识别 task_id、状态和服务返回的上传及下载链接。"""
|
|
|
|
def __init__(self, files: list[InputFile]) -> None:
|
|
"""读取连接与轮询预算，不依赖 MinerU SDK 或云存储 SDK。"""
|
|
self.base = os.environ.get("MINERU_API_URL", "").rstrip("/")
|
|
if not self.base:
|
|
raise ClientError("Set MINERU_API_URL to the API base URL")
|
|
self.origin = _origin(self.base)
|
|
if urlsplit(self.base).query or urlsplit(self.base).fragment:
|
|
raise ClientError("MINERU_API_URL must not contain a query or fragment")
|
|
self.key = os.environ.get("MINERU_API_KEY", "")
|
|
self.files = files
|
|
self.by_id = {item.client_file_id: item for item in files}
|
|
self.connect_timeout = _integer("CONNECT_TIMEOUT", 10)
|
|
self.timeout = _integer("REQUEST_TIMEOUT", 600)
|
|
self.max_polls = _integer("MAX_POLLS", 300, minimum=0)
|
|
try:
|
|
self.poll_interval = float(os.environ.get("POLL_INTERVAL", "1"))
|
|
except ValueError as exc:
|
|
raise ClientError("POLL_INTERVAL must be a non-negative number") from exc
|
|
if not 0 <= self.poll_interval < float("inf"):
|
|
raise ClientError("POLL_INTERVAL must be a non-negative finite number")
|
|
self.task_id: str | None = None
|
|
self.status_url: str | None = None
|
|
self.saved_ids: set[str] = set()
|
|
|
|
def resolve(self, url: str) -> str:
|
|
"""解析服务返回的绝对或相对链接，保留 OSS 签名查询串。"""
|
|
resolved = urljoin(self.base + "/", url)
|
|
_origin(resolved)
|
|
return resolved
|
|
|
|
def _headers(self, url: str, headers: Mapping[str, str] | None) -> dict[str, str]:
|
|
"""只向 API 同源地址附加 Key，同时保留明确给出的存储服务请求头。"""
|
|
result = dict(headers or {})
|
|
for name, value in result.items():
|
|
if not isinstance(name, str) or not isinstance(value, str) or "\r" in name + value or "\n" in name + value:
|
|
raise ClientError("Invalid service-provided HTTP header")
|
|
same_origin = _origin(url) == self.origin
|
|
if self.key and same_origin:
|
|
if not any(name.lower() == "authorization" for name in result):
|
|
result["Authorization"] = f"Bearer {self.key}"
|
|
elif self.key:
|
|
for name, value in result.items():
|
|
if name.lower() == "authorization" and value == f"Bearer {self.key}":
|
|
raise ClientError("Refusing to forward the API Key to another origin")
|
|
return result
|
|
|
|
@contextmanager
|
|
def request(
|
|
self,
|
|
method: str,
|
|
url: str,
|
|
*,
|
|
headers: Mapping[str, str] | None = None,
|
|
body: BodyFactory | None = None,
|
|
length: int | None = None,
|
|
) -> Iterator[http.client.HTTPResponse]:
|
|
"""流式发送与读取一次请求；创建任务不进行隐式重试或 POST 重定向。"""
|
|
resolved = self.resolve(url)
|
|
parts = urlsplit(resolved)
|
|
connection_type = http.client.HTTPSConnection if parts.scheme == "https" else http.client.HTTPConnection
|
|
connection = connection_type(parts.hostname, parts.port, timeout=self.connect_timeout)
|
|
outgoing = self._headers(resolved, headers)
|
|
if length is not None:
|
|
outgoing = {
|
|
name: value for name, value in outgoing.items() if name.lower() not in {"content-length", "transfer-encoding"}
|
|
}
|
|
outgoing["Content-Length"] = str(length)
|
|
iterator = body() if body is not None else None
|
|
try:
|
|
try:
|
|
connection.connect()
|
|
if connection.sock is not None:
|
|
connection.sock.settimeout(self.timeout)
|
|
target = parts.path or "/"
|
|
if parts.query:
|
|
target += "?" + parts.query
|
|
connection.request(method, target, body=iterator, headers=outgoing, encode_chunked=False)
|
|
response = connection.getresponse()
|
|
except (OSError, http.client.HTTPException, socket.timeout) as exc:
|
|
raise ClientError(f"{method} connection failed: {exc}", retryable=True) from exc
|
|
try:
|
|
yield response
|
|
finally:
|
|
response.close()
|
|
finally:
|
|
if iterator is not None:
|
|
iterator.close()
|
|
connection.close()
|
|
|
|
def _read_response(self, response: http.client.HTTPResponse, size: int) -> bytes:
|
|
"""只把响应读取异常标记为网络错误，不能把本地写盘错误错误地加入重试。"""
|
|
try:
|
|
return response.read(size)
|
|
except (OSError, http.client.HTTPException, socket.timeout) as exc:
|
|
raise ClientError(f"Response interrupted: {exc}", retryable=True) from exc
|
|
|
|
def _http_error(self, response: http.client.HTTPResponse) -> ClientError:
|
|
"""提取小型错误消息，不在日志中输出带签名的下载地址或认证信息。"""
|
|
content = self._read_response(response, MAX_JSON_BYTES + 1)
|
|
message = f"HTTP {response.status}"
|
|
try:
|
|
value = json.loads(content)
|
|
detail = value.get("error") if isinstance(value, dict) else None
|
|
if isinstance(detail, dict):
|
|
message += f": {detail.get('code', 'request_failed')}: {detail.get('message', '')}"
|
|
except (ValueError, UnicodeError):
|
|
pass
|
|
return ClientError(message, retryable=response.status in RETRYABLE_HTTP, status=response.status)
|
|
|
|
def json_request(
|
|
self,
|
|
method: str,
|
|
url: str,
|
|
*,
|
|
data: bytes | None = None,
|
|
body: BodyFactory | None = None,
|
|
length: int | None = None,
|
|
headers: Mapping[str, str] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""读取有界的任务元数据，HTTP 200 的非法 JSON 也立即报告错误。"""
|
|
if data is not None:
|
|
|
|
def data_body() -> Generator[bytes, None, None]:
|
|
"""闭包固定本次 JSON 元数据字节，不把源文件主体放入创建请求。"""
|
|
yield from _bytes_body(data)
|
|
|
|
body, length = data_body, len(data)
|
|
with self.request(method, url, body=body, length=length, headers=headers) as response:
|
|
if response.status not in {200, 202}:
|
|
raise self._http_error(response)
|
|
content = self._read_response(response, MAX_JSON_BYTES + 1)
|
|
if len(content) > MAX_JSON_BYTES:
|
|
raise ClientError("Task metadata response is too large")
|
|
try:
|
|
payload = json.loads(content)
|
|
except (ValueError, UnicodeError) as exc:
|
|
raise ClientError("Task response is not valid JSON") from exc
|
|
if not isinstance(payload, dict):
|
|
raise ClientError("Task response must be a JSON object")
|
|
return self.validate_task(payload)
|
|
|
|
def validate_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""校验任务身份和文件映射，恢复请求不能误把文件上传给另一项。"""
|
|
task_id = payload.get("task_id")
|
|
if not isinstance(task_id, str) or not task_id or (self.task_id is not None and task_id != self.task_id):
|
|
raise ClientError("Missing or inconsistent task_id")
|
|
if payload.get("status") not in STATUSES:
|
|
raise ClientError("Unknown task status")
|
|
if not isinstance(payload.get("status_url"), str) or not payload["status_url"]:
|
|
raise ClientError("Missing status_url")
|
|
files = payload.get("files")
|
|
if not isinstance(files, list) or len(files) != len(self.files):
|
|
raise ClientError("Task file count does not match local inputs")
|
|
seen: set[str] = set()
|
|
for item in files:
|
|
if not isinstance(item, dict):
|
|
raise ClientError("Invalid task file entry")
|
|
client_id = item.get("client_file_id")
|
|
if not isinstance(client_id, str) or client_id not in self.by_id or client_id in seen:
|
|
raise ClientError("Invalid or duplicate client_file_id")
|
|
seen.add(client_id)
|
|
local = self.by_id[client_id]
|
|
if item.get("file_name") != local.path.name or item.get("bytes") != local.size:
|
|
raise ClientError("Task files differ from local inputs; use the original file list and order")
|
|
if local.sha256sum is not None and item.get("sha256sum") not in {None, local.sha256sum}:
|
|
raise ClientError("Task hash differs from the local file")
|
|
if item.get("upload_status") not in {"pending", "ready", "expired"}:
|
|
raise ClientError("Unknown upload_status")
|
|
uploads = payload.get("uploads")
|
|
if not isinstance(uploads, list):
|
|
raise ClientError("Missing uploads list")
|
|
seen_uploads: set[str] = set()
|
|
pending = {item["client_file_id"] for item in files if item["upload_status"] == "pending"}
|
|
for link in uploads:
|
|
if not isinstance(link, dict):
|
|
raise ClientError("Invalid upload descriptor")
|
|
client_id = link.get("client_file_id")
|
|
if not isinstance(client_id, str) or client_id not in self.by_id or client_id in seen_uploads:
|
|
raise ClientError("Invalid or duplicate upload client_file_id")
|
|
seen_uploads.add(client_id)
|
|
if link.get("upload_method") != "PUT" or not isinstance(link.get("upload_url"), str) or not link["upload_url"]:
|
|
raise ClientError("Invalid upload URL or method")
|
|
if not isinstance(link.get("upload_headers"), dict):
|
|
raise ClientError("Invalid upload headers")
|
|
if not seen_uploads.issubset(pending):
|
|
raise ClientError("Upload descriptors do not match pending files")
|
|
if payload["status"] in TERMINAL and seen_uploads:
|
|
raise ClientError("Terminal task still returns upload URLs")
|
|
formats = payload.get("output_formats")
|
|
if (
|
|
not isinstance(formats, list)
|
|
or not formats
|
|
or any(value not in {"default", "docx", "html", "latex"} for value in formats)
|
|
):
|
|
raise ClientError("Invalid task output_formats")
|
|
for item in files:
|
|
if item.get("status") not in {"waiting_upload", "queued", "running", "completed", "failed", "canceled"}:
|
|
raise ClientError("Unknown file status")
|
|
if item["status"] == "completed":
|
|
if not isinstance(item.get("result_url"), str) or not item["result_url"]:
|
|
raise ClientError("Completed file has no download URL")
|
|
elif item.get("result_url") is not None:
|
|
raise ClientError("Unfinished file has a download URL")
|
|
self.task_id = task_id
|
|
self.status_url = self.resolve(payload["status_url"])
|
|
return payload
|
|
|
|
def create(self) -> dict[str, Any]:
|
|
"""用一次 JSON 请求创建清单，文件主体只发送到服务返回的上传链接。"""
|
|
formats = [value.strip() for value in os.environ.get("OUTPUT_FORMATS", "default").split(",")]
|
|
if not formats or any(value not in {"default", "docx", "html", "latex"} for value in formats):
|
|
raise ClientError("OUTPUT_FORMATS must contain default/docx/html/latex")
|
|
body: dict[str, Any] = {
|
|
"files": [file.task_file() for file in self.files],
|
|
"ocr_mode": os.environ.get("OCR_MODE", "auto"),
|
|
"page_range": os.environ.get("PAGE_RANGE", ""),
|
|
"output_formats": list(dict.fromkeys(formats)),
|
|
}
|
|
model_output = os.environ.get("INCLUDE_MODEL_OUTPUT", "0")
|
|
if model_output not in {"0", "1"}:
|
|
raise ClientError("INCLUDE_MODEL_OUTPUT must be 0 or 1")
|
|
if model_output == "1":
|
|
if "default" not in formats:
|
|
raise ClientError("INCLUDE_MODEL_OUTPUT requires default output format")
|
|
body["include_model_output"] = True
|
|
if os.environ.get("MINERU_TIER"):
|
|
body["tier"] = os.environ["MINERU_TIER"]
|
|
return self.json_request(
|
|
"POST",
|
|
self.base + "/v1/tasks",
|
|
data=json.dumps(body, ensure_ascii=False).encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
|
|
def refresh(self) -> dict[str, Any]:
|
|
"""只查询已有任务，任何查询失败都不能触发创建重试。"""
|
|
if self.status_url is None:
|
|
raise ClientError("No task status URL")
|
|
return self.json_request("GET", self.status_url)
|
|
|
|
def refresh_after_transfer_error(self) -> dict[str, Any]:
|
|
"""传输结果不确定时先恢复状态查询，无法确认时不盲目重复发送文件。"""
|
|
for attempt in range(3):
|
|
try:
|
|
return self.refresh()
|
|
except ClientError as exc:
|
|
if not exc.retryable or attempt == 2:
|
|
raise
|
|
time.sleep(attempt + 1)
|
|
raise AssertionError("unreachable")
|
|
|
|
def upload_missing(self, task: dict[str, Any]) -> dict[str, Any]:
|
|
"""逐文件上传与恢复，已经 ready 的文件永远不会重新发送。"""
|
|
for original in task["uploads"]:
|
|
client_id = original["client_file_id"]
|
|
local = self.by_id[client_id]
|
|
link = original
|
|
for attempt in range(3):
|
|
|
|
def upload_body() -> Generator[bytes, None, None]:
|
|
"""每次重试重新打开同一个文件，单文件从头重传。"""
|
|
yield from _file_body(local.path)
|
|
|
|
try:
|
|
with self.request(
|
|
"PUT",
|
|
link["upload_url"],
|
|
headers=link["upload_headers"],
|
|
body=upload_body,
|
|
length=local.size,
|
|
) as response:
|
|
if not 200 <= response.status < 300:
|
|
raise self._http_error(response)
|
|
break
|
|
except ClientError as exc:
|
|
if not exc.retryable and exc.status not in {403, 410}:
|
|
raise
|
|
# 请求可能已经落盘但响应丢失；先看状态，再决定是否重传。
|
|
current = self.refresh_after_transfer_error()
|
|
remote = next(item for item in current["files"] if item["client_file_id"] == client_id)
|
|
if remote["upload_status"] == "ready":
|
|
break
|
|
if remote["status"] in {"failed", "canceled"}:
|
|
break
|
|
replacement = next((item for item in current["uploads"] if item["client_file_id"] == client_id), None)
|
|
if replacement is None:
|
|
raise ClientError("Pending file has no upload URL") from exc
|
|
if exc.status == 403 and replacement["upload_url"] == link["upload_url"]:
|
|
raise
|
|
if attempt == 2:
|
|
raise ClientError(f"Upload failed after three attempts: {local.path.name}") from exc
|
|
link = replacement
|
|
time.sleep(attempt + 1)
|
|
return self.refresh()
|
|
|
|
def download(self, task: dict[str, Any], file: dict[str, Any], output: Path) -> None:
|
|
"""下载一个已完成文件，按所选格式校验，并安全刷新过期下载链接。"""
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
descriptor, name = tempfile.mkstemp(prefix=output.name + ".", suffix=".part", dir=output.parent)
|
|
os.close(descriptor)
|
|
temporary = Path(name)
|
|
url = self.resolve(file["result_url"])
|
|
try:
|
|
for attempt in range(3):
|
|
try:
|
|
for _ in range(6):
|
|
with self.request("GET", url) as response:
|
|
if response.status in REDIRECTS:
|
|
location = response.getheader("Location")
|
|
if not location:
|
|
raise ClientError("Download redirect is missing Location")
|
|
url = self.resolve(urljoin(url, location))
|
|
continue
|
|
if response.status != 200:
|
|
raise self._http_error(response)
|
|
with temporary.open("wb") as target:
|
|
while chunk := self._read_response(response, CHUNK_BYTES):
|
|
try:
|
|
target.write(chunk)
|
|
except OSError as exc:
|
|
raise ClientError(f"Cannot save result: {exc}") from exc
|
|
expected = response.getheader("Content-Length")
|
|
if expected is not None and temporary.stat().st_size != int(expected):
|
|
raise ClientError("Incomplete result download", retryable=True)
|
|
break
|
|
else:
|
|
raise ClientError("Too many download redirects")
|
|
formats = list(dict.fromkeys(task["output_formats"]))
|
|
if formats == ["html"]:
|
|
if "<html" not in temporary.read_text(encoding="utf-8").lower():
|
|
raise ClientError("Downloaded result is not an HTML document")
|
|
else:
|
|
if not zipfile.is_zipfile(temporary):
|
|
raise ClientError("Downloaded result is not a ZIP archive")
|
|
with zipfile.ZipFile(temporary) as archive:
|
|
if archive.testzip() is not None:
|
|
raise ClientError("Downloaded ZIP is corrupt")
|
|
if formats == ["docx"] and "word/document.xml" not in archive.namelist():
|
|
raise ClientError("Downloaded result is not a DOCX document")
|
|
temporary.replace(output)
|
|
return
|
|
except ClientError as exc:
|
|
if attempt == 2 or (not exc.retryable and exc.status not in {401, 403, 404}):
|
|
raise
|
|
current = self.refresh_after_transfer_error()
|
|
remote = next(item for item in current["files"] if item["client_file_id"] == file["client_file_id"])
|
|
if remote["status"] != "completed":
|
|
raise ClientError("File no longer has a downloadable result") from exc
|
|
replacement = self.resolve(remote["result_url"])
|
|
if exc.status in {401, 403, 404} and replacement == url:
|
|
raise
|
|
url = replacement
|
|
time.sleep(attempt + 1)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
|
|
def download_completed(self, task: dict[str, Any], output: Path) -> None:
|
|
"""按 client_file_id 保存独立结果，同名输入和多次轮询不会相互覆盖。"""
|
|
formats = list(dict.fromkeys(task["output_formats"]))
|
|
extension = formats[0] if len(formats) == 1 and formats[0] in {"html", "docx"} else "zip"
|
|
for file in task["files"]:
|
|
client_id = file["client_file_id"]
|
|
if file["status"] != "completed" or client_id in self.saved_ids:
|
|
continue
|
|
destination = output / client_id / (self.by_id[client_id].path.stem + "." + extension)
|
|
self.download(task, file, destination)
|
|
self.saved_ids.add(client_id)
|
|
print(f"Saved {destination}")
|
|
|
|
def run(self, output: Path) -> int:
|
|
"""上传缺失文件并逐个下载成功结果，失败与取消也保留已经完成的产物。"""
|
|
resume = os.environ.get("TASK_ID", "")
|
|
if resume:
|
|
self.task_id = resume
|
|
self.status_url = self.base + "/v1/tasks/" + quote(resume, safe="")
|
|
task = self.refresh()
|
|
else:
|
|
task = self.create()
|
|
print(f"task_id={self.task_id}", file=sys.stderr)
|
|
self.download_completed(task, output)
|
|
for link in list(task["uploads"]):
|
|
current_links = [item for item in task["uploads"] if item["client_file_id"] == link["client_file_id"]]
|
|
if not current_links:
|
|
continue
|
|
task = self.upload_missing({**task, "uploads": current_links})
|
|
self.download_completed(task, output)
|
|
polls = 0
|
|
while task["status"] not in TERMINAL:
|
|
if polls >= self.max_polls:
|
|
print(f"Polling budget exhausted; task_id={self.task_id} continues running", file=sys.stderr)
|
|
return 124
|
|
polls += 1
|
|
time.sleep(self.poll_interval)
|
|
try:
|
|
task = self.refresh()
|
|
except ClientError as exc:
|
|
if not exc.retryable:
|
|
raise
|
|
continue
|
|
self.download_completed(task, output)
|
|
if task["status"] in {"failed", "canceled"}:
|
|
print(
|
|
json.dumps({"status": task["status"], "error": task.get("error"), "files": task["files"]}, ensure_ascii=False),
|
|
file=sys.stderr,
|
|
)
|
|
return 1
|
|
return 2 if task["status"] == "partial" else 0
|
|
|
|
|
|
def main(arguments: list[str]) -> int:
|
|
"""构造真实文件清单并报告可恢复的失败，所有文件读取都采用流式处理。"""
|
|
client: TaskClient | None = None
|
|
try:
|
|
if not arguments:
|
|
raise ClientError("usage: bash scripts/http_task_example.sh <file> [file ...]")
|
|
if os.environ.get("UPLOAD_MODE") or os.environ.get("RESULT_FILE"):
|
|
raise ClientError("UPLOAD_MODE and RESULT_FILE were removed; use JSON preflight and RESULT_DIR")
|
|
send_hash = os.environ.get("SEND_HASH", "1")
|
|
if send_hash not in {"0", "1"}:
|
|
raise ClientError("SEND_HASH must be 0 or 1")
|
|
inputs: list[InputFile] = []
|
|
for index, value in enumerate(arguments):
|
|
path = Path(value).resolve()
|
|
if not path.is_file() or path.stat().st_size == 0:
|
|
raise ClientError(f"Input must be a non-empty file: {value}")
|
|
hasher = hashlib.sha256() if send_hash == "1" else None
|
|
if hasher is not None:
|
|
for chunk in _file_body(path):
|
|
hasher.update(chunk)
|
|
inputs.append(InputFile(path, f"f{index + 1}", path.stat().st_size, hasher.hexdigest() if hasher else None))
|
|
client = TaskClient(inputs)
|
|
return client.run(Path(os.environ.get("RESULT_DIR", "results")))
|
|
except (ClientError, OSError, ValueError) as exc:
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
if client is not None and client.task_id:
|
|
print(f"Resume with TASK_ID={client.task_id} and the original file list", file=sys.stderr)
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main(sys.argv[1:]))
|