1
/*
2
 *  This file is part of CoCalc: Copyright © 2023 Sagemath, Inc.
3
 *  License: MS-RSL – see LICENSE.md for details
4
 */
5

6
/*
7
Launching and managing Jupyter kernels in a pool for
8
performance.
9
*/
10

11
import { reuseInFlight } from "@cocalc/util/reuse-in-flight";
12
import { delay } from "awaiting";
13
import json from "json-stable-stringify";
14
import nodeCleanup from "node-cleanup";
15
import { unlinkSync } from "node:fs";
16
import { mkdir, readFile, writeFile } from "node:fs/promises";
17
import getLogger from "@cocalc/backend/logger";
18
import createChdirCommand from "@cocalc/util/jupyter-api/chdir-commands";
19
import createSetenvCommand from "@cocalc/util/jupyter-api/setenv-commands";
20
import { exists, unlink } from "@cocalc/backend/misc/async-utils-node";
21
import { getLanguage } from "@cocalc/jupyter/kernel/kernel-data";
22
import launchJupyterKernelNoPool, {
23
  LaunchJupyterOpts,
24
  SpawnedKernel,
25
} from "@cocalc/jupyter/kernel/launch-kernel";
26
import {
27
  getConfig,
28
  getConfigDir,
29
  getLaunchDelayMS,
30
  getSize,
31
  getTimeoutS,
32
} from "./pool-params";
33
import { getAbsolutePathFromHome } from "@cocalc/jupyter/util/fs";
34

35
// any kernel name whose lowercase representation contains one of these strings
36
// will never use the pool. See https://github.com/sagemathinc/cocalc/issues/7041
37
const BLACKLIST = ["julia"];
38

39
function isBlacklisted(kernel: string) {
40
  const s = kernel.toLowerCase();
41
  for (const n of BLACKLIST) {
42
    if (s.includes(n)) {
43
      return true;
44
    }
45
  }
46
  return false;
47
}
48

49
export type { LaunchJupyterOpts, SpawnedKernel };
50

51
const log = getLogger("jupyter:pool");
52

53
async function writeConfig(content: string): Promise<void> {
54
  try {
55
    // harmless to call if dir already exists
56
    await mkdir(getConfigDir(), { recursive: true });
57
    await writeFile(getConfig(), content);
58
  } catch (error) {
59
    log.debug("Error writeConfig -- ", error);
60
  }
61
}
62

63
async function readConfig(): Promise<string> {
64
  return (await readFile(getConfig())).toString();
65
}
66

67
const POOL: { [key: string]: SpawnedKernel[] } = {};
68
const EXPIRE: { [key: string]: number } = {};
69

70
// Make key for cache that describes this kernel.  We explicitly omit
71
// the parameters that aren't generic and would make it not possible to
72
// put this in a pool:
73
//   - opts.cwd : current working directory
74
function makeKey({ name, opts }) {
75
  // Copy of opts but delete opts.cwd and opts.env.COCALC_JUPYTER_FILENAME.
76
  // We don't change opts though!
77
  const opts0 = { ...opts };
78
  delete opts0.cwd;
79
  opts0.env = { ...opts.env };
80
  delete opts0.env.COCALC_JUPYTER_FILENAME;
81
  return json({ name, opts: opts0 })!;
82
}
83

84
export default async function launchJupyterKernel(
85
  name: string, // name of the kernel
86
  opts: LaunchJupyterOpts,
87
  size_arg?: number, // min number of these in the pool
88
  timeout_s_arg?: number,
89
): Promise<SpawnedKernel> {
90
  const size: number = size_arg ?? getSize();
91
  const timeout_s: number = timeout_s_arg ?? getTimeoutS();
92
  if (isBlacklisted(name)) {
93
    log.debug(`not using kernel pool for ${name} because it is blacklisted`);
94
    return await launchJupyterKernelNoPool(name, opts);
95
  }
96
  let language;
97
  try {
98
    language = await getLanguage(name);
99
  } catch (error) {
100
    log.error("Failed to get language of kernel -- not using pool", error);
101
    return await launchJupyterKernelNoPool(name, opts);
102
  }
103

104
  let initCode: string[] = [];
105
  if (opts.cwd) {
106
    try {
107
      const absPath = getAbsolutePathFromHome(opts.cwd);
108
      initCode.push(createChdirCommand(language, absPath));
109
    } catch (error) {
110
      log.error("Failed to get chdir command -- not using pool", error);
111
      return await launchJupyterKernelNoPool(name, opts);
112
    }
113
  }
114
  if (opts.env?.COCALC_JUPYTER_FILENAME) {
115
    try {
116
      initCode.push(
117
        createSetenvCommand(
118
          language,
119
          "COCALC_JUPYTER_FILENAME",
120
          opts.env.COCALC_JUPYTER_FILENAME,
121
        ),
122
      );
123
    } catch (error) {
124
      log.error("Failed to get setenv command -- not using pool", error);
125
      return await launchJupyterKernelNoPool(name, opts);
126
    }
127
  }
128

129
  const key = makeKey({ name, opts });
130
  log.debug("launchJupyterKernel", key.slice(0, 30));
131
  try {
132
    if (POOL[key] == null) {
133
      POOL[key] = [];
134
    }
135
    if (POOL[key].length > 0) {
136
      const kernel = POOL[key].shift();
137
      replenishPool(key, size, timeout_s);
138
      return { ...(kernel as SpawnedKernel), initCode };
139
    }
140
    const kernel = await launchJupyterKernelNoPool(name, opts);
141

142
    // we don't start replenishing the pool until the kernel is initialized,
143
    // since we don't want to slow down creating the kernel itself!
144
    replenishPool(key, size, timeout_s);
145

146
    // we do NOT include the initCode here; it's not needed since this kernel
147
    // isn't from the pool.
148
    return kernel;
149
  } catch (error) {
150
    log.error("Failed to launch Jupyter kernel", error);
151
    throw error;
152
  }
153
}
154

155
// Don't replenish pool for same key twice at same time, or
156
// pool could end up a little too big.
157
const replenishPool = reuseInFlight(
158
  async (key: string, size_arg?: number, timeout_s_arg?: number) => {
159
    const { name, opts } = JSON.parse(key);
160
    if (isBlacklisted(name)) {
161
      log.debug(
162
        "replenishPool",
163
        key.slice(0, 30),
164
        ` -- skipping since ${name} is blacklisted`,
165
      );
166
      return;
167
    }
168
    const size: number = size_arg ?? getSize();
169
    const timeout_s: number = timeout_s_arg ?? getTimeoutS();
170
    log.debug("replenishPool", key.slice(0, 30), { size, timeout_s });
171
    try {
172
      if (POOL[key] == null) {
173
        POOL[key] = [];
174
      }
175
      const pool = POOL[key];
176
      while (pool.length < size) {
177
        log.debug("replenishPool - creating a kernel", key.slice(0, 30));
178
        writeConfig(key);
179
        await delay(getLaunchDelayMS());
180
        const kernel = await launchJupyterKernelNoPool(name, opts);
181
        pool.push(kernel);
182
        EXPIRE[key] = Math.max(EXPIRE[key] ?? 0, Date.now() + 1000 * timeout_s);
183
      }
184
    } catch (error) {
185
      log.error(
186
        "Failed to replenish Jupyter kernel pool",
187
        key.slice(0, 30),
188
        error,
189
      );
190
      throw error;
191
    }
192
  },
193
  {
194
    createKey: (args) => args[0],
195
  },
196
);
197

198
/*
199
If there is nothing in the pool, find the newest non-hidden ipynb files in
200
the current directory or in any immediate subdirectory.  It is a JSON file,
201
and we parse the
202

203
*/
204
async function fillWhenEmpty() {
205
  for (const key in POOL) {
206
    if (POOL[key].length > 0) {
207
      // nothing to do
208
      return;
209
    }
210
  }
211
  // pool is empty, so possibly put something in it.
212
  try {
213
    // this can throw, e.g., a corrupt file
214
    const key = await readConfig();
215
    if (key) {
216
      // this can definitely throw, e.g., change image and then available kernels change.  No need to crash the entire project in that case!
217
      await replenishPool(key);
218
    }
219
  } catch {}
220
}
221

222
async function maintainPool() {
223
  log.debug("maintainPool");
224
  const now = Date.now();
225
  for (const key in EXPIRE) {
226
    if (EXPIRE[key] < now) {
227
      log.debug("maintainPool -- expiring key=", key.slice(0, 30));
228
      const pool = POOL[key] ?? [];
229
      while (pool.length > 0) {
230
        const kernel = pool.shift() as SpawnedKernel;
231
        try {
232
          await killKernel(kernel);
233
        } catch (error) {
234
          // won't happen
235
          log.error("Failed to kill Jupyter kernel", error);
236
        }
237
      }
238
    }
239
  }
240
  fillWhenEmpty();
241
}
242

243
export function init() {
244
  // DO NOT create the pool if we're running under jest testing, since
245
  // then tests don't exit cleanly.
246
  if (process.env.NODE_ENV != "test") {
247
    setInterval(maintainPool, 30 * 1000);
248
    maintainPool();
249
  }
250
}
251

252
nodeCleanup(() => {
253
  for (const key in POOL) {
254
    for (const kernel of POOL[key]) {
255
      try {
256
        process.kill(-kernel.spawn.pid, "SIGTERM");
257
        unlinkSync(kernel.connectionFile);
258
      } catch (_) {}
259
    }
260
  }
261
});
262

263
export async function killKernel(kernel: SpawnedKernel) {
264
  kernel.spawn?.removeAllListeners();
265
  try {
266
    if (kernel.spawn?.pid) {
267
      log.debug("killKernel pid=", kernel.spawn.pid);
268
      try {
269
        process.kill(-kernel.spawn.pid, "SIGTERM");
270
      } catch (error) {
271
        log.error("Failed to send SIGTERM to Jupyter kernel", error);
272
      }
273
    }
274
    kernel.spawn?.close?.();
275
    if (await exists(kernel.connectionFile)) {
276
      try {
277
        await unlink(kernel.connectionFile);
278
      } catch (error) {
279
        log.error(
280
          `Failed to delete Jupyter kernel connection file ${kernel.connectionFile}`,
281
          error,
282
        );
283
      }
284
    }
285
  } catch (error) {
286
    log.error("Failed to kill Jupyter kernel", error);
287
  }
288
}
289

290