194 lines
6.2 KiB
Python
194 lines
6.2 KiB
Python
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
import time
|
||
from typing import Any
|
||
|
||
from core.config import DEPLOY_REGION, TAT_METRIC_TIMEOUT_SECONDS, TENCENT_SECRET_ID, TENCENT_SECRET_KEY, utc_now
|
||
|
||
|
||
def decode_task_output(output: str) -> str:
|
||
# 没有输出时直接返回空串。
|
||
if not output:
|
||
return ""
|
||
|
||
# 把 TAT 返回的 base64 文本解码成普通字符串。
|
||
return base64.b64decode(output).decode("utf-8", errors="replace")
|
||
|
||
|
||
def build_tat_client() -> tuple[Any | None, str]:
|
||
# TAT 运行需要完整的腾讯云密钥和区域。
|
||
if not TENCENT_SECRET_ID or not TENCENT_SECRET_KEY or not DEPLOY_REGION:
|
||
return None, "missing TAT credentials in .env"
|
||
|
||
try:
|
||
# 运行时才导入 SDK，避免没装依赖时整个服务启动失败。
|
||
from tencentcloud.common import credential
|
||
# 腾讯云异常类型用于重试。
|
||
from tencentcloud.common.profile.client_profile import ClientProfile
|
||
# 指定 TAT API endpoint。
|
||
from tencentcloud.common.profile.http_profile import HttpProfile
|
||
# TAT 客户端。
|
||
from tencentcloud.tat.v20201028 import tat_client
|
||
except Exception as exc: # noqa: BLE001
|
||
# SDK 不存在时直接返回错误。
|
||
return None, f"tencentcloud sdk unavailable: {exc}"
|
||
|
||
# 创建凭据对象。
|
||
cred = credential.Credential(TENCENT_SECRET_ID, TENCENT_SECRET_KEY)
|
||
|
||
# 创建 TAT 客户端。
|
||
client = tat_client.TatClient(
|
||
cred,
|
||
DEPLOY_REGION,
|
||
ClientProfile(httpProfile=HttpProfile(endpoint="tat.tencentcloudapi.com")),
|
||
)
|
||
|
||
# 成功时返回客户端。
|
||
return client, ""
|
||
|
||
|
||
def invoke_tat(action: Any, label: str, retries: int = 3, delay_seconds: int = 2) -> Any:
|
||
# 这里延迟导入异常类型，避免 SDK 缺失时模块导入失败。
|
||
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
|
||
|
||
# 记录最后一次异常。
|
||
last_error: Exception | None = None
|
||
|
||
# 简单重试几次，防止瞬时 API 抖动。
|
||
for attempt in range(1, retries + 1):
|
||
try:
|
||
# 成功时直接返回结果。
|
||
return action()
|
||
except TencentCloudSDKException as exc:
|
||
# 保存最后一次异常。
|
||
last_error = exc
|
||
|
||
# 最后一次重试失败时不再等待。
|
||
if attempt == retries:
|
||
break
|
||
|
||
# 做一个线性退避。
|
||
time.sleep(delay_seconds * attempt)
|
||
|
||
# 有异常时抛出最后一次异常。
|
||
if last_error is not None:
|
||
raise last_error
|
||
|
||
# 理论上不会走到这里，这里只是兜底。
|
||
raise RuntimeError(f"{label} failed without error")
|
||
|
||
|
||
def run_shell_command(
|
||
instance_ids: list[str],
|
||
script: str,
|
||
*,
|
||
command_name: str,
|
||
timeout_seconds: int = TAT_METRIC_TIMEOUT_SECONDS,
|
||
working_directory: str = "/root",
|
||
) -> dict[str, dict[str, Any]]:
|
||
# 先初始化 TAT 客户端。
|
||
client, error = build_tat_client()
|
||
|
||
# 客户端不可用时直接抛错。
|
||
if client is None:
|
||
raise RuntimeError(error)
|
||
|
||
# 这里延迟导入 TAT models。
|
||
from tencentcloud.tat.v20201028 import models as tat_models
|
||
|
||
# 创建命令执行请求。
|
||
req = tat_models.RunCommandRequest()
|
||
|
||
# 组装批量执行命令。
|
||
req.from_json_string(
|
||
json.dumps(
|
||
{
|
||
"CommandName": command_name,
|
||
"CommandType": "SHELL",
|
||
"Content": base64.b64encode(script.encode("utf-8")).decode("ascii"),
|
||
"InstanceIds": instance_ids,
|
||
"Timeout": timeout_seconds,
|
||
"WorkingDirectory": working_directory,
|
||
}
|
||
)
|
||
)
|
||
|
||
# 发起远端命令执行。
|
||
response = json.loads(invoke_tat(lambda: client.RunCommand(req), "RunCommand").to_json_string())
|
||
# 拿到 invocation id，后面轮询执行结果。
|
||
invocation_id = response["InvocationId"]
|
||
# 设置轮询截止时间。
|
||
deadline = time.time() + timeout_seconds + 60
|
||
# 这里保存各实例的最终结果。
|
||
results: dict[str, dict[str, Any]] = {}
|
||
|
||
# 开始轮询任务结果。
|
||
while time.time() < deadline:
|
||
# 轮询间隔不必太短。
|
||
time.sleep(2)
|
||
|
||
# 创建查询结果请求。
|
||
query = tat_models.DescribeInvocationTasksRequest()
|
||
# 按 invocation id 查询。
|
||
query.from_json_string(
|
||
json.dumps(
|
||
{
|
||
"HideOutput": False,
|
||
"Limit": 50,
|
||
"Filters": [{"Name": "invocation-id", "Values": [invocation_id]}],
|
||
}
|
||
)
|
||
)
|
||
|
||
# 拉取当前命令所有实例的执行状态。
|
||
payload = json.loads(
|
||
invoke_tat(lambda: client.DescribeInvocationTasks(query), "DescribeInvocationTasks").to_json_string()
|
||
)
|
||
# 任务列表可能暂时为空。
|
||
tasks = payload.get("InvocationTaskSet") or []
|
||
|
||
# 没任务时继续等待。
|
||
if not tasks:
|
||
continue
|
||
|
||
# 遍历当前拿到的实例结果。
|
||
for task in tasks:
|
||
# 读取实例 ID。
|
||
instance_id = task.get("InstanceId", "")
|
||
# 读取任务状态。
|
||
status = task.get("TaskStatus", "")
|
||
|
||
# 只处理终态。
|
||
if status not in {"SUCCESS", "FAILED", "TIMEOUT", "CANCELLED", "START_FAILED"}:
|
||
continue
|
||
|
||
# 解码远端输出。
|
||
output = decode_task_output(task.get("TaskResult", {}).get("Output", ""))
|
||
|
||
# 统一存成结构化结果。
|
||
results[instance_id] = {
|
||
"status": status,
|
||
"output": output,
|
||
"updatedAt": utc_now(),
|
||
}
|
||
|
||
# 收齐全部实例结果后结束轮询。
|
||
if len(results) >= len(instance_ids):
|
||
break
|
||
|
||
# 为缺失结果的实例补超时占位。
|
||
for instance_id in instance_ids:
|
||
results.setdefault(
|
||
instance_id,
|
||
{
|
||
"status": "TIMEOUT",
|
||
"output": "",
|
||
"updatedAt": utc_now(),
|
||
},
|
||
)
|
||
|
||
# 返回按实例 ID 索引的结果。
|
||
return results
|