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()
|