gpu-lock

1import { atom, read, update } from 'claude-code'
2import type { Hook, Register } from 'claude-code'
3
4import { runCommand } from '../src/commands'
5import { createGuard } from '../src/guard'
6import type { Io } from '../src/io'
7import { readOptions } from '../src/options'
8import type { Mode } from '../src/options'
9import { createRuntime } from '../src/runtime'
10import type { Runtime } from '../src/runtime'
11import type { Charts } from '../src/history'
12import { bandView, paneView, withBelow } from '../src/ui'
13import { bandModel, paneModel } from '../src/view'
14
15const snapshotAtom = atom({ plugin: 'gpu-lock', key: 'snapshot' } as const, null)
16const logAtom = atom({ plugin: 'gpu-lock', key: 'log' } as const, [])
17const waitingAtom = atom({ plugin: 'gpu-lock', key: 'waiting' } as const, null)
18const historyAtom = atom({ plugin: 'gpu-lock', key: 'history' } as const, [])
19const enabledAtom = atom({ plugin: 'gpu-lock', key: 'enabled' } as const, true)
20
21const PANE_ID = 'gpu-lock'
22const DEFAULT_COLUMNS = 100
23const PANE_CHROME_ROWS = 6
24// Pane padding on both sides, the percent labels and the space after them.
25const CHART_SIDE_COLUMNS = 7
26
27// Rows the charts take: each card's title, rows and summary, the axis, the legend, and the gap above.
28const chartHeight = (charts: Charts | null): number =>
29  charts === null ? 0 : charts.cards.reduce((n, c) => n + c.rows.length + 2, 0) + 1 + charts.legend.length + 1
30
31const TAIL =
32  'To say how much GPU memory a command needs, end it with `# gpu 8g`. To mark a command that does not use the GPU, end it with `# nogpu`.'
33
34const rulesFor = (mode: Mode, percent: number): string =>
35  mode === 'shared'
36    ? `GPU lock (gpu-lock plugin): this machine shares the GPU between CUDA jobs, up to ${percent}% of its memory together. A Bash command that uses the GPU is checked against that limit first. If it is refused, wait for a running job to end, then try again; do not work around the limit. ${TAIL}`
37    : `GPU lock (gpu-lock plugin): this machine runs one CUDA job at a time. A Bash command that uses the GPU is run under a lock and checked against free GPU memory first. If it is refused, wait for the running job to end; do not start a second one. ${TAIL}`
38
39let io: Io | null = null
40let rt: Runtime | null = null
41
42const startHook: Hook<'session.start'> = async ($, e, next) => {
43    io = {
44      process: { run: (argv, init) => $.process.run(argv, init) },
45      fs: { read: path => $.fs.read(path), exists: path => $.fs.exists(path) },
46      clock: { now: () => $.clock.now(), every: (ms, fn) => $.clock.every(ms, fn) },
47      session: { id: () => $.session.id(), cwd: () => $.session.cwd() },
48      store: { get: key => $.store.get(key), set: (key, value) => $.store.set(key, value) },
49      ui: { toast: text => $.ui.toast(text), status: text => $.ui.status(text) },
50      state: {
51        snapshot: { get: () => read($, snapshotAtom), set: async fn => void (await update($, snapshotAtom, fn)) },
52        log: { get: () => read($, logAtom), set: async fn => void (await update($, logAtom, fn)) },
53        waiting: { get: () => read($, waitingAtom), set: async fn => void (await update($, waitingAtom, fn)) },
54        history: { get: () => read($, historyAtom), set: async fn => void (await update($, historyAtom, fn)) },
55        enabled: { get: () => read($, enabledAtom), set: async fn => void (await update($, enabledAtom, fn)) },
56      },
57    }
58    try {
59      await $.command.register({
60        name: 'gpu',
61        description: 'GPU lock: open or close the pane, status, on, off, check <command>',
62        argumentHint: '[close|status|on|off|check <command>]',
63        immediate: true,
64      })
65    } catch {
66      $.ui.toast('gpu-lock: the /gpu command could not be registered; the lock still works')
67    }
68    await rt?.start(io, e.isInteractive)
69    return next(e)
70}
71
72export const register: Register = (on, options) => {
73  const config = readOptions(options)
74  const runtime = createRuntime(config)
75  const guard = createGuard(config, runtime)
76  rt = runtime
77
78  on('session.start', startHook)
79
80  on('session.end', async (_$, e, next) => {
81    if (io !== null && e.reason === 'clear') await runtime.resetSession(io)
82    else runtime.stopPolling()
83    return next(e)
84  })
85
86  on('tool.call', { tool: 'Bash' }, (_$, e, next) => {
87    if (io === null) return next(e)
88    return guard(io, {
89      command: e.command,
90      agentId: e.agentId,
91      background: e.run_in_background === true,
92      run: command => next({ ...e, command }),
93    })
94  })
95
96  on('prompt.compose', async ($, e, next) => {
97    const composed = await next(e)
98    if (!(await read($, enabledAtom))) return composed
99    const text = rulesFor(config.mode, config.maxVramPercent)
100    return { sections: [...composed.sections, { id: 'gpu-lock:rules', text, scope: 'session' }] }
101  })
102
103  on('command.run', { command: 'gpu' }, async ($, e) => {
104    if (io === null) return { text: 'gpu-lock has not started yet.' }
105    const args = e.args ?? ''
106    const word = (args.trim().split(/\s+/)[0] ?? '').toLowerCase()
107    const enabled = await read($, enabledAtom)
108    const wantsFresh = enabled && (word === '' || word === 'status')
109    const isPaneOpen = (await $.ui.panes()).some(pane => pane.id === PANE_ID)
110    const out = runCommand(args, {
111      config,
112      snapshot: wantsFresh ? await runtime.refresh(io) : await read($, snapshotAtom),
113      enabled,
114      isPaneOpen,
115      now: await $.clock.now(),
116      session: await $.session.id(),
117    })
118    if (out.action === 'on') await runtime.setEnabled(io, true)
119    if (out.action === 'off') await runtime.setEnabled(io, false)
120    if (out.action === 'pane') await $.ui.open({ id: PANE_ID, title: 'GPU lock', closeOnEscape: true })
121    if (out.action === 'close') await $.ui.close({ id: PANE_ID })
122    return { text: out.text }
123  })
124
125  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
126    if (e.props.hasSurvey) return next(e)
127    const model = bandModel({
128      snapshot: await read($, snapshotAtom),
129      waiting: await read($, waitingAtom),
130      enabled: await read($, enabledAtom),
131      mode: config.band,
132      now: await $.clock.now(),
133      session: await $.session.id(),
134      width: e.viewport?.columns ?? DEFAULT_COLUMNS,
135    })
136    const below = await next(e)
137    if (model === null) return below
138    const els = $.ui.resolve(e)
139    return withBelow(els, below, bandView(els, model))
140  })
141
142  on('ui.render', { component: 'Pane', requestId: PANE_ID }, async ($, e) => {
143    const enabled = await read($, enabledAtom)
144    const model = paneModel({
145      snapshot: await read($, snapshotAtom),
146      waiting: await read($, waitingAtom),
147      enabled,
148      log: await read($, logAtom),
149      bad: config.bad,
150      now: await $.clock.now(),
151      session: await $.session.id(),
152      maxPercent: config.mode === 'shared' ? config.maxVramPercent : null,
153      history: await read($, historyAtom),
154      columns: (e.props.bodyColumns ?? e.viewport?.columns ?? DEFAULT_COLUMNS) - CHART_SIDE_COLUMNS,
155    })
156    const room = Math.max(1, (e.viewport?.rows ?? 24) - PANE_CHROME_ROWS - chartHeight(model.chart))
157    return paneView($.ui.resolve(e), { ...model, log: model.log.slice(0, room) }, {
158      onRefresh: () => {
159        if (io !== null) void runtime.refresh(io)
160      },
161      onToggle: () => {
162        if (io !== null) void runtime.setEnabled(io, !enabled)
163      },
164      onClose: () => $.ui.close({ id: PANE_ID }),
165    })
166  })
167}
168

1import type { Reading, Snapshot } from '../types'
2import { classify } from './classify'
3import { gib } from './format'
4import { waitersOf } from './lock'
5import { describeHolder, describeHolders } from './messages'
6import type { Config } from './options'
7
8export type Context = {
9  config: Config
10  snapshot: Snapshot | null
11  enabled: boolean
12  isPaneOpen: boolean
13  now: number
14  session: string
15}
16
17export type Outcome = { text: string; action?: 'pane' | 'close' | 'on' | 'off' }
18
19const PANE_CLOSED = 'GPU lock pane closed.'
20
21const HELP = [
22  '/gpu                    open the GPU pane, or close it if it is open',
23  '/gpu close              close the GPU pane',
24  '/gpu status             print the lock and memory state',
25  '/gpu on | off           turn the lock on or off',
26  '/gpu check <command>    say whether a command counts as a GPU job',
27  'Mark a command by ending it with # gpu 8g (needs 8 GiB), # gpu (force the lock) or # nogpu (skip it).',
28].join('\n')
29
30const apps = (gpu: Reading): string => {
31  if (gpu.kind !== 'ok') return ''
32  const users = gpu.apps
33    .filter(a => a.usedMiB > 0)
34    .sort((a, b) => b.usedMiB - a.usedMiB)
35    .map(a => `${a.name} (${gib(a.usedMiB)})`)
36  return users.length === 0 ? '' : `Other GPU users: ${users.join(', ')}.`
37}
38
39const gpuLines = (gpu: Reading): string[] => {
40  if (gpu.kind === 'none') return ['GPU: no NVIDIA GPU or nvidia-smi found.']
41  if (gpu.kind === 'error') return [`GPU: nvidia-smi failed: ${gpu.message}.`]
42  return gpu.gpus.map(g => {
43    const parts = [
44      g.name,
45      `${(g.usedMiB / 1024).toFixed(1)} of ${gib(g.totalMiB)} used`,
46      g.util === null ? '' : `${g.util}% busy`,
47      g.tempC === null ? '' : `${g.tempC}°C`,
48    ]
49    return `${gpu.gpus.length === 1 ? 'GPU' : `GPU ${g.index}`}: ${parts.filter(p => p !== '').join(', ')}.`
50  })
51}
52
53const lockLine = (s: Snapshot, c: Context): string => {
54  if (s.lock.kind === 'free') return 'Lock: free.'
55  if (s.lock.kind === 'unavailable') return `Lock: unavailable (${s.lock.reason}).`
56  if (s.lock.kind === 'shared') {
57    return `Lock: shared by ${describeHolders(s.lock.holders, c.now, c.session)}.`
58  }
59  return `Lock: held by ${describeHolder(s.lock.holder, c.now, c.session)}.`
60}
61
62const queueLine = (s: Snapshot, c: Context): string => {
63  const waiters = waitersOf(s)
64  return waiters.length === 0 ? '' : `Waiting: ${describeHolders(waiters, c.now, c.session)}.`
65}
66
67const rulesLine = (config: Config): string => {
68  const mode = config.mode === 'shared' ? `shared up to ${config.maxVramPercent}% of memory` : 'one job at a time'
69  return `Rules: ${mode}, wait ${config.waitSeconds}s, least free ${gib(config.minFreeMiB)}, band ${config.band}, lock file ${config.lockPath}.`
70}
71
72const status = (c: Context): string => {
73  const { config, snapshot } = c
74  const lines = [`gpu-lock is ${c.enabled ? 'on' : 'off'}.`]
75  if (snapshot === null) {
76    lines.push('Lock: not read yet.')
77  } else {
78    lines.push(lockLine(snapshot, c), queueLine(snapshot, c), ...gpuLines(snapshot.gpu), apps(snapshot.gpu))
79  }
80  lines.push(rulesLine(config))
81  if (config.bad.length > 0) lines.push(`Ignored patterns (not valid): ${config.bad.join(', ')}`)
82  return lines.filter(l => l !== '').join('\n')
83}
84
85const check = (command: string, c: Context): string => {
86  if (command === '') return 'Usage: /gpu check <command>'
87  const v = classify(command, { extra: c.config.extra, skip: c.config.skip, lockPath: c.config.lockPath })
88  if (!v.gpu) return 'That does not count as a GPU job. End a command with `# gpu` to force the lock.'
89  const floor = c.config.minFreeMiB
90  const need =
91    v.needMiB !== null
92      ? `It needs ${gib(v.needMiB)} free.`
93      : floor > 0
94        ? `It needs ${gib(floor)} free (the least-free rule).`
95        : 'It has no memory floor.'
96  const devices = v.devices === null ? '' : ` Devices: ${v.devices.join(', ')}.`
97  return `That counts as a GPU job (${v.why}). ${need}${devices}`
98}
99
100export const runCommand = (args: string, c: Context): Outcome => {
101  const trimmed = args.trim()
102  const [word = '', ...rest] = trimmed.split(/\s+/)
103  const sub = word.toLowerCase()
104  if (sub === '') {
105    return c.isPaneOpen ? { text: PANE_CLOSED, action: 'close' } : { text: 'GPU lock pane opened.', action: 'pane' }
106  }
107  if (sub === 'close') return c.isPaneOpen ? { text: PANE_CLOSED, action: 'close' } : { text: 'The GPU lock pane is not open.' }
108  if (sub === 'on') return { text: 'gpu-lock is on.', action: 'on' }
109  if (sub === 'off') return { text: 'gpu-lock is off. GPU commands run unchecked.', action: 'off' }
110  if (sub === 'status') return { text: status(c) }
111  if (sub === 'check') return { text: check(trimmed.slice(word.length).trim(), c) }
112  if (sub === 'help') return { text: HELP }
113  return { text: `Unknown option "${word}".\n${HELP}` }
114}
115

1import type { ToolCallResult } from 'claude-code'
2
3import { classify } from './classify'
4import type { Verdict } from './classify'
5import { clean } from './format'
6import type { Io } from './io'
7import { gateOf, jobsDirOf, ownerPathOf, waitingDirOf } from './lock'
8import { acquiredNote, refusedBusy, refusedMemory, refusedQueue, refusedShare, whoLabel } from './messages'
9import type { Config } from './options'
10import { decide } from './policy'
11import type { Runtime } from './runtime'
12import { ensureLock, readGpu, readLock } from './sys'
13import { wrap } from './wrap'
14import type { ShareSpec } from './wrap'
15
16// The wrapper's own refusal: a message of ours that says the command did not run.
17const isLost = (text: string): boolean => text.includes('gpu-lock: another job took the GPU') && text.includes('so this command did not run')
18const LOG_CMD_CHARS = 120
19const TOAST_CMD_CHARS = 50
20
21export type Guarded = {
22  command: string
23  agentId: string | undefined
24  background: boolean
25  run: (command: string) => Promise<ToolCallResult>
26}
27
28const shareOf = (config: Config, verdict: Extract<Verdict, { gpu: true }>): ShareSpec | null =>
29  config.mode === 'shared'
30    ? {
31        jobsDir: jobsDirOf(config.lockPath),
32        gatePath: gateOf(config.lockPath),
33        maxPercent: config.maxVramPercent,
34        needMiB: verdict.needMiB ?? config.minFreeMiB,
35        devices: verdict.devices,
36      }
37    : null
38
39export const createGuard = (config: Config, rt: Runtime) => {
40  const options = { extra: config.extra, skip: config.skip, lockPath: config.lockPath }
41
42  const refuse = async (io: Io, command: string, now: number, note: string, toast: string) => {
43    await rt.record(io, { at: now, kind: 'refused', cmd: clean(command, LOG_CMD_CHARS), ms: null, note })
44    io.ui.toast(toast)
45  }
46
47  const lockBroken = (io: Io, reason: string) => {
48    rt.setFault(reason)
49    if (rt.takeFaultNotice()) io.ui.toast(`gpu-lock is off here: ${reason}`)
50  }
51
52  const runWrapped = async (io: Io, g: Guarded, wrapped: string, queued: boolean, now: number) => {
53    if (queued) await io.state.waiting.set(() => ({ cmd: clean(g.command, LOG_CMD_CHARS), since: now }))
54    try {
55      return await g.run(wrapped)
56    } finally {
57      if (queued) await io.state.waiting.set(() => null)
58    }
59  }
60
61  return async (io: Io, g: Guarded): Promise<ToolCallResult> => {
62    if (!(await io.state.enabled.get())) return g.run(g.command)
63    const verdict = classify(g.command, options)
64    if (!verdict.gpu) return g.run(g.command)
65
66    const failure = await ensureLock(io, config.lockPath)
67    if (failure !== null) {
68      lockBroken(io, failure)
69      return g.run(g.command)
70    }
71    rt.setFault(null)
72
73    const [locking, gpu, now, session, cwd] = await Promise.all([
74      readLock(io, config.lockPath),
75      readGpu(io),
76      io.clock.now(),
77      io.session.id(),
78      io.session.cwd(),
79    ])
80    const { lock } = locking
81    await rt.publish(io, { at: now, lock, gpu, waiters: locking.waiters })
82    const d = decide({
83      verdict,
84      lock,
85      gpu,
86      minFreeMiB: config.minFreeMiB,
87      waitSeconds: config.waitSeconds,
88      mode: config.mode,
89      maxPercent: config.maxVramPercent,
90      now,
91      waiters: locking.waiters,
92    })
93    const short = clean(g.command, TOAST_CMD_CHARS)
94
95    if (d.kind === 'pass') return g.run(g.command)
96    if (d.kind === 'refuse-busy') {
97      const who = d.holders.map(h => whoLabel(h, session)).join(', ') || whoLabel(null, session)
98      await refuse(io, g.command, now, `locked by ${who}`, `GPU locked · refused: ${short}`)
99      return { deny: refusedBusy(d.holders, now, session, gpu) }
100    }
101    if (d.kind === 'refuse-queue') {
102      await refuse(io, g.command, now, `${d.waiters.length} waiting ahead`, `GPU queue · refused: ${short}`)
103      return { deny: refusedQueue(d.waiters, now, session) }
104    }
105    if (d.kind === 'refuse-memory') {
106      await refuse(io, g.command, now, 'not enough free memory', `GPU memory low · refused: ${short}`)
107      return { deny: refusedMemory(d.needMiB, d.freeMiB, d.totalMiB, gpu) }
108    }
109    if (d.kind === 'refuse-share') {
110      await refuse(io, g.command, now, `over the ${d.percent}% limit`, `GPU memory at ${d.percent}% · refused: ${short}`)
111      return { deny: refusedShare(d, d.holders, now, session) }
112    }
113
114    const share = shareOf(config, verdict)
115    const wrapped = wrap(g.command, {
116      lockPath: config.lockPath,
117      ownerPath: ownerPathOf(config.lockPath),
118      waitingDir: waitingDirOf(config.lockPath),
119      waitSeconds: d.waits ? config.waitSeconds : 0,
120      owner: { session, agent: g.agentId ?? null, cmd: g.command, cwd, why: verdict.why, needMiB: verdict.needMiB },
121      share,
122    })
123    const ran = await runWrapped(io, g, wrapped, d.waits, now)
124    if (ran.deny !== undefined) return ran
125
126    const ended = await io.clock.now()
127    if (ran.isError === true && isLost(ran.text ?? '')) {
128      const why = share === null ? 'another job took the lock first' : 'another job took the GPU memory first'
129      await refuse(io, g.command, ended, why, `GPU locked · refused: ${short} (${why})`)
130      rt.refreshQuietly(io)
131      return ran
132    }
133    await rt.record(io, {
134      at: ended,
135      kind: ran.isError === true ? 'failed' : 'ran',
136      cmd: clean(g.command, LOG_CMD_CHARS),
137      ms: g.background ? null : ended - now,
138      note: g.background ? 'background' : '',
139    })
140    if (g.background) rt.backgroundStarted(ended)
141    rt.refreshQuietly(io)
142    const note = acquiredNote({
143      freeMiB: d.freeMiB,
144      totalMiB: d.totalMiB,
145      reading: gpu,
146      background: g.background,
147      ...(share === null ? {} : { sharedPercent: share.maxPercent }),
148    })
149    return { ...ran, context: [...(ran.context ?? []), note] } as ToolCallResult
150  }
151}
152

1import type { Entry, Sample, Snapshot, Waiting } from '../types'
2
3export type Cell<T> = {
4  get: () => Promise<T>
5  set: (fn: (value: T) => T) => Promise<void>
6}
7
8export type Io = {
9  process: {
10    run: (
11      argv: readonly string[],
12      init?: { timeoutMs?: number },
13    ) => Promise<{ exitCode: number; stdout: string; stderr: string }>
14  }
15  fs: {
16    read: (path: string) => Promise<string>
17    exists: (path: string) => Promise<boolean>
18  }
19  clock: {
20    now: () => Promise<number>
21    every: (ms: number, fn: () => void) => { cancel: () => void }
22  }
23  session: {
24    id: () => Promise<string>
25    cwd: () => Promise<string>
26  }
27  store: {
28    get: (key: string) => Promise<unknown>
29    set: (key: string, value: unknown) => Promise<void>
30  }
31  ui: {
32    toast: (text: string) => void
33    status: (text: string | undefined) => void
34  }
35  state: {
36    snapshot: Cell<Snapshot | null>
37    log: Cell<Entry[]>
38    waiting: Cell<Waiting | null>
39    history: Cell<Sample[]>
40    enabled: Cell<boolean>
41  }
42}
43

1import { compilePatterns } from './classify'
2
3export const DEFAULT_LOCK_PATH = '/tmp/gpu-lock/lock'
4
5export type Band = 'auto' | 'always' | 'never'
6
7export type Mode = 'exclusive' | 'shared'
8
9export type Config = {
10  mode: Mode
11  maxVramPercent: number
12  waitSeconds: number
13  minFreeMiB: number
14  band: Band
15  lockPath: string
16  extra: RegExp[]
17  skip: RegExp[]
18  bad: string[]
19}
20
21const whole = (value: unknown, min: number, max: number, fallback: number): number => {
22  if (typeof value !== 'number' || !Number.isFinite(value)) return fallback
23  return Math.min(max, Math.max(min, Math.floor(value)))
24}
25
26const lines = (value: unknown): string => (typeof value === 'string' ? value : '')
27
28const isLockPath = (value: unknown): value is string =>
29  typeof value === 'string' && /^\/[^\0\r\n]*[^/\0\r\n]$/.test(value)
30
31export const readOptions = (o: Readonly<Record<string, unknown>>): Config => {
32  const extra = compilePatterns(lines(o['extraPatterns']))
33  const skip = compilePatterns(lines(o['skipPatterns']))
34  const band = o['band']
35  return {
36    mode: o['mode'] === 'shared' ? 'shared' : 'exclusive',
37    maxVramPercent: whole(o['maxVramPercent'], 10, 100, 75),
38    waitSeconds: whole(o['waitSeconds'], 0, 600, 0),
39    minFreeMiB: whole(o['minFreeMiB'], 0, 1048576, 1024),
40    band: band === 'always' || band === 'never' ? band : 'auto',
41    lockPath: isLockPath(o['lockPath']) ? o['lockPath'] : DEFAULT_LOCK_PATH,
42    extra: extra.regexes,
43    skip: skip.regexes,
44    bad: [...extra.bad, ...skip.bad],
45  }
46}
47

1import type { Entry, Snapshot } from '../types'
2import { appendSample, readSamples } from './histfile'
3import { MIN_GAP_MS, needsTree, sampleOf } from './history'
4import type { Io } from './io'
5import type { Config } from './options'
6import { newMemo, publishSnapshot } from './publish'
7import { PS_ARGV, parseParents } from './procs'
8import { ensureLock, takeSnapshot } from './sys'
9
10const POLL_MS = 3000
11const PS_TIMEOUT_MS = 3000
12const MAX_LOG = 20
13
14export type Runtime = {
15  refresh: (io: Io) => Promise<Snapshot>
16  publish: (io: Io, snap: Snapshot) => Promise<void>
17  record: (io: Io, entry: Entry) => Promise<void>
18  start: (io: Io, isInteractive: boolean) => Promise<void>
19  stopPolling: () => void
20  setEnabled: (io: Io, enabled: boolean) => Promise<void>
21  resetSession: (io: Io) => Promise<void>
22  refreshQuietly: (io: Io) => void
23  setFault: (reason: string | null) => void
24  takeFaultNotice: () => boolean
25  backgroundStarted: (at: number) => void
26}
27
28export const createRuntime = (config: Config): Runtime => {
29  const memo = newMemo()
30  let timer: { cancel: () => void } | null = null
31  let isReading = false
32  let fault: string | null = null
33  let isFaultShown = false
34
35  const publish = (io: Io, snap: Snapshot) => publishSnapshot(io, config, memo, snap)
36
37  const record = async (io: Io, entry: Entry) => {
38    await io.state.log.set(list => [...list, entry].slice(-MAX_LOG))
39  }
40
41  const readParents = async (io: Io): Promise<Map<number, number> | null> => {
42    try {
43      const r = await io.process.run(PS_ARGV, { timeoutMs: PS_TIMEOUT_MS })
44      return r.exitCode === 0 ? parseParents(r.stdout) : null
45    } catch {
46      return null
47    }
48  }
49
50  // The file every session shares is the history: add this reading when the newest is old enough, then read it all.
51  const trackHistory = async (io: Io, snap: Snapshot) => {
52    try {
53      const samples = await readSamples(io, config.lockPath, snap.at)
54      const newest = samples[samples.length - 1]?.at ?? 0
55      if (snap.at - newest >= MIN_GAP_MS) {
56        const sample = sampleOf(snap, needsTree(snap) ? await readParents(io) : null)
57        if (sample !== null && (await appendSample(io, config.lockPath, sample))) samples.push(sample)
58      }
59      await io.state.history.set(() => samples)
60    } catch {
61      return
62    }
63  }
64
65  const refresh = async (io: Io): Promise<Snapshot> => {
66    const now = await io.clock.now()
67    if (fault !== null) fault = await ensureLock(io, config.lockPath)
68    const raw = await takeSnapshot(io, config.lockPath, now)
69    const snap: Snapshot = fault === null ? raw : { ...raw, lock: { kind: 'unavailable', reason: fault } }
70    await trackHistory(io, snap)
71    await publish(io, snap)
72    return snap
73  }
74
75  const refreshQuietly = (io: Io) => {
76    if (isReading) return
77    isReading = true
78    refresh(io)
79      .catch(() => undefined)
80      .finally(() => {
81        isReading = false
82      })
83  }
84
85  const startPolling = (io: Io) => {
86    timer?.cancel()
87    timer = io.clock.every(POLL_MS, () => refreshQuietly(io))
88    refreshQuietly(io)
89  }
90
91  const stopPolling = () => {
92    timer?.cancel()
93    timer = null
94  }
95
96  const start = async (io: Io, isInteractive: boolean) => {
97    const enabled = (await io.store.get('enabled')) !== false
98    await io.state.enabled.set(() => enabled)
99    if (!enabled) return
100    fault = await ensureLock(io, config.lockPath)
101    if (isInteractive) startPolling(io)
102  }
103
104  const setEnabled = async (io: Io, enabled: boolean) => {
105    await io.store.set('enabled', enabled)
106    await io.state.enabled.set(() => enabled)
107    if (enabled) {
108      fault = await ensureLock(io, config.lockPath)
109      startPolling(io)
110      return
111    }
112    stopPolling()
113    await io.state.waiting.set(() => null)
114    memo.lastStatus = undefined
115    io.ui.status(undefined)
116  }
117
118  const resetSession = async (io: Io) => {
119    await io.state.log.set(() => [])
120    await io.state.waiting.set(() => null)
121  }
122
123  return {
124    refresh,
125    publish,
126    record,
127    start,
128    stopPolling,
129    setEnabled,
130    resetSession,
131    refreshQuietly,
132    setFault: reason => {
133      fault = reason
134    },
135    takeFaultNotice: () => {
136      const first = !isFaultShown
137      isFaultShown = true
138      return first
139    },
140    backgroundStarted: at => {
141      memo.backgroundSince = at
142    },
143  }
144}
145

1import type { CardSample, Holder, LockState, Sample, Snapshot } from '../types'
2import { clean, gib, ratio, shortId } from './format'
3import { jobOf } from './procs'
4
5export type { Sample } from '../types'
6
7export const SPAN_MS = 600_000
8export const MAX_SAMPLES = 400
9export const PALETTE_SIZE = 5
10// Samples closer together than this are one: sessions that poll at once write one between them.
11export const MIN_GAP_MS = 2000
12
13const MIN_COLUMNS = 8
14const LEVELS = 8
15const BLOCKS = ' ▁▂▃▄▅▆▇█'
16const CAP_MARK = '╌'
17const LABEL_WIDTH = 4
18const MAX_CARDS = 16
19const MAX_JOBS = 32
20const SESSION_CHARS = 64
21const CMD_CHARS = 60
22const FUTURE_MS = 60_000
23const JOB_KEY = /^\d{1,10}$/
24
25const holdersOf = (lock: LockState): Holder[] => {
26  if (lock.kind === 'shared') return lock.holders
27  return lock.kind === 'held' && lock.holder !== null ? [lock.holder] : []
28}
29
30// A process counts on the card its uuid names; with one card and no uuid it can only be that one.
31const cardOf = (app: { gpu?: string }, gpus: readonly { uuid?: string }[]): number | null => {
32  if (app.gpu === undefined) return gpus.length === 1 ? 0 : null
33  const i = gpus.findIndex(g => g.uuid === app.gpu)
34  return i < 0 ? null : i
35}
36
37// The process tree is only worth a `ps` while a job holds the lock and the card has processes to place.
38export const needsTree = (snap: Snapshot): boolean =>
39  snap.gpu.kind === 'ok' && snap.gpu.apps.length > 0 && holdersOf(snap.lock).length > 0
40
41const sum = (jobs: Record<string, number>): number => Object.values(jobs).reduce((a, b) => a + b, 0)
42
43// Every card as one point in time, and the memory each running job's processes hold on it.
44export const sampleOf = (snap: Snapshot, parents: ReadonlyMap<number, number> | null): Sample | null => {
45  if (snap.gpu.kind !== 'ok' || snap.gpu.gpus.length === 0) return null
46  const { gpus, apps } = snap.gpu
47  const holders = holdersOf(snap.lock)
48  const pids = new Set(holders.map(h => h.pid))
49  const cards: CardSample[] = gpus.map(g => ({ usedMiB: g.usedMiB, totalMiB: g.totalMiB, jobs: {} }))
50  const meta: Sample['jobs'] = {}
51  if (parents !== null && holders.length > 0) {
52    for (const app of apps) {
53      const card = cards[cardOf(app, gpus) ?? -1]
54      const pid = jobOf(app.pid, pids, parents)
55      if (card === undefined || pid === null || app.usedMiB <= 0) continue
56      const add = Math.min(app.usedMiB, Math.max(0, card.usedMiB - sum(card.jobs)))
57      if (add <= 0) continue
58      const key = String(pid)
59      card.jobs[key] = (card.jobs[key] ?? 0) + add
60      const h = holders.find(x => x.pid === pid)
61      if (h !== undefined) meta[key] = { session: clean(h.session, SESSION_CHARS), cmd: clean(h.cmd, CMD_CHARS) }
62    }
63  }
64  return { at: snap.at, gpus: cards, jobs: meta }
65}
66
67export const sampleLine = (s: Sample): string => JSON.stringify(s)
68
69const isObject = (v: unknown): v is Record<string, unknown> => typeof v === 'object' && v !== null && !Array.isArray(v)
70
71const count = (v: unknown): number | null => (typeof v === 'number' && Number.isFinite(v) && v >= 0 ? v : null)
72
73const cardFrom = (raw: unknown): CardSample | null => {
74  if (!isObject(raw)) return null
75  const used = count(raw['usedMiB'])
76  const total = count(raw['totalMiB'])
77  if (used === null || total === null || total <= 0) return null
78  const jobs: Record<string, number> = {}
79  if (isObject(raw['jobs'])) {
80    for (const [key, v] of Object.entries(raw['jobs']).slice(0, MAX_JOBS)) {
81      const mib = count(v)
82      if (JOB_KEY.test(key) && mib !== null) jobs[key] = mib
83    }
84  }
85  return { usedMiB: used, totalMiB: total, jobs }
86}
87
88const sampleFrom = (raw: unknown): Sample | null => {
89  if (!isObject(raw)) return null
90  const at = raw['at']
91  if (typeof at !== 'number' || !Number.isFinite(at) || at <= 0) return null
92  if (!Array.isArray(raw['gpus']) || raw['gpus'].length === 0 || raw['gpus'].length > MAX_CARDS) return null
93  const gpus: CardSample[] = []
94  for (const c of raw['gpus']) {
95    const card = cardFrom(c)
96    if (card === null) return null
97    gpus.push(card)
98  }
99  const jobs: Sample['jobs'] = {}
100  if (isObject(raw['jobs'])) {
101    for (const [key, v] of Object.entries(raw['jobs']).slice(0, MAX_JOBS)) {
102      if (!JOB_KEY.test(key) || !isObject(v)) continue
103      const session = typeof v['session'] === 'string' ? clean(v['session'], SESSION_CHARS) : ''
104      const cmd = typeof v['cmd'] === 'string' ? clean(v['cmd'], CMD_CHARS) : ''
105      jobs[key] = { session, cmd }
106    }
107  }
108  return { at, gpus, jobs }
109}
110
111// The shared file may hold anything, so each line is checked; a line that is not a sample is skipped.
112export const parseSamples = (text: string, now: number): Sample[] => {
113  const out: Sample[] = []
114  for (const line of text.split('\n')) {
115    if (!line.startsWith('{')) continue
116    try {
117      const s = sampleFrom(JSON.parse(line))
118      if (s !== null && s.at >= now - SPAN_MS && s.at <= now + FUTURE_MS) out.push(s)
119    } catch {
120      continue
121    }
122  }
123  return out.sort((a, b) => a.at - b.at).slice(-MAX_SAMPLES)
124}
125
126// What a cell is drawn in: a job's palette slot, or one of the plain kinds.
127export type Paint = 'blank' | 'cap' | 'other' | number
128export type ChartRun = { text: string; fg: Paint; bg: Paint | null }
129export type ChartCard = { title: string; summary: string; rows: { label: string; runs: ChartRun[] }[] }
130export type LegendItem = { paint: Paint; label: string; mem: string }
131export type Charts = { cards: ChartCard[]; axis: string; legend: LegendItem[] }
132
133export type ChartInput = {
134  history: readonly Sample[] | undefined
135  now: number
136  columns: number
137  rows: number
138  capPercent: number | null
139  session: string
140  titles: readonly string[]
141}
142
143type Segment = { paint: Paint; upTo: number }
144
145const axisOf = (columns: number): string => {
146  const left = `${SPAN_MS / 60_000}m ago`
147  const right = 'now'
148  const gap = columns - left.length - right.length
149  const body = gap >= 1 ? `${left}${' '.repeat(gap)}${right}` : right.padStart(columns)
150  return ' '.repeat(LABEL_WIDTH + 1) + body
151}
152
153// Jobs take palette slots in the order they first show memory, so a job keeps its colour along the span.
154const rankJobs = (history: readonly Sample[]): Map<string, number> => {
155  const rank = new Map<string, number>()
156  for (const s of history) {
157    for (const card of s.gpus) {
158      for (const key of Object.keys(card.jobs).sort((a, b) => Number(a) - Number(b))) {
159        if ((card.jobs[key] ?? 0) > 0 && !rank.has(key)) rank.set(key, rank.size)
160      }
161    }
162  }
163  return rank
164}
165
166const stack = (card: CardSample, levels: number, rank: ReadonlyMap<string, number>): { used: number; segs: Segment[] } => {
167  const share = (mib: number) => Math.round(Math.min(1, mib / card.totalMiB) * levels)
168  const used = card.usedMiB <= 0 ? 0 : Math.max(1, share(card.usedMiB))
169  const segs: Segment[] = []
170  let covered = 0
171  let mib = 0
172  const keys = Object.keys(card.jobs)
173    .filter(k => (card.jobs[k] ?? 0) > 0 && rank.has(k))
174    .sort((a, b) => (rank.get(a) ?? 0) - (rank.get(b) ?? 0))
175  for (const key of keys) {
176    mib += card.jobs[key] ?? 0
177    const upTo = Math.min(used, share(mib))
178    if (upTo > covered) segs.push({ paint: (rank.get(key) ?? 0) % PALETTE_SIZE, upTo })
179    covered = Math.max(covered, upTo)
180  }
181  if (used > covered) segs.push({ paint: 'other', upTo: used })
182  return { used, segs }
183}
184
185// One character cell of a stacked column. A cell holding two jobs takes the lower as its glyph and the upper as its background.
186const cellOf = (segs: readonly Segment[], used: number, lo: number): { text: string; fg: Paint; bg: Paint | null } | null => {
187  if (used <= lo) return null
188  const at = (level: number) => segs.find(s => s.upTo > level)?.paint ?? 'other'
189  const top = Math.min(used, lo + LEVELS)
190  const first = segs.find(s => s.upTo > lo)
191  const edge = Math.min(first?.upTo ?? top, top)
192  const lower = at(lo)
193  if (edge >= top) return { text: BLOCKS[top - lo] ?? ' ', fg: lower, bg: null }
194  if (top < lo + LEVELS) {
195    const paint = edge - lo >= top - edge ? lower : at(top - 1)
196    return { text: BLOCKS[top - lo] ?? ' ', fg: paint, bg: null }
197  }
198  return { text: BLOCKS[edge - lo] ?? ' ', fg: lower, bg: at(top - 1) }
199}
200
201const labelOf = (meta: { session: string; cmd: string } | undefined, session: string): string => {
202  const who = meta === undefined ? 'session ?' : meta.session === session ? 'this session' : `session ${shortId(meta.session) || '?'}`
203  return [who, meta?.cmd ?? ''].filter(p => p !== '').join(' · ')
204}
205
206const legendOf = (latest: Sample, rank: ReadonlyMap<string, number>, session: string): LegendItem[] => {
207  const mib = new Map<string, number>()
208  for (const card of latest.gpus) {
209    for (const [key, v] of Object.entries(card.jobs)) if (v > 0 && rank.has(key)) mib.set(key, (mib.get(key) ?? 0) + v)
210  }
211  const items = [...mib.entries()]
212    .sort((a, b) => (rank.get(a[0]) ?? 0) - (rank.get(b[0]) ?? 0))
213    .map(([key, v]) => ({
214      paint: ((rank.get(key) ?? 0) % PALETTE_SIZE) as Paint,
215      label: labelOf(latest.jobs[key], session),
216      mem: gib(v),
217    }))
218  if (items.length === 0) return []
219  const used = latest.gpus.reduce((s, c) => s + c.usedMiB, 0)
220  const rest = used - [...mib.values()].reduce((s, v) => s + v, 0)
221  return rest > 0 ? [...items, { paint: 'other', label: 'other processes', mem: gib(rest) }] : items
222}
223
224export const chartsOf = (i: ChartInput): Charts | null => {
225  if (i.columns < MIN_COLUMNS) return null
226  const inSpan = (i.history ?? []).filter(s => s.at >= i.now - SPAN_MS)
227  const latest = inSpan[inSpan.length - 1]
228  if (latest === undefined || latest.gpus.length === 0) return null
229
230  const rank = rankJobs(inSpan)
231  const levels = i.rows * LEVELS
232  const capLevel = i.capPercent === null ? null : Math.max(1, Math.round((i.capPercent / 100) * levels))
233  const capRow = capLevel === null ? -1 : i.rows - 1 - Math.floor((capLevel - 1) / LEVELS)
234  const start = i.now - SPAN_MS
235
236  const cards = latest.gpus.map((_, c): ChartCard => {
237    // Each column shows the fullest sample of its slice of time, with that sample's own stacking.
238    const picks: (CardSample | null)[] = Array.from({ length: i.columns }, () => null)
239    let peak = 0
240    let newest: CardSample | null = null
241    for (const s of inSpan) {
242      const card = s.gpus[c]
243      if (card === undefined) continue
244      peak = Math.max(peak, card.usedMiB)
245      newest = card
246      const col = Math.min(i.columns - 1, Math.floor(((s.at - start) / SPAN_MS) * i.columns))
247      const held = picks[col]
248      if (held === null || held === undefined || card.usedMiB / card.totalMiB > held.usedMiB / held.totalMiB) picks[col] = card
249    }
250    const stacks = picks.map(p => (p === null || p === undefined ? null : stack(p, levels, rank)))
251
252    const rows = Array.from({ length: i.rows }, (_, r) => {
253      const lo = (i.rows - 1 - r) * LEVELS
254      const runs: ChartRun[] = []
255      for (const st of stacks) {
256        const filled = st === null ? null : cellOf(st.segs, st.used, lo)
257        const cell: ChartRun =
258          filled ?? (r === capRow ? { text: CAP_MARK, fg: 'cap', bg: null } : { text: ' ', fg: 'blank', bg: null })
259        const last = runs[runs.length - 1]
260        if (last !== undefined && last.fg === cell.fg && last.bg === cell.bg) last.text += cell.text
261        else runs.push(cell)
262      }
263      const label = r === 0 ? '100%' : r === capRow ? `${i.capPercent}%` : r === i.rows - 1 ? '0%' : ''
264      return { label: label.padStart(LABEL_WIDTH), runs }
265    })
266
267    return {
268      title: latest.gpus.length === 1 ? 'GPU memory' : (i.titles[c] ?? `GPU ${c}`),
269      summary: newest === null ? '' : `now ${ratio(newest.usedMiB, newest.totalMiB)} · peak ${gib(peak)}`,
270      rows,
271    }
272  })
273
274  return { cards, axis: axisOf(i.columns), legend: legendOf(latest, rank, i.session) }
275}
276

1import type { Elements, RenderElement, RenderSurface } from 'claude-code'
2
3import type { ChartCard, Charts, LegendItem, Paint } from './history'
4import { columnWidths } from './view'
5import type { BandModel, PaneJob, PaneModel, Tone } from './view'
6
7export type Els = Elements[RenderSurface]
8
9const COLOR: Record<Tone, string | undefined> = {
10  mine: 'success',
11  other: 'warning',
12  waiting: 'warning',
13  free: 'success',
14  off: undefined,
15}
16
17const DOT: Record<Tone, string> = { mine: '●', other: '●', waiting: '◐', free: '○', off: '○' }
18
19const ICON_COLOR: Record<string, string | undefined> = { '✓': 'success', '✗': 'error', '⊘': 'warning' }
20
21// The band is one tree, so gpu-lock draws under whatever the plugins beneath it drew.
22export const withBelow = (els: Els, below: RenderElement | null | undefined, mine: RenderElement): RenderElement => {
23  const { Box } = els
24  return below == null ? (
25    mine
26  ) : (
27    <Box flexDirection="column">
28      {below}
29      {mine}
30    </Box>
31  )
32}
33
34// The filled part of a meter takes the tone colour; the empty part recedes.
35const meterBar = (els: Els, bar: string, color: string | undefined) => {
36  const { Box, Text } = els
37  const filled = bar.replace(/░/g, '')
38  return (
39    <Box>
40      <Text color={color}>{filled}</Text>
41      <Text dimColor>{bar.slice(filled.length)}</Text>
42    </Box>
43  )
44}
45
46// One line per session. Columns share widths across rows, and no cell wraps: the command gives way first.
47export const bandView = (els: Els, m: BandModel) => {
48  const { Box, Text } = els
49  const w = columnWidths(m.rows)
50  const meterWidth = m.meter === null ? 0 : m.meter.bar.length + 1 + m.meter.text.length
51  return (
52    <Box flexDirection="column">
53      {m.rows.map((r, i) => {
54        const color = COLOR[r.tone]
55        const quiet = r.tone === 'free' || r.tone === 'off'
56        return (
57          <Box key={`${r.title}-${i}`} columnGap={2}>
58            <Box width={1} flexShrink={0}>
59              <Text color={color} dimColor={quiet}>{DOT[r.tone]}</Text>
60            </Box>
61            <Box width={w.title} flexShrink={0}>
62              <Text bold dimColor={quiet} wrap="truncate-end">{r.title}</Text>
63            </Box>
64            {w.who > 0 && (
65              <Box width={w.who} flexShrink={0}>
66                <Text dimColor wrap="truncate-end">{r.who}</Text>
67              </Box>
68            )}
69            <Box flexGrow={1} flexShrink={1}>
70              <Text wrap="truncate-end">{r.cmd}</Text>
71            </Box>
72            {w.elapsed > 0 && (
73              <Box width={w.elapsed} flexShrink={0} justifyContent="flex-end">
74                <Text dimColor wrap="truncate-end">{r.elapsed}</Text>
75              </Box>
76            )}
77            {m.meter !== null && (
78              <Box width={meterWidth} flexShrink={0} columnGap={1}>
79                {i === 0 && meterBar(els, m.meter.bar, color)}
80                {i === 0 && <Text dimColor>{m.meter.text}</Text>}
81              </Box>
82            )}
83          </Box>
84        )
85      })}
86    </Box>
87  )
88}
89
90const jobList = (els: Els, title: string, jobs: readonly PaneJob[], color: string) => {
91  const { Box, Text } = els
92  const who = Math.max(...jobs.map(j => (j.place > 0 ? `${j.place}. ${j.who}` : j.who).length))
93  const elapsed = Math.max(...jobs.map(j => j.elapsed.length))
94  return (
95    <Box flexDirection="column">
96      <Text bold dimColor>{title}</Text>
97      {jobs.map(j => (
98        <Box key={`${j.who}-${j.cmd}-${j.elapsed}`} columnGap={2}>
99          <Box width={who} flexShrink={0}>
100            <Text color={j.mine ? color : undefined} wrap="truncate-end">{j.place > 0 ? `${j.place}. ${j.who}` : j.who}</Text>
101          </Box>
102          <Box flexGrow={1} flexShrink={1}>
103            <Text wrap="truncate-end">{j.cmd}</Text>
104          </Box>
105          {j.need !== '' && <Text dimColor>{`needs ${j.need}`}</Text>}
106          <Box width={elapsed} flexShrink={0} justifyContent="flex-end">
107            <Text dimColor>{j.elapsed}</Text>
108          </Box>
109        </Box>
110      ))}
111    </Box>
112  )
113}
114
115// Terminal palette names, so a job keeps a distinct colour on any theme. Memory no job owns is grey.
116const JOB_COLORS = ['cyan', 'magenta', 'blue', 'yellow', 'green']
117
118const paintColor = (paint: Paint | null): string | undefined => {
119  if (paint === null || paint === 'blank' || paint === 'cap') return undefined
120  return paint === 'other' ? 'gray' : JOB_COLORS[paint % JOB_COLORS.length]
121}
122
123const legendRow = (els: Els, item: LegendItem) => {
124  const { Box, Text } = els
125  return (
126    <Box key={`${item.label}-${item.mem}`} columnGap={1}>
127      <Text color={paintColor(item.paint)}>■</Text>
128      <Box flexGrow={1} flexShrink={1}>
129        <Text wrap="truncate-end">{item.label}</Text>
130      </Box>
131      <Box flexShrink={0}>
132        <Text dimColor>{item.mem}</Text>
133      </Box>
134    </Box>
135  )
136}
137
138const chartRow = (els: Els, row: ChartCard['rows'][number], i: number) => {
139  const { Box, Text } = els
140  return (
141    <Box key={`row-${i}`} columnGap={1}>
142      <Text dimColor>{row.label}</Text>
143      <Box>
144        {row.runs.map((run, n) => (
145          <Text key={`${i}-${n}`} color={paintColor(run.fg)} backgroundColor={paintColor(run.bg)} dimColor={run.fg === 'cap'}>
146            {run.text}
147          </Text>
148        ))}
149      </Box>
150    </Box>
151  )
152}
153
154const chartCard = (els: Els, card: ChartCard, c: number) => {
155  const { Box, Text } = els
156  return (
157    <Box key={`card-${c}`} flexDirection="column">
158      <Text bold dimColor wrap="truncate-end">{card.title}</Text>
159      {card.rows.map((r, i) => chartRow(els, r, i))}
160      <Text dimColor wrap="truncate-end">{card.summary}</Text>
161    </Box>
162  )
163}
164
165// Memory over the last minutes, one chart per card: a column per slice of time, stacked by job, the cap a dashed line.
166const chartView = (els: Els, charts: Charts) => {
167  const { Box, Text } = els
168  return (
169    <Box flexDirection="column">
170      {charts.cards.map((card, c) => chartCard(els, card, c))}
171      <Text dimColor>{charts.axis}</Text>
172      {charts.legend.map(item => legendRow(els, item))}
173    </Box>
174  )
175}
176
177export type PaneActions = { onRefresh: () => void; onToggle: () => void; onClose: () => void }
178
179export const paneView = (els: Els, m: PaneModel, actions: PaneActions) => {
180  const { Box, Text, Button } = els
181  const color = COLOR[m.tone === 'unknown' ? 'off' : m.tone]
182  const dot = m.tone === 'unknown' ? '○' : DOT[m.tone]
183  const width = Math.max(4, ...m.apps.map(a => a.name.length))
184  const memWidth = Math.max(...m.apps.map(a => a.mem.length), 0)
185  return (
186    <Box flexDirection="column" paddingX={1} gap={1}>
187      <Box flexDirection="column">
188        <Text color={color} bold>{`${dot} ${m.headline}`}</Text>
189        {m.cmd !== '' && <Text>{m.cmd}</Text>}
190        {m.detail !== '' && <Text dimColor>{m.detail}</Text>}
191      </Box>
192
193      {m.gpus.map(g => (
194        <Box flexDirection="column">
195          <Text bold>{g.title}</Text>
196          <Box>
197            {meterBar(els, g.bar, color)}
198            <Text>{` ${g.used}`}</Text>
199            {g.stats !== '' && <Text dimColor>{`  ${g.stats}`}</Text>}
200          </Box>
201        </Box>
202      ))}
203      {m.note !== '' && <Text color="warning">{m.note}</Text>}
204      {m.chart !== null && chartView(els, m.chart)}
205
206      {m.running.length > 0 && jobList(els, 'Running', m.running, 'success')}
207      {m.queue.length > 0 && jobList(els, 'Waiting for a turn', m.queue, 'warning')}
208
209      {m.apps.length > 0 && (
210        <Box flexDirection="column">
211          <Text bold dimColor>Using the GPU</Text>
212          {m.apps.map(a => (
213            <Text>{`${a.name.padEnd(width)}  ${a.mem.padStart(memWidth)}`}</Text>
214          ))}
215        </Box>
216      )}
217
218      {m.log.length > 0 && (
219        <Box flexDirection="column">
220          <Text bold dimColor>This session</Text>
221          {m.log.map(l => (
222            <Box>
223              <Text color={ICON_COLOR[l.icon]}>{`${l.icon} `}</Text>
224              <Text>{l.text}</Text>
225              <Text dimColor>{`  ${l.when}`}</Text>
226            </Box>
227          ))}
228        </Box>
229      )}
230
231      {m.bad.length > 0 && <Text color="warning">{`Ignored patterns (not valid): ${m.bad.join(', ')}`}</Text>}
232
233      <Box gap={2}>
234        <Button key="refresh" label="Refresh" hotkey="r" onPress={actions.onRefresh} />
235        <Button key="toggle" label={m.enabled ? 'Turn off' : 'Turn on'} hotkey="t" onPress={actions.onToggle} />
236        <Button key="close" label="Close" hotkey="c" onPress={actions.onClose} />
237      </Box>
238    </Box>
239  )
240}
241

1import type { Entry, Holder, Reading, Sample, Snapshot, Waiting } from '../types'
2import type { Band } from './options'
3import { bar, clean, duration, gib, ratio } from './format'
4import { chartsOf } from './history'
5import type { Charts } from './history'
6import { waitersOf } from './lock'
7import { isMine, whoLabel } from './messages'
8
9export type Tone = 'mine' | 'other' | 'waiting' | 'free' | 'off'
10
11export type BandRow = { tone: Tone; title: string; who: string; cmd: string; elapsed: string }
12
13export type BandModel = {
14  rows: BandRow[]
15  meter: { bar: string; text: string } | null
16}
17
18export type BandInput = {
19  snapshot: Snapshot | null
20  waiting: Waiting | null
21  enabled: boolean
22  mode: Band
23  now: number
24  session: string
25  width: number
26}
27
28const METER_WIDTH = 10
29const MIN_CMD = 8
30
31const meterOf = (gpu: Reading): { bar: string; text: string } | null => {
32  if (gpu.kind !== 'ok' || gpu.gpus.length === 0) return null
33  const used = gpu.gpus.reduce((sum, g) => sum + g.usedMiB, 0)
34  const total = gpu.gpus.reduce((sum, g) => sum + g.totalMiB, 0)
35  return { bar: bar(used, total, METER_WIDTH), text: `${(used / 1024).toFixed(1)} / ${gib(total)}` }
36}
37
38const whoOf = (lock: Snapshot['lock'], session: string): { who: string; mine: boolean } => {
39  const holder = lock.kind === 'held' ? lock.holder : null
40  return { who: whoLabel(holder, session), mine: isMine(holder, session) }
41}
42
43const COLUMN_GAP = 2
44const MAX_ROWS = 5
45
46export const columnWidths = (rows: readonly BandRow[]) => ({
47  title: Math.max(...rows.map(r => r.title.length)),
48  who: Math.max(...rows.map(r => r.who.length)),
49  elapsed: Math.max(...rows.map(r => r.elapsed.length)),
50})
51
52// Columns across a row: dot, title, who, cmd, elapsed, meter. The dot is 1 cell; empty columns drop out.
53const fit = (rows: BandRow[], meter: BandModel['meter'], width: number): BandModel => {
54  const w = columnWidths(rows)
55  const room = (withMeter: boolean): number => {
56    const meterWidth = withMeter && meter !== null ? METER_WIDTH + 1 + meter.text.length : 0
57    const fixed = [1, w.title, w.who, w.elapsed, meterWidth].filter(c => c > 0)
58    return width - fixed.reduce((a, b) => a + b, 0) - COLUMN_GAP * fixed.length
59  }
60  let kept = meter
61  let budget = room(true)
62  if (budget < MIN_CMD && kept !== null) {
63    kept = null
64    budget = room(false)
65  }
66  return { rows: rows.map(r => ({ ...r, cmd: budget < MIN_CMD ? '' : clean(r.cmd, budget) })), meter: kept }
67}
68
69const holderRow = (h: Holder | null, title: string, i: BandInput): BandRow => ({
70  tone: isMine(h, i.session) ? 'mine' : 'other',
71  title,
72  who: whoLabel(h, i.session),
73  cmd: h?.cmd ?? '',
74  elapsed: h === null ? '' : duration(i.now - h.since),
75})
76
77const waiterRow = (h: Holder, place: number, i: BandInput): BandRow => ({
78  tone: 'waiting',
79  title: `GPU waiting #${place}`,
80  who: whoLabel(h, i.session),
81  cmd: h.cmd,
82  elapsed: duration(i.now - h.since),
83})
84
85// One row per running job, then one per waiting job from every session.
86// This session's own wait shows at once, before the next reading lists it.
87const sessionRows = (i: BandInput, snap: Snapshot): BandRow[] => {
88  const { lock } = snap
89  const rows: BandRow[] = []
90  if (lock.kind === 'held') rows.push(holderRow(lock.holder, 'GPU locked', i))
91  if (lock.kind === 'shared') {
92    if (lock.holders.length === 0) rows.push(holderRow(null, 'GPU shared', i))
93    for (const h of lock.holders) rows.push(holderRow(h, 'GPU shared', i))
94  }
95  const waiters = waitersOf(snap)
96  waiters.forEach((h, n) => rows.push(waiterRow(h, n + 1, i)))
97  if (i.waiting !== null && !waiters.some(h => isMine(h, i.session))) {
98    rows.push({
99      tone: 'waiting',
100      title: `GPU waiting #${waiters.length + 1}`,
101      who: 'this session',
102      cmd: i.waiting.cmd,
103      elapsed: duration(i.now - i.waiting.since),
104    })
105  }
106  return rows
107}
108
109const limitRows = (rows: BandRow[]): BandRow[] =>
110  rows.length <= MAX_ROWS
111    ? rows
112    : [...rows.slice(0, MAX_ROWS - 1), { tone: 'off', title: '', who: '', cmd: `+${rows.length - MAX_ROWS + 1} more`, elapsed: '' }]
113
114export const bandModel = (i: BandInput): BandModel | null => {
115  if (!i.enabled || i.mode === 'never' || i.snapshot === null) return null
116  const { lock, gpu } = i.snapshot
117  const meter = meterOf(gpu)
118
119  const rows = limitRows(sessionRows(i, i.snapshot))
120  if (rows.length > 0) return fit(rows, meter, i.width)
121  if (i.mode !== 'always') return null
122  if (lock.kind === 'unavailable') {
123    return fit([{ tone: 'off', title: 'GPU lock off', who: '', cmd: lock.reason, elapsed: '' }], null, i.width)
124  }
125  return fit([{ tone: 'free', title: 'GPU free', who: '', cmd: '', elapsed: '' }], meter, i.width)
126}
127
128export const statusLine = (model: BandModel | null): string | undefined => {
129  const row = model?.rows[0]
130  if (model === null || row === undefined) return undefined
131  const more = model.rows.length > 1 ? `+${model.rows.length - 1} more` : ''
132  return [row.title, row.who, row.cmd, row.elapsed, more].filter(p => p !== '').join(' · ')
133}
134
135// place is 0 for a running job, and the 1-based position in line for a waiting one.
136export type PaneJob = { who: string; cmd: string; elapsed: string; need: string; mine: boolean; place: number }
137
138export type PaneModel = {
139  headline: string
140  tone: Tone | 'unknown'
141  cmd: string
142  detail: string
143  gpus: { title: string; bar: string; used: string; stats: string }[]
144  apps: { name: string; mem: string }[]
145  running: PaneJob[]
146  queue: PaneJob[]
147  note: string
148  chart: Charts | null
149  log: { icon: string; text: string; when: string }[]
150  bad: string[]
151  enabled: boolean
152}
153
154export type PaneInput = {
155  snapshot: Snapshot | null
156  waiting: Waiting | null
157  enabled: boolean
158  log: readonly Entry[]
159  bad: readonly string[]
160  now: number
161  session: string
162  maxPercent: number | null
163  history?: readonly Sample[]
164  columns?: number
165}
166
167// A lone card gets a taller chart; with several, each takes fewer rows so they all fit.
168const CHART_ROWS = 4
169const CARD_CHART_ROWS = 3
170
171const paneJob = (h: Holder, i: PaneInput, place = 0): PaneJob => ({
172  place,
173  who: whoLabel(h, i.session),
174  cmd: h.cmd,
175  elapsed: duration(i.now - h.since),
176  need: h.needMiB === null ? '' : gib(h.needMiB),
177  mine: isMine(h, i.session),
178})
179
180const LOG_ROWS = 8
181const ICONS = { ran: '✓', failed: '✗', refused: '⊘' } as const
182
183const lockPart = (i: PaneInput): Pick<PaneModel, 'headline' | 'tone' | 'cmd' | 'detail'> => {
184  const none = { cmd: '', detail: '' }
185  if (!i.enabled) return { headline: 'Turned off', tone: 'off', ...none }
186  if (i.waiting !== null) {
187    return { headline: 'Waiting for the lock', tone: 'waiting', cmd: i.waiting.cmd, detail: duration(i.now - i.waiting.since) }
188  }
189  if (i.snapshot === null) return { headline: 'Reading the GPU…', tone: 'unknown', ...none }
190  const { lock } = i.snapshot
191  if (lock.kind === 'free') return { headline: 'Free', tone: 'free', ...none }
192  if (lock.kind === 'unavailable') return { headline: 'Lock unavailable', tone: 'off', cmd: '', detail: lock.reason }
193  if (lock.kind === 'shared') {
194    const n = lock.holders.length
195    const limit = i.maxPercent === null ? '' : `jobs may use ${i.maxPercent}% of memory together`
196    return {
197      headline: n === 0 ? 'Shared by an unknown process' : `Shared by ${n} ${n === 1 ? 'job' : 'jobs'}`,
198      tone: lock.holders.some(h => isMine(h, i.session)) ? 'mine' : 'other',
199      cmd: '',
200      detail: limit,
201    }
202  }
203  const { who, mine } = whoOf(lock, i.session)
204  const h = lock.holder
205  const detail = h === null ? '' : [h.cwd, h.why, duration(i.now - h.since)].filter(p => p !== '').join(' · ')
206  return { headline: `Locked by ${who}`, tone: mine ? 'mine' : 'other', cmd: h?.cmd ?? '', detail }
207}
208
209const notePart = (gpu: Reading | undefined): string => {
210  if (gpu === undefined || gpu.kind === 'ok') return ''
211  return gpu.kind === 'none' ? 'No NVIDIA GPU or nvidia-smi found.' : `nvidia-smi failed: ${gpu.message}`
212}
213
214export const paneModel = (i: PaneInput): PaneModel => {
215  const gpu = i.snapshot?.gpu
216  const ok = gpu?.kind === 'ok' ? gpu : null
217  return {
218    ...lockPart(i),
219    enabled: i.enabled,
220    gpus: (ok?.gpus ?? []).map(g => ({
221      title: `GPU ${g.index}  ${g.name}`,
222      bar: bar(g.usedMiB, g.totalMiB, METER_WIDTH),
223      used: ratio(g.usedMiB, g.totalMiB),
224      stats: [g.util === null ? '' : `${g.util}% busy`, g.tempC === null ? '' : `${g.tempC}°C`]
225        .filter(s => s !== '')
226        .join(' · '),
227    })),
228    apps: (ok?.apps ?? [])
229      .filter(a => a.usedMiB > 0)
230      .sort((a, b) => b.usedMiB - a.usedMiB)
231      .map(a => ({ name: a.name, mem: gib(a.usedMiB) })),
232    running: i.snapshot?.lock.kind === 'shared' ? i.snapshot.lock.holders.map(h => paneJob(h, i)) : [],
233    queue: (i.snapshot === null ? [] : waitersOf(i.snapshot)).map((h, n) => paneJob(h, i, n + 1)),
234    note: notePart(gpu),
235    chart: chartsOf({
236      history: i.history,
237      now: i.now,
238      columns: i.columns ?? 0,
239      rows: (ok?.gpus.length ?? 1) > 1 ? CARD_CHART_ROWS : CHART_ROWS,
240      capPercent: i.maxPercent,
241      session: i.session,
242      titles: (ok?.gpus ?? []).map(g => `GPU ${g.index}  ${g.name}`),
243    }),
244    log: [...i.log]
245      .reverse()
246      .slice(0, LOG_ROWS)
247      .map(e => ({
248        icon: ICONS[e.kind],
249        text: [e.cmd, e.ms === null ? '' : duration(e.ms), e.note].filter(p => p !== '').join(' · '),
250        when: `${duration(i.now - e.at)} ago`,
251      })),
252    bad: [...i.bad],
253  }
254}
255

1import { scan } from './shell'
2import type { Segment } from './shell'
3
4export type Patterns = {
5  extra: readonly RegExp[]
6  skip: readonly RegExp[]
7  lockPath?: string
8}
9
10export type Verdict =
11  | { gpu: false }
12  | { gpu: true; why: string; needMiB: number | null; devices: number[] | null }
13
14const MAX_HEAD = 20000
15
16const HINTS = [
17  'cuda', 'cudnn', 'gpu', 'torch', 'pytorch', 'tensorflow', 'jax', 'triton', 'tensorrt', 'onnxruntime', 'vllm',
18  'transformers', 'diffusers', 'bitsandbytes', 'flash-attn', 'flash_attn', 'xformers', 'cupy', 'numba', 'nccl',
19  'nvenc', 'nvdec', 'cuvid',
20]
21const CARGO_HINTS = ['cuda', 'cudarc', 'gpu', 'candle', 'wgpu', 'torch', 'tch', 'onnx', 'ort', 'tensorrt', 'burn']
22
23const wordRegex = (word: string): RegExp => new RegExp(`(?<![a-z])${word.replace(/[-_]/g, '\\$&')}(?![a-z])`, 'i')
24const HINT_TESTS = HINTS.map(w => [w, wordRegex(w)] as const)
25const CARGO_TESTS = CARGO_HINTS.map(w => [w, wordRegex(w)] as const)
26
27const firstHint = (tests: readonly (readonly [string, RegExp])[], text: string): string | null =>
28  tests.find(([, re]) => re.test(text))?.[0] ?? null
29const SCRIPT = /(?<![a-z])((?:train|finetune|fine_tune|infer|inference|bench)\w*\.py)/i
30const CARGO_JOB = /(?<![a-z])(train|finetune|infer|inference|bench|benchmark)\w*/i
31const STRONG = /(--device[ =]cuda|cuda:\d|--gpus|CUDA_VISIBLE_DEVICES)/
32const ASSIGN = /^([A-Za-z_][A-Za-z0-9_]*)=([\s\S]*)$/
33const NEED = /^(\d+(?:\.\d+)?)\s*(g|gb|gib|m|mb|mib)?$/i
34
35const EXEMPT = new Set([
36  'nvidia-smi', 'nvtop', 'gpustat', 'rg', 'rkgrep', 'grep', 'egrep', 'fgrep', 'cat', 'ls', 'head', 'tail', 'less',
37  'more', 'git', 'gh', 'sed', 'awk', 'echo', 'printf', 'wc', 'find', 'fd', 'cp', 'mv', 'rm', 'mkdir', 'rmdir',
38  'touch', 'chmod', 'chown', 'cd', 'pwd', 'which', 'whereis', 'type', 'jq', 'sqlite3', 'curl', 'wget', 'tar',
39  'unzip', 'zip', 'gzip', 'diff', 'sort', 'uniq', 'tr', 'cut', 'tee', 'kill', 'pkill', 'killall', 'ps', 'pgrep',
40  'lsof', 'fuser', 'df', 'du', 'free', 'top', 'htop', 'sleep', 'date', 'export', 'source', 'true', 'false', 'test',
41  '[', 'ssh', 'scp', 'rsync', 'stat', 'file', 'ln', 'realpath', 'basename', 'dirname', 'tree', 'bat', 'code',
42  'pip', 'npm', 'pnpm', 'yarn', 'apt', 'apt-get', 'brew', 'dpkg', 'uv', 'poetry', 'pipenv', 'pdm', 'rye', 'hatch',
43  'conda', 'mamba', 'micromamba', 'claude', 'man', 'help', 'nproc', 'uname', 'whoami', 'id', 'env', 'ulimit', 'xdg-open',
44])
45
46const ALWAYS = new Set([
47  'torchrun', 'deepspeed', 'vllm', 'llama-server', 'llama-cli', 'llama-bench', 'llama-perplexity', 'ncu', 'nsys',
48  'trtexec',
49])
50
51const INTERPRETERS = new Set(['python', 'pytest', 'py.test', 'ipython', 'jupyter', 'node', 'bun', 'deno', 'julia', 'ffmpeg'])
52
53const PIP_LIKE = new Set(['pip', 'venv', 'ensurepip', 'build', 'http.server', 'json.tool'])
54
55const CARGO_RUN = new Set(['run', 'r', 'test', 't', 'bench', 'nextest'])
56
57const RUNNERS = new Set(['uv', 'poetry', 'pipenv', 'pdm', 'rye', 'hatch', 'conda', 'mamba', 'micromamba'])
58const RUNNER_VALUED = new Set(['-n', '-p', '--with', '--python', '--project', '--directory', '--env-file', '--group', '--extra', '--name', '--prefix'])
59
60const skipFlags = (words: readonly string[], from: number, valued: ReadonlySet<string>): number => {
61  let i = from
62  while (i < words.length && (words[i] ?? '').startsWith('-')) {
63    i += valued.has(words[i] ?? '') ? 2 : 1
64  }
65  return i
66}
67
68const NO_VALUE: ReadonlySet<string> = new Set()
69
70const WRAPPERS = new Map<string, (words: readonly string[], from: number) => number>([
71  ['sudo', (w, i) => skipFlags(w, i, new Set(['-u', '-g', '-h', '-p', '-C', '-D', '-R', '-T', '-U']))],
72  ['nohup', (_w, i) => i],
73  ['time', (w, i) => skipFlags(w, i, NO_VALUE)],
74  ['command', (w, i) => skipFlags(w, i, NO_VALUE)],
75  ['exec', (w, i) => skipFlags(w, i, NO_VALUE)],
76  ['setsid', (w, i) => skipFlags(w, i, NO_VALUE)],
77  ['env', (w, i) => skipFlags(w, i, new Set(['-u', '-C', '-S']))],
78  ['nice', (w, i) => skipFlags(w, i, new Set(['-n']))],
79  ['ionice', (w, i) => skipFlags(w, i, new Set(['-c', '-n', '-p']))],
80  ['stdbuf', (w, i) => skipFlags(w, i, NO_VALUE)],
81  ['npx', (w, i) => skipFlags(w, i, NO_VALUE)],
82  ['bunx', (w, i) => skipFlags(w, i, NO_VALUE)],
83  ['timeout', (w, i) => skipFlags(w, i, new Set(['-s', '-k'])) + 1],
84])
85
86type Unwrapped = { env: Record<string, string>; program: string | undefined; args: string[] }
87
88const baseName = (word: string): string => {
89  const name = (word.split('/').pop() ?? '').toLowerCase()
90  if (/^python[\d.]*$/.test(name)) return 'python'
91  if (/^pip[\d.]*$/.test(name)) return 'pip'
92  return name
93}
94
95const unwrap = (words: readonly string[]): Unwrapped => {
96  const env: Record<string, string> = {}
97  let i = 0
98  while (i < words.length) {
99    const word = words[i] ?? ''
100    const assigned = ASSIGN.exec(word)
101    if (assigned !== null) {
102      env[assigned[1] ?? ''] = assigned[2] ?? ''
103      i += 1
104      continue
105    }
106    const base = baseName(word)
107    const wrap = WRAPPERS.get(base)
108    if (wrap !== undefined) {
109      i = wrap(words, i + 1)
110      continue
111    }
112    const next = words[i + 1]
113    if (RUNNERS.has(base) && (next === 'run' || next === 'exec')) {
114      i = skipFlags(words, i + 2, RUNNER_VALUED)
115      continue
116    }
117    return { env, program: word, args: words.slice(i + 1) }
118  }
119  return { env, program: undefined, args: [] }
120}
121
122const parseNeed = (text: string): number | null => {
123  const m = NEED.exec(text.trim())
124  if (m === null) return null
125  const n = Number(m[1])
126  const unit = (m[2] ?? 'm').toLowerCase()
127  return Math.round(unit.startsWith('g') ? n * 1024 : n)
128}
129
130const parseDevices = (value: string): number[] | null => {
131  const ids = value.split(',').map(s => s.trim())
132  const nums = ids.filter(s => /^\d+$/.test(s)).map(Number)
133  return nums.length > 0 && nums.length === ids.length ? nums : null
134}
135
136type Hit = { why: string; devices: number[] | null }
137
138const devicesOf = (env: Record<string, string>, text: string): number[] | null => {
139  const fromEnv = env['CUDA_VISIBLE_DEVICES']
140  if (fromEnv !== undefined) return parseDevices(fromEnv)
141  const m = /cuda:(\d+)/.exec(text)
142  return m === null ? null : [Number(m[1])]
143}
144
145const isCpuOnly = (value: string): boolean => value.trim() === '' || value.trim() === '-1' || value.trim() === 'NoDevFiles'
146
147const judge = (segment: Segment): Hit | null => {
148  const { env, program, args } = unwrap(segment.words)
149  if (program === undefined) return null
150  const base = baseName(program)
151  const text = segment.text
152  const visible = env['CUDA_VISIBLE_DEVICES']
153  if (visible !== undefined && isCpuOnly(visible)) return null
154  if (visible !== undefined && !EXEMPT.has(base)) return { why: 'CUDA_VISIBLE_DEVICES', devices: parseDevices(visible) }
155  if (EXEMPT.has(base)) return null
156  const devices = devicesOf(env, text)
157
158  if (ALWAYS.has(base)) return { why: base, devices }
159  if (base === 'accelerate') return args.includes('launch') ? { why: 'accelerate launch', devices } : null
160  if (base === 'ollama') {
161    const sub = args[0]
162    return sub === 'run' || sub === 'serve' ? { why: `ollama ${sub}`, devices } : null
163  }
164  if (base === 'docker' || base === 'podman') {
165    return /--gpus|--runtime[ =]nvidia/.test(text) ? { why: `${base} + gpus`, devices } : null
166  }
167  if (base === 'cargo') {
168    const sub = args.find(a => !a.startsWith('-') && !a.startsWith('+'))
169    if (sub === undefined || !CARGO_RUN.has(sub)) return null
170    const hint = firstHint(CARGO_TESTS, text) ?? CARGO_JOB.exec(text)?.[1]?.toLowerCase() ?? null
171    return hint === null ? null : { why: `cargo + ${hint}`, devices }
172  }
173  if (INTERPRETERS.has(base)) {
174    if (base === 'python' && args[0] === '-m' && PIP_LIKE.has(args[1] ?? '')) return null
175    const hint = firstHint(HINT_TESTS, text)
176    if (hint !== null) return { why: `${base} + ${hint}`, devices }
177    const script = SCRIPT.exec(text)
178    return script === null ? null : { why: `${base} + ${(script[1] ?? '').toLowerCase()}`, devices }
179  }
180  const strong = STRONG.exec(text)
181  return strong === null ? null : { why: `${base} + ${strong[1] ?? ''}`, devices }
182}
183
184type Marker = { kind: 'none' } | { kind: 'skip' } | { kind: 'force'; needMiB: number | null }
185
186const markerOf = (comments: readonly string[]): Marker => {
187  let found: Marker = { kind: 'none' }
188  for (const comment of comments) {
189    if (/^nogpu\b/i.test(comment)) return { kind: 'skip' }
190    const m = /^gpu\b\s*(.*)$/i.exec(comment)
191    if (m !== null) found = { kind: 'force', needMiB: parseNeed(m[1] ?? '') }
192  }
193  return found
194}
195
196export const compilePatterns = (text: string): { regexes: RegExp[]; bad: string[] } => {
197  const regexes: RegExp[] = []
198  const bad: string[] = []
199  for (const line of text.split('\n').map(l => l.trim()).filter(l => l !== '')) {
200    try {
201      regexes.push(new RegExp(line, 'i'))
202    } catch {
203      bad.push(line)
204    }
205  }
206  return { regexes, bad }
207}
208
209export const classify = (command: string, patterns: Patterns): Verdict => {
210  const head = command.slice(0, MAX_HEAD)
211  if (head.trim() === '') return { gpu: false }
212  if (patterns.lockPath !== undefined && patterns.lockPath !== '' && head.includes(patterns.lockPath)) {
213    return { gpu: false }
214  }
215  const { segments, comments } = scan(head)
216  const marker = markerOf(comments)
217  if (marker.kind === 'skip') return { gpu: false }
218  if (patterns.skip.some(r => r.test(head))) return { gpu: false }
219
220  const needMiB = marker.kind === 'force' ? marker.needMiB : null
221  for (const segment of segments) {
222    const hit = judge(segment)
223    if (hit !== null) return { gpu: true, why: hit.why, needMiB, devices: hit.devices }
224  }
225  if (marker.kind === 'force') {
226    const first = segments.map(s => unwrap(s.words).env).find(e => e['CUDA_VISIBLE_DEVICES'] !== undefined)
227    return { gpu: true, why: 'marker # gpu', needMiB, devices: devicesOf(first ?? {}, head) }
228  }
229  if (patterns.extra.some(r => r.test(head))) return { gpu: true, why: 'extra pattern', needMiB, devices: null }
230  return { gpu: false }
231}
232

1export const gib = (mib: number): string => `${(mib / 1024).toFixed(1)} GiB`
2
3export const ratio = (partMiB: number, totalMiB: number): string =>
4  `${(partMiB / 1024).toFixed(1)} of ${gib(totalMiB)}`
5
6export const duration = (ms: number): string => {
7  const total = Math.max(0, Math.floor(ms / 1000))
8  const days = Math.floor(total / 86400)
9  const hours = Math.floor((total % 86400) / 3600)
10  const minutes = Math.floor((total % 3600) / 60)
11  const seconds = total % 60
12  if (days > 0) return `${days}d ${hours}h`
13  if (hours > 0) return `${hours}h ${String(minutes).padStart(2, '0')}m`
14  if (minutes > 0) return `${minutes}m ${String(seconds).padStart(2, '0')}s`
15  return `${seconds}s`
16}
17
18export const bar = (used: number, total: number, width: number): string => {
19  const share = total > 0 ? Math.min(1, Math.max(0, used / total)) : 0
20  const filled = Math.round(share * width)
21  return '█'.repeat(filled) + '░'.repeat(width - filled)
22}
23
24export const clean = (text: string, max: number): string => {
25  const flat = text
26    .replace(/[\t\n\r]+/g, ' ')
27    .replace(/[\u0000-\u001f\u007f-\u009f]/g, '')
28    .replace(/ {2,}/g, ' ')
29    .trim()
30  return flat.length > max ? `${flat.slice(0, max - 1)}…` : flat
31}
32
33export const shortId = (id: string): string => id.slice(0, 4)
34

1import type { Holder, LockState } from '../types'
2import { clean } from './format'
3import { RAMP_SECONDS } from './share'
4
5export const ownerPathOf = (lockPath: string): string => {
6  const cut = lockPath.lastIndexOf('/')
7  return `${cut < 0 ? '.' : lockPath.slice(0, cut) || ''}/owner.json`
8}
9
10export const dirOf = (lockPath: string): string => {
11  const cut = lockPath.lastIndexOf('/')
12  return cut < 0 ? '.' : lockPath.slice(0, cut) || ''
13}
14
15export const jobsDirOf = (lockPath: string): string => `${dirOf(lockPath)}/jobs`
16export const waitingDirOf = (lockPath: string): string => `${dirOf(lockPath)}/waiting`
17export const gateOf = (lockPath: string): string => `${dirOf(lockPath)}/gate`
18
19// Exit 0: held exclusively. 1: free. 4: held only by shared holders. 3: could not tell.
20// Stdout lists the live slot files, one per line: `jobs <json>` or `waiting <json>`.
21// A slot whose pid is gone is skipped once it is older than the ramp, so a job that
22// has just started is not dropped before its shell is visible.
23const PROBE = `f=$1
24d=$(dirname "$f")
25list() {
26  now=$(date +%s)
27  for k in jobs waiting; do
28    for p in "$d/$k"/*.json; do
29      [ -f "$p" ] || continue
30      b=$(head -c 2000 "$p" | tr '\\n\\r' '  ')
31      pid=$(printf '%s' "$b" | sed -n 's/.*"pid": *\\([0-9][0-9]*\\).*/\\1/p')
32      since=$(printf '%s' "$b" | sed -n 's/.*"since": *\\([0-9][0-9]*\\).*/\\1/p')
33      if [ -d /proc/self ] && [ -n "$pid" ] && [ ! -d "/proc/$pid" ] && [ $((now - \${since:-0})) -ge ${RAMP_SECONDS} ]; then continue; fi
34      printf '%s %s\\n' "$k" "$b"
35    done
36  done
37}
38[ -e "$f" ] || exit 1
39if [ -r /proc/locks ]; then
40  dev=$(stat -c %d "$f" 2>/dev/null) || exit 3
41  ino=$(stat -c %i "$f" 2>/dev/null) || exit 3
42  awk -v d="$dev" -v i="$ino" '
43    BEGIN { x = int(d / 4096); want = sprintf("%02x:%02x:%d", int(d / 256) % 4096, d % 256 + x - x % 256, i) }
44    $2 == "FLOCK" && $0 !~ /->/ && $6 == want { if ($4 == "WRITE") w = 1; else r = 1 }
45    END { exit w ? 0 : (r ? 4 : 1) }
46  ' /proc/locks
47  s=$?
48else
49  flock -w 1 -E 9 "$f" true
50  case $? in 0) s=1 ;; 9) s=0 ;; *) s=3 ;; esac
51fi
52list
53exit $s`
54
55const ENSURE = `f=$1
56d=$(dirname "$f")
57mkdir -p -m 777 "$d" 2>/dev/null
58for s in jobs waiting; do mkdir -p -m 777 "$d/$s" 2>/dev/null; done
59[ -e "$f" ] || { : > "$f" && chmod 666 "$f"; } 2>/dev/null
60for s in gate history.jsonl history.lock; do [ -e "$d/$s" ] || { : > "$d/$s" && chmod 666 "$d/$s"; } 2>/dev/null; done
61[ -r "$f" ] || exit 3
62command -v flock >/dev/null 2>&1 || exit 4
63exit 0`
64
65export const probeArgv = (lockPath: string): string[] => ['sh', '-c', PROBE, 'gpu-lock-probe', lockPath]
66
67export const ensureArgv = (lockPath: string): string[] => ['sh', '-c', ENSURE, 'gpu-lock-ensure', lockPath]
68
69export const ensureReason = (exitCode: number): string | null => {
70  if (exitCode === 0) return null
71  if (exitCode === 4) return 'flock is not installed'
72  return 'the lock file cannot be created or read'
73}
74
75const text = (v: unknown, max: number): string => (typeof v === 'string' ? clean(v, max) : '')
76
77export const parseHolder = (raw: string): Holder | null => {
78  let data: unknown
79  try {
80    data = JSON.parse(raw)
81  } catch {
82    return null
83  }
84  if (typeof data !== 'object' || data === null || Array.isArray(data)) return null
85  const o = data as Record<string, unknown>
86  const pid = o['pid']
87  const since = o['since']
88  if (typeof pid !== 'number' || !Number.isInteger(pid) || pid <= 0) return null
89  if (typeof since !== 'number' || !Number.isFinite(since) || since <= 0) return null
90  const need = o['needMiB']
91  const agent = text(o['agent'], 64)
92  return {
93    session: text(o['session'], 64),
94    agent: agent === '' ? null : agent,
95    cmd: text(o['cmd'], 200),
96    cwd: text(o['cwd'], 200),
97    why: text(o['why'], 80),
98    needMiB: typeof need === 'number' && Number.isFinite(need) && need > 0 ? need : null,
99    pid,
100    since: since * 1000,
101  }
102}
103
104// State outlives a reload, so a stored snapshot may predate the waiters field.
105export const waitersOf = (snap: { waiters?: Holder[] }): Holder[] => snap.waiters ?? []
106
107export const parseListing = (stdout: string): { jobs: Holder[]; waiters: Holder[] } => {
108  const jobs: Holder[] = []
109  const waiters: Holder[] = []
110  for (const line of stdout.split('\n')) {
111    const cut = line.indexOf(' ')
112    const kind = line.slice(0, cut)
113    if (kind !== 'jobs' && kind !== 'waiting') continue
114    const h = parseHolder(line.slice(cut + 1))
115    if (h !== null) (kind === 'jobs' ? jobs : waiters).push(h)
116  }
117  return { jobs: jobs.sort((a, b) => a.since - b.since), waiters: waiters.sort((a, b) => a.since - b.since) }
118}
119
120export const lockState = (probeExit: number, ownerText: string | null, ownerAlive: boolean): LockState => {
121  if (probeExit === 1) return { kind: 'free' }
122  if (probeExit !== 0) return { kind: 'unavailable', reason: `the lock probe failed (exit ${probeExit})` }
123  const holder = ownerText === null ? null : parseHolder(ownerText)
124  return { kind: 'held', holder: holder !== null && ownerAlive ? holder : null }
125}
126