440 lines
13 KiB
Python
440 lines
13 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import sys
|
|||
|
|
import time
|
|||
|
|
import argparse
|
|||
|
|
import threading
|
|||
|
|
import subprocess
|
|||
|
|
import shutil
|
|||
|
|
import logging
|
|||
|
|
from datetime import datetime
|
|||
|
|
import queue
|
|||
|
|
import re
|
|||
|
|
|
|||
|
|
from prometheus_client import start_http_server, Counter, Gauge, Summary, Histogram
|
|||
|
|
from flask import Flask, request, jsonify
|
|||
|
|
|
|||
|
|
# —— 常量 & 路径 —— #
|
|||
|
|
BASE = os.getcwd()
|
|||
|
|
INPUT_ROOT = os.path.join(BASE, "input")
|
|||
|
|
OUTPUT_ROOT = os.path.join(BASE, "output")
|
|||
|
|
EMPTY_DIR = os.path.join(BASE, "empty")
|
|||
|
|
LOG_DIR = os.path.join(BASE, "logs")
|
|||
|
|
|
|||
|
|
COS_BUCKET = "mb_raw_rosbag_decode_dirs"
|
|||
|
|
DOCKER_IMAGE = (
|
|||
|
|
"artifact.swfcn.i.mercedes-benz.com/swfcn_docker/perception-3d/mmtbag_decoder:v6.6"
|
|||
|
|
)
|
|||
|
|
DOCKER_CMD_TEMPLATE = [
|
|||
|
|
"docker",
|
|||
|
|
"run",
|
|||
|
|
"--rm",
|
|||
|
|
"-v",
|
|||
|
|
"{in_dir}:/input",
|
|||
|
|
"-v",
|
|||
|
|
"{out_dir}:/output",
|
|||
|
|
DOCKER_IMAGE,
|
|||
|
|
"bash",
|
|||
|
|
"-c",
|
|||
|
|
"source /opt/ros/noetic/setup.bash && "
|
|||
|
|
"/opt/perception-3d/scripts/tools/"
|
|||
|
|
"mmt_bag_decoder_scripts/decoded-bag.sh /input /output 3 1",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
BATCH_SIZE = 50
|
|||
|
|
MAX_LOCAL = 100
|
|||
|
|
MAX_RETRIES = 3
|
|||
|
|
RETRY_DELAY_S = 2
|
|||
|
|
METRICS_PORT = 8000
|
|||
|
|
|
|||
|
|
SENTINEL = (None, None)
|
|||
|
|
|
|||
|
|
# —— 日志配置 —— #
|
|||
|
|
os.makedirs(LOG_DIR, exist_ok=True)
|
|||
|
|
logger = logging.getLogger("pipeline")
|
|||
|
|
logger.setLevel(logging.INFO)
|
|||
|
|
h_info = logging.FileHandler(os.path.join(LOG_DIR, "pipeline.log"), encoding="utf-8")
|
|||
|
|
h_err = logging.FileHandler(os.path.join(LOG_DIR, "error_tasks.log"), encoding="utf-8")
|
|||
|
|
fmt = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
|
|||
|
|
h_info.setFormatter(fmt)
|
|||
|
|
h_err.setFormatter(fmt)
|
|||
|
|
h_err.setLevel(logging.ERROR)
|
|||
|
|
logger.addHandler(h_info)
|
|||
|
|
logger.addHandler(h_err)
|
|||
|
|
|
|||
|
|
# —— Prometheus 指标 —— #
|
|||
|
|
DL_TOTAL = Counter("pipeline_download_total", "下载尝试总数")
|
|||
|
|
DL_FAIL = Counter("pipeline_download_failures", "下载失败总数")
|
|||
|
|
DL_RETRY = Counter("pipeline_download_retries", "下载重试总数")
|
|||
|
|
PR_TOTAL = Counter("pipeline_process_total", "处理尝试总数")
|
|||
|
|
PR_FAIL = Counter("pipeline_process_failures", "处理失败总数")
|
|||
|
|
PR_RETRY = Counter("pipeline_process_retries", "处理重试总数")
|
|||
|
|
UP_TOTAL = Counter("pipeline_upload_total", "上传尝试总数")
|
|||
|
|
UP_FAIL = Counter("pipeline_upload_failures", "上传失败总数")
|
|||
|
|
UP_RETRY = Counter("pipeline_upload_retries", "上传重试总数")
|
|||
|
|
|
|||
|
|
DL_DUR = Summary("pipeline_download_duration_seconds", "单批下载耗时秒")
|
|||
|
|
PR_DUR = Summary("pipeline_process_duration_seconds", "单批处理耗时秒")
|
|||
|
|
UP_DUR = Summary("pipeline_upload_duration_seconds", "单批上传耗时秒")
|
|||
|
|
|
|||
|
|
BATCH_SIZE_HIST = Histogram(
|
|||
|
|
"pipeline_batch_size",
|
|||
|
|
"单批任务中文件数量分布",
|
|||
|
|
buckets=[1, 10, 20, 50, 100, 200, 500],
|
|||
|
|
)
|
|||
|
|
FILE_DL_DUR = Histogram("pipeline_file_download_duration_seconds", "单文件下载耗时分布")
|
|||
|
|
BATCH_OUT_FILES = Gauge("pipeline_batch_output_file_count", "单批处理后输出文件数")
|
|||
|
|
|
|||
|
|
Q_BATCH = Gauge("pipeline_queue_batches", "待下载批次数")
|
|||
|
|
Q_PROC = Gauge("pipeline_queue_processing", "待处理批次数")
|
|||
|
|
Q_UP = Gauge("pipeline_queue_uploading", "待上传批次数")
|
|||
|
|
LOCAL_FILES = Gauge("pipeline_local_file_count", "本地 input 文件总数")
|
|||
|
|
|
|||
|
|
# —— 全局队列 & inflight 计数 —— #
|
|||
|
|
batch_q = queue.Queue()
|
|||
|
|
proc_q = queue.Queue()
|
|||
|
|
up_q = queue.Queue()
|
|||
|
|
|
|||
|
|
inflight = 0
|
|||
|
|
inflight_lock = threading.Lock()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# —— 辅助函数 —— #
|
|||
|
|
def count_local_files():
|
|||
|
|
return sum(len(files) for _, _, files in os.walk(INPUT_ROOT))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run(cmd, timeout=None):
|
|||
|
|
logger.info("CMD: %s", " ".join(cmd))
|
|||
|
|
try:
|
|||
|
|
p = subprocess.Popen(
|
|||
|
|
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
|
|||
|
|
)
|
|||
|
|
start = time.time()
|
|||
|
|
out_lines = []
|
|||
|
|
timed_out = False
|
|||
|
|
for line in p.stdout:
|
|||
|
|
out_lines.append(line)
|
|||
|
|
if timeout and (time.time() - start) > timeout:
|
|||
|
|
p.kill()
|
|||
|
|
timed_out = True
|
|||
|
|
break
|
|||
|
|
code = p.wait()
|
|||
|
|
return code, timed_out, "".join(out_lines)
|
|||
|
|
except Exception:
|
|||
|
|
logger.exception("CMD 执行异常")
|
|||
|
|
return -1, False, ""
|
|||
|
|
|
|||
|
|
|
|||
|
|
def with_retry(tag, fn, *args):
|
|||
|
|
for i in range(1, MAX_RETRIES + 1):
|
|||
|
|
code, timed_out, _ = fn(*args)
|
|||
|
|
if code == 0:
|
|||
|
|
return True
|
|||
|
|
if timed_out:
|
|||
|
|
logger.error("%s 阶段超时，不再重试", tag)
|
|||
|
|
break
|
|||
|
|
# 计数重试
|
|||
|
|
if tag.startswith("DL["):
|
|||
|
|
DL_RETRY.inc()
|
|||
|
|
if tag.startswith("PR["):
|
|||
|
|
PR_RETRY.inc()
|
|||
|
|
if tag.startswith("UP["):
|
|||
|
|
UP_RETRY.inc()
|
|||
|
|
logger.warning("%s 重试 %d/%d", tag, i, MAX_RETRIES)
|
|||
|
|
time.sleep(RETRY_DELAY_S)
|
|||
|
|
logger.error("%s 最终失败", tag)
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
|
|||
|
|
# —— 下载 —— #
|
|||
|
|
@DL_DUR.time()
|
|||
|
|
def do_download(batch_id, paths, batch_timeout):
|
|||
|
|
if batch_id is None:
|
|||
|
|
proc_q.put(SENTINEL)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
DL_TOTAL.inc()
|
|||
|
|
start = time.time()
|
|||
|
|
in_dir = os.path.join(INPUT_ROOT, batch_id)
|
|||
|
|
os.makedirs(in_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 限制本地文件数
|
|||
|
|
while count_local_files() >= MAX_LOCAL:
|
|||
|
|
logger.warning("本地文件过多，暂停下载5分钟")
|
|||
|
|
time.sleep(300)
|
|||
|
|
|
|||
|
|
for p in paths:
|
|||
|
|
if time.time() - start > batch_timeout:
|
|||
|
|
logger.error("DL[%s] 下载阶段超时，跳过剩余", batch_id)
|
|||
|
|
DL_FAIL.inc()
|
|||
|
|
break
|
|||
|
|
dst = os.path.join(in_dir, os.path.basename(p))
|
|||
|
|
f_start = time.time()
|
|||
|
|
ok = with_retry(
|
|||
|
|
f"DL[{batch_id}]",
|
|||
|
|
lambda s, d: run(
|
|||
|
|
["coscmd", "-s", "download", s, d],
|
|||
|
|
timeout=batch_timeout - (time.time() - start),
|
|||
|
|
),
|
|||
|
|
p,
|
|||
|
|
dst,
|
|||
|
|
)
|
|||
|
|
FILE_DL_DUR.observe(time.time() - f_start)
|
|||
|
|
if not ok:
|
|||
|
|
DL_FAIL.inc()
|
|||
|
|
|
|||
|
|
proc_q.put((batch_id, in_dir))
|
|||
|
|
|
|||
|
|
|
|||
|
|
# —— 处理 —— #
|
|||
|
|
@PR_DUR.time()
|
|||
|
|
def do_process(batch_id, in_dir, batch_timeout):
|
|||
|
|
if batch_id is None:
|
|||
|
|
up_q.put(SENTINEL)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
PR_TOTAL.inc()
|
|||
|
|
start = time.time()
|
|||
|
|
out_dir = os.path.join(OUTPUT_ROOT, batch_id)
|
|||
|
|
os.makedirs(out_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
cmd = [c.format(in_dir=in_dir, out_dir=out_dir) for c in DOCKER_CMD_TEMPLATE]
|
|||
|
|
|
|||
|
|
def run_pr(command):
|
|||
|
|
p = subprocess.Popen(
|
|||
|
|
command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
|
|||
|
|
)
|
|||
|
|
for line in p.stdout:
|
|||
|
|
logger.info("[PR %s] %s", batch_id, line.rstrip())
|
|||
|
|
if time.time() - start > batch_timeout:
|
|||
|
|
p.kill()
|
|||
|
|
return p.wait(), True, ""
|
|||
|
|
return p.wait(), False, ""
|
|||
|
|
|
|||
|
|
ok = with_retry(f"PR[{batch_id}]", run_pr, cmd)
|
|||
|
|
if not ok:
|
|||
|
|
PR_FAIL.inc()
|
|||
|
|
|
|||
|
|
# 统计输出文件
|
|||
|
|
files = []
|
|||
|
|
for r, _, fs in os.walk(out_dir):
|
|||
|
|
for fn in fs:
|
|||
|
|
files.append(os.path.relpath(os.path.join(r, fn), out_dir))
|
|||
|
|
BATCH_OUT_FILES.set(len(files))
|
|||
|
|
|
|||
|
|
shutil.rmtree(in_dir, ignore_errors=True)
|
|||
|
|
up_q.put((batch_id, out_dir))
|
|||
|
|
|
|||
|
|
|
|||
|
|
# —— 上传 —— #
|
|||
|
|
@UP_DUR.time()
|
|||
|
|
def do_upload(batch_id, out_dir, batch_timeout):
|
|||
|
|
global inflight
|
|||
|
|
if batch_id is None:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
UP_TOTAL.inc()
|
|||
|
|
ok = with_retry(
|
|||
|
|
f"UP[{batch_id}]",
|
|||
|
|
lambda d: run(
|
|||
|
|
["coscmd", "-s", "upload", "-r", d, COS_BUCKET], timeout=batch_timeout
|
|||
|
|
),
|
|||
|
|
out_dir,
|
|||
|
|
)
|
|||
|
|
if not ok:
|
|||
|
|
UP_FAIL.inc()
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 删除目录结构
|
|||
|
|
for cmd in [
|
|||
|
|
["sudo", "rsync", "-av", "--delete", f"{EMPTY_DIR}/", f"{out_dir}/"],
|
|||
|
|
["sudo", "rm", "-rf", out_dir],
|
|||
|
|
]:
|
|||
|
|
run(cmd, timeout=60)
|
|||
|
|
|
|||
|
|
logger.info("UP[%s] 完成", batch_id)
|
|||
|
|
finally:
|
|||
|
|
# 无论成功失败，任务算完成，inflight-1
|
|||
|
|
with inflight_lock:
|
|||
|
|
inflight -= 1
|
|||
|
|
|
|||
|
|
|
|||
|
|
# —— Worker 模板 —— #
|
|||
|
|
def worker(q, fn, timeout):
|
|||
|
|
while True:
|
|||
|
|
bid, data = q.get()
|
|||
|
|
fn(bid, data, timeout)
|
|||
|
|
q.task_done()
|
|||
|
|
if bid is None:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
|
|||
|
|
# —— Service HTTP —— #
|
|||
|
|
app = Flask(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@app.route("/ready", methods=["GET"])
|
|||
|
|
def api_ready():
|
|||
|
|
with inflight_lock:
|
|||
|
|
busy = inflight > 0
|
|||
|
|
return jsonify(ready=not busy)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@app.route("/notify", methods=["POST"])
|
|||
|
|
def api_notify():
|
|||
|
|
global inflight
|
|||
|
|
data = request.get_json(force=True)
|
|||
|
|
if not isinstance(data, list):
|
|||
|
|
return jsonify(error="Expect JSON list"), 400
|
|||
|
|
|
|||
|
|
# 兼容 bag-checker，只发 name 时补前缀
|
|||
|
|
paths = []
|
|||
|
|
for item in data:
|
|||
|
|
if not isinstance(item, str):
|
|||
|
|
continue
|
|||
|
|
if item.startswith("mb_cuct_data_collection/"):
|
|||
|
|
paths.append(item)
|
|||
|
|
else:
|
|||
|
|
paths.append("mb_cuct_data_collection/" + item)
|
|||
|
|
|
|||
|
|
TIME_RE = re.compile(r"_(\d{8})-(\d{6})_") # 匹配 20230803-160828
|
|||
|
|
|
|||
|
|
def extract_ts(p: str) -> datetime:
|
|||
|
|
m = TIME_RE.search(os.path.basename(p))
|
|||
|
|
if not m:
|
|||
|
|
return datetime.min # 无法解析的放最后
|
|||
|
|
date_part, time_part = m.groups()
|
|||
|
|
ts_str = f"{date_part}{time_part}"
|
|||
|
|
return datetime.strptime(ts_str, "%Y%m%d%H%M%S")
|
|||
|
|
|
|||
|
|
paths.sort(key=extract_ts, reverse=True)
|
|||
|
|
|
|||
|
|
with inflight_lock:
|
|||
|
|
inflight += 1
|
|||
|
|
|
|||
|
|
for idx in range(0, len(paths), BATCH_SIZE):
|
|||
|
|
blk = paths[idx : idx + BATCH_SIZE]
|
|||
|
|
BATCH_SIZE_HIST.observe(len(blk))
|
|||
|
|
bid = datetime.utcnow().strftime("%Y%m%dT%H%M%SZ") + f"_{idx // BATCH_SIZE + 1}"
|
|||
|
|
batch_q.put((bid, blk))
|
|||
|
|
batch_q.put(SENTINEL)
|
|||
|
|
|
|||
|
|
# batch_q.put((bid, paths))
|
|||
|
|
return jsonify(status="accepted", batch_size=BATCH_SIZE), 202
|
|||
|
|
|
|||
|
|
|
|||
|
|
def start_metric_updater():
|
|||
|
|
def loop():
|
|||
|
|
while True:
|
|||
|
|
Q_BATCH.set(batch_q.qsize())
|
|||
|
|
Q_PROC.set(proc_q.qsize())
|
|||
|
|
Q_UP.set(up_q.qsize())
|
|||
|
|
LOCAL_FILES.set(count_local_files())
|
|||
|
|
time.sleep(1)
|
|||
|
|
|
|||
|
|
t = threading.Thread(target=loop, daemon=True)
|
|||
|
|
t.start()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# —— 两种模式的入口 —— #
|
|||
|
|
def file_mode(args):
|
|||
|
|
# 读 tasks-file，分批入队，放入 sentinel，然后启动处理
|
|||
|
|
lines = [
|
|||
|
|
line.strip() for line in open(args.tasks_file, encoding="utf-8") if line.strip()
|
|||
|
|
]
|
|||
|
|
for idx in range(0, len(lines), args.batch_size):
|
|||
|
|
blk = lines[idx : idx + args.batch_size]
|
|||
|
|
BATCH_SIZE_HIST.observe(len(blk))
|
|||
|
|
bid = (
|
|||
|
|
datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
|
|||
|
|
+ f"_{idx // args.batch_size + 1}"
|
|||
|
|
)
|
|||
|
|
batch_q.put((bid, blk))
|
|||
|
|
batch_q.put(SENTINEL)
|
|||
|
|
|
|||
|
|
start_http_server(METRICS_PORT)
|
|||
|
|
logger.info("Metrics HTTP 启动，端口 %d", METRICS_PORT)
|
|||
|
|
|
|||
|
|
threads = [
|
|||
|
|
threading.Thread(
|
|||
|
|
target=worker, args=(batch_q, do_download, args.batch_timeout)
|
|||
|
|
),
|
|||
|
|
threading.Thread(target=worker, args=(proc_q, do_process, args.batch_timeout)),
|
|||
|
|
threading.Thread(target=worker, args=(up_q, do_upload, args.batch_timeout)),
|
|||
|
|
]
|
|||
|
|
for t in threads:
|
|||
|
|
t.start()
|
|||
|
|
|
|||
|
|
# 更新指标 & 等待完成
|
|||
|
|
while batch_q.unfinished_tasks or proc_q.unfinished_tasks or up_q.unfinished_tasks:
|
|||
|
|
Q_BATCH.set(batch_q.qsize())
|
|||
|
|
Q_PROC.set(proc_q.qsize())
|
|||
|
|
Q_UP.set(up_q.qsize())
|
|||
|
|
LOCAL_FILES.set(count_local_files())
|
|||
|
|
time.sleep(1)
|
|||
|
|
|
|||
|
|
for t in threads:
|
|||
|
|
t.join()
|
|||
|
|
|
|||
|
|
logger.info("文件模式处理完成，退出。")
|
|||
|
|
sys.exit(0)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def service_mode(args):
|
|||
|
|
# 确保目录存在
|
|||
|
|
for d in (INPUT_ROOT, OUTPUT_ROOT, EMPTY_DIR, LOG_DIR):
|
|||
|
|
os.makedirs(d, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 启动 Prometheus 和指标更新
|
|||
|
|
start_http_server(METRICS_PORT)
|
|||
|
|
logger.info("Metrics HTTP 启动，端口 %d", METRICS_PORT)
|
|||
|
|
start_metric_updater()
|
|||
|
|
|
|||
|
|
# 启动后台 worker
|
|||
|
|
threads = [
|
|||
|
|
threading.Thread(
|
|||
|
|
target=worker, args=(batch_q, do_download, args.batch_timeout), daemon=True
|
|||
|
|
),
|
|||
|
|
threading.Thread(
|
|||
|
|
target=worker, args=(proc_q, do_process, args.batch_timeout), daemon=True
|
|||
|
|
),
|
|||
|
|
threading.Thread(
|
|||
|
|
target=worker, args=(up_q, do_upload, args.batch_timeout), daemon=True
|
|||
|
|
),
|
|||
|
|
]
|
|||
|
|
for t in threads:
|
|||
|
|
t.start()
|
|||
|
|
|
|||
|
|
# 启动 Flask
|
|||
|
|
logger.info("Decode Service 启动 HTTP on %s:%d", args.host, args.port)
|
|||
|
|
app.run(host=args.host, port=args.port, threaded=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
p = argparse.ArgumentParser()
|
|||
|
|
sub = p.add_subparsers(dest="mode", required=True)
|
|||
|
|
|
|||
|
|
f = sub.add_parser("file", help="文件模式：--tasks-file")
|
|||
|
|
f.add_argument("--tasks-file", required=True)
|
|||
|
|
f.add_argument("--batch-size", type=int, default=BATCH_SIZE)
|
|||
|
|
f.add_argument("--batch-timeout", type=int, default=3600)
|
|||
|
|
|
|||
|
|
s = sub.add_parser("service", help="服务模式：启动 HTTP ready/notify")
|
|||
|
|
s.add_argument("--batch-timeout", type=int, default=3600)
|
|||
|
|
s.add_argument("--host", default="0.0.0.0")
|
|||
|
|
s.add_argument("--port", type=int, default=5000)
|
|||
|
|
|
|||
|
|
args = p.parse_args()
|
|||
|
|
if args.mode == "file":
|
|||
|
|
file_mode(args)
|
|||
|
|
else:
|
|||
|
|
service_mode(args)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|