758 lines
27 KiB
Python
758 lines
27 KiB
Python
import atexit
|
||
import concurrent.futures
|
||
import fcntl
|
||
import gc
|
||
import glob
|
||
import json
|
||
import os
|
||
import random
|
||
import signal
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import zipfile
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
|
||
import yaml
|
||
from fabric import Connection
|
||
from vmplatform import VMOS, Client, VMDataDisk
|
||
|
||
from schemas.context import ASRContext
|
||
from utils.client_async import ClientAsync
|
||
from utils.evaluator import BaseEvaluator
|
||
from utils.logger import logger
|
||
from utils.service import register_sut
|
||
|
||
IN_TEST = os.getenv("SUBMIT_CONFIG_FILEPATH", None) is None
|
||
UNIT_TEST = os.getenv("UNIT_TEST", 0)
|
||
|
||
DATASET_NUM = os.getenv("DATASET_NUM")
|
||
|
||
# vm榜单参数
|
||
SUT_TYPE = os.getenv("SUT_TYPE", "kubernetes")
|
||
SHARE_SUT = os.getenv("SHARE_SUT", "true") == "true"
|
||
VM_ID = 0
|
||
VM_IP = ""
|
||
do_deploy_chart = True
|
||
VM_CPU = int(os.getenv("VM_CPU", "2"))
|
||
VM_MEM = int(os.getenv("VM_MEM", "4096"))
|
||
MODEL_BASEPATH = os.getenv("MODEL_BASEPATH", "/tmp/customer/leaderboard/pc_asr")
|
||
MODEL_MAPPING = json.loads(os.getenv("MODEL_MAPPING", "{}"))
|
||
SSH_KEY_DIR = os.getenv("SSH_KEY_DIR", "/workspace")
|
||
SSH_PUBLIC_KEY_FILE = os.path.join(SSH_KEY_DIR, "ssh-key-ecdsa.pub")
|
||
SSH_KEY_FILE = os.path.join(SSH_KEY_DIR, "ssh-key-ecdsa")
|
||
|
||
CONNECT_KWARGS = {"key_filename": SSH_KEY_FILE}
|
||
|
||
# 共享sut参数
|
||
JOB_ID = os.getenv("JOB_ID")
|
||
dirname = "/tmp/submit_private/sut_share"
|
||
os.makedirs(dirname, exist_ok=True)
|
||
SUT_SHARE_LOCK = os.path.join(dirname, "lock.lock")
|
||
SUT_SHARE_USE_LOCK = os.path.join(dirname, "use.lock")
|
||
SUT_SHARE_STATUS = os.path.join(dirname, "status.json")
|
||
SUT_SHARE_JOB_STATUS = os.path.join(dirname, f"job_status.{JOB_ID}")
|
||
SUT_SHARE_PUBLIC_FAIL = os.path.join(dirname, "one_job_failed")
|
||
fd_lock = open(SUT_SHARE_USE_LOCK, "a")
|
||
|
||
|
||
def clean_vm_atexit():
|
||
global VM_ID, do_deploy_chart
|
||
if not VM_ID:
|
||
return
|
||
if not do_deploy_chart:
|
||
return
|
||
logger.info("删除vm")
|
||
vmclient = Client()
|
||
err_msg = vmclient.delete_vm(VM_ID)
|
||
if err_msg:
|
||
logger.warning(f"删除vm失败: {err_msg}")
|
||
|
||
|
||
def put_file_to_vm(c: Connection, local_path: str, remote_path: str):
|
||
logger.info(f"uploading file {local_path} to {remote_path}")
|
||
result = c.put(local_path, remote_path)
|
||
logger.info("uploaded {0.local} to {0.remote}".format(result))
|
||
|
||
|
||
def deploy_windows_sut():
|
||
global VM_ID
|
||
global VM_IP
|
||
|
||
submit_config_filepath = os.getenv("SUBMIT_CONFIG_FILEPATH", "")
|
||
with open(submit_config_filepath, "r") as fp:
|
||
st_config = yaml.safe_load(fp)
|
||
assert "model" in st_config, "未配置model"
|
||
assert "model_key" in st_config, "未配置model_key"
|
||
assert "config.json" in st_config, "未配置config.json"
|
||
nfs = st_config.get("leaderboard_options", {}).get("nfs", [])
|
||
assert len(nfs) > 0, "未配置nfs"
|
||
assert st_config["model"] in MODEL_MAPPING, "提交模型不在可用模型范围内"
|
||
|
||
model = st_config["model"]
|
||
model_key = st_config["model_key"]
|
||
model_path = ""
|
||
config = st_config["config.json"]
|
||
exist = False
|
||
for nfs_item in nfs:
|
||
if nfs_item["name"] == model_key:
|
||
exist = True
|
||
if nfs_item["source"] == "ceph_customer":
|
||
model_path = os.path.join(
|
||
"/tmp/customer",
|
||
nfs_item["srcRelativePath"],
|
||
)
|
||
else:
|
||
model_path = os.path.join(
|
||
"/tmp/juicefs",
|
||
nfs_item["srcRelativePath"],
|
||
)
|
||
break
|
||
if not exist:
|
||
raise RuntimeError(f"未找到nfs配置项 name={model_key}")
|
||
config_path = os.path.join(tempfile.mkdtemp(), "config.json")
|
||
model_dir = os.path.basename(model_path).split(".")[0]
|
||
config["model_path"] = f"E:\\model\\{model_dir}"
|
||
with open(config_path, "w") as fp:
|
||
json.dump(config, fp, ensure_ascii=False, indent=4)
|
||
|
||
vmclient = Client()
|
||
with open(SSH_PUBLIC_KEY_FILE, "r") as fp:
|
||
sshpublickey = fp.read().rstrip()
|
||
VM_ID = vmclient.create_vm(
|
||
"amd64",
|
||
VMOS.windows10,
|
||
VM_CPU,
|
||
VM_MEM,
|
||
"leaderboard-%s-submit-%s-job-%s"
|
||
% (
|
||
os.getenv("BENCHMARK_NAME"),
|
||
os.getenv("SUBMIT_ID"),
|
||
os.getenv("JOB_ID"),
|
||
),
|
||
sshpublickey,
|
||
datadisks=[
|
||
VMDataDisk(
|
||
size=50,
|
||
disk_type="ssd",
|
||
mount_path="/",
|
||
filesystem="NTFS",
|
||
)
|
||
],
|
||
)
|
||
atexit.register(clean_vm_atexit)
|
||
signal.signal(signal.SIGTERM, lambda signum, _: sys.exit(signum))
|
||
VM_IP = vmclient.wait_until_vm_running(VM_ID)
|
||
logger.info("vm created successfully, vm_ip: %s", VM_IP)
|
||
|
||
def sut_startup():
|
||
with Connection(
|
||
VM_IP,
|
||
"administrator",
|
||
connect_kwargs=CONNECT_KWARGS,
|
||
) as c:
|
||
script_path = "E:\\base\\asr\\faster-whisper\\server"
|
||
script_path = "E:\\install\\asr\\sensevoice\\server"
|
||
bat_filepath = f"{script_path}\\start.bat"
|
||
config_filepath = "E:\\submit\\config.json"
|
||
result = c.run("")
|
||
assert result.ok
|
||
c.run(
|
||
f'cd /d {script_path} & set "EDGE_ML_ENV_HOME=E:\\install" & {bat_filepath} {config_filepath}',
|
||
warn=True,
|
||
)
|
||
|
||
with Connection(
|
||
VM_IP,
|
||
"administrator",
|
||
connect_kwargs=CONNECT_KWARGS,
|
||
) as c:
|
||
model_filepath = os.path.join(MODEL_BASEPATH, MODEL_MAPPING[model])
|
||
filename = os.path.basename(model_filepath)
|
||
put_file_to_vm(c, model_filepath, "/E:/")
|
||
|
||
result = c.run("mkdir E:\\base")
|
||
assert result.ok
|
||
result = c.run("mkdir E:\\model")
|
||
assert result.ok
|
||
result = c.run("mkdir E:\\submit")
|
||
assert result.ok
|
||
|
||
result = c.run(
|
||
f"tar zxvf E:\\{filename} -C E:\\base --strip-components 1"
|
||
)
|
||
assert result.ok
|
||
|
||
result = c.run("E:\\base\\setup-win.bat E:\\install")
|
||
assert result.ok
|
||
|
||
put_file_to_vm(c, config_path, "/E:/submit")
|
||
put_file_to_vm(c, model_path, "/E:/model")
|
||
result = c.run(
|
||
f"tar zxvf E:\\model\\{os.path.basename(model_path)} -C E:\\model"
|
||
)
|
||
assert result.ok
|
||
threading.Thread(target=sut_startup, daemon=True).start()
|
||
time.sleep(60)
|
||
|
||
return f"ws://{VM_IP}:{config['port']}"
|
||
|
||
|
||
def deploy_macos_sut():
|
||
global VM_ID
|
||
global VM_IP
|
||
|
||
submit_config_filepath = os.getenv("SUBMIT_CONFIG_FILEPATH", "")
|
||
with open(submit_config_filepath, "r") as fp:
|
||
st_config = yaml.safe_load(fp)
|
||
assert "model" in st_config, "未配置model"
|
||
assert "model_key" in st_config, "未配置model_key"
|
||
assert "config.json" in st_config, "未配置config.json"
|
||
nfs = st_config.get("leaderboard_options", {}).get("nfs", [])
|
||
assert len(nfs) > 0, "未配置nfs"
|
||
assert st_config["model"] in MODEL_MAPPING, "提交模型不在可用模型范围内"
|
||
|
||
model = st_config["model"]
|
||
model_key = st_config["model_key"]
|
||
model_path = ""
|
||
config = st_config["config.json"]
|
||
exist = False
|
||
for nfs_item in nfs:
|
||
if nfs_item["name"] == model_key:
|
||
exist = True
|
||
if nfs_item["source"] == "ceph_customer":
|
||
model_path = os.path.join(
|
||
"/tmp/customer",
|
||
nfs_item["srcRelativePath"],
|
||
)
|
||
else:
|
||
model_path = os.path.join(
|
||
"/tmp/juicefs",
|
||
nfs_item["srcRelativePath"],
|
||
)
|
||
break
|
||
if not exist:
|
||
raise RuntimeError(f"未找到nfs配置项 name={model_key}")
|
||
config_path = os.path.join(tempfile.mkdtemp(), "config.json")
|
||
model_dir = os.path.basename(model_path).split(".")[0]
|
||
|
||
vmclient = Client()
|
||
with open(SSH_PUBLIC_KEY_FILE, "r") as fp:
|
||
sshpublickey = fp.read().rstrip()
|
||
VM_ID = vmclient.create_vm(
|
||
"amd64",
|
||
VMOS.macos12,
|
||
VM_CPU,
|
||
VM_MEM,
|
||
"leaderboard-%s-submit-%s-job-%s"
|
||
% (
|
||
os.getenv("BENCHMARK_NAME"),
|
||
os.getenv("SUBMIT_ID"),
|
||
os.getenv("JOB_ID"),
|
||
),
|
||
sshpublickey,
|
||
datadisks=[
|
||
VMDataDisk(
|
||
size=50,
|
||
disk_type="ssd",
|
||
mount_path="/",
|
||
filesystem="apfs",
|
||
)
|
||
],
|
||
)
|
||
atexit.register(clean_vm_atexit)
|
||
signal.signal(signal.SIGTERM, lambda signum, _: sys.exit(signum))
|
||
VM_IP = vmclient.wait_until_vm_running(VM_ID)
|
||
logger.info("vm created successfully, vm_ip: %s", VM_IP)
|
||
|
||
with Connection(
|
||
VM_IP,
|
||
"admin",
|
||
connect_kwargs=CONNECT_KWARGS,
|
||
) as c:
|
||
result = c.run("ls -d /Volumes/data*")
|
||
assert result.ok
|
||
volume_path = result.stdout.strip()
|
||
|
||
config["model_path"] = f"{volume_path}/model/{model_dir}"
|
||
with open(config_path, "w") as fp:
|
||
json.dump(config, fp, ensure_ascii=False, indent=4)
|
||
|
||
def sut_startup():
|
||
with Connection(
|
||
VM_IP,
|
||
"admin",
|
||
connect_kwargs=CONNECT_KWARGS,
|
||
) as c:
|
||
script_path = f"{volume_path}/install/asr/sensevoice/server"
|
||
startsh = f"{script_path}/start.sh"
|
||
config_filepath = f"{volume_path}/submit/config.json"
|
||
c.run(
|
||
f"cd {script_path} && sh {startsh} {config_filepath}",
|
||
warn=True,
|
||
)
|
||
|
||
with Connection(
|
||
VM_IP,
|
||
"admin",
|
||
connect_kwargs=CONNECT_KWARGS,
|
||
) as c:
|
||
model_filepath = os.path.join(MODEL_BASEPATH, MODEL_MAPPING[model])
|
||
filename = os.path.basename(model_filepath)
|
||
put_file_to_vm(c, model_filepath, f"{volume_path}")
|
||
|
||
result = c.run(f"mkdir {volume_path}/base")
|
||
assert result.ok
|
||
result = c.run(f"mkdir {volume_path}/model")
|
||
assert result.ok
|
||
result = c.run(f"mkdir {volume_path}/submit")
|
||
assert result.ok
|
||
|
||
result = c.run(
|
||
f"tar zxvf {volume_path}/{filename} -C {volume_path}/base --strip-components 1" # noqa: E501
|
||
)
|
||
assert result.ok
|
||
|
||
result = c.run(
|
||
f"sh {volume_path}/base/setup-mac.sh {volume_path}/install x64"
|
||
)
|
||
assert result.ok
|
||
|
||
put_file_to_vm(c, config_path, f"{volume_path}/submit")
|
||
put_file_to_vm(c, model_path, f"{volume_path}/model")
|
||
result = c.run(
|
||
f"tar zxvf {volume_path}/model/{os.path.basename(model_path)} -C {volume_path}/model" # noqa: E501
|
||
)
|
||
assert result.ok
|
||
threading.Thread(target=sut_startup, daemon=True).start()
|
||
time.sleep(60)
|
||
|
||
return f"ws://{VM_IP}:{config['port']}"
|
||
|
||
|
||
def get_sut_url_vm(vm_type: str):
|
||
global VM_ID
|
||
global VM_IP
|
||
global do_deploy_chart
|
||
|
||
do_deploy_chart = True
|
||
# 拉起SUT
|
||
|
||
def check_job_failed():
|
||
while True:
|
||
time.sleep(30)
|
||
if os.path.exists(SUT_SHARE_PUBLIC_FAIL):
|
||
logger.error("there is a job failed in current submit")
|
||
sys.exit(1)
|
||
|
||
sut_url = ""
|
||
threading.Thread(target=check_job_failed, daemon=True).start()
|
||
if SHARE_SUT:
|
||
|
||
time.sleep(10 * random.random())
|
||
try:
|
||
open(SUT_SHARE_LOCK, "x").close()
|
||
except Exception:
|
||
do_deploy_chart = False
|
||
|
||
start_at = time.time()
|
||
|
||
def file_last_updated_at(file: str):
|
||
return os.stat(file).st_mtime if os.path.exists(file) else start_at
|
||
|
||
if not do_deploy_chart:
|
||
with open(SUT_SHARE_JOB_STATUS, "w") as f:
|
||
f.write("waiting")
|
||
while (
|
||
time.time() - file_last_updated_at(SUT_SHARE_STATUS)
|
||
<= 60 * 60 * 24
|
||
):
|
||
logger.info(
|
||
"Waiting sut application to be deployed by another job"
|
||
)
|
||
time.sleep(10 + random.random())
|
||
if os.path.exists(SUT_SHARE_STATUS):
|
||
get_status = False
|
||
for _ in range(10):
|
||
try:
|
||
with open(SUT_SHARE_STATUS, "r") as f:
|
||
status = json.load(f)
|
||
get_status = True
|
||
break
|
||
except Exception:
|
||
time.sleep(1 + random.random())
|
||
continue
|
||
if not get_status:
|
||
raise RuntimeError(
|
||
"Failed to get status of sut application"
|
||
)
|
||
assert (
|
||
status.get("status") != "failed"
|
||
), "Failed to deploy sut application, \
|
||
please check other job logs"
|
||
if status.get("status") == "running":
|
||
VM_ID = status.get("vmid")
|
||
VM_IP = status.get("vmip")
|
||
sut_url = status.get("sut_url")
|
||
with open(SSH_PUBLIC_KEY_FILE, "w") as fp:
|
||
fp.write(status.get("pubkey"))
|
||
with open(SSH_KEY_FILE, "w") as fp:
|
||
fp.write(status.get("prikey"))
|
||
logger.info("Successfully get deployed sut application")
|
||
break
|
||
|
||
if do_deploy_chart:
|
||
try:
|
||
fcntl.flock(fd_lock, fcntl.LOCK_EX)
|
||
with open(SUT_SHARE_JOB_STATUS, "w") as f:
|
||
f.write("waiting")
|
||
pending = True
|
||
|
||
def update_status():
|
||
while pending:
|
||
time.sleep(30)
|
||
if not pending:
|
||
break
|
||
with open(SUT_SHARE_STATUS, "w") as f:
|
||
json.dump({"status": "pending"}, f)
|
||
|
||
threading.Thread(target=update_status, daemon=True).start()
|
||
if vm_type == "windows":
|
||
sut_url = deploy_windows_sut()
|
||
else:
|
||
sut_url = deploy_macos_sut()
|
||
except Exception:
|
||
open(SUT_SHARE_PUBLIC_FAIL, "w").close()
|
||
with open(SUT_SHARE_STATUS, "w") as f:
|
||
json.dump({"status": "failed"}, f)
|
||
raise
|
||
finally:
|
||
pending = False
|
||
with open(SUT_SHARE_STATUS, "w") as f:
|
||
pubkey = ""
|
||
with open(SSH_PUBLIC_KEY_FILE, "r") as fp:
|
||
pubkey = fp.read().rstrip()
|
||
prikey = ""
|
||
with open(SSH_KEY_FILE, "r") as fp:
|
||
prikey = fp.read()
|
||
json.dump(
|
||
{
|
||
"status": "running",
|
||
"vmid": VM_ID,
|
||
"vmip": VM_IP,
|
||
"pubkey": pubkey,
|
||
"sut_url": sut_url,
|
||
"prikey": prikey,
|
||
},
|
||
f,
|
||
)
|
||
else:
|
||
while True:
|
||
time.sleep(5 + random.random())
|
||
try:
|
||
fcntl.flock(fd_lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||
break
|
||
except Exception:
|
||
logger.info("尝试抢占调用sut失败，继续等待 5s ...")
|
||
|
||
with open(SUT_SHARE_JOB_STATUS, "w") as f:
|
||
f.write("running")
|
||
|
||
return sut_url
|
||
|
||
|
||
def get_sut_url():
|
||
if SUT_TYPE in ("windows", "macos"):
|
||
return get_sut_url_vm(SUT_TYPE)
|
||
|
||
submit_config_filepath = os.getenv(
|
||
"SUBMIT_CONFIG_FILEPATH", "./tests/resources/submit_config"
|
||
)
|
||
CPU = os.getenv("SUT_CPU", "2")
|
||
MEMORY = os.getenv("SUT_MEMORY", "4Gi")
|
||
resource_name = os.getenv("BENCHMARK_NAME")
|
||
|
||
# 任务信息
|
||
# 斯拉夫语族：俄语、波兰语
|
||
# 日耳曼语族：英语、德语、荷兰语
|
||
# 拉丁语族（罗曼语族）：西班牙语、葡萄牙语、法国语、意大利语
|
||
# 闪米特语族：阿拉伯语、希伯来语
|
||
|
||
# 提交配置 & 启动被测服务
|
||
if os.getenv("DATASET_FILEPATH", ""):
|
||
with open(submit_config_filepath, "r") as fp:
|
||
st_config = yaml.safe_load(fp)
|
||
if "values" not in st_config:
|
||
st_config["values"] = {}
|
||
st_config["values"]["resources"] = {}
|
||
st_config["values"]["resources"]["limits"] = {}
|
||
st_config["values"]["resources"]["limits"]["cpu"] = CPU
|
||
st_config["values"]["resources"]["limits"]["memory"] = MEMORY
|
||
# st_config["values"]['resources']['limits']['nvidia.com/gpu'] = '1'
|
||
# st_config["values"]['resources']['limits']['nvidia.com/gpumem'] = "1843"
|
||
# st_config["values"]['resources']['limits']['nvidia.com/gpucores'] = "8"
|
||
st_config["values"]["resources"]["requests"] = {}
|
||
st_config["values"]["resources"]["requests"]["cpu"] = CPU
|
||
st_config["values"]["resources"]["requests"]["memory"] = MEMORY
|
||
# st_config["values"]['resources']['requests']['nvidia.com/gpu'] = '1'
|
||
# st_config["values"]['resources']['requests']['nvidia.com/gpumem'] = "1843"
|
||
# st_config["values"]['resources']['requests']['nvidia.com/gpucores'] = "8"
|
||
# st_config['values']['nodeSelector'] = {}
|
||
# st_config["values"]["nodeSelector"][
|
||
# "contest.4pd.io/accelerator"
|
||
# ] = "A10vgpu"
|
||
# st_config['values']['tolerations'] = []
|
||
# toleration_item = {}
|
||
# toleration_item['key'] = 'hosttype'
|
||
# toleration_item['operator'] = 'Equal'
|
||
# toleration_item['value'] = 'vgpu'
|
||
# toleration_item['effect'] = 'NoSchedule'
|
||
# st_config['values']['tolerations'].append(toleration_item)
|
||
if os.getenv("RESOURCE_TYPE", "cpu") == "cpu":
|
||
values = st_config["values"]
|
||
limits = values.get("resources", {}).get("limits", {})
|
||
requests = values.get("resources", {}).get("requests", {})
|
||
if (
|
||
"nvidia.com/gpu" in limits
|
||
or "nvidia.com/gpumem" in limits
|
||
or "nvidia.com/gpucores" in limits
|
||
or "nvidia.com/gpu" in requests
|
||
or "nvidia.com/gpumem" in requests
|
||
or "nvidia.com/gpucores" in requests
|
||
):
|
||
raise Exception("禁止使用GPU!")
|
||
else:
|
||
vgpu_num = int(os.getenv("SUT_VGPU", "3"))
|
||
st_config["values"]["resources"]["limits"]["nvidia.com/gpu"] = (
|
||
str(vgpu_num)
|
||
)
|
||
st_config["values"]["resources"]["limits"][
|
||
"nvidia.com/gpumem"
|
||
] = str(1843 * vgpu_num)
|
||
st_config["values"]["resources"]["limits"][
|
||
"nvidia.com/gpucores"
|
||
] = str(8 * vgpu_num)
|
||
st_config["values"]["resources"]["requests"][
|
||
"nvidia.com/gpu"
|
||
] = str(vgpu_num)
|
||
st_config["values"]["resources"]["requests"][
|
||
"nvidia.com/gpumem"
|
||
] = str(1843 * vgpu_num)
|
||
st_config["values"]["resources"]["requests"][
|
||
"nvidia.com/gpucores"
|
||
] = str(8 * vgpu_num)
|
||
st_config["values"]["nodeSelector"] = {}
|
||
st_config["values"]["nodeSelector"][
|
||
"contest.4pd.io/accelerator"
|
||
] = "A10vgpu"
|
||
st_config["values"]["tolerations"] = []
|
||
toleration_item = {}
|
||
toleration_item["key"] = "hosttype"
|
||
toleration_item["operator"] = "Equal"
|
||
toleration_item["value"] = "vgpu"
|
||
toleration_item["effect"] = "NoSchedule"
|
||
st_config["values"]["tolerations"].append(toleration_item)
|
||
if "docker_images" in st_config:
|
||
sut_url = "ws://172.26.1.75:9827"
|
||
os.environ["test"] = "1"
|
||
elif "docker_image" in st_config:
|
||
sut_url = register_sut(st_config, resource_name)
|
||
elif UNIT_TEST:
|
||
sut_url = "ws://172.27.231.36:80"
|
||
else:
|
||
logger.error("config 配置错误，没有 docker_image")
|
||
os._exit(1)
|
||
return sut_url
|
||
else:
|
||
os.environ["test"] = "1"
|
||
sut_url = "ws://172.27.231.36:80"
|
||
sut_url = "ws://172.26.1.75:9827"
|
||
return sut_url
|
||
|
||
|
||
def load_merge_dataset(dataset_filepath: str) -> dict:
|
||
local_dataset_path = "./dataset"
|
||
os.makedirs(local_dataset_path, exist_ok=True)
|
||
with zipfile.ZipFile(dataset_filepath) as zf:
|
||
zf.extractall(local_dataset_path)
|
||
|
||
config = {}
|
||
sub_datasets = os.listdir(local_dataset_path)
|
||
for sub_dataset in sub_datasets:
|
||
if sub_dataset.startswith("asr."):
|
||
lang = sub_dataset[4:]
|
||
lang_path = os.path.join(local_dataset_path, lang)
|
||
os.makedirs(lang_path, exist_ok=True)
|
||
with zipfile.ZipFile(
|
||
os.path.join(local_dataset_path, sub_dataset)
|
||
) as zf:
|
||
zf.extractall(lang_path)
|
||
lang_config_path = os.path.join(lang_path, "data.yaml")
|
||
with open(lang_config_path, "r") as fp:
|
||
lang_config = yaml.safe_load(fp)
|
||
audio_lengths = {}
|
||
for query_item in lang_config.get("query_data", []):
|
||
audio_path = os.path.join(
|
||
lang_path,
|
||
query_item["file"],
|
||
)
|
||
query_item["file"] = audio_path
|
||
audio_lengths[query_item["file"]] = os.path.getsize(
|
||
audio_path,
|
||
)
|
||
lang_config["query_data"] = sorted(
|
||
lang_config.get("query_data", []),
|
||
key=lambda x: audio_lengths[x["file"]],
|
||
reverse=True,
|
||
)
|
||
|
||
idx = 0
|
||
length = 0.0
|
||
for query_item in lang_config["query_data"]:
|
||
audio_length = audio_lengths[query_item["file"]]
|
||
length += audio_length / 32000
|
||
idx += 1
|
||
# 每个语言限制半个小时长度
|
||
if length >= 30 * 60:
|
||
break
|
||
|
||
lang_config["query_data"] = lang_config["query_data"][:idx]
|
||
config[lang] = lang_config
|
||
|
||
config["query_data"] = []
|
||
for lang, lang_config in config.items():
|
||
if lang == "query_data":
|
||
continue
|
||
for query_item in lang_config["query_data"]:
|
||
config["query_data"].append(
|
||
{
|
||
**query_item,
|
||
"lang": lang,
|
||
}
|
||
)
|
||
random.Random(0).shuffle(config["query_data"])
|
||
|
||
return config
|
||
|
||
|
||
def postprocess_failed():
|
||
open(SUT_SHARE_PUBLIC_FAIL, "w").close()
|
||
|
||
|
||
def main():
|
||
dataset_filepath = os.getenv(
|
||
"DATASET_FILEPATH",
|
||
"/Users/4paradigm/Projects/dataset/asr/de.zip",
|
||
# "./tests/resources/en.zip",
|
||
)
|
||
result_filepath = os.getenv("RESULT_FILEPATH", "./out/result")
|
||
bad_cases_filepath = os.getenv("BAD_CASES_FILEPATH", "./out/badcase")
|
||
detail_cases_filepath = os.getenv(
|
||
"DETAILED_CASES_FILEPATH", "./out/detailcase.jsonl"
|
||
)
|
||
thread_num = int(os.getenv("THREAD_NUM", "1"))
|
||
|
||
# 数据集处理
|
||
config = {}
|
||
if os.getenv("MERGE_DATASET", "1"):
|
||
config = load_merge_dataset(dataset_filepath)
|
||
dataset_query = config["query_data"]
|
||
else:
|
||
local_dataset_path = "./dataset"
|
||
os.makedirs(local_dataset_path, exist_ok=True)
|
||
with zipfile.ZipFile(dataset_filepath) as zf:
|
||
zf.extractall(local_dataset_path)
|
||
config_path = os.path.join(local_dataset_path, "data.yaml")
|
||
with open(config_path, "r") as fp:
|
||
dataset_config = yaml.safe_load(fp)
|
||
# 读取所有的音频，进而获得音频的总长度，最后按照音频长度对 query_data 进行降序排序
|
||
lang = os.getenv("lang")
|
||
if lang is None:
|
||
lang = dataset_config.get("global", {}).get("lang", "en")
|
||
audio_lengths = []
|
||
for query_item in dataset_config.get("query_data", []):
|
||
query_item["lang"] = lang
|
||
audio_path = os.path.join(local_dataset_path, query_item["file"])
|
||
query_item["file"] = audio_path
|
||
audio_lengths.append(os.path.getsize(audio_path) / 1024 / 1024)
|
||
dataset_config["query_data"] = sorted(
|
||
dataset_config.get("query_data", []),
|
||
key=lambda x: audio_lengths[dataset_config["query_data"].index(x)],
|
||
reverse=True,
|
||
)
|
||
# 数据集信息
|
||
# dataset_global_config = dataset_config.get("global", {})
|
||
dataset_query = dataset_config.get("query_data", {})
|
||
config[lang] = dataset_config
|
||
|
||
# sut url
|
||
sut_url = get_sut_url()
|
||
|
||
try:
|
||
# 开始测试
|
||
logger.info("开始执行")
|
||
evaluator = BaseEvaluator()
|
||
future_list = []
|
||
with ThreadPoolExecutor(max_workers=thread_num) as executor:
|
||
for idx, query_item in enumerate(dataset_query):
|
||
context = ASRContext(
|
||
**config[query_item["lang"]].get("global", {}),
|
||
)
|
||
context.lang = query_item["lang"]
|
||
context.file_path = query_item["file"]
|
||
context.append_labels(query_item["voice"])
|
||
future = executor.submit(
|
||
ClientAsync(sut_url, context, idx).action
|
||
)
|
||
future_list.append(future)
|
||
for future in concurrent.futures.as_completed(future_list):
|
||
context = future.result()
|
||
evaluator.evaluate(context)
|
||
detail_case = evaluator.gen_detail_case()
|
||
with open(detail_cases_filepath, "a") as fp:
|
||
fp.write(
|
||
json.dumps(
|
||
detail_case.to_dict(),
|
||
ensure_ascii=False,
|
||
)
|
||
+ "\n",
|
||
)
|
||
del context
|
||
gc.collect()
|
||
|
||
evaluator.post_evaluate()
|
||
output_result = evaluator.gen_result()
|
||
logger.info("执行完成")
|
||
|
||
with open(result_filepath, "w") as fp:
|
||
json.dump(output_result, fp, indent=2, ensure_ascii=False)
|
||
with open(bad_cases_filepath, "w") as fp:
|
||
fp.write("当前榜单不存在 Bad Case\n")
|
||
|
||
if SHARE_SUT:
|
||
with open(SUT_SHARE_JOB_STATUS, "w") as f:
|
||
f.write("success")
|
||
|
||
fcntl.flock(fd_lock, fcntl.LOCK_UN)
|
||
fd_lock.close()
|
||
while SHARE_SUT and do_deploy_chart:
|
||
time.sleep(30)
|
||
success_num = 0
|
||
for job_status_file in glob.glob(dirname + "/job_status.*"):
|
||
with open(job_status_file, "r") as f:
|
||
job_status = f.read()
|
||
success_num += job_status == "success"
|
||
if success_num == int(DATASET_NUM):
|
||
break
|
||
logger.info("Waiting for all jobs to complete")
|
||
except Exception:
|
||
if SHARE_SUT:
|
||
postprocess_failed()
|
||
raise
|
||
sys.exit(0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|