Skip to content

Commit

Permalink
change to use auto config
Browse files Browse the repository at this point in the history
  • Loading branch information
elazarcoh committed Sep 13, 2021
1 parent 039edea commit e48b308
Show file tree
Hide file tree
Showing 5 changed files with 310 additions and 188 deletions.
53 changes: 53 additions & 0 deletions src/config.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
import { configUtils } from "vscode-extensions-json-generator/utils";

type ItemOption =
| "product_name"
| "gpu_temp"
| "gpu_util"
| "fan_speed"
| "memory_util"
| "memory_total"
| "memory_free"
| "memory_used"
| "memory_used_percent";


export interface Config {
/**
* @title nvidia-smi Path
* @default "nvidia-smi"
*/
executablePath: string;

/**
* @title Auto Refresh
* @default true
*/
"refresh.autoRefresh": boolean;

/**
* @asType integer
* @title Time Interval (sec)
* @default 3
* @minimum 1
* @maximum 3600
* @markdownDescription "Controls the refresh time interval in seconds. Only applies when `#nvidia-smi-plus.refresh.autoRefresh#` is set to `true`."
*/
"refresh.timeInterval": number;

/**
* @default "product_name"
*/
"view.gpuMainDescription": ItemOption;

/**
* @title GPU Information
* @default ["memory_total", "memory_free", "memory_used_percent"]
* @description Configure which information will be shown on the view. see https://github.com/RSIP-Vision/vscode-nividia-smi-plus#available-information-fields for the complete list.
*/
"view.gpuItems": ItemOption[];
}

const configSection = "nvidia-smi-plus";
export const getConfiguration =
configUtils.ConfigurationGetter<Config>(configSection);
227 changes: 120 additions & 107 deletions src/gpu-info-service.ts
Original file line number Diff line number Diff line change
@@ -1,144 +1,157 @@
import * as vscode from 'vscode';

var parser = require('fast-xml-parser');
import { spawn } from 'child_process';
import { CronJob } from 'cron';
import { shallowEqual } from './utils';
import { NVIDIA_SMI_FIELDS, resolveGpuInfoField } from './nvidia-smi-fields';
import * as vscode from "vscode";

import parser = require("fast-xml-parser");
import { spawn } from "child_process";
import { CronJob } from "cron";
import { shallowEqual } from "./utils";
import { NVIDIA_SMI_FIELDS, resolveGpuInfoField } from "./nvidia-smi-fields";
import { getConfiguration } from "./config";

export type GpuInfo = {
id: number;
[key: string]: any;
id: number;
[key: string]: any;
};

export type NvidiaSmiInfo = {
gpus: GpuInfo[];
gpus: GpuInfo[];
};


export interface NvidiaSmiEvent {
info: NvidiaSmiInfo
info: NvidiaSmiInfo;
}

function asCronTime(seconds: number) {
if (seconds < 60) {
return `*/${seconds} * * * * *`;
} else if (seconds < 3600) {
let minutes = Math.trunc(seconds / 60);
seconds = seconds % 60;
return `*/${seconds} */${minutes} * * * *`;
} else {
return `* * 1 * * *`;
}
if (seconds < 60) {
return `*/${seconds} * * * * *`;
} else if (seconds < 3600) {
const minutes = Math.trunc(seconds / 60);
seconds = seconds % 60;
return `*/${seconds} */${minutes} * * * *`;
} else {
return `* * 1 * * *`;
}
}

type RefreshConfig = {
autoRefresh: boolean
refreshInterval: number
autoRefresh: boolean;
refreshInterval: number;
};
function refreshConfiguration(): RefreshConfig {
const extConfig = vscode.workspace.getConfiguration('nvidia-smi-plus');

const autoRefresh = extConfig.get<boolean>('refresh.autoRefresh');
const seconds = extConfig.get<number>('refresh.timeInterval');
return {
autoRefresh: Boolean(autoRefresh),
refreshInterval: seconds ? seconds : 0,
};

const autoRefresh = getConfiguration("refresh.autoRefresh");
const seconds = getConfiguration("refresh.timeInterval");
return {
autoRefresh: Boolean(autoRefresh),
refreshInterval: seconds ? seconds : 0,
};
}

export class NvidiaSmiService implements vscode.Disposable {
private _updateInfoJob: CronJob | undefined;
private _currentRefreshSettings: RefreshConfig | undefined;

constructor() {
this.setAutoUpdate();
}

setAutoUpdate() {
const currentConfig = refreshConfiguration();
if (!this._currentRefreshSettings || !shallowEqual(this._currentRefreshSettings, currentConfig)) {
this._currentRefreshSettings = currentConfig;
this._updateInfoJob?.stop();
if (this._currentRefreshSettings.autoRefresh) {
const seconds = this._currentRefreshSettings.refreshInterval;
this._updateInfoJob = new CronJob(asCronTime(seconds), () => this.update());
if (!this._updateInfoJob.running) {
this._updateInfoJob.start();
}
}
private _updateInfoJob: CronJob | undefined;
private _currentRefreshSettings: RefreshConfig | undefined;

constructor() {
this.setAutoUpdate();
}

setAutoUpdate() {
const currentConfig = refreshConfiguration();
if (
!this._currentRefreshSettings ||
!shallowEqual(this._currentRefreshSettings, currentConfig)
) {
this._currentRefreshSettings = currentConfig;
this._updateInfoJob?.stop();
if (this._currentRefreshSettings.autoRefresh) {
const seconds = this._currentRefreshSettings.refreshInterval;
this._updateInfoJob = new CronJob(asCronTime(seconds), () =>
this.update()
);
if (!this._updateInfoJob.running) {
this._updateInfoJob.start();
}
}
}
}

private readonly _onDidInfoAcquired = new vscode.EventEmitter<NvidiaSmiEvent>();
readonly onDidInfoAcquired = this._onDidInfoAcquired.event;
private readonly _onDidInfoAcquired =
new vscode.EventEmitter<NvidiaSmiEvent>();
readonly onDidInfoAcquired = this._onDidInfoAcquired.event;

private _currentState: NvidiaSmiInfo | undefined;
private _currentState: NvidiaSmiInfo | undefined;

async update(): Promise<void> {
this._currentState = await this.currentNvidiaStatus();
if (this._currentState) {
this._onDidInfoAcquired.fire({ info: this._currentState });
}
async update(): Promise<void> {
this._currentState = await this.currentNvidiaStatus();
if (this._currentState) {
this._onDidInfoAcquired.fire({ info: this._currentState });
}

async currentNvidiaStatus(): Promise<NvidiaSmiInfo | undefined> {
try {
const jsonObj = await nvidiaSmiAsJsonObject();
const gpus: GpuInfo[] = [];
for (const [gpuId, gpuInfo] of jsonObj.nvidia_smi_log.gpu.entries()) {
const gpuInfoFields: Record<string, any> = {};
for (const [name, field] of Object.entries(NVIDIA_SMI_FIELDS)) {
gpuInfoFields[name] = resolveGpuInfoField(gpuInfo, field, gpuInfoFields);
}
gpus.push({
id: gpuId,
...gpuInfoFields,
});
}
return {
gpus: gpus,
};
} catch (error) {
console.log(error.message);
}

async currentNvidiaStatus(): Promise<NvidiaSmiInfo | undefined> {
try {
const jsonObj = await nvidiaSmiAsJsonObject();
const gpus: GpuInfo[] = [];
for (const [gpuId, gpuInfo] of jsonObj.nvidia_smi_log.gpu.entries()) {
const gpuInfoFields: Record<string, any> = {};
for (const [name, field] of Object.entries(NVIDIA_SMI_FIELDS)) {
gpuInfoFields[name] = resolveGpuInfoField(
gpuInfo,
field,
gpuInfoFields
);
}
gpus.push({
id: gpuId,
...gpuInfoFields,
});
}
return {
gpus: gpus,
};
} catch (error) {
console.error(error);
}
}

dispose() {
if (this._updateInfoJob?.running) {
this._updateInfoJob.stop();
}
dispose() {
if (this._updateInfoJob?.running) {
this._updateInfoJob.stop();
}
}
}

export async function nvidiaSmiAsJsonObject() {
const extConfig = vscode.workspace.getConfiguration('nvidia-smi-plus');
const exec = extConfig.get<string>("executablePath") ?? "nvidia-smi";

const child = spawn(exec, ["-q", "-x"]);
let xmlData = '';
for await (const data of child.stdout) {
xmlData += data.toString();
};
const jsonObj = parser.parse(xmlData, {}, true);

// for a system with a single GPU
if (!Array.isArray(jsonObj.nvidia_smi_log.gpu)) {
jsonObj.nvidia_smi_log.gpu = [jsonObj.nvidia_smi_log.gpu];
}
const exec = getConfiguration("executablePath") ?? "nvidia-smi";

return jsonObj;
}
const child = spawn(exec, ["-q", "-x"]);
let xmlData = "";
for await (const data of child.stdout) {
xmlData += data.toString();
}
const jsonObj = parser.parse(xmlData, {}, true);

export async function openAsJsonFile() {
const jsonObj = await nvidiaSmiAsJsonObject();
// for a system with a single GPU
if (!Array.isArray(jsonObj.nvidia_smi_log.gpu)) {
jsonObj.nvidia_smi_log.gpu = [jsonObj.nvidia_smi_log.gpu];
}

const fileName = 'nvidia-smi.json';
const newUri = vscode.Uri.file(fileName).with({ scheme: 'untitled', path: fileName });
return jsonObj;
}

const document = await vscode.workspace.openTextDocument(newUri);
const textEdit = await vscode.window.showTextDocument(document);
await textEdit.edit(edit => edit.insert(new vscode.Position(0, 0), JSON.stringify(jsonObj, undefined, 4)));
}
export async function openAsJsonFile() {
const jsonObj = await nvidiaSmiAsJsonObject();

const fileName = "nvidia-smi.json";
const newUri = vscode.Uri.file(fileName).with({
scheme: "untitled",
path: fileName,
});

const document = await vscode.workspace.openTextDocument(newUri);
const textEdit = await vscode.window.showTextDocument(document);
await textEdit.edit((edit) =>
edit.insert(
new vscode.Position(0, 0),
JSON.stringify(jsonObj, undefined, 4)
)
);
}
Loading

0 comments on commit e48b308

Please sign in to comment.