Skip to content

Commit 5302e4b

Browse files
committed
feat: add Active Runs support to tray menu
This PR sets up the tray menu to watch for active runs. Upon selection, it will either open the dashboard, if we have already attached, or attach an aggregator and then open the dashboard. This PR sneaks in some minor bits to add a PyTorch Profiler tray menu option.
1 parent 91b03a9 commit 5302e4b

File tree

15 files changed

+308
-11
lines changed

15 files changed

+308
-11
lines changed

package-lock.json

Lines changed: 21 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@
9292
"@kui-shell/webpack": "11.5.0-dev-20220801-133414",
9393
"@playwright/test": "^1.24.2",
9494
"@types/debug": "^4.1.7",
95+
"@types/needle": "^2.5.3",
9596
"@types/node": "14.11.8",
9697
"@types/react": "17.0.39",
9798
"@types/react-dom": "17.0.11",
1.28 KB
Loading
1.88 KB
Loading

plugins/plugin-codeflare/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
"@patternfly/react-core": "^4.224.1",
2929
"asciinema-player": "^3.0.1",
3030
"chokidar": "^3.5.3",
31+
"needle": "^3.1.0",
3132
"open": "^8.4.0",
3233
"pretty-bytes": "^6.0.0",
3334
"pretty-ms": "8.0.0",

plugins/plugin-codeflare/src/tray/icons.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
/** Icon set for the tray menu items */
1818

19+
import ray from "@kui-shell/client/icons/png/rayTemplate.png"
1920
import profile from "@kui-shell/client/icons/png/profileTemplate.png"
2021
import bug from "@kui-shell/client/icons/png/bugTemplate.png"
2122
import powerOff from "@kui-shell/client/icons/png/powerOffTemplate.png"
@@ -32,6 +33,7 @@ function iconFor(filepath: string) {
3233
return join(iconHome, filepath)
3334
}
3435

36+
export const rayIcon = iconFor(ray)
3537
export const profileIcon = iconFor(profile)
3638
export const bugIcon = iconFor(bug)
3739
export const powerOffIcon = iconFor(powerOff)

plugins/plugin-codeflare/src/tray/main.ts

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,17 @@ class LiveMenu {
3939
this.render()
4040
}
4141

42-
public async render() {
43-
this.tray.setToolTip(productName)
44-
this.tray.setContextMenu(await buildContextMenu(this.createWindow, this.render.bind(this)))
42+
/** Avoid a flurry of re-renders */
43+
private debounce: null | ReturnType<typeof setTimeout> = null
44+
45+
public render() {
46+
if (this.debounce != null) {
47+
clearTimeout(this.debounce)
48+
}
49+
this.debounce = setTimeout(async () => {
50+
this.tray.setToolTip(productName)
51+
this.tray.setContextMenu(await buildContextMenu(this.createWindow, this.render.bind(this)))
52+
}, 200)
4553
}
4654
}
4755

plugins/plugin-codeflare/src/tray/menus/index.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ export default async function buildContextMenu(
3535
{ type: "separator" },
3636
{ label: `Codeflare ${version}`, enabled: false },
3737
{
38-
label: `CodeFlare Explorer`,
38+
label: `Explore CodeFlare`,
3939
icon: gettingStartedIcon,
4040
click: () => createWindow([], { width: 1200, height: 800 }),
4141
},
Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
/*
2+
* Copyright 2022 The Kubernetes Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import { Profiles } from "madwizard"
18+
import { basename, join } from "path"
19+
import { cli } from "madwizard/dist/fe/cli"
20+
import { MenuItemConstructorOptions } from "electron"
21+
import { CreateWindowFunction } from "@kui-shell/core"
22+
23+
import { rayIcon } from "../../icons"
24+
import UpdateFunction from "../../update"
25+
import ProfileActiveRunWatcher from "../../watchers/profile/active-runs"
26+
27+
import { runMenuItems } from "./runs"
28+
29+
/** active run watcher per profile */
30+
const watchers: Record<string, ProfileActiveRunWatcher> = {}
31+
32+
/** This is the utility function that will open a CodeFlare Dashboard window, pointing to the given local filesystem `logdir` */
33+
async function openDashboard(createWindow: CreateWindowFunction, logdir: string, follow = true) {
34+
await createWindow(["codeflare", "dashboard", follow ? "-f" : "", logdir], {
35+
title: "CodeFlare Dashboard - " + basename(logdir),
36+
})
37+
}
38+
39+
/** This is the click handler for a an active run menu item */
40+
async function openMenuItem(this: CreateWindowFunction, profile: string, runId: string) {
41+
// check to see if we already have a log aggregator's capture
42+
const logdir = join(Profiles.guidebookJobDataPath({ profile }), runId)
43+
if (
44+
await import("fs/promises")
45+
.then((_) => _.access(logdir))
46+
.then(() => true)
47+
.catch(() => false)
48+
) {
49+
// yup, so we can just open up what we are capturing/have
50+
// already captured
51+
return openDashboard(this, logdir)
52+
}
53+
54+
// otherwise, we will need to start of a log aggregator
55+
const guidebook = "ml/ray/aggregator/with-jobid"
56+
const rayAddress = await watchers[profile].rayAddress
57+
if (rayAddress) {
58+
process.env.JOB_ID = runId
59+
process.env.RAY_ADDRESS = rayAddress
60+
process.env.NO_WAIT = "true" // don't wait for job termination
61+
process.env.QUIET_CONSOLE = "true" // don't tee logs to the console
62+
const resp = await cli(["madwizard", "guide", guidebook], undefined, {
63+
profile,
64+
clean: false /* don't kill the port-forward subprocess! we'll manage that */,
65+
interactive: false,
66+
store: process.env.GUIDEBOOK_STORE,
67+
})
68+
69+
if (resp) {
70+
if (!resp.env.LOGDIR_STAGE) {
71+
console.error("Failed to attach to job", runId)
72+
} else {
73+
await openDashboard(this, resp.env.LOGDIR_STAGE)
74+
75+
// now the window has closed, so we can clean up any
76+
// subprocesses spawned by the guidebook
77+
if (typeof resp.cleanExit === "function") {
78+
resp.cleanExit()
79+
process.on("exit", () => resp.cleanExit())
80+
}
81+
}
82+
}
83+
}
84+
}
85+
86+
/** @return menu items that allow attaching to an active run in the given `profileName` */
87+
export default function activeRuns(
88+
profile: string,
89+
createWindow: CreateWindowFunction,
90+
updateFn: UpdateFunction
91+
): MenuItemConstructorOptions[] {
92+
if (!watchers[profile]) {
93+
watchers[profile] = new ProfileActiveRunWatcher(updateFn, profile)
94+
}
95+
96+
const runs = watchers[profile].runs
97+
if (runs.length > 0) {
98+
return runMenuItems(
99+
profile,
100+
openMenuItem.bind(createWindow),
101+
runs.map((_) => Object.assign(_, { icon: rayIcon }))
102+
)
103+
} else {
104+
return [{ label: "No active runs", enabled: false }]
105+
}
106+
}

plugins/plugin-codeflare/src/tray/menus/profiles/dashboards/index.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,20 @@ export default async function dashboards(
2828
): Promise<MenuItemConstructorOptions[]> {
2929
const mlflow = { name: "MLFlow", portEnv: "MLFLOW_PORT" }
3030
const tensorboard = { name: "Tensorboard", portEnv: "TENSORBOARD_PORT" }
31+
const pytorchProfiler = {
32+
name: "PyTorch Profiler",
33+
nameForGuidebook: "Tensorboard",
34+
portEnv: "TENSORBOARD_PORT",
35+
path: "#pytorch_profiler",
36+
}
3137

3238
return [
3339
{ label: "CodeFlare", submenu: await codeflare(profile, createWindow, updateFn) },
3440
{ label: "MLFlow", click: () => import("./open").then((_) => _.default(mlflow, profile, createWindow)) },
3541
{ label: "Tensorboard", click: () => import("./open").then((_) => _.default(tensorboard, profile, createWindow)) },
42+
{
43+
label: "PyTorch Profiler",
44+
click: () => import("./open").then((_) => _.default(pytorchProfiler, profile, createWindow)),
45+
},
3646
]
3747
}

0 commit comments

Comments
 (0)