diff --git a/.github/CI_DOCKER.md b/.github/CI_DOCKER.md index ad1cdf2..0faa8bc 100644 --- a/.github/CI_DOCKER.md +++ b/.github/CI_DOCKER.md @@ -1,4 +1,4 @@ # Start ci container ``` docker container create --name ci-mtml-py --env MTHREADS_VISIBLE_DEVICES=all --env MTHREADS_DRIVER_CAPABILITIES=all -v $PWD:/app/test -it registry.cn-hangzhou.aliyuncs.com/fishingfly/mthreads-ml-py-ci:latest bash -``` \ No newline at end of file +``` diff --git a/.github/Dockerfile b/.github/Dockerfile index 554b334..8692adc 100644 --- a/.github/Dockerfile +++ b/.github/Dockerfile @@ -45,4 +45,4 @@ RUN chown -R mccxadmin:mccxadmin /app USER mccxadmin -CMD ["/bin/bash"] \ No newline at end of file +CMD ["/bin/bash"] diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..2b796e9 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,48 @@ +# Hooks every change to this repository has to pass. +# +# pre-commit install # run them on every commit +# pre-commit run --all-files # check the whole tree (formatting hooks fix) +# +# The checks mirror what the Makefile used to call directly (flake8 with the +# same limits, isort, black); pre-commit provisions each tool in its own +# environment, so no global install of them is needed. +default_stages: [pre-commit, manual] + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v6.0.0 + hooks: + - id: check-ast + - id: check-yaml + - id: end-of-file-fixer + exclude: ^mtml_2\.2\.0\.h$ + - id: trailing-whitespace + exclude: ^mtml_2\.2\.0\.h$ + - id: check-merge-conflict + - id: detect-private-key + - id: debug-statements + - repo: https://github.com/PyCQA/isort + rev: 7.0.0 + hooks: + - id: isort + - repo: https://github.com/psf/black + rev: 26.1.0 + hooks: + - id: black + - repo: https://github.com/pycqa/flake8 + rev: 7.1.1 + hooks: + - id: flake8 + # The Makefile's limits, extended rather than replaced so flake8's own + # defaults keep working. E203 is what black produces; E266 allows the + # `##` banner comments this file has always used; F403/F405 cover the + # `from pymtml import *` that the examples and tests are meant to use. + args: ["--max-line-length=120", "--extend-ignore=E203,E501,W503,E266,F403,F405"] + - repo: https://github.com/codespell-project/codespell + rev: v2.4.1 + hooks: + - id: codespell + # musa/mtml/pymtml and renderD (from the /dev/dri/renderD128 path) are + # names, not misspellings. The vendored C header keeps its own text. + exclude: ^mtml_2\.2\.0\.h$ + args: ['-L', 'musa,musas,mtml,musa_,pymtml,renderD'] diff --git a/AGENTS.md b/AGENTS.md index 5cb77d0..0ef55b9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -136,4 +136,3 @@ mtmlDeviceFreeMemory(memory) - `nvmlDeviceGetCudaComputeCapability()` returns `(0, 0)` unless torch_musa is available - Use `patch_torch_c_for_musa()` to patch torch._C with functions from torch_musa._MUSAC - PCI busId field may be empty from driver; library auto-fills from sbdf if needed - diff --git a/Makefile b/Makefile index 01a838d..5c64471 100644 --- a/Makefile +++ b/Makefile @@ -3,27 +3,32 @@ # Default Python interpreter PYTHON ?= python3 +# The lint and format tools are pinned in .pre-commit-config.yaml; pre-commit +# provisions each one in its own environment, so no global install is needed. +PRECOMMIT ?= pre-commit + # Package name PACKAGE = pymtml help: @echo "Available targets:" - @echo " format - Format code with isort + black" - @echo " lint - Run linter (flake8)" + @echo " format - Format the tree with the pinned pre-commit hooks" + @echo " lint - Run the read-only checks (flake8, codespell, ast)" @echo " test - Run tests" @echo " build - Build wheel package" @echo " clean - Clean build artifacts" @echo " publish - Upload wheel to PyPI" @echo " all - Format, lint, test, build" -# Format code using isort + black +# Format the tree (isort, black, whitespace and end-of-file hooks) format: - $(PYTHON) -m isort $(PACKAGE).py example.py test_*.py - $(PYTHON) -m black $(PACKAGE).py example.py test_*.py + $(PRECOMMIT) run --all-files -# Lint code using flake8 +# Read-only checks: lint, spelling and syntax, the same hooks the gate runs lint: - $(PYTHON) -m flake8 $(PACKAGE).py example.py --max-line-length=120 --ignore=E501,W503 + $(PRECOMMIT) run flake8 --all-files + $(PRECOMMIT) run codespell --all-files + $(PRECOMMIT) run check-ast --all-files # Run tests test: @@ -45,4 +50,3 @@ publish: build # Run all checks and build all: format lint test build - diff --git a/PUBLISH.md b/PUBLISH.md index 48e6e00..a364bc8 100644 --- a/PUBLISH.md +++ b/PUBLISH.md @@ -5,4 +5,4 @@ python setup.py bdist_wheel # Upload the packages to PyPI python -m twine upload --repository pypi dist/* -``` \ No newline at end of file +``` diff --git a/examples/02_device_info.py b/examples/02_device_info.py index 8adb084..030141b 100644 --- a/examples/02_device_info.py +++ b/examples/02_device_info.py @@ -18,7 +18,9 @@ print(f" 名称: {mtmlDeviceGetName(device)}") print(f" 索引: {mtmlDeviceGetIndex(device)}") print(f" UUID: {mtmlDeviceGetUUID(device)}") - print(f" 品牌: {'MTT' if mtmlDeviceGetBrand(device) == MTML_BRAND_MTT else 'Unknown'}") + print( + f" 品牌: {'MTT' if mtmlDeviceGetBrand(device) == MTML_BRAND_MTT else 'Unknown'}" + ) try: print(f" 序列号: {mtmlDeviceGetSerialNumber(device)}") @@ -33,8 +35,10 @@ print(f" MtBIOS 版本: [不可用: {e}]") print(f" GPU 核心数: {mtmlDeviceCountGpuCores(device)}") - print(f" 功耗: {mtmlDeviceGetPowerUsage(device)} mW " - f"({mtmlDeviceGetPowerUsage(device) / 1000:.1f} W)") + print( + f" 功耗: {mtmlDeviceGetPowerUsage(device)} mW " + f"({mtmlDeviceGetPowerUsage(device) / 1000:.1f} W)" + ) # PCI 信息 pci = mtmlDeviceGetPciInfo(device) @@ -43,7 +47,9 @@ print(f" PCI 设备ID: {pci.pciDeviceId:#010x}") print(f" PCIe 代数: 当前 Gen{pci.pciCurGen} / 最大 Gen{pci.pciMaxGen}") print(f" PCIe 宽度: 当前 x{pci.pciCurWidth} / 最大 x{pci.pciMaxWidth}") - print(f" PCIe 速率: 当前 {pci.pciCurSpeed:.1f} GT/s / 最大 {pci.pciMaxSpeed:.1f} GT/s") + print( + f" PCIe 速率: 当前 {pci.pciCurSpeed:.1f} GT/s / 最大 {pci.pciMaxSpeed:.1f} GT/s" + ) # PCIe 插槽信息 try: @@ -54,7 +60,11 @@ # 设备属性 prop = mtmlDeviceGetProperty(device) - virt_cap = "支持" if prop.virtCapability == MTML_DEVICE_SUPPORT_VIRTUALIZATION else "不支持" + virt_cap = ( + "支持" + if prop.virtCapability == MTML_DEVICE_SUPPORT_VIRTUALIZATION + else "不支持" + ) mpc_cap = "支持" if prop.mpcCapability == MTML_DEVICE_SUPPORT_MPC else "不支持" print(f" 虚拟化: {virt_cap}") print(f" MPC: {mpc_cap}") diff --git a/examples/03_gpu_monitoring.py b/examples/03_gpu_monitoring.py index d94e240..c7538ba 100644 --- a/examples/03_gpu_monitoring.py +++ b/examples/03_gpu_monitoring.py @@ -35,7 +35,7 @@ print(f" GPU 时钟: {clock} / {max_clock} MHz") # 各引擎利用率 - print(f" 引擎利用率:") + print(" 引擎利用率:") for engine_id, engine_name in engine_names.items(): try: engine_util = mtmlGpuGetEngineUtilization(gpu, engine_id) diff --git a/examples/05_vpu_monitoring.py b/examples/05_vpu_monitoring.py index 2bf8924..a68e3de 100644 --- a/examples/05_vpu_monitoring.py +++ b/examples/05_vpu_monitoring.py @@ -36,30 +36,38 @@ # 编码器会话状态 try: enc_states = mtmlVpuGetEncoderSessionStates(vpu, 8) - active_enc = sum(1 for s in enc_states if s.state == MTML_CODEC_SESSION_STATE_ACTIVE) + active_enc = sum( + 1 for s in enc_states if s.state == MTML_CODEC_SESSION_STATE_ACTIVE + ) print(f" 活跃编码会话: {active_enc}") for s in enc_states: if s.state == MTML_CODEC_SESSION_STATE_ACTIVE: metrics = mtmlVpuGetEncoderSessionMetrics(vpu, s.sessionId) - print(f" 编码会话 {s.sessionId}: {metrics.width}x{metrics.height}, " - f"codec={metrics.codecType}, fps={metrics.fps}") + print( + f" 编码会话 {s.sessionId}: {metrics.width}x{metrics.height}, " + f"codec={metrics.codecType}, fps={metrics.fps}" + ) break except MTMLError: - print(f" 活跃编码会话: [不可用]") + print(" 活跃编码会话: [不可用]") # 解码器会话状态 try: dec_states = mtmlVpuGetDecoderSessionStates(vpu, 8) - active_dec = sum(1 for s in dec_states if s.state == MTML_CODEC_SESSION_STATE_ACTIVE) + active_dec = sum( + 1 for s in dec_states if s.state == MTML_CODEC_SESSION_STATE_ACTIVE + ) print(f" 活跃解码会话: {active_dec}") for s in dec_states: if s.state == MTML_CODEC_SESSION_STATE_ACTIVE: metrics = mtmlVpuGetDecoderSessionMetrics(vpu, s.sessionId) - print(f" 解码会话 {s.sessionId}: {metrics.width}x{metrics.height}, " - f"codec={metrics.codecType}, fps={metrics.fps}") + print( + f" 解码会话 {s.sessionId}: {metrics.width}x{metrics.height}, " + f"codec={metrics.codecType}, fps={metrics.fps}" + ) break except MTMLError: - print(f" 活跃解码会话: [不可用]") + print(" 活跃解码会话: [不可用]") except MTMLError as e: print(f" VPU 不可用: {e}") diff --git a/examples/07_ecc_errors.py b/examples/07_ecc_errors.py index 9ddb629..859d021 100644 --- a/examples/07_ecc_errors.py +++ b/examples/07_ecc_errors.py @@ -28,8 +28,11 @@ # 已纠正错误 (Volatile) try: corrected_vol = mtmlMemoryGetEccErrorCounter( - memory, MTML_MEMORY_ERROR_TYPE_CORRECTED, - MTML_VOLATILE_ECC, MTML_MEMORY_LOCATION_DRAM) + memory, + MTML_MEMORY_ERROR_TYPE_CORRECTED, + MTML_VOLATILE_ECC, + MTML_MEMORY_LOCATION_DRAM, + ) print(f" 已纠正错误 (当前会话): {corrected_vol}") except MTMLError as e: print(f" 已纠正错误 (当前会话): [不可用: {e}]") @@ -37,8 +40,11 @@ # 未纠正错误 (Volatile) try: uncorrected_vol = mtmlMemoryGetEccErrorCounter( - memory, MTML_MEMORY_ERROR_TYPE_UNCORRECTED, - MTML_VOLATILE_ECC, MTML_MEMORY_LOCATION_DRAM) + memory, + MTML_MEMORY_ERROR_TYPE_UNCORRECTED, + MTML_VOLATILE_ECC, + MTML_MEMORY_LOCATION_DRAM, + ) print(f" 未纠正错误 (当前会话): {uncorrected_vol}") except MTMLError as e: print(f" 未纠正错误 (当前会话): [不可用: {e}]") @@ -46,8 +52,11 @@ # 已纠正错误 (Aggregate) try: corrected_agg = mtmlMemoryGetEccErrorCounter( - memory, MTML_MEMORY_ERROR_TYPE_CORRECTED, - MTML_AGGREGATE_ECC, MTML_MEMORY_LOCATION_DRAM) + memory, + MTML_MEMORY_ERROR_TYPE_CORRECTED, + MTML_AGGREGATE_ECC, + MTML_MEMORY_LOCATION_DRAM, + ) print(f" 已纠正错误 (累计): {corrected_agg}") except MTMLError as e: print(f" 已纠正错误 (累计): [不可用: {e}]") @@ -55,8 +64,11 @@ # 未纠正错误 (Aggregate) try: uncorrected_agg = mtmlMemoryGetEccErrorCounter( - memory, MTML_MEMORY_ERROR_TYPE_UNCORRECTED, - MTML_AGGREGATE_ECC, MTML_MEMORY_LOCATION_DRAM) + memory, + MTML_MEMORY_ERROR_TYPE_UNCORRECTED, + MTML_AGGREGATE_ECC, + MTML_MEMORY_LOCATION_DRAM, + ) print(f" 未纠正错误 (累计): {uncorrected_agg}") except MTMLError as e: print(f" 未纠正错误 (累计): [不可用: {e}]") diff --git a/examples/08_topology_and_mtlink.py b/examples/08_topology_and_mtlink.py index 42f42e3..3caeba3 100644 --- a/examples/08_topology_and_mtlink.py +++ b/examples/08_topology_and_mtlink.py @@ -23,7 +23,7 @@ print(f" 设备 {i}: {name}") # MtLink 信息 -print(f"\n--- MtLink 信息 ---") +print("\n--- MtLink 信息 ---") mtlink_state_names = { MTML_MTLINK_STATE_DOWN: "断开", MTML_MTLINK_STATE_UP: "连接", @@ -75,7 +75,7 @@ MTML_TOPOLOGY_SYSTEM: "跨NUMA节点", } - print(f"\n--- 拓扑关系矩阵 ---") + print("\n--- 拓扑关系矩阵 ---") for i in range(device_count): for j in range(i + 1, device_count): try: @@ -86,36 +86,54 @@ print(f" 设备 {i} <-> 设备 {j}: [不可用: {e}]") # P2P 状态 - print(f"\n--- P2P 状态 ---") + print("\n--- P2P 状态 ---") for i in range(device_count): for j in range(i + 1, device_count): try: - read_status = mtmlDeviceGetP2PStatus(devices[i], devices[j], MTML_P2P_CAPS_READ) - write_status = mtmlDeviceGetP2PStatus(devices[i], devices[j], MTML_P2P_CAPS_WRITE) - read_ok = "OK" if read_status == MTML_P2P_STATUS_OK else f"Status={read_status}" - write_ok = "OK" if write_status == MTML_P2P_STATUS_OK else f"Status={write_status}" + read_status = mtmlDeviceGetP2PStatus( + devices[i], devices[j], MTML_P2P_CAPS_READ + ) + write_status = mtmlDeviceGetP2PStatus( + devices[i], devices[j], MTML_P2P_CAPS_WRITE + ) + read_ok = ( + "OK" + if read_status == MTML_P2P_STATUS_OK + else f"Status={read_status}" + ) + write_ok = ( + "OK" + if write_status == MTML_P2P_STATUS_OK + else f"Status={write_status}" + ) print(f" 设备 {i} <-> 设备 {j}: 读={read_ok}, 写={write_ok}") except MTMLError as e: print(f" 设备 {i} <-> 设备 {j}: [不可用: {e}]") # MtLink 布局 - print(f"\n--- MtLink 布局 ---") + print("\n--- MtLink 布局 ---") for i in range(device_count): for j in range(i + 1, device_count): try: link_count = mtmlDeviceCountMtLinkLayouts(devices[i], devices[j]) print(f" 设备 {i} <-> 设备 {j}: {link_count} 条链路") if link_count > 0: - layouts = mtmlDeviceGetMtLinkLayouts(devices[i], devices[j], link_count) + layouts = mtmlDeviceGetMtLinkLayouts( + devices[i], devices[j], link_count + ) for layout in layouts: - print(f" 本地链路ID={layout.localLinkId} <-> 远程链路ID={layout.remoteLinkId}") + print( + f" 本地链路ID={layout.localLinkId} <-> 远程链路ID={layout.remoteLinkId}" + ) except MTMLError as e: print(f" 设备 {i} <-> 设备 {j}: [不可用: {e}]") # MtLink 最短路径 (两设备间) - print(f"\n--- MtLink 最短路径 ---") + print("\n--- MtLink 最短路径 ---") try: - path_count, path_length = mtmlDeviceCountMtLinkShortestPaths(devices[0], devices[1]) + path_count, path_length = mtmlDeviceCountMtLinkShortestPaths( + devices[0], devices[1] + ) print(f" 设备 0 -> 设备 1: {path_count} 条路径, 长度 {path_length}") if path_count > 0 and path_length > 0: paths = mtmlDeviceGetMtLinkShortestPaths( @@ -127,12 +145,14 @@ print(f" [不可用: {e}]") # 按拓扑层级统计设备 / 获取同层级设备 - print(f"\n--- 按拓扑层级设备 ---") + print("\n--- 按拓扑层级设备 ---") for level in range(MTML_TOPOLOGY_SYSTEM + 1): try: count = mtmlDeviceCountDeviceByTopologyLevel(devices[0], level) if count > 0: - level_devs = mtmlDeviceGetDeviceByTopologyLevel(devices[0], level, count) + level_devs = mtmlDeviceGetDeviceByTopologyLevel( + devices[0], level, count + ) print(f" 层级 {level}: {count} 个设备 (GetDeviceByTopologyLevel)") except MTMLError: pass diff --git a/examples/09_nvml_compatibility.py b/examples/09_nvml_compatibility.py index caa2f9d..5e38807 100644 --- a/examples/09_nvml_compatibility.py +++ b/examples/09_nvml_compatibility.py @@ -89,22 +89,24 @@ # ECC ecc_current, ecc_pending = pynvml.nvmlDeviceGetEccMode(handle) - print(f" ECC 模式: 当前={'启用' if ecc_current else '禁用'}, " - f"待定={'启用' if ecc_pending else '禁用'}") + print( + f" ECC 模式: 当前={'启用' if ecc_current else '禁用'}, " + f"待定={'启用' if ecc_pending else '禁用'}" + ) # MUSA 计算能力 (需要 torch_musa) major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle) if major > 0: print(f" MUSA 计算能力: {major}.{minor}") else: - print(f" MUSA 计算能力: 不可用 (需要 torch_musa)") + print(" MUSA 计算能力: 不可用 (需要 torch_musa)") # NVLink/MtLink 状态 (链路 0) try: nvlink_state = pynvml.nvmlDeviceGetNvLinkState(handle, 0) print(f" NvLink/MtLink 链路0: {'UP' if nvlink_state else 'DOWN'}") except pynvml.NVMLError: - print(f" NvLink/MtLink 链路0: [不可用]") + print(" NvLink/MtLink 链路0: [不可用]") try: cap = pynvml.nvmlDeviceGetNvLinkCapability(handle, 0, 0) print(f" NvLink/MtLink 链路0 能力(0): {cap}") @@ -119,13 +121,19 @@ # 多 GPU 时: P2P 与拓扑 (NVML 接口) if device_count >= 2: - print(f"\n--- 多 GPU: P2P 与拓扑 (NVML) ---") + print("\n--- 多 GPU: P2P 与拓扑 (NVML) ---") h0 = pynvml.nvmlDeviceGetHandleByIndex(0) h1 = pynvml.nvmlDeviceGetHandleByIndex(1) try: - p2p_read = pynvml.nvmlDeviceGetP2PStatus(h0, h1, pynvml.NVML_P2P_CAPS_INDEX_READ) - p2p_write = pynvml.nvmlDeviceGetP2PStatus(h0, h1, pynvml.NVML_P2P_CAPS_INDEX_WRITE) - p2p_nvlink = pynvml.nvmlDeviceGetP2PStatus(h0, h1, pynvml.NVML_P2P_CAPS_INDEX_NVLINK) + p2p_read = pynvml.nvmlDeviceGetP2PStatus( + h0, h1, pynvml.NVML_P2P_CAPS_INDEX_READ + ) + p2p_write = pynvml.nvmlDeviceGetP2PStatus( + h0, h1, pynvml.NVML_P2P_CAPS_INDEX_WRITE + ) + p2p_nvlink = pynvml.nvmlDeviceGetP2PStatus( + h0, h1, pynvml.NVML_P2P_CAPS_INDEX_NVLINK + ) print(f" P2P 读: {p2p_read}, 写: {p2p_write}, NvLink/MtLink: {p2p_nvlink}") except pynvml.NVMLError as e: print(f" P2P: [不可用: {e}]") @@ -142,5 +150,5 @@ # 关闭 pynvml.nvmlShutdown() -print(f"\n[4] nvmlShutdown() 成功") +print("\n[4] nvmlShutdown() 成功") print("\n完成!") diff --git a/examples/10_device_paths.py b/examples/10_device_paths.py index 7b6ab7a..86db32e 100644 --- a/examples/10_device_paths.py +++ b/examples/10_device_paths.py @@ -37,8 +37,10 @@ try: spec = mtmlDeviceGetDisplayInterfaceSpec(device, d) type_name = disp_type_names.get(spec.type, f"Unknown({spec.type})") - print(f" 接口 {d}: {type_name}, " - f"最大分辨率 {spec.maxResWidth}x{spec.maxResHeight}") + print( + f" 接口 {d}: {type_name}, " + f"最大分辨率 {spec.maxResWidth}x{spec.maxResHeight}" + ) except MTMLError as e: print(f" 接口 {d}: [不可用: {e}]") except MTMLError as e: diff --git a/examples/11_mpc_and_virtualization.py b/examples/11_mpc_and_virtualization.py index 1699863..9fbdec2 100644 --- a/examples/11_mpc_and_virtualization.py +++ b/examples/11_mpc_and_virtualization.py @@ -28,7 +28,9 @@ MTML_MPC_TYPE_INSTANCE: "MPC 实例", } print(f" MPC 支持: {'是' if mpc_cap else '否'}") - print(f" MPC 类型: {mpc_type_names.get(prop.mpcType, f'Unknown({prop.mpcType})')}") + print( + f" MPC 类型: {mpc_type_names.get(prop.mpcType, f'Unknown({prop.mpcType})')}" + ) try: mode = mtmlDeviceGetMpcMode(device) @@ -42,8 +44,10 @@ if profile_count > 0: profiles = mtmlDeviceGetSupportedMpcProfiles(device, profile_count) for p in profiles: - print(f" ID={p.profileId}, 名称={p.name}, " - f"显存={p.memSize / 1024**3:.1f}GB, 核心数={p.gpuCores}") + print( + f" ID={p.profileId}, 名称={p.name}, " + f"显存={p.memSize / 1024**3:.1f}GB, 核心数={p.gpuCores}" + ) except MTMLError as e: print(f" MPC 配置文件: [不可用: {e}]") @@ -67,7 +71,9 @@ MTML_VIRT_ROLE_GUEST_VIRTDEVICE: "虚拟机虚拟设备", } print(f"\n 虚拟化支持: {'是' if virt_cap else '否'}") - print(f" 虚拟化角色: {virt_role_names.get(prop.virtRole, f'Unknown({prop.virtRole})')}") + print( + f" 虚拟化角色: {virt_role_names.get(prop.virtRole, f'Unknown({prop.virtRole})')}" + ) try: supported = mtmlDeviceCountSupportedVirtTypes(device) @@ -76,8 +82,10 @@ virt_types = mtmlDeviceGetSupportedVirtTypes(device, supported) for vt in virt_types: print(f" ID={vt.id}, 名称={vt.name}, 类别={vt.deviceClass}") - print(f" 最大实例数={vt.maxInstances}, 显存={vt.memSize / 1024**3:.1f}GB, " - f"核心数={vt.gpuCores}") + print( + f" 最大实例数={vt.maxInstances}, 显存={vt.memSize / 1024**3:.1f}GB, " + f"核心数={vt.gpuCores}" + ) except MTMLError as e: print(f" 支持的虚拟化类型: [不可用: {e}]") diff --git a/examples/12_affinity_and_log.py b/examples/12_affinity_and_log.py index cb72715..3df185f 100644 --- a/examples/12_affinity_and_log.py +++ b/examples/12_affinity_and_log.py @@ -29,7 +29,9 @@ cores.append(word_idx * 64 + bit) if cores: if len(cores) > 16: - print(f" CPU 核心列表: {cores[:8]}...{cores[-8:]} (共 {len(cores)} 核心)") + print( + f" CPU 核心列表: {cores[:8]}...{cores[-8:]} (共 {len(cores)} 核心)" + ) else: print(f" CPU 核心列表: {cores}") except MTMLError as e: @@ -50,7 +52,7 @@ print(f" 内存亲和性: [不可用: {e}]") # 日志配置 -print(f"\n--- 日志配置 ---") +print("\n--- 日志配置 ---") try: log_config = mtmlLogGetConfiguration() log_level_names = { @@ -62,7 +64,9 @@ } print(f" 日志文件: {log_config.filePath}") print(f" 最大大小: {log_config.maxSize}") - print(f" 日志级别: {log_level_names.get(log_config.logLevel, f'Unknown({log_config.logLevel})')}") + print( + f" 日志级别: {log_level_names.get(log_config.logLevel, f'Unknown({log_config.logLevel})')}" + ) except MTMLError as e: print(f" [不可用: {e}]") diff --git a/examples/13_comprehensive_report.py b/examples/13_comprehensive_report.py index 13d1a22..35ade12 100644 --- a/examples/13_comprehensive_report.py +++ b/examples/13_comprehensive_report.py @@ -65,7 +65,11 @@ def _safe(fn, default=None, *args, **kwargs): pass gpu_util_str = "{}%".format(gpu_util) if gpu_util is not None else "N/A" gpu_temp_str = "{}°C".format(gpu_temp) if gpu_temp is not None else "N/A" - gpu_clock_str = "{}/{} MHz".format(gpu_clock, gpu_max_clock) if (gpu_clock is not None and gpu_max_clock is not None) else "N/A" + gpu_clock_str = ( + "{}/{} MHz".format(gpu_clock, gpu_max_clock) + if (gpu_clock is not None and gpu_max_clock is not None) + else "N/A" + ) mem_total = mem_used = mem_util = mem_clock = mem_max_clock = None mem_ctx = _safe(mtmlMemoryContext, None, device) @@ -84,7 +88,9 @@ def _safe(fn, default=None, *args, **kwargs): if mem_util is not None: mem_str += " ({}%)".format(mem_util) mem_str += " 时钟 " + ( - "{}/{} MHz".format(mem_clock, mem_max_clock) if (mem_clock is not None and mem_max_clock is not None) else "N/A" + "{}/{} MHz".format(mem_clock, mem_max_clock) + if (mem_clock is not None and mem_max_clock is not None) + else "N/A" ) else: mem_str = "N/A" @@ -101,7 +107,8 @@ def _safe(fn, default=None, *args, **kwargs): else: pci_str = "N/A" - print(""" + print( + """ GPU {}: {} ├─ UUID: {} ├─ PCI: {} @@ -109,8 +116,18 @@ def _safe(fn, default=None, *args, **kwargs): ├─ 温度: {} 风扇: {} ├─ GPU: 利用率 {} 时钟 {} └─ 显存: {}""".format( - i, name, uuid, pci_str, power_str, gpu_temp_str, fan_str, gpu_util_str, gpu_clock_str, mem_str - )) + i, + name, + uuid, + pci_str, + power_str, + gpu_temp_str, + fan_str, + gpu_util_str, + gpu_clock_str, + mem_str, + ) + ) print() print("=" * 72) diff --git a/pymtml.py b/pymtml.py index 1c89b17..3cfbc15 100644 --- a/pymtml.py +++ b/pymtml.py @@ -7,11 +7,11 @@ import string import sys import threading +from contextlib import contextmanager from ctypes import * from dataclasses import dataclass from functools import wraps from typing import TYPE_CHECKING as _TYPE_CHECKING -from contextlib import contextmanager if _TYPE_CHECKING: from typing_extensions import TypeAlias as _TypeAlias # Python 3.10+ @@ -863,7 +863,7 @@ def _mtmlGetFunctionPointer(name): libLoadLock.acquire() try: # ensure library was loaded - if mtmlLib == None: + if mtmlLib is None: raise MTMLError(MTML_ERROR_FUNCTION_NOT_FOUND) try: _mtmlGetFunctionPointer_cache[name] = getattr(mtmlLib, name) @@ -904,19 +904,20 @@ def wrapper(*args, **kwargs): return wrapper return func + def _LoadWindowsLibrary(): """ Load the MTML library on Windows platform """ lib_name = "mtml.dll" lib_loader = WinDLL - + # Try loading the library with different names try: return lib_loader(lib_name) - except OSError as ose: + except OSError: pass - + # If all attempts failed, try search PATH for p in os.environ.get("PATH", "").split(os.pathsep): if not p: @@ -926,9 +927,10 @@ def _LoadWindowsLibrary(): if os.path.isfile(candidate): os.add_dll_directory(p) return WinDLL(candidate) - + raise OSError(f"Failed to load MTML library on Windows. Tried: {lib_name}") + def _LoadLinuxLibrary(): """ Load the MTML library on Linux/Unix platform @@ -947,13 +949,13 @@ def _LoadLinuxLibrary(): load_errors.append(f"{lib_name}: {ose}") error_msg = ( - "Failed to load MTML library on Linux. " - f"Tried: {', '.join(lib_names)}" + "Failed to load MTML library on Linux. " f"Tried: {', '.join(lib_names)}" ) if load_errors: error_msg += "\nLoad errors:\n" + "\n".join(load_errors) raise OSError(error_msg) + ## C function wrappers ## def _LoadMtmlLibrary(): """ @@ -961,16 +963,16 @@ def _LoadMtmlLibrary(): """ global mtmlLib - if mtmlLib == None: + if mtmlLib is None: # lock to ensure only one caller loads the library libLoadLock.acquire() try: # ensure the library still isn't loaded - if mtmlLib == None: + if mtmlLib is None: # Platform-specific library loading platform = sys.platform - + if platform.startswith("win32") or platform.startswith("cygwin"): # Windows platform mtmlLib = _LoadWindowsLibrary() @@ -978,8 +980,10 @@ def _LoadMtmlLibrary(): # Linux/Unix platform mtmlLib = _LoadLinuxLibrary() - if mtmlLib == None: - raise OSError("Failed to load MTML library on platform: " + platform) + if mtmlLib is None: + raise OSError( + "Failed to load MTML library on platform: " + platform + ) finally: # lock is always freed libLoadLock.release() @@ -1017,9 +1021,7 @@ def mtmlLibraryShutDown(): _mtmlCheckReturn(ret) _mtmlLib_handles.pop() _mtmlLib_refcount = len(_mtmlLib_handles) - libHandle = ( - _mtmlLib_handles[0] if _mtmlLib_handles else c_mtmlLibrary_t() - ) + libHandle = _mtmlLib_handles[0] if _mtmlLib_handles else c_mtmlLibrary_t() return None @@ -1135,6 +1137,7 @@ def mtmlDeviceGetPciInfo(device): _mtmlCheckReturn(ret) return c_pciinfo + @convertStrBytes def mtmlDeviceGetSerialNumber(device): global libHandle @@ -1285,6 +1288,7 @@ def mtmlVpuGetMaxClock(vpu): _mtmlCheckReturn(ret) return c_clock.value + @convertStrBytes def mtmlSystemGetDriverVersion(system): c_version = create_string_buffer(MTML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) @@ -1602,6 +1606,7 @@ def mtmlGpuGetEngineUtilization(gpu, engine): _mtmlCheckReturn(ret) return c_util.value + @contextmanager def mtmlGpuContext(device): c_gpu = None @@ -1662,6 +1667,7 @@ def mtmlMemoryGetType(memory): _mtmlCheckReturn(ret) return c_type.value + @contextmanager def mtmlMemoryContext(device): c_memory = None @@ -1672,6 +1678,7 @@ def mtmlMemoryContext(device): if c_memory is not None: mtmlDeviceFreeMemory(c_memory) + ## VPU APIs def mtmlVpuGetUtilization(vpu): c_util = c_mtmlCodecUtil_t() @@ -1727,6 +1734,7 @@ def mtmlVpuGetDecoderSessionMetrics(vpu, sessionId): _mtmlCheckReturn(ret) return c_metrics + @contextmanager def mtmlVpuContext(device): c_vpu = None @@ -1737,6 +1745,7 @@ def mtmlVpuContext(device): if c_vpu is not None: mtmlDeviceFreeVpu(c_vpu) + ## Log configuration APIs def mtmlLogSetConfiguration(configuration): fn = _mtmlGetFunctionPointer("mtmlLogSetConfiguration") @@ -2070,7 +2079,7 @@ def mtmlLibrarySetMpcConfigurationInBatch(devices, mpcConfigIds): _nvmlComputeMode_t = c_uint NVML_COMPUTEMODE_DEFAULT = 0 -NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed +NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 # Support Removed NVML_COMPUTEMODE_PROHIBITED = 2 NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3 NVML_COMPUTEMODE_COUNT = 4 @@ -2281,7 +2290,7 @@ def nvmlDeviceGetPowerUsage(device): return mtmlDeviceGetPowerUsage(device) -# cannot expose this function directly since _nvmlGetFunctionPointer will retrive the function pointer from the mtml library directly. +# cannot expose this function directly since _nvmlGetFunctionPointer will retrieve the function pointer from the mtml library directly. # it will cause a MTMLError_FunctionNotFound exception when we get nvml function pointer from the mtml library. # def _nvmlGetFunctionPointer(name): # return _mtmlGetFunctionPointer(name) @@ -2394,7 +2403,7 @@ def nvmlDeviceGetCudaComputeCapability(device): """ try: import torch - import torch_musa + import torch_musa # noqa: F401 - registers the MUSA backend for torch # Get device index from the device handle # We need to find which device index this handle corresponds to diff --git a/setup.py b/setup.py index 7b22b9e..08ae6a3 100644 --- a/setup.py +++ b/setup.py @@ -5,15 +5,14 @@ from distutils.core import setup from distutils.dist import DistributionMetadata -from sys import version -from sys import exit +from sys import exit, version # earlier versions don't support all classifiers -if version < '2.2.3': +if version < "2.2.3": DistributionMetadata.classifiers = None DistributionMetadata.download_url = None -_package_name='mthreads-ml-py' +_package_name = "mthreads-ml-py" long_description = None try: @@ -21,14 +20,15 @@ # This will fail with an import error on earlier version of python3 and python2 # Code is from https://packaging.python.org/en/latest/guides/making-a-pypi-friendly-readme/ from pathlib import Path + this_dir = Path(__file__).parent - long_description = (this_dir / 'README.md').read_text() + long_description = (this_dir / "README.md").read_text() except ImportError: # If pathlib can not find the readme, try reading it from the current dir. # This should always work if the readme is in the parent directory. # This will also work with python2. try: - with open('README.md', 'r') as readme: + with open("README.md", "r") as readme: long_description = readme.read() except IOError: # The python2 error for files that don't exist. @@ -41,32 +41,33 @@ finally: # Exit if there is no long description. if long_description is None: - print('Unable to load README.md as long description.') - print('Make sure it is in the correct directory.') + print("Unable to load README.md as long description.") + print("Make sure it is in the correct directory.") exit(1) -setup(name=_package_name, - version='2.4.3', - description='Python Bindings for the Moore Threads GPU Management Library', - long_description=long_description, - long_description_content_type='text/markdown', - py_modules=['pymtml', 'example'], - package_data={_package_name: ['Example.txt']}, - license='MIT', - url='https://developer.mthreads.com', - author='Moore Threads Corporation', - author_email='mt-sw-cloudcomputing@mthreads.com', - classifiers=[ - 'Development Status :: 5 - Production/Stable', - 'Intended Audience :: Developers', - 'Intended Audience :: System Administrators', - 'License :: OSI Approved :: MIT License', - 'Operating System :: Microsoft :: Windows', - 'Operating System :: POSIX :: Linux', - 'Programming Language :: Python', - 'Topic :: Software Development :: Libraries :: Python Modules', - 'Topic :: System :: Hardware', - 'Topic :: System :: Systems Administration', - ], - ) +setup( + name=_package_name, + version="2.4.3", + description="Python Bindings for the Moore Threads GPU Management Library", + long_description=long_description, + long_description_content_type="text/markdown", + py_modules=["pymtml", "example"], + package_data={_package_name: ["Example.txt"]}, + license="MIT", + url="https://developer.mthreads.com", + author="Moore Threads Corporation", + author_email="mt-sw-cloudcomputing@mthreads.com", + classifiers=[ + "Development Status :: 5 - Production/Stable", + "Intended Audience :: Developers", + "Intended Audience :: System Administrators", + "License :: OSI Approved :: MIT License", + "Operating System :: Microsoft :: Windows", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python", + "Topic :: Software Development :: Libraries :: Python Modules", + "Topic :: System :: Hardware", + "Topic :: System :: Systems Administration", + ], +) diff --git a/test_pymtml.py b/test_pymtml.py index 412bbc2..c4368aa 100644 --- a/test_pymtml.py +++ b/test_pymtml.py @@ -4,9 +4,9 @@ Run with: python test_pymtml.py """ +import os import sys import traceback -import os from pymtml import * @@ -408,7 +408,7 @@ def main(): try: mtmlLibraryShutDown() print("\nMTML library shutdown complete.") - except: + except Exception: pass # Test init/shutdown cycle (after main tests complete) diff --git a/test_pynvml.py b/test_pynvml.py index 88247b0..9a7fcef 100644 --- a/test_pynvml.py +++ b/test_pynvml.py @@ -64,7 +64,9 @@ def test_device_handle_apis(self): return None # Get UUID for testing GetHandleByUuid - _, uuid = test_api("nvmlDeviceGetUUID", lambda: pynvml.nvmlDeviceGetUUID(handle)) + _, uuid = test_api( + "nvmlDeviceGetUUID", lambda: pynvml.nvmlDeviceGetUUID(handle) + ) if uuid: test_api( "nvmlDeviceGetHandleByUuid", @@ -92,7 +94,8 @@ def test_device_info_apis(self, device): test_api("nvmlDeviceGetSerial", lambda: pynvml.nvmlDeviceGetSerial(device)) test_api("nvmlDeviceGetBrand", lambda: pynvml.nvmlDeviceGetBrand(device)) test_api( - "nvmlDeviceGetVbiosVersion", lambda: pynvml.nvmlDeviceGetVbiosVersion(device) + "nvmlDeviceGetVbiosVersion", + lambda: pynvml.nvmlDeviceGetVbiosVersion(device), ) test_api( "nvmlDeviceGetMinorNumber", lambda: pynvml.nvmlDeviceGetMinorNumber(device) @@ -201,7 +204,10 @@ def test_pci_apis(self, device): if pci_info.busId and pci_info.busId[0].isalnum(): print_result("nvmlDeviceGetPciInfo busId populated", pci_info.busId) else: - print_result("nvmlDeviceGetPciInfo busId populated", "[FAIL: busId is empty or invalid]") + print_result( + "nvmlDeviceGetPciInfo busId populated", + "[FAIL: busId is empty or invalid]", + ) test_api( "nvmlDeviceGetPcieThroughput (TX)", lambda: pynvml.nvmlDeviceGetPcieThroughput( @@ -419,4 +425,3 @@ def main(): if __name__ == "__main__": sys.exit(main()) - diff --git a/test_sglang_compat.py b/test_sglang_compat.py index 3c9a857..aa030a3 100644 --- a/test_sglang_compat.py +++ b/test_sglang_compat.py @@ -333,7 +333,7 @@ def main(): try: pynvml.nvmlShutdown() print("\nLibrary shutdown complete (via pynvml.nvmlShutdown()).") - except: + except Exception: pass return 0