query_dna_jobs.py 5.35 KB
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""查询阿里云 DNA 作业状态。

用法:
    # 查询指定 JobId
    python scripts/aliyun_dna/query_dna_jobs.py --job-ids job_id_1,job_id_2

    # 查看最近入库状态(从 state-file 读取 JobId)
    python scripts/aliyun_dna/query_dna_jobs.py --state-file upload_dna_state.json

    # 只显示失败的作业
    python scripts/aliyun_dna/query_dna_jobs.py --state-file upload_dna_state.json --filter failed
"""

import argparse
import json
import logging
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

from dotenv import load_dotenv
load_dotenv(Path(__file__).resolve().parent.parent.parent / ".env")

from alibabacloud_ice20201109.client import Client as ICEClient
from alibabacloud_tea_openapi.models import Config as OpenApiConfig
from alibabacloud_ice20201109 import models as ice_models

logger = logging.getLogger(__name__)

STATUS_EMOJI = {
    "Success": "✅",
    "Fail": "❌",
    "Queuing": "⏳",
    "Analysing": "🔄",
}


def _get_ice_client() -> ICEClient:
    access_key_id = os.environ["ALIYUN_ICE_ACCESS_KEY_ID"]
    access_key_secret = os.environ["ALIYUN_ICE_ACCESS_KEY_SECRET"]
    region = os.environ.get("ALIYUN_ICE_REGION", "cn-hangzhou")

    config = OpenApiConfig(
        access_key_id=access_key_id,
        access_key_secret=access_key_secret,
        endpoint=f"ice.{region}.aliyuncs.com",
        region_id=region,
    )
    return ICEClient(config)


def load_jobs_from_state(state_file: str) -> dict[str, str]:
    """从 state-file 提取 job_id 映射 {audio_path: job_id}。"""
    if not Path(state_file).exists():
        logger.error("状态文件不存在: %s", state_file)
        return {}
    with open(state_file, encoding="utf-8") as f:
        state = json.load(f)
    return {k: v for k, v in state.items() if k.endswith("::job_id")}


def query_jobs(ice_client: ICEClient, job_ids: list[str]) -> list[dict]:
    """批量查询 DNA 作业状态。"""
    request = ice_models.QueryDNAJobListRequest(job_ids=",".join(job_ids))
    response = ice_client.query_dnajob_list(request)

    results = []
    for job in response.body.job_list or []:
        results.append({
            "job_id": job.id,
            "status": job.status,
            "primary_key": job.primary_key,
            "creation_time": job.creation_time,
            "finish_time": job.finish_time,
            "message": job.message or "",
            "dna_result_url": job.dnaresult or "",
            "config": job.config or "",
        })
    return results


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
    )

    parser = argparse.ArgumentParser(description="查询阿里云 DNA 作业状态")
    parser.add_argument("--job-ids", help="JobId 列表,逗号分隔")
    parser.add_argument("--state-file", help="从上传状态文件中提取 JobId")
    parser.add_argument("--filter", choices=["all", "success", "failed", "pending"],
                        default="all", help="过滤作业状态")
    args = parser.parse_args()

    # 验证配置
    required_env = ["ALIYUN_ICE_ACCESS_KEY_ID", "ALIYUN_ICE_ACCESS_KEY_SECRET", "ALIYUN_DNA_DB_ID"]
    for key in required_env:
        if not os.environ.get(key):
            logger.error("缺少环境变量: %s", key)
            sys.exit(1)

    # 收集 JobId
    job_ids = []
    if args.job_ids:
        job_ids = [j.strip() for j in args.job_ids.split(",") if j.strip()]
    elif args.state_file:
        state_jobs = load_jobs_from_state(args.state_file)
        job_ids = list(state_jobs.values())
        logger.info("从状态文件提取到 %d 个 JobId", len(job_ids))
    else:
        logger.error("请指定 --job-ids 或 --state-file")
        sys.exit(1)

    if not job_ids:
        logger.info("没有找到 JobId")
        return

    # 分批查询(每次最多 10 个)
    ice_client = _get_ice_client()
    all_results = []

    for batch_start in range(0, len(job_ids), 10):
        batch = job_ids[batch_start:batch_start + 10]
        results = query_jobs(ice_client, batch)
        all_results.extend(results)

    # 过滤
    if args.filter == "success":
        all_results = [r for r in all_results if r["status"] == "Success"]
    elif args.filter == "failed":
        all_results = [r for r in all_results if r["status"] == "Fail"]
    elif args.filter == "pending":
        all_results = [r for r in all_results if r["status"] in ("Queuing", "Analysing")]

    # 输出
    status_count = {}
    for r in all_results:
        s = r["status"]
        status_count[s] = status_count.get(s, 0) + 1

    print(f"\n作业状态汇总 (共 {len(all_results)} 个):")
    for status, count in sorted(status_count.items()):
        emoji = STATUS_EMOJI.get(status, "❓")
        print(f"  {emoji} {status}: {count}")

    print(f"\n{'JobId':<40} {'状态':<12} {'PrimaryKey':<20} {'创建时间':<22} {'错误信息'}")
    print("-" * 140)
    for r in sorted(all_results, key=lambda x: x.get("creation_time", "")):
        emoji = STATUS_EMOJI.get(r["status"], "❓")
        print(f"{emoji} {r['job_id']:<38} {r['status']:<12} {r['primary_key']:<20} "
              f"{r['creation_time'] or 'N/A':<22} {r['message']}")

    # 汇总输出
    print(f"\n总计: {len(all_results)} 个作业")


if __name__ == "__main__":
    main()