264 lines
10 KiB
Python
264 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""七牛云对象存储的轻量封装（零 SDK 依赖）。
|
||
|
||
仅用 Python 标准库（hmac/hashlib/base64/json）+ curl 实现七牛云的上传/
|
||
下载/列举。worker 端无需 pip 装任何东西。
|
||
|
||
七牛上传签名机制：
|
||
PutPolicy = {"scope":"<bucket>:<key>", "deadline":<unix秒>}
|
||
编码后用 AccessKey/SecretKey 做 HMAC-SHA1 签名，得到 uploadToken。
|
||
|
||
用途（本分布式网格）：
|
||
- 七牛云当"种子库"：每个收敛模型的 .7 大气（393KB）上传为
|
||
<seed_prefix>/<model_name>.7，供异地 worker 在 seed_step 回退时拉取。
|
||
- 仅 master 持有 AK/SK 生成 uptoken；worker 只拿 token 上传（无需密钥）。
|
||
- 大产物（.spec 等 13MB）不上传，留各机本地。
|
||
|
||
环境变量：
|
||
QINIU_ACCESS_KEY / QINIU_SECRET_KEY / QINIU_BUCKET / QINIU_DOMAIN
|
||
（可被 dist_config.yaml 的 qiniu 段覆盖；master 在派发时把 token 传给 worker）
|
||
"""
|
||
import base64
|
||
import hashlib
|
||
import hmac
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import time
|
||
|
||
try:
|
||
import urllib.request as urlreq
|
||
except ImportError: # py2 兜底（不会触发）
|
||
urlreq = None
|
||
|
||
|
||
class QiniuStore:
|
||
"""七牛云存储客户端。master 与 worker 共用；worker 只用 download/upload(token)。"""
|
||
|
||
def __init__(self, access_key=None, secret_key=None, bucket=None,
|
||
domain=None, seed_prefix="seeds"):
|
||
self.ak = access_key or os.environ.get("QINIU_ACCESS_KEY", "")
|
||
self.sk = secret_key or os.environ.get("QINIU_SECRET_KEY", "")
|
||
self.bucket = bucket or os.environ.get("QINIU_BUCKET", "")
|
||
# 绑定域名（如 http:// 或 https:// + 域名，无尾斜杠）。公开读下载用。
|
||
self.domain = (domain or os.environ.get("QINIU_DOMAIN", "")).rstrip("/")
|
||
self.seed_prefix = seed_prefix
|
||
|
||
# ---- base64 url-safe（七牛要求）----
|
||
@staticmethod
|
||
def _b64u(data):
|
||
"""bytes -> url-safe base64 字符串（去 padding 也可，七牛两边都接受）。"""
|
||
if isinstance(data, str):
|
||
data = data.encode("utf-8")
|
||
return base64.urlsafe_b64encode(data).decode("ascii")
|
||
|
||
@staticmethod
|
||
def _b64u_json(obj):
|
||
return QiniuStore._b64u(json.dumps(obj, separators=(",", ":")))
|
||
|
||
def _sign(self, data):
|
||
"""HMAC-SHA1(data, sk) -> url-safe base64。"""
|
||
if isinstance(data, str):
|
||
data = data.encode("utf-8")
|
||
return self._b64u(hmac.new(self.sk.encode("utf-8"), data,
|
||
hashlib.sha1).digest())
|
||
|
||
# ---- uptoken 生成（master 持有 AK/SK 时调用）----
|
||
def gen_uptoken(self, key, expires=3600):
|
||
"""为指定 key 生成上传 token。
|
||
|
||
key: 对象名（不含 bucket 前缀），如 "seeds/t30000_g5.0_he0_c-1_n-1_o-1.7"
|
||
返回 uploadToken 字符串，交给 worker 用 curl 上传。
|
||
"""
|
||
scope = "{}:{}".format(self.bucket, key) if key else self.bucket
|
||
policy = {"scope": scope, "deadline": int(time.time()) + expires}
|
||
encoded = self._b64u_json(policy)
|
||
sign = self._sign(encoded)
|
||
return "{}:{}".format(self.ak, sign) + ":" + encoded
|
||
|
||
def gen_private_url(self, key, expires=3600):
|
||
"""生成私有空间的临时下载 URL（带 e/token 签名）。
|
||
|
||
若 bucket 为公开读，直接用 domain/key 即可，不必调这个。
|
||
"""
|
||
if not self.domain:
|
||
raise ValueError("QINIU_DOMAIN 未配置")
|
||
url = "{}/{}".format(self.domain, key)
|
||
e = int(time.time()) + expires
|
||
to_sign = "{}?e={}".format(url, e)
|
||
token = "{}:{}".format(self.ak, self._sign(to_sign))
|
||
return "{}?e={}&token={}".format(url, e, token)
|
||
|
||
def public_url(self, key):
|
||
"""公开空间的下载 URL（无签名）。"""
|
||
if not self.domain:
|
||
raise ValueError("QINIU_DOMAIN 未配置")
|
||
return "{}/{}".format(self.domain, key)
|
||
|
||
# ---- 上传 / 下载（通过 curl，worker 无需 pip）----
|
||
def upload(self, local_path, key, uptoken=None):
|
||
"""上传本地文件到七牛。uptoken 可由 master 预生成传入。
|
||
|
||
用 curl 表单 POST（七牛上传接口）：
|
||
POST http://upload.qiniup.com multipart/form-data
|
||
field "token"=uptoken, "key"=key, "file"=@local_path
|
||
返回 (ok:bool, resp_text)。
|
||
"""
|
||
if uptoken is None:
|
||
uptoken = self.gen_uptoken(key)
|
||
# 七牛上传域名：华东 upload.qiniup.com，其它区域用对应域名。
|
||
upload_host = os.environ.get("QINIU_UPLOAD_HOST", "upload.qiniup.com")
|
||
url = "http://{}/".format(upload_host)
|
||
cmd = ["curl", "-sS", "-m", "300", "-X", "POST", url,
|
||
"-F", "token={}".format(uptoken),
|
||
"-F", "key={}".format(key),
|
||
"-F", "file=@{}".format(local_path)]
|
||
try:
|
||
out = subprocess.run(cmd, capture_output=True, text=True,
|
||
timeout=320)
|
||
except subprocess.TimeoutExpired:
|
||
return False, "upload timeout (>300s)"
|
||
if out.returncode != 0:
|
||
return False, out.stderr.strip() or out.stdout.strip()
|
||
try:
|
||
resp = json.loads(out.stdout)
|
||
except ValueError:
|
||
return False, out.stdout.strip()
|
||
# 成功响应含 "key" 与 "hash"；失败含 "error"
|
||
if "error" in resp:
|
||
return False, resp["error"]
|
||
return True, resp.get("key", "")
|
||
|
||
def download(self, key, local_path, timeout=120):
|
||
"""下载对象到本地。公开空间用 public_url，私有空间用 private_url。
|
||
|
||
返回 True/False。
|
||
"""
|
||
if self.ak and self.sk:
|
||
url = self.gen_private_url(key)
|
||
else:
|
||
url = self.public_url(key)
|
||
cmd = ["curl", "-sS", "-m", str(timeout), "-f", "-o", local_path, url]
|
||
try:
|
||
r = subprocess.run(cmd, capture_output=True, timeout=timeout + 10)
|
||
return r.returncode == 0
|
||
except subprocess.TimeoutExpired:
|
||
return False
|
||
|
||
# ---- 种子库语义封装 ----
|
||
def seed_key(self, model_name):
|
||
"""规范种子对象名：seeds/<model>.7"""
|
||
return "{}/{}.7".format(self.seed_prefix, model_name)
|
||
|
||
def upload_seed(self, local_path, model_name, uptoken=None):
|
||
"""上传某模型的最终大气 .7 作为种子。返回 (ok, msg)。"""
|
||
return self.upload(local_path, self.seed_key(model_name), uptoken)
|
||
|
||
def download_seed(self, model_name, local_path):
|
||
"""下载某模型的种子 .7 到 local_path。返回 True/False。"""
|
||
return self.download(self.seed_key(model_name), local_path)
|
||
|
||
# ---- 列举（用于 find_seed 跨网查最近邻种子）----
|
||
def list_seeds(self):
|
||
"""返回七牛种子库中所有 <model_name>.7 的模型名列表。
|
||
|
||
调用七牛 RS API: GET /list?bucket=<b>&prefix=<seed_prefix>
|
||
（需 AK/SK 签名）。返回 [] 当未配置或失败。
|
||
结果缓存到本地文件避免每次都拉（见 _seeds_cache_path）。
|
||
"""
|
||
if not (self.ak and self.sk and self.bucket):
|
||
return []
|
||
cache = self._seeds_cache_path()
|
||
# 缓存 5 分钟
|
||
if os.path.isfile(cache) and time.time() - os.path.getmtime(cache) < 300:
|
||
try:
|
||
return json.load(open(cache))
|
||
except Exception:
|
||
pass
|
||
names = self._fetch_seed_list()
|
||
try:
|
||
with open(cache, "w") as f:
|
||
json.dump(names, f)
|
||
except Exception:
|
||
pass
|
||
return names
|
||
|
||
def _seeds_cache_path(self):
|
||
return os.path.join(os.path.dirname(os.path.abspath(__file__)),
|
||
"..", ".qiniu_seeds_cache.json")
|
||
|
||
def _fetch_seed_list(self):
|
||
"""分页拉取七牛种子列表，返回 model_name 列表。"""
|
||
rs_host = os.environ.get("QINIU_RS_HOST", "http://rsf.qiniuapi.com")
|
||
names = []
|
||
marker = ""
|
||
prefix = self.seed_prefix + "/"
|
||
while True:
|
||
path = "/list?bucket={}&prefix={}".format(self.bucket, prefix)
|
||
if marker:
|
||
path += "&marker={}".format(marker)
|
||
# 签名：AccessToken = AK:urlsafe_b64(hmac(sk, path)):<path>
|
||
encoded_path = path # 七牛签名用原始 path
|
||
sign = self._sign(encoded_path)
|
||
access_token = "{}:{}".format(self.ak, sign) + encoded_path
|
||
url = rs_host + path
|
||
cmd = ["curl", "-sS", "-m", "60",
|
||
"-H", "Authorization: QBox " + access_token, url]
|
||
try:
|
||
r = subprocess.run(cmd, capture_output=True, text=True,
|
||
timeout=70)
|
||
if r.returncode != 0:
|
||
break
|
||
data = json.loads(r.stdout)
|
||
except (ValueError, subprocess.TimeoutExpired):
|
||
break
|
||
for item in data.get("items", []):
|
||
k = item.get("key", "")
|
||
# seeds/<model>.7 -> <model>
|
||
if k.startswith(prefix) and k.endswith(".7"):
|
||
names.append(k[len(prefix):-2])
|
||
marker = data.get("marker", "")
|
||
if not marker:
|
||
break
|
||
return names
|
||
|
||
|
||
def load_from_dist_config(cfg):
|
||
"""从 dist_config.yaml 的 qiniu 段构造 QiniuStore。"""
|
||
q = (cfg or {}).get("qiniu", {}) if isinstance(cfg, dict) else {}
|
||
return QiniuStore(
|
||
access_key=q.get("access_key"),
|
||
secret_key=q.get("secret_key"),
|
||
bucket=q.get("bucket"),
|
||
domain=q.get("domain"),
|
||
seed_prefix=q.get("seed_prefix", "seeds"),
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 自检：打印 token（需配置 AK/SK/Bucket）
|
||
import argparse
|
||
ap = argparse.ArgumentParser(description="qiniu store 自检")
|
||
ap.add_argument("--gen-token", help="为某 key 生成 uptoken")
|
||
ap.add_argument("--list", action="store_true", help="列出种子")
|
||
ap.add_argument("--config", help="dist_config.yaml 路径")
|
||
args = ap.parse_args()
|
||
if args.config:
|
||
try:
|
||
import yaml
|
||
cfg = yaml.safe_load(open(args.config))
|
||
except Exception:
|
||
# 兜底：不装 yaml 时只从环境变量读
|
||
cfg = {"qiniu": {}}
|
||
store = load_from_dist_config(cfg)
|
||
else:
|
||
store = QiniuStore()
|
||
if args.gen_token:
|
||
print(store.gen_uptoken(args.gen_token))
|
||
elif args.list:
|
||
for n in store.list_seeds():
|
||
print(n)
|
||
else:
|
||
print("AK set:", bool(store.ak), "Bucket:", store.bucket,
|
||
"Domain:", store.domain or "(未配置)")
|