Skip to content

Commit

Permalink
feat: 安装 Agent 将上报的 CPU 架构信息更新到 CMDB (closed TencentBlueKing#1708)
Browse files Browse the repository at this point in the history
# Reviewed, transaction id: 3346
  • Loading branch information
Huayeaaa committed Mar 5, 2024
1 parent 3879585 commit af052ad
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 3 deletions.
15 changes: 13 additions & 2 deletions apps/backend/components/collections/agent_new/install.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@
from apps.prometheus.helper import SetupObserve
from apps.utils import concurrent, sync
from apps.utils.exc import ExceptionHandler
from common.api import JobApi
from common.api import CCApi, JobApi
from common.log import logger
from pipeline.core.flow import Service, StaticIntervalGenerator

Expand Down Expand Up @@ -823,13 +823,23 @@ def _schedule(self, data, parent_data, callback_data=None):
cpu_arch__host_id_map = defaultdict(list)
os_version__host_id_map = defaultdict(list)
host_id__agent_id_map: Dict[int, str] = {}
update_list = []
for result in results:
# 对于未完成的实例,记录下来到下一次schedule中继续检查
if not result["is_finished"]:
left_scheduling_sub_inst_ids.append(result["sub_inst_id"])
# 按 CPU 架构对主机进行分组
bk_host_id = common_data.sub_inst_id__host_id_map.get(result["sub_inst_id"])
cpu_arch__host_id_map[result["cpu_arch"]].append(bk_host_id)
cpu_arch = result["cpu_arch"]
cpu_arch__host_id_map[cpu_arch].append(bk_host_id)
update_params: Dict[str, Any] = {
"bk_host_id": bk_host_id,
"properties": {
"bk_cpu_architecture": constants.CMDB_CPU_ARCH_MAP.get(cpu_arch, "x86"),
"bk_os_bit": constants.OsBitType.cpu_type__bit_map().get(cpu_arch),
},
}
update_list.append(update_params)
# 记录不为空的 agent_id 和 bk_host_id 的对应关系
agent_id: str = result.get("agent_id") or ""
agent_id = agent_id.split(":")[-1].strip()
Expand Down Expand Up @@ -866,6 +876,7 @@ def _schedule(self, data, parent_data, callback_data=None):
data.outputs.scheduling_sub_inst_ids = left_scheduling_sub_inst_ids
if not left_scheduling_sub_inst_ids:
self.finish_schedule()
CCApi.batch_update_host({"update": update_list})
return True

polling_time = data.get_one_of_outputs("polling_time")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ class InstallBaseTestCase(utils.AgentServiceBaseTestCase):
NODE_TYPE = constants.NodeType.AGENT
DOWNLOAD_PATH = "/tmp/data/bkee/public/bknodeman/download"
JOB_API_MOCK_PATH = "apps.backend.components.collections.agent_new.install.JobApi"
CMDB_API_MOCK_PATH = "apps.backend.components.collections.agent_new.install.CCApi"
EXECUTE_CMD_MOCK_PATH = "apps.backend.components.collections.agent_new.install.execute_cmd"
PUT_FILE_MOCK_PATH = "apps.backend.components.collections.agent_new.install.put_file"
CUSTOM_DATAIPC_DIR = "/var/run/gse_test"
Expand All @@ -70,6 +71,12 @@ def init_mock_clients(self):
return_type=mock_data_utils.MockReturnType.RETURN_VALUE.value, return_obj={"job_instance_id": 1}
),
)
self.cmdb_mock_client = api_mkd.cmdb.utils.CCApiMockClient(
batch_update_host=mock_data_utils.MockReturn(
return_type=mock_data_utils.MockReturnType.RETURN_VALUE.value, return_obj={"message": "success"}
),
)


def init_redis_data(self):
# 初始化redis数据,用于schedule时读取解析
Expand Down Expand Up @@ -147,6 +154,7 @@ def update_common_inputs(self):

def start_patch(self):
mock.patch(self.JOB_API_MOCK_PATH, self.job_mock_client).start()
mock.patch(self.CMDB_API_MOCK_PATH, self.cmdb_mock_client).start()
mock.patch(target=self.EXECUTE_CMD_MOCK_PATH, return_value="").start()
mock.patch(target=self.PUT_FILE_MOCK_PATH, return_value="").start()
base.get_asyncssh_connect_mock_patch().start()
Expand All @@ -158,7 +166,8 @@ def start_patch(self):
with open(
os.path.join(self.DOWNLOAD_PATH, constants.SetupScriptFileName.SETUP_PAGENT_PY.value), mode="w+"
) as fs:
fs.write("哈哈哈113343ddfd🐒")
fs.write("哈哈哈113343ddfd")
# fs.write("哈哈哈113343ddfd🐒")

def setUp(self) -> None:
self.update_callback_url()
Expand Down Expand Up @@ -1031,3 +1040,14 @@ def _test_shell_solution(self, validate_encrypted_password: bool):
run_cmd,
],
)


class ReportCpuArchTestCase(LinuxInstallTestCase):
def tearDown(self) -> None:
mock_call_obj = self.cmdb_mock_client.batch_update_host.call_args
if mock_call_obj:
call_args = mock_call_obj[0][0]
self.assertEqual(call_args["update"][0]["bk_host_id"], self.obj_factory.bk_host_ids[0])
self.assertEqual(call_args["update"][0]["properties"]["bk_cpu_architecture"], "arm")
self.assertEqual(call_args["update"][0]["properties"]["bk_os_bit"], "arm-64bit")
super().tearDown()
27 changes: 27 additions & 0 deletions apps/node_man/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -503,6 +503,11 @@ def _get_member__alias_map(cls) -> Dict[Enum, str]:
}
CMDB_CPU_MAP = {"x86": CpuType.x86, "arm": CpuType.aarch64}

CMDB_CPU_ARCH = ("x86", "arm")
CPU_ARCH_CHOICES = tuple_choices(CMDB_CPU_ARCH)
CpuArchType = choices_to_namedtuple(CPU_ARCH_CHOICES)
CMDB_CPU_ARCH_MAP = {"x86": CpuArchType.x86, "x86_64": CpuArchType.x86, "aarch64": CpuArchType.arm}

PACKAGE_PATH_RE = re.compile(
f"(?P<is_external>external_)?plugins_(?P<os>({'|'.join(map(str, PLUGIN_OS_TUPLE))}))"
f"_(?P<cpu_arch>({'|'.join(map(str, CPU_TUPLE))})?$)"
Expand Down Expand Up @@ -1137,3 +1142,25 @@ class CommonExecutionSolutionStepType(EnhanceEnum):
@classmethod
def _get_member__alias_map(cls) -> Dict[Enum, str]:
return {cls.DEPENDENCIES: _("依赖文件"), cls.COMMANDS: _("命令")}


class OsBitType(EnhanceEnum):
BIT32 = "32-bit"
BIT64 = "64-bit"
ARM64 = "arm-64bit"

@classmethod
def _get_member__alias_map(cls) -> Dict[Enum, str]:
return {
cls.BIT32: _("操作系统位数:32-bit"),
cls.BIT64: _("操作系统位数:64-bit"),
cls.ARM64: _("操作系统位数:arm-64bit")
}

@classmethod
def cpu_type__bit_map(cls):
return {
CpuType.x86: cls.BIT32.value,
CpuType.x86_64: cls.BIT64.value,
CpuType.aarch64: cls.ARM64.value
}

0 comments on commit af052ad

Please sign in to comment.