import os
import re
import subprocess
import tempfile
import threading
from itertools import chain
from shutil import copyfile
from typing import Optional
import backoff
import yaml
from leaderboard.be import LOG, exc
from leaderboard.be.config import BaseConfig, K8sConfig
from leaderboard.be.typing import CmdBase, CmdFull
from . import HELM_PATTERN
K8S_NAMESPACE = K8sConfig.NAMESPACE
KUBE_CONFIG_PATH = BaseConfig.KUBE_CONFIG_PATH
def sut_name_prune(resource_name: str) -> str:
"""将sut名字改为helm可用的名字"""
resource_name = resource_name.lower()
resource_name = re.sub("_", "-", resource_name)
resource_name = re.sub("[^-a-z0-9]", "", resource_name)
return resource_name
def get_helm_value(chart_name: str, chart_version: Optional[str] = None) -> dict:
"""获取helm value的内容
Args:
chart_name (str): chart名
chart_version (str): chart版本
Returns:
Any: helm value, 获取失败为空
"""
ok, values, msg = Helm.helm_show_value(chart_name, chart_version)
if not ok:
LOG.error(msg)
return {}
return yaml.full_load(values)
def chart2info(chart: str) -> tuple[str, Optional[str], dict]:
"""将chart地址为chart的详细信息
Args:
chart (str): helm chart地址
Raises:
exc.ArgumentError
Returns:
tuple[str, Optional[str], dict]: helm名, helm版本, value信息
"""
if "/" not in chart:
raise exc.ArgumentError("chart地址格式错误")
ok, chart_repo, msg = Helm.helm_add_repo(chart[: chart.rindex("/")])
if not ok:
raise exc.ArgumentError("解析repo失败: %s" % msg)
ok, _, msg = Helm.helm_repo_update()
if not ok:
raise exc.ArgumentError("更新repo失败: %s" % msg)
repo_chart = chart_repo + chart[chart.rindex("/") :]
if ":" in repo_chart:
chart_name, chart_version = repo_chart.rsplit(":", 1)
else:
chart_name, chart_version = repo_chart, None
if not Helm.validate_helm(chart_name, chart_version):
raise exc.ArgumentError("chart地址不存在")
chart_values = get_helm_value(chart_name, chart_version)
return chart_name, chart_version, chart_values
def chart2images(chart_name: str, chart_version: Optional[str], value_path: Optional[str]) -> set[str]:
ok, workloads, msg = Helm.helm_template(chart_name, chart_version, value_path)
if not ok:
LOG.error(msg)
return set()
def pod_spec2images(pod_spec) -> list[str]:
return [container["image"] for container in pod_spec["containers"]]
def pod_template_spec2images(pod_template_spec) -> list[str]:
return pod_spec2images(pod_template_spec["template"]["spec"])
def object2image(obj) -> str:
return obj["image"]
def dagtask2image(dagtask) -> list[str]:
if "inline" in dagtask:
return list(chain(*[template2images(template) for template in dagtask["inline"]]))
return []
def template2images(template) -> list[str]:
images = []
if "container" in template:
images.append(object2image(template["container"]))
if "script" in template:
images.append(object2image(template["script"]))
if "initContainers" in template:
images.extend(object2image(obj) for obj in template["initContainers"])
if "sidecars" in template:
images.extend(object2image(obj) for obj in template["sidecars"])
if "containerSet" in template:
containerSet = template["containerSet"]
if "containers" in containerSet:
images.extend(object2image(obj) for obj in containerSet["containers"])
if "dag" in template:
dag = template["dag"]
if "tasks" in dag:
for task in dag["tasks"]:
images.extend(dagtask2image(task))
if "steps" in template:
for step in template["steps"]:
for workflowstep in step:
images.extend(workflowstep2images(workflowstep))
return images
def workflowstep2images(workflowstep) -> list[str]:
if "inline" in workflowstep:
return list(chain(*[template2images(template) for template in workflowstep["inline"]]))
return []
def workflowspec2images(workflowspec) -> list[str]:
images = []
if "templates" in workflowspec:
for template in workflowspec["templates"]:
images.extend(template2images(template))
if "templateDefaults" in workflowspec:
for template in workflowspec["templates"]:
images.extend(template2images(template))
return images
def workflowstatus2images(workflowstatus) -> list[str]:
images = []
if "storedTemplates" in workflowstatus:
for template in workflowstatus["storedTemplates"]:
images.extend(template2images(template))
if "storedWorkflowTemplateSpec" in workflowstatus:
images.extend(workflowspec2images(workflowstatus["storedWorkflowTemplateSpec"]))
return images
def workflow2images(workflow) -> list[str]:
images = []
if "status" in workflow:
images.extend(workflowstatus2images(workflow["status"]))
if "spec" in workflow:
images.extend(workflowspec2images(workflow["spec"]))
return images
def workflowtemplate2images(workflowtemplate) -> list[str]:
if "spec" in workflowtemplate:
return workflowspec2images(workflowtemplate["spec"])
return []
def cronworkflowspec2images(cronworkflowspec) -> list[str]:
if "workflowSpec" in cronworkflowspec:
return workflowspec2images(cronworkflowspec["workflowSpec"])
return []
def cronworkflow2images(cronworkflow) -> list[str]:
if "spec" in cronworkflow:
return cronworkflowspec2images(cronworkflow["spec"])
return []
images = []
for workload in workloads:
if "kind" not in workload:
continue
kind = workload["kind"]
if kind == "Pod":
images.extend(pod_spec2images(workload["spec"]))
elif kind == "PodTemplate":
images.extend(pod_template_spec2images(workload))
elif kind in (
"ReplicationController",
"ReplicaSet",
"Deployment",
"StatefulSet",
"ControllerRevision",
"DaemonSet",
"Job",
"CronJob",
):
images.extend(pod_template_spec2images(workload["spec"]))
elif kind == "Workflow":
images.extend(workflow2images(workload))
elif kind == "WorkflowTemplate":
images.extend(workflowtemplate2images(workload))
elif kind == "CronWorkflow":
images.extend(cronworkflow2images(workload))
return set(images) - set(["busybox"])
def helm_predicate(*ret) -> bool: # pragma: no cover
SPEC_MSG = "Error: INSTALLATION FAILED: cannot re-use a name that is still in use"
return not ret[0] and not re.match(f".*{SPEC_MSG}.*", ret[1])
class Helm: # pragma: no cover
mutex = threading.Lock()
@staticmethod
@backoff.on_predicate(backoff.constant, predicate=helm_predicate, max_tries=3, interval=1, logger=LOG)
def helm_install(
name: str, chart_name: str, chart_version: Optional[str] = None, value_path: Optional[str] = None
, k8s_namespace: Optional[str] = K8S_NAMESPACE, kube_config_path: Optional[str] = KUBE_CONFIG_PATH
) -> CmdBase:
"""helm install
Args:
name (str): 名称
chart_name (str): chart URL 或 repo/chartname
chart_version (str): chart版本
value_path (str, optional): 指定value.yaml路径. Defaults to None.
k8s_namespace (str, optional): 指定k8s namespace. Defaults to K8S_NAMESPACE.
kube_config_path (str, optional): 指定kube_config_path路径. Defaults to KUBE_CONFIG_PATH.
Returns:
CmdBase: 是否安装成功, 错误信息
"""
with Helm.mutex:
cmd = ["helm", "install", name, chart_name, "--namespace", k8s_namespace, "--kubeconfig", kube_config_path]
if chart_version is not None:
cmd.extend(["--version", chart_version])
if value_path is not None:
cmd.extend(["-f", value_path])
proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
return proc.returncode == 0, proc.stderr.decode()
@staticmethod
def helm_template(
chart_name: str, chart_version: Optional[str] = None, value_path: Optional[str] = None
) -> tuple[bool, list[dict], str]:
"""helm install
Args:
chart_name (str): chart URL 或 repo/chartname
chart_version (str): chart版本
value_path (str, optional): 指定value.yaml路径. Defaults to None.
Returns:
tuple[bool, list[dict], str]: tuple[是否成功, 渲染结果, 错误信息]
"""
with Helm.mutex:
cmd = ["helm", "template", "test", chart_name]
if chart_version is not None:
cmd.extend(["--version", chart_version])
if value_path is not None:
cmd.extend(["-f", value_path])
proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
return proc.returncode == 0, list(yaml.full_load_all(proc.stdout.decode())), proc.stderr.decode()
@staticmethod
@backoff.on_predicate(backoff.constant, predicate=helm_predicate, max_tries=3, interval=1, logger=LOG)
def helm_check_gpu_var(
chart_name: str, chart_version: Optional[str] = None, value_path: Optional[str] = None
) -> CmdBase:
"""helm check gpu varible
Args:
chart_name (str): chart URL 或 repo/chartname
chart_version (str): chart版本
value_path (str, optional): 指定value.yaml路径. Defaults to None.
Returns:
CmdBase: 检查安装成功, 错误信息
"""
with Helm.mutex:
script_dir = os.path.dirname(os.path.abspath(__file__))
if not os.access(script_dir + "/../cmds/checkgpuvarforhelm", os.X_OK):
os.chmod(script_dir + "/../cmds/checkgpuvarforhelm", 0o755)
if not os.access(script_dir + "/../cmds/skopeo", os.X_OK):
os.chmod(script_dir + "/../cmds/skopeo", 0o755)
cmd = [script_dir + "/../cmds/checkgpuvarforhelm", "-helmChartPath", chart_name]
if chart_version is not None:
cmd.extend(["-version", chart_version])
if value_path is not None:
cmd.extend(["-valuesYamlPath", value_path])
else:
cmd.extend(["-valuesYamlPath", chart_name + "/values.yaml"])
proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
return proc.returncode == 0, proc.stderr.decode()
@staticmethod
@backoff.on_predicate(backoff.constant, predicate=helm_predicate, max_tries=3, interval=1, logger=LOG)
def helm_uninstall(name: str) -> CmdBase:
"""helm uninstall
Args:
name (str): 名称
Returns:
CmdBase: 是否卸载成功, 错误信息
"""
cmd = ["helm", "uninstall", name, "--namespace", K8S_NAMESPACE, "--kubeconfig", KUBE_CONFIG_PATH]
proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
return proc.returncode == 0, proc.stderr.decode()
@staticmethod
def list_helms(name_pattern: Optional[str] = HELM_PATTERN) -> list[str]:
"""列出helm charts
Args:
name_pattern (str, optional): 对名称的限制要求. Defaults to None.
Returns:
list[str]: 符合要求的helm的名字
"""
with Helm.mutex:
helms = (
subprocess.run(
f"helm list --max 20000 --namespace {K8S_NAMESPACE} "
+ f"--kubeconfig {KUBE_CONFIG_PATH} | awk '{{print $1}}'",
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
shell=True,
)
.stdout.decode()
.split(os.linesep)[1:-1]
)
if name_pattern is not None:
helms = [helm for helm in helms if re.match(name_pattern, helm)]
return helms
@staticmethod
def helm_add_repo(url: str) -> CmdFull:
"""helm add repo, repo已存在则直接返回
Args:
url (str): helm repo路径
Returns:
CmdFull: 是否获取成功, repo名, 错误信息
"""
with Helm.mutex:
cmd = ["helm repo list | awk '{print $1, $2}'"]
proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
if proc.returncode == 0:
for repo in proc.stdout.decode().splitlines()[1:]:
repo_name, repo_url = repo.split(" ", 1)
if repo_url == repo[1]:
return True, repo_name, ""
name = url.replace("://", "-").replace("/", "-")
proc = subprocess.run(["helm", "repo", "add", name, url], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
return proc.returncode == 0, name, proc.stderr.decode()
@staticmethod
def helm_repo_update() -> CmdFull:
"""helm repo update
Returns:
CmdFull: 是否成功, 正确返回, 错误信息
"""
with Helm.mutex:
proc = subprocess.run(["helm", "repo", "update"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
return proc.returncode == 0, proc.stdout.decode(), proc.stderr.decode()
@staticmethod
def helm_show_value(chart_name: str, chart_version: Optional[str] = None) -> CmdFull:
"""helm show value
Args:
chart_name (str): chart URL 或 repo/chartname
chart_version (str): chart版本
Returns:
CmdFull: 是否获取成功, 正确返回, 错误信息
"""
with Helm.mutex:
cmd = ["helm", "show", "values", chart_name]
if chart_version is not None:
cmd.extend(["--version", chart_version])
proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
return proc.returncode == 0, proc.stdout.decode(), proc.stderr.decode()
@staticmethod
def helm_pull(destination: str, chart_name: str, chart_version: Optional[str] = None) -> bool:
"""helm pull
Args:
destination (str): helm包存放位置
chart_name (str): chart URL 或 repo/chartname
chart_version (Optional[str], optional): chart版本. Defaults to None.
Returns:
bool: 是否成功
"""
Helm.helm_repo_update()
with Helm.mutex:
with tempfile.TemporaryDirectory() as temp_dir:
cmd = ["helm", "pull", chart_name, "--destination", temp_dir]
if chart_version is not None:
cmd.extend(["--version", chart_version])
proc = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
if len(os.listdir(temp_dir)) != 1:
return False
copyfile(os.path.join(temp_dir, os.listdir(temp_dir)[0]), destination)
return proc.returncode == 0
@staticmethod
def validate_helm(chart_name: str, chart_version: Optional[str] = None) -> bool:
"""检查helm是否存在
Args:
chart_name (str): chart URL 或 repo/chartname
chart_version (Optional[str], optional): chart版本. Defaults to None.
Returns:
bool: helm是否存在
"""
with Helm.mutex:
cmd = ["helm", "show", "chart", chart_name]
if chart_version is not None:
cmd.extend(["--version", chart_version])
proc = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
return proc.returncode == 0
__all__ = [
"K8S_NAMESPACE",
"KUBE_CONFIG_PATH",
"sut_name_prune",
"get_helm_value",
"chart2info",
"chart2images",
"helm_predicate",
"Helm",
]