Merge pull request #14 from maleksan85/amd_gpu_monitoring

AMD GPUs support with rocm-smi

package.json

@@ -16,7 +16,7 @@
  "url": "https://github.com/RSIP-Vision/vscode-nividia-smi-plus"
  },
  "description": "Watch Nvidia GPUs state on your machine",
- "version": "1.0.0",
+ "version": "1.0.1",
  "engines": {
  "vscode": "^1.55.0"
  },

src/extension.ts

@@ -1,25 +1,29 @@
 import * as vscode from "vscode";
 import { configurations } from "./config";
-
-import { NvidiaSmiService, openAsJsonFile } from "./gpu-info-service";
+import { SmiService, NvidiaSmiService, RocmSmiService, openAsJsonFile } from "./gpu-info-service";
 import { GPUInfoProvider } from "./gpu-treeview";
 
-let nvidiaSmiService: NvidiaSmiService | undefined;
+let smiService: SmiService | undefined;
 
 export function updateNvidiaInfo(): void {
- nvidiaSmiService?.update();
+ smiService?.update();
 }
 
 export async function setAutoRefresh(value: boolean): Promise<void> {
  await configurations.update("refresh.autoRefresh", value);
- nvidiaSmiService?.setAutoUpdate();
+ smiService?.setAutoUpdate();
 }
 
 export async function activate(
  context: vscode.ExtensionContext
 ): Promise<void> {
- nvidiaSmiService = new NvidiaSmiService();
- context.subscriptions.push(nvidiaSmiService);
+ const exec = configurations.get("executablePath", undefined, "");
+ if (exec.includes("rocm")) {
+ smiService = new RocmSmiService();
+ } else {
+ smiService = new NvidiaSmiService();
+ }
+ context.subscriptions.push(smiService);
 
  const nvidiaRefreshCmd = vscode.commands.registerCommand(
  "nvidia-smi-plus.refresh",
@@ -45,15 +49,15 @@ export async function activate(
  context.subscriptions.push(nvidiaOpenJsonCmd);
 
  const gpuInfoProvider = new GPUInfoProvider();
- nvidiaSmiService.onDidInfoAcquired(gpuInfoProvider.refresh, gpuInfoProvider);
+ smiService.onDidInfoAcquired(gpuInfoProvider.refresh, gpuInfoProvider);
  vscode.window.registerTreeDataProvider("nvidia-gpus", gpuInfoProvider);
 
  const configChange = vscode.workspace.onDidChangeConfiguration((event) => {
  if (event.affectsConfiguration("nvidia-smi-plus.refresh")) {
- nvidiaSmiService?.setAutoUpdate();
+ smiService?.setAutoUpdate();
  }
  if (event.affectsConfiguration("nvidia-smi-plus")) {
- nvidiaSmiService?.update();
+ smiService?.update();
  }
  });
  context.subscriptions.push(configChange);

src/gpu-info-service.ts

@@ -1,12 +1,12 @@
 import * as vscode from "vscode";
-
-import parser = require("fast-xml-parser");
 import { spawn } from "child_process";
 import { CronJob } from "cron";
 import { json, shallowEqual } from "./utils";
-import { NVIDIA_SMI_FIELDS, resolveGpuInfoField } from "./nvidia-smi-fields";
+import { NVIDIA_SMI_FIELDS, ROCM_SMI_FIELDS, resolveGpuInfoField } from "./nvidia-smi-fields";
 import { configurations } from "./config";
 
+/* eslint-disable */
+const { XMLParser } = require("fast-xml-parser");
 
 type NvidiaSmiInfoJson = {
  nvidia_smi_log: {
@@ -15,16 +15,16 @@ type NvidiaSmiInfoJson = {
 }
 
 export type GpuInfo = {
- id: number;
+ id: string | number;
  [key: string]: string | number;
 };
 
-export type NvidiaSmiInfo = {
+export type SmiInfo = {
  gpus: GpuInfo[];
 };
 
-export interface NvidiaSmiEvent {
- info: NvidiaSmiInfo;
+export interface SmiEvent {
+ info: SmiInfo;
 }
 
 function asCronTime(seconds: number) {
@@ -52,14 +52,20 @@ function refreshConfiguration(): RefreshConfig {
  };
 }
 
-export class NvidiaSmiService implements vscode.Disposable {
+export abstract class SmiService implements vscode.Disposable {
  private _updateInfoJob: CronJob | undefined;
  private _currentRefreshSettings: RefreshConfig | undefined;
 
  constructor() {
  this.setAutoUpdate();
  }
 
+ protected readonly _onDidInfoAcquired =
+ new vscode.EventEmitter<SmiEvent>();
+ readonly onDidInfoAcquired = this._onDidInfoAcquired.event;
+
+ abstract update(): void;
+
  setAutoUpdate(): void {
  const currentConfig = refreshConfiguration();
  if (
@@ -80,11 +86,59 @@ export class NvidiaSmiService implements vscode.Disposable {
  }
  }
 
- private readonly _onDidInfoAcquired =
- new vscode.EventEmitter<NvidiaSmiEvent>();
- readonly onDidInfoAcquired = this._onDidInfoAcquired.event;
+ dispose(): void {
+ if (this._updateInfoJob?.running) {
+ this._updateInfoJob.stop();
+ }
+ }
+}
+
+export class RocmSmiService extends SmiService {
+ constructor() {
+ super()
+ }
+
+ private _currentState: SmiInfo | undefined;
+
+ async update(): Promise<void> {
+ this._currentState = await this.currentRocmStatus();
+ if (this._currentState) {
+ this._onDidInfoAcquired.fire({ info: this._currentState });
+ }
+ }
+
+ async currentRocmStatus(): Promise<SmiInfo | undefined> {
+ try {
+ const jsonObj: json = await rocmSmiAsJsonObject();
+ const gpus: GpuInfo[] = [];
+ for (let [gpuId, gpuInfo] of Object.entries(jsonObj!)) {
+ if (!gpuId.includes("card"))
+ continue
+
+ const gpuInfoFields: Record<string, number | string> = {};
+ for (const [name, field] of Object.entries(ROCM_SMI_FIELDS)) {
+ gpuInfoFields[name] = resolveGpuInfoField(
+ gpuInfo,
+ field,
+ gpuInfoFields
+ ) ?? "null";
+ }
+ gpus.push({
+ id: gpuId,
+ ...gpuInfoFields,
+ });
+ }
+ return {
+ gpus: gpus,
+ }
+ } catch (error) {
+ console.error(error);
+ }
+ }
+}
 
- private _currentState: NvidiaSmiInfo | undefined;
+export class NvidiaSmiService extends SmiService {
+ private _currentState: SmiInfo | undefined;
 
  async update(): Promise<void> {
  this._currentState = await this.currentNvidiaStatus();
@@ -93,11 +147,11 @@ export class NvidiaSmiService implements vscode.Disposable {
  }
  }
 
- async currentNvidiaStatus(): Promise<NvidiaSmiInfo | undefined> {
+ async currentNvidiaStatus(): Promise<SmiInfo | undefined> {
  try {
  const jsonObj = await nvidiaSmiAsJsonObject();
  const gpus: GpuInfo[] = [];
- for (const [gpuId, gpuInfo] of jsonObj.nvidia_smi_log.gpu.entries()) {
+ for (const [gpuId, gpuInfo] of jsonObj?.nvidia_smi_log.gpu.entries()) {
  const gpuInfoFields: Record<string, number | string> = {};
  for (const [name, field] of Object.entries(NVIDIA_SMI_FIELDS)) {
  gpuInfoFields[name] = resolveGpuInfoField(
@@ -118,12 +172,19 @@ export class NvidiaSmiService implements vscode.Disposable {
  console.error(error);
  }
  }
+}
 
- dispose(): void {
- if (this._updateInfoJob?.running) {
- this._updateInfoJob.stop();
- }
+export async function rocmSmiAsJsonObject(): Promise<json> {
+ const exec = configurations.get("executablePath", undefined, "rocm-smi");
+
+ const child = spawn(exec, ["--showallinfo", "--showfan", "--showmeminfo", "VRAM", "--json"]);
+ let jsonData = "";
+ for await (const data of child.stdout) {
+ jsonData += data.toString();
  }
+ const jsonObj: json = JSON.parse(jsonData)
+
+ return jsonObj;
 }
 
 export async function nvidiaSmiAsJsonObject(): Promise<NvidiaSmiInfoJson> {
@@ -134,6 +195,7 @@ export async function nvidiaSmiAsJsonObject(): Promise<NvidiaSmiInfoJson> {
  for await (const data of child.stdout) {
  xmlData += data.toString();
  }
+ const parser = new XMLParser();
  const jsonObj: NvidiaSmiInfoJson = parser.parse(xmlData, {}, true);
 
  // for a system with a single GPU
@@ -145,9 +207,18 @@ export async function nvidiaSmiAsJsonObject(): Promise<NvidiaSmiInfoJson> {
 }
 
 export async function openAsJsonFile(): Promise<void> {
- const jsonObj = await nvidiaSmiAsJsonObject();
+ const exec = configurations.get("executablePath", undefined, "");
+ let jsonObj: json;
+ let fileName;
+ if (exec.includes("rocm")) {
+ jsonObj = await rocmSmiAsJsonObject()
+ fileName = "rocm-smi.json";
+ }
+ else {
+ jsonObj = await nvidiaSmiAsJsonObject();
+ fileName = "nvidia-smi.json";
+ }
 
- const fileName = "nvidia-smi.json";
  const newUri = vscode.Uri.file(fileName).with({
  scheme: "untitled",
  path: fileName,

src/gpu-treeview.ts

@@ -1,7 +1,7 @@
 import * as vscode from "vscode";
 import { configurations } from "./config";
-import { NvidiaSmiInfo, NvidiaSmiEvent, GpuInfo } from "./gpu-info-service";
-import { NVIDIA_SMI_FIELDS } from "./nvidia-smi-fields";
+import { SmiInfo, SmiEvent, GpuInfo } from "./gpu-info-service";
+import { NVIDIA_SMI_FIELDS, ROCM_SMI_FIELDS } from "./nvidia-smi-fields";
 
 enum GPUTreeItemType {
  gpuItem = "GPUItem",
@@ -23,8 +23,14 @@ class GPUItem extends vscode.TreeItem {
 }
 
 function itemLabel(itemId: string, _itemValue: string | number): string {
- const field = NVIDIA_SMI_FIELDS[itemId];
- return `${field.label}`;
+ const exec = configurations.get("executablePath", undefined, "");
+ if (exec.includes("rocm")) {
+ const field = ROCM_SMI_FIELDS[itemId];
+ return `${field.label}`;
+ } else {
+ const field = NVIDIA_SMI_FIELDS[itemId];
+ return `${field.label}`;
+ }
 }
 
 function itemDescription(itemId: string, itemValue: string | number): string {
@@ -33,27 +39,34 @@ function itemDescription(itemId: string, itemValue: string | number): string {
 
 function gpuInfoItems(gpu: GpuInfo): GPUItem[] {
  const infoItemsToShow = configurations.get("view.gpuItems");
+ const exec = configurations.get("executablePath", undefined, "");
 
  if (!infoItemsToShow) {
  return [];
  } else {
  const items = [];
  for (const infoId of infoItemsToShow) {
+ let iconPath: string = ""
+ if (exec.includes("rocm")) {
+ iconPath = ROCM_SMI_FIELDS[infoId].iconPath!
+ } else {
+ iconPath = NVIDIA_SMI_FIELDS[infoId].iconPath!
+ }
  items.push(
  new GPUItem(
  itemLabel(infoId, gpu[infoId]),
  vscode.TreeItemCollapsibleState.None,
  GPUTreeItemType.gpuInfoItem,
  itemDescription(infoId, gpu[infoId]),
- NVIDIA_SMI_FIELDS[infoId].iconPath
+ iconPath
  )
  );
  }
  return items;
  }
 }
 
-function gpusInfo(info: NvidiaSmiInfo): GPUItem[] {
+function gpusInfo(info: SmiInfo): GPUItem[] {
  const gpuMainDescription = configurations.get("view.gpuMainDescription");
 
  return info.gpus.map(
@@ -75,9 +88,9 @@ export class GPUInfoProvider implements vscode.TreeDataProvider<GPUItem> {
  >();
  readonly onDidChangeTreeData = this._onDidChangeTreeData.event;
 
- private _currentInfo: NvidiaSmiInfo | undefined;
+ private _currentInfo: SmiInfo | undefined;
 
- refresh(event: NvidiaSmiEvent): void {
+ refresh(event: SmiEvent): void {
  this._currentInfo = event.info;
  this._onDidChangeTreeData.fire(undefined);
  }

src/nvidia-smi-fields.ts

@@ -1,6 +1,5 @@
 // keep the NVIDIA_SMI_FIELDS records in snake_case to match the configuration style
 /* eslint-disable @typescript-eslint/naming-convention */
-import * as vscode from 'vscode';
 import { json, replaceAll } from './utils';
 
 enum InfoFieldType {
@@ -14,7 +13,7 @@ export class InfoField {
  public readonly label: string,
  public readonly accessor: string[] | string,
  public readonly depends?: string[],
- public readonly iconPath?: vscode.Uri | string,
+ public readonly iconPath?: string,
  ) { }
 }
 
@@ -37,6 +36,22 @@ export const NVIDIA_SMI_FIELDS: Record<string, InfoField> = {
 };
 
 
+// for `Expr` Field, the order is significant. you can't reference a field that has not beed calculated yet.
+export const ROCM_SMI_FIELDS: Record<string, InfoField> = {
+ gpu_temp: new InfoField(InfoFieldType.value, 'GPU Temperature', ['Temperature (Sensor edge) (C)']),
+ gpu_util: new InfoField(InfoFieldType.value, 'GPU Utilization', ['GPU use (%)']),
+ memory_util: new InfoField(InfoFieldType.value, 'Memory Utilization', ['GPU memory use (%)']),
+ memory_total: new InfoField(InfoFieldType.value, 'Total memory', ['VRAM Total Memory (B)']),
+ memory_used: new InfoField(InfoFieldType.value, 'Used memory', ['VRAM Total Used Memory (B)']),
+ memory_used_percent: new InfoField(
+ InfoFieldType.expr,
+ "Memory Used (%)",
+ '`${Math.trunc(("{memory_used}") / ("{memory_total}") * 100)} %`',
+ ['memory_used', 'memory_total']),
+ fan_speed: new InfoField(InfoFieldType.value, "Fan speed", ['Fan speed (%)']),
+ product_name: new InfoField(InfoFieldType.value, "Product Name", ['Card series']),
+};
+
 export function resolveGpuInfoField(gpuInfo: json, field: InfoField, values: Record<string, string | number>): string | number | undefined {
  switch (field.type) {
  case InfoFieldType.value: