cluster-context

1import type { EngineInterface, Register } from 'claude-code'
2
3import { downloadTarget, fmtLeft, gpuModelFrom, isAllocation, isOffScratch, nodeFrom, secondsLeft, sectionText } from './facts'
4import type { Node } from './facts'
5
6const LAUNCH = 'Launch new allocation'
7const HERE = 'Run on this node'
8const ANYWAY = 'Download there anyway'
9const SCRATCH = 'Use scratch instead'
10
11export const register: Register = (on, options) => {
12  const configuredScratch = typeof options.scratchRoot === 'string' ? options.scratchRoot.trim() : ''
13  // Empty setting means /scratch/$USER, resolved at session start.
14  let scratchRoot = configuredScratch || '/scratch'
15  const minMinutesLeft = typeof options.minMinutesLeft === 'number' ? options.minMinutesLeft : 30
16
17  // Read once per load: the allocation does not change under a running session.
18  let node: Node | null = null
19  let hfCache: string | null = null
20  let home: string | null = null
21  let isInteractive = false
22
23  on('session.start', async ($, e, next) => {
24    const started = await next(e)
25    isInteractive = e.isInteractive
26    home = (await $.env.get('HOME')) ?? null
27    const user = (await $.env.get('USER')) ?? (await $.env.get('LOGNAME'))
28    if (!configuredScratch && user) scratchRoot = `/scratch/${user}`
29    node = nodeFrom({
30      SLURM_JOB_ID: await $.env.get('SLURM_JOB_ID'),
31      SLURM_JOB_NAME: await $.env.get('SLURM_JOB_NAME'),
32      SLURM_JOB_PARTITION: await $.env.get('SLURM_JOB_PARTITION'),
33      SLURMD_NODENAME: await $.env.get('SLURMD_NODENAME'),
34      SLURM_CPUS_ON_NODE: await $.env.get('SLURM_CPUS_ON_NODE'),
35      SLURM_MEM_PER_NODE: await $.env.get('SLURM_MEM_PER_NODE'),
36      SLURM_JOB_GPUS: await $.env.get('SLURM_JOB_GPUS'),
37      SLURM_STEP_GPUS: await $.env.get('SLURM_STEP_GPUS'),
38      CUDA_VISIBLE_DEVICES: await $.env.get('CUDA_VISIBLE_DEVICES'),
39      SLURM_GPUS_ON_NODE: await $.env.get('SLURM_GPUS_ON_NODE'),
40      SLURM_JOB_END_TIME: await $.env.get('SLURM_JOB_END_TIME'),
41    })
42    if (node && node.gpus > 0) {
43      try {
44        const r = await $.process.run(['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader'], { timeoutMs: 10_000 })
45        if (r.exitCode === 0) node = { ...node, gpuModel: gpuModelFrom(r.stdout) }
46      } catch {
47        // No nvidia-smi on this node: the count from Slurm stands.
48      }
49    }
50    // Only the caches move: HF_HOME stays put, since the login token lives at $HF_HOME/token
51    // (~/.cache/huggingface/token) and moving it would log huggingface_hub out.
52    const hfHome = await $.env.get('HF_HOME')
53    hfCache = (await $.env.get('HF_HUB_CACHE')) ?? null
54    if (!hfHome && !hfCache && scratchRoot !== '/scratch') {
55      const root = `${scratchRoot.replace(/\/$/, '')}/hf_cache`
56      hfCache = `${root}/hub`
57      await $.env.set('HF_HUB_CACHE', hfCache)
58      if (!(await $.env.get('HF_XET_CACHE'))) await $.env.set('HF_XET_CACHE', `${root}/xet`)
59    } else if (!hfCache && hfHome) {
60      hfCache = `${hfHome}/hub`
61    }
62    return started
63  })
64
65  on('prompt.compose', async ($, e, next) => {
66    const composed = await next(e)
67    return {
68      sections: [...composed.sections, { id: 'cluster-context:node', text: sectionText(node, { scratchRoot, hfCache }), scope: 'session' }],
69    }
70  })
71
72  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
73    if (!isInteractive) {
74      $.ui.status('cluster-context: guards off (session not interactive)')
75      return next(e)
76    }
77    const cwd = (await $.session.cwd()) ?? ''
78
79    if (node && node.gpus > 0 && isAllocation(e.command)) {
80      const left = secondsLeft(node, await $.clock.now())
81      if (left === null || left >= minMinutesLeft * 60) {
82        const gpus = `${node.gpus} GPU${node.gpus > 1 ? 's' : ''}${node.gpuModel ? ` (${node.gpuModel})` : ''}`
83        const time = left === null ? '' : ` and ${fmtLeft(left)} left`
84        const answer = await askOrNull($, `Claude wants a new Slurm allocation (${e.command.slice(0, 100)}), but this node${node.host ? ` (${node.host})` : ''} already has ${gpus}${time}. Launch it anyway?`, 'New node', [LAUNCH, HERE])
85        if (answer === null) return { deny: 'cluster-context: the user dismissed the question about starting a new allocation. Ask them how to proceed.' }
86        if (answer === HERE) return { deny: `cluster-context: the user wants this run on the current node (${gpus}${time}), not in a new allocation. Run the work directly in this shell without srun/sbatch/salloc.` }
87        if (answer !== LAUNCH) return { deny: `cluster-context: the user answered: ${answer}` }
88      }
89    }
90
91    const target = downloadTarget(e.command, cwd, home)
92    if (target && isOffScratch(target, scratchRoot) && isOffScratch(await realpath($, target), scratchRoot)) {
93      const answer = await askOrNull($, `This download writes to ${target}, outside ${scratchRoot}. Download there anyway?`, 'Data path', [ANYWAY, SCRATCH])
94      if (answer === null) return { deny: 'cluster-context: the user dismissed the question about the download location. Ask them where it should go.' }
95      if (answer === SCRATCH) return { deny: `cluster-context: downloads, weights, datasets and caches go under ${scratchRoot} (e.g. ${scratchRoot}/datasets/<name>). Rerun with the target there.` }
96      if (answer !== ANYWAY) return { deny: `cluster-context: the user answered: ${answer}` }
97    }
98
99    return next(e)
100  }).catch(($, e, next) => {
101    // A guard bug never blocks work, but it shows: the status line names the error.
102    try {
103      $.ui.status(`cluster-context: guard failed: ${String(next.error).slice(0, 160)}`)
104    } catch {
105      // Re-entry: this $ cannot draw.
106    }
107    return next(e)
108  })
109}
110
111async function askOrNull($: EngineInterface, question: string, header: string, options: readonly string[]): Promise<string | null> {
112  try {
113    return await $.ui.ask(question, { header, options })
114  } catch {
115    // Dismissed, or nobody to ask.
116    return null
117  }
118}
119
120/** The path with symlinks followed (T4 under /projects is a link into scratch); the path itself when that fails. */
121async function realpath($: EngineInterface, path: string): Promise<string> {
122  try {
123    const r = await $.process.run(['realpath', '-m', path], { timeoutMs: 5_000 })
124    return r.exitCode === 0 && r.stdout.trim() ? r.stdout.trim() : path
125  } catch {
126    return path
127  }
128}
129

1/** What this shell's Slurm allocation holds, read from its environment only (never squeue). */
2export type Node = {
3  job: string
4  name: string | null
5  partition: string | null
6  host: string | null
7  cpus: string | null
8  /** Memory as Slurm gives it in MB, shown as G. */
9  memMb: number | null
10  gpus: number
11  gpuModel: string | null
12  /** Epoch ms; null when Slurm did not say. */
13  endsAt: number | null
14}
15
16export type SlurmEnv = {
17  SLURM_JOB_ID?: string
18  SLURM_JOB_NAME?: string
19  SLURM_JOB_PARTITION?: string
20  SLURMD_NODENAME?: string
21  SLURM_CPUS_ON_NODE?: string
22  SLURM_MEM_PER_NODE?: string
23  SLURM_JOB_GPUS?: string
24  SLURM_STEP_GPUS?: string
25  CUDA_VISIBLE_DEVICES?: string
26  SLURM_GPUS_ON_NODE?: string
27  SLURM_JOB_END_TIME?: string
28}
29
30/** Same reading as pace-line's slurmFrom, plus the facts the prompt section needs. */
31export function nodeFrom(env: SlurmEnv): Node | null {
32  if (!env.SLURM_JOB_ID) return null
33  const ids = (env.SLURM_JOB_GPUS || env.SLURM_STEP_GPUS || env.CUDA_VISIBLE_DEVICES || '')
34    .split(',')
35    .filter(s => s.trim() !== '' && s !== 'NoDevFiles')
36  const onNode = Number(env.SLURM_GPUS_ON_NODE)
37  const mem = Number(env.SLURM_MEM_PER_NODE)
38  // Epoch seconds, set by Slurm 23.02+ when the job starts.
39  const end = Number(env.SLURM_JOB_END_TIME)
40  return {
41    job: env.SLURM_JOB_ID,
42    name: env.SLURM_JOB_NAME || null,
43    partition: env.SLURM_JOB_PARTITION || null,
44    host: env.SLURMD_NODENAME || null,
45    cpus: env.SLURM_CPUS_ON_NODE || null,
46    memMb: mem > 0 ? mem : null,
47    gpus: ids.length || (onNode > 0 ? onNode : 0),
48    gpuModel: null,
49    endsAt: end > 0 ? end * 1000 : null,
50  }
51}
52
53/** First line of `nvidia-smi --query-gpu=name,memory.total --format=csv,noheader`: "NVIDIA A100-SXM4-80GB, 81920 MiB". */
54export function gpuModelFrom(stdout: string): string | null {
55  const first = stdout.split('\n').find(l => l.trim() !== '')
56  if (!first) return null
57  const [name, mem] = first.split(',').map(s => s.trim())
58  const gib = Math.round(Number.parseInt(mem ?? '', 10) / 1024)
59  return name ? (gib > 0 ? `${name.replace(/^NVIDIA /, '')} ${gib}GB` : name.replace(/^NVIDIA /, '')) : null
60}
61
62export function fmtLeft(seconds: number): string {
63  const h = Math.floor(seconds / 3600)
64  const m = Math.floor((seconds % 3600) / 60)
65  return h > 0 ? `${h}h${String(m).padStart(2, '0')}m` : `${m}m`
66}
67
68/** Seconds until the job ends; null when unknown. */
69export function secondsLeft(node: Node, nowMs: number): number | null {
70  return node.endsAt === null ? null : Math.max(0, Math.floor((node.endsAt - nowMs) / 1000))
71}
72
73export function sectionText(node: Node | null, o: { scratchRoot: string; hfCache: string | null }): string {
74  const data = `Downloads, model weights, datasets and caches go under ${o.scratchRoot}, never /home or /projects.` +
75    (o.hfCache ? ` Hugging Face models and datasets cache in ${o.hfCache}; leave HF_HOME alone (the login token lives there).` : '')
76  if (!node) return `# Where you are running (cluster-context)\nNot inside a Slurm allocation.\n${data}`
77  const where = [
78    `Slurm job ${node.job}`,
79    node.name || node.partition ? ` (${[node.name, node.partition && `partition ${node.partition}`].filter(Boolean).join(', ')})` : '',
80    node.host ? ` on node ${node.host}` : '',
81  ].join('')
82  const res = [
83    node.cpus && `${node.cpus} CPUs`,
84    node.memMb && `${Math.round(node.memMb / 1024)}G RAM`,
85    node.gpus > 0 ? `${node.gpus} GPU${node.gpus > 1 ? 's' : ''}${node.gpuModel ? ` (${node.gpuModel})` : ''}` : 'NO GPU',
86  ].filter(Boolean).join(', ')
87  const ends = node.endsAt === null ? '' : ` The job ends at ${new Date(node.endsAt).toISOString().slice(0, 16).replace('T', ' ')} UTC (compare with \`date -u\`).`
88  const rule = node.gpus > 0
89    ? 'This node has a GPU: run short GPU work (model tests, small evals, downloads) right here, and do not start a new allocation (srun/sbatch/salloc) for it. Use sbatch only for work longer than the time left, or that must survive this session.'
90    : 'There is no GPU here: anything needing CUDA goes through sbatch. Downloads and light CPU work can run here; no new node is needed for them.'
91  return `# Where you are running (cluster-context)\n${where}: ${res}.${ends}\n${rule}\n${data}\nThese facts are read from this session's environment and override any GPU/CPU counts written in CLAUDE.md.`
92}
93
94const ALLOC = /(^|[\s;&|(`$])(srun|sbatch|salloc)(?=\s|$)/
95/** A command that asks Slurm for a new allocation; an srun step inside this job (--jobid/--overlap) is not one. */
96export function isAllocation(cmd: string): boolean {
97  return ALLOC.test(cmd) && !/--overlap\b|--jobid[= ]/.test(cmd)
98}
99
100const DOWNLOADERS: readonly RegExp[] = [
101  /(^|[\s;&|(])(hf|huggingface-cli)\s+download\b/,
102  /\bsnapshot_download\s*\(/,
103  /\bhf_hub_download\s*\(/,
104  /(^|[\s;&|(])wget\s/,
105  /(^|[\s;&|(])curl\b[^|]*\s(-[a-zA-Z]*[oO]\b|--output\b|--remote-name\b)/,
106  /(^|[\s;&|(])aria2c\s/,
107  /(^|[\s;&|(])git\s+lfs\s+(pull|fetch|clone)\b/,
108  /(^|[\s;&|(])git\s+clone\s/,
109  /(^|[\s;&|(])(rsync|scp)\s[^;&|]*\s[\w.-]+@?[\w.-]*:[^\s]/,
110]
111
112function unquote(s: string): string {
113  return s.replace(/^['"]|['"]$/g, '')
114}
115
116function resolve(p: string, cwd: string, home: string | null): string {
117  if (p.startsWith('~') && home) return home + p.slice(1)
118  if (p.startsWith('/')) return p
119  return `${cwd.replace(/\/$/, '')}/${p.replace(/^\.\//, '')}`
120}
121
122/**
123 * Where a download command writes, or null when it is not a download or the
124 * target is the HF cache (HF_HOME already decides that). A leading `cd X &&`
125 * moves the cwd the same way the shell would.
126 */
127export function downloadTarget(cmd: string, cwd: string, home: string | null): string | null {
128  if (!DOWNLOADERS.some(r => r.test(cmd))) return null
129  const cdMatch = cmd.match(/(?:^|[;&|]\s*)cd\s+("[^"]+"|'[^']+'|\S+)\s*&&/)
130  const here = cdMatch?.[1] ? resolve(unquote(cdMatch[1]), cwd, home) : cwd
131  const flag = cmd.match(/(?:--local-dir|--local_dir|--cache-dir|--directory-prefix|--output|--dir|-P|-O|-o|-d)[=\s]+("[^"]+"|'[^']+'|[^\s;&|]+)/)
132  const kw = cmd.match(/\b(?:local_dir|cache_dir)\s*=\s*("[^"]+"|'[^']+'|[^\s,)]+)/)
133  const explicit = flag?.[1] ?? kw?.[1]
134  if (explicit) {
135    const p = unquote(explicit)
136    return p === '-' ? null : resolve(p, here, home)
137  }
138  // hf / snapshot_download with no target write to the HF cache, which HF_HOME places.
139  if (/(hf|huggingface-cli)\s+download\b|snapshot_download|hf_hub_download/.test(cmd)) return null
140  const scp = cmd.match(/(?:rsync|scp)\s.*\s("[^"]+"|'[^']+'|[^\s;&|]+)\s*$/)
141  if (scp?.[1] && !/:/.test(scp[1])) return resolve(unquote(scp[1]), here, home)
142  const clone = cmd.match(/git\s+(?:lfs\s+)?clone\s+(?:-\S+\s+)*\S+\s+("[^"]+"|'[^']+'|[^\s;&|-][^\s;&|]*)/)
143  if (clone?.[1]) return resolve(unquote(clone[1]), here, home)
144  return here
145}
146
147/** True when `path` belongs outside scratch: anywhere but scratch, /tmp, or /dev. */
148export function isOffScratch(path: string, scratchRoot: string): boolean {
149  const root = scratchRoot.replace(/\/$/, '')
150  return !(path === root || path.startsWith(`${root}/`) || path.startsWith('/scratch/') || path.startsWith('/tmp/') || path.startsWith('/dev/'))
151}
152