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.

Author: starpit

package-lock.json

@@ -92,6 +92,7 @@
     "@kui-shell/webpack": "11.5.0-dev-20220801-133414",
     "@playwright/test": "^1.24.2",
     "@types/debug": "^4.1.7",
+    "@types/needle": "^2.5.3",
     "@types/node": "14.11.8",
     "@types/react": "17.0.39",
     "@types/react-dom": "17.0.11",

package.json

@@ -28,6 +28,7 @@
     "@patternfly/react-core": "^4.224.1",
     "asciinema-player": "^3.0.1",
     "chokidar": "^3.5.3",
+    "needle": "^3.1.0",
     "open": "^8.4.0",
     "pretty-bytes": "^6.0.0",
     "pretty-ms": "8.0.0",

plugins/plugin-client-default/icons/png/rayTemplate.png

@@ -16,6 +16,7 @@
 
 /** Icon set for the tray menu items */
 
+import ray from "@kui-shell/client/icons/png/rayTemplate.png"
 import profile from "@kui-shell/client/icons/png/profileTemplate.png"
 import bug from "@kui-shell/client/icons/png/bugTemplate.png"
 import powerOff from "@kui-shell/client/icons/png/powerOffTemplate.png"
@@ -32,6 +33,7 @@ function iconFor(filepath: string) {
   return join(iconHome, filepath)
 }
 
+export const rayIcon = iconFor(ray)
 export const profileIcon = iconFor(profile)
 export const bugIcon = iconFor(bug)
 export const powerOffIcon = iconFor(powerOff)

plugins/plugin-client-default/icons/png/rayTemplate@2x.png

@@ -39,9 +39,17 @@ class LiveMenu {
     this.render()
   }
 
-  public async render() {
-    this.tray.setToolTip(productName)
-    this.tray.setContextMenu(await buildContextMenu(this.createWindow, this.render.bind(this)))
+  /** Avoid a flurry of re-renders */
+  private debounce: null | ReturnType<typeof setTimeout> = null
+
+  public render() {
+    if (this.debounce != null) {
+      clearTimeout(this.debounce)
+    }
+    this.debounce = setTimeout(async () => {
+      this.tray.setToolTip(productName)
+      this.tray.setContextMenu(await buildContextMenu(this.createWindow, this.render.bind(this)))
+    }, 200)
   }
 }
 

plugins/plugin-codeflare/package.json

@@ -35,7 +35,7 @@ export default async function buildContextMenu(
     { type: "separator" },
     { label: `Codeflare ${version}`, enabled: false },
     {
-      label: `CodeFlare Explorer`,
+      label: `Explore CodeFlare`,
       icon: gettingStartedIcon,
       click: () => createWindow([], { width: 1200, height: 800 }),
     },

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

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2022 The Kubernetes Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { Profiles } from "madwizard"
+import { basename, join } from "path"
+import { cli } from "madwizard/dist/fe/cli"
+import { MenuItemConstructorOptions } from "electron"
+import { CreateWindowFunction } from "@kui-shell/core"
+
+import { rayIcon } from "../../icons"
+import UpdateFunction from "../../update"
+import ProfileActiveRunWatcher from "../../watchers/profile/active-runs"
+
+import { runMenuItems } from "./runs"
+
+/** active run watcher per profile */
+const watchers: Record<string, ProfileActiveRunWatcher> = {}
+
+/** This is the utility function that will open a CodeFlare Dashboard window, pointing to the given local filesystem `logdir` */
+async function openDashboard(createWindow: CreateWindowFunction, logdir: string, follow = true) {
+  await createWindow(["codeflare", "dashboard", follow ? "-f" : "", logdir], {
+    title: "CodeFlare Dashboard - " + basename(logdir),
+  })
+}
+
+/** This is the click handler for a an active run menu item */
+async function openMenuItem(this: CreateWindowFunction, profile: string, runId: string) {
+  // check to see if we already have a log aggregator's capture
+  const logdir = join(Profiles.guidebookJobDataPath({ profile }), runId)
+  if (
+    await import("fs/promises")
+      .then((_) => _.access(logdir))
+      .then(() => true)
+      .catch(() => false)
+  ) {
+    // yup, so we can just open up what we are capturing/have
+    // already captured
+    return openDashboard(this, logdir)
+  }
+
+  // otherwise, we will need to start of a log aggregator
+  const guidebook = "ml/ray/aggregator/with-jobid"
+  const rayAddress = await watchers[profile].rayAddress
+  if (rayAddress) {
+    process.env.JOB_ID = runId
+    process.env.RAY_ADDRESS = rayAddress
+    process.env.NO_WAIT = "true" // don't wait for job termination
+    process.env.QUIET_CONSOLE = "true" // don't tee logs to the console
+    const resp = await cli(["madwizard", "guide", guidebook], undefined, {
+      profile,
+      clean: false /* don't kill the port-forward subprocess! we'll manage that */,
+      interactive: false,
+      store: process.env.GUIDEBOOK_STORE,
+    })
+
+    if (resp) {
+      if (!resp.env.LOGDIR_STAGE) {
+        console.error("Failed to attach to job", runId)
+      } else {
+        await openDashboard(this, resp.env.LOGDIR_STAGE)
+
+        // now the window has closed, so we can clean up any
+        // subprocesses spawned by the guidebook
+        if (typeof resp.cleanExit === "function") {
+          resp.cleanExit()
+          process.on("exit", () => resp.cleanExit())
+        }
+      }
+    }
+  }
+}
+
+/** @return menu items that allow attaching to an active run in the given `profileName` */
+export default function activeRuns(
+  profile: string,
+  createWindow: CreateWindowFunction,
+  updateFn: UpdateFunction
+): MenuItemConstructorOptions[] {
+  if (!watchers[profile]) {
+    watchers[profile] = new ProfileActiveRunWatcher(updateFn, profile)
+  }
+
+  const runs = watchers[profile].runs
+  if (runs.length > 0) {
+    return runMenuItems(
+      profile,
+      openMenuItem.bind(createWindow),
+      runs.map((_) => Object.assign(_, { icon: rayIcon }))
+    )
+  } else {
+    return [{ label: "No active runs", enabled: false }]
+  }
+}

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

@@ -28,10 +28,20 @@ export default async function dashboards(
 ): Promise<MenuItemConstructorOptions[]> {
   const mlflow = { name: "MLFlow", portEnv: "MLFLOW_PORT" }
   const tensorboard = { name: "Tensorboard", portEnv: "TENSORBOARD_PORT" }
+  const pytorchProfiler = {
+    name: "PyTorch Profiler",
+    nameForGuidebook: "Tensorboard",
+    portEnv: "TENSORBOARD_PORT",
+    path: "#pytorch_profiler",
+  }
 
   return [
     { label: "CodeFlare", submenu: await codeflare(profile, createWindow, updateFn) },
     { label: "MLFlow", click: () => import("./open").then((_) => _.default(mlflow, profile, createWindow)) },
     { label: "Tensorboard", click: () => import("./open").then((_) => _.default(tensorboard, profile, createWindow)) },
+    {
+      label: "PyTorch Profiler",
+      click: () => import("./open").then((_) => _.default(pytorchProfiler, profile, createWindow)),
+    },
   ]
 }