🌍↔️🐍 Limit the number of simultaneous executions

.changeset/cute-aliens-follow.md

@@ -0,0 +1,5 @@
+---
+'myst-cli': patch
+---
+
+Enhancement: Making it possible to configure the maximum number of simultaneous executions

docs/execute-notebooks.md

@@ -39,6 +39,14 @@ If you enable execution with the `--execute` flag as above, the following conten
 In order to execute your MyST content, you must install a Jupyter Server and the kernel needed to execute your code (e.g., the [IPython kernel](https://ipython.readthedocs.io/en/stable/), the [Xeus Python kernel](https://github.com/jupyter-xeus/xeus-python), or the [IRKernel](https://irkernel.github.io/).)
 :::
 
+## Limiting simultaneous executions
+
+By default, up to {math}`N-1` executable files are run concurrently, where {math}`N` is the number of available CPUs. 
+
+You can change this by using the `--execute-parallel <n>` option in your build command, where `<n>` sets the maximum number of executable documents that can run at the same time. This option is useful when your project includes executable content that is resource intensive or starts multiple subprocesses that might interfere with one another when run in parallel.
+
+For example, setting `--execute-parallel 1` will execute the documents sequentially.
+
 ## Show raw Python objects like modules and classes
 
 By default, MyST will suppress outputs from cells that return **raw** Python objects - like modules and classes - that don't have a string representation. For example with regular Python, you would observe this:

packages/myst-cli-utils/src/exec.ts

@@ -5,7 +5,7 @@ import type { Logger } from './types.js';
 
 function execWrapper(
   command: string,
-  options?: { cwd?: string },
+  options?: child_process.ExecOptionsWithStringEncoding,
   callback?: (error: child_process.ExecException | null, stdout: string, stderr: string) => void,
 ) {
   const childProcess = child_process.exec(command, options ?? {}, callback);

packages/myst-cli/package.json

@@ -45,6 +45,7 @@
     "@jupyterlab/services": "^7.3.0",
     "@reduxjs/toolkit": "^2.1.0",
     "adm-zip": "^0.5.10",
+    "async-mutex": "^0.5.0",
     "boxen": "^7.1.1",
     "cffjs": "^0.0.2",
     "chalk": "^5.2.0",

packages/myst-cli/src/cli/build.ts

@@ -17,6 +17,7 @@ import {
   makeWatchOption,
   makeCIOption,
   makeExecuteOption,
+  makeExecuteParallelOption,
   makeMaxSizeWebpOption,
   makeDOIBibOption,
   makeCffOption,
@@ -29,6 +30,7 @@ export function makeBuildCommand() {
     .description('Build PDF, LaTeX, Word and website exports from MyST files')
     .argument('[files...]', 'list of files to export')
     .addOption(makeExecuteOption('Execute Notebooks'))
+    .addOption(makeExecuteParallelOption())
     .addOption(makePdfOption('Build PDF output'))
     .addOption(makeTexOption('Build LaTeX outputs'))
     .addOption(makeTypstOption('Build Typst outputs'))

packages/myst-cli/src/cli/options.ts

@@ -1,4 +1,5 @@
 import { InvalidArgumentError, Option } from 'commander';
+import { cpus } from 'node:os';
 
 export const MYST_DOI_BIB_FILE = 'myst.doi.bib';
 
@@ -60,6 +61,19 @@ export function makeExecuteOption(description: string) {
   return new Option('--execute', description).default(false);
 }
 
+export function makeExecuteParallelOption() {
+  const defaultParallelism = Math.max(1, cpus().length - 1);
+  return new Option('--execute-parallel <n>', `Maximum number of notebooks to execute in parallel`)
+    .argParser((value) => {
+      const parsedValue = parseInt(value);
+      if (parsedValue < 1) {
+        throw new InvalidArgumentError('Must be an integer greater than or equal to 1.');
+      }
+      return parsedValue;
+    })
+    .default(defaultParallelism);
+}
+
 export function makeAllOption(description: string) {
   return new Option('-a, --all', description).default(false);
 }

packages/myst-cli/src/process/mdast.ts

@@ -184,14 +184,20 @@ export async function transformMdast(
 
   if (execute && !frontmatter.skip_execution) {
     const cachePath = path.join(session.buildPath(), 'execute');
-    await kernelExecutionTransform(mdast, vfile, {
-      basePath: session.sourcePath(),
-      cache: new LocalDiskCache<(IExpressionResult | IOutput[])[]>(cachePath),
-      sessionFactory: () => session.jupyterSessionManager(),
-      frontmatter: frontmatter,
-      ignoreCache: false,
-      errorIsFatal: false,
-      log: session.log,
+    const fileName = path.basename(file);
+    session.log.debug(`⏳ Waiting for execution slot: ${fileName}`);
+    await session.executionSemaphore.runExclusive(async () => {
+      session.log.debug(`▶️  Executing: ${fileName}`);
+      await kernelExecutionTransform(mdast, vfile, {
+        basePath: session.sourcePath(),
+        cache: new LocalDiskCache<(IExpressionResult | IOutput[])[]>(cachePath),
+        sessionFactory: () => session.jupyterSessionManager(),
+        frontmatter: frontmatter,
+        ignoreCache: false,
+        errorIsFatal: false,
+        log: session.log,
+      });
+      session.log.debug(`✅ Completed execution: ${fileName}`);
     });
   }
 

packages/myst-cli/src/process/site.ts

@@ -46,6 +46,7 @@ import type { TransformFn } from './mdast.js';
 import { finalizeMdast, postProcessMdast, transformMdast } from './mdast.js';
 import { toSectionedParts, buildHierarchy, sectionToHeadingLevel } from './search.js';
 import { SPEC_VERSION } from '../spec-version.js';
+import { cpus } from 'node:os';
 
 const WEB_IMAGE_EXTENSIONS = [
   ImageExtensions.mp4,
@@ -596,6 +597,7 @@ export async function processProject(
       }),
     ),
   );
+
   const pageReferenceStates = selectPageReferenceStates(session, pagesToTransform);
   // Handle all cross references
   await Promise.all(

packages/myst-cli/src/session/session.ts

@@ -10,6 +10,8 @@ import chalk from 'chalk';
 import { HttpsProxyAgent } from 'https-proxy-agent';
 import pLimit from 'p-limit';
 import type { Limit } from 'p-limit';
+import { Semaphore } from 'async-mutex';
+import { cpus } from 'node:os';
 import {
   findCurrentProjectAndLoad,
   findCurrentSiteAndLoad,
@@ -87,6 +89,7 @@ export class Session implements ISession {
   store: Store<RootState>;
   $logger: Logger;
   doiLimiter: Limit;
+  executionSemaphore: Semaphore;
 
   proxyAgent?: HttpsProxyAgent<string>;
   _shownUpgrade = false;
@@ -97,11 +100,20 @@ export class Session implements ISession {
     return this.$logger;
   }
 
-  constructor(opts: { logger?: Logger; doiLimiter?: Limit; configFiles?: string[] } = {}) {
+  constructor(
+    opts: {
+      logger?: Logger;
+      doiLimiter?: Limit;
+      executionSemaphore?: Semaphore;
+      configFiles?: string[];
+    } = {},
+  ) {
     this.API_URL = API_URL;
     this.configFiles = (opts.configFiles ? opts.configFiles : CONFIG_FILES).slice();
     this.$logger = opts.logger ?? chalkLogger(LogLevel.info, process.cwd());
     this.doiLimiter = opts.doiLimiter ?? pLimit(3);
+    this.executionSemaphore =
+      opts.executionSemaphore ?? new Semaphore(Math.max(1, cpus().length - 1));
     const proxyUrl = process.env.HTTPS_PROXY;
     if (proxyUrl) this.proxyAgent = new HttpsProxyAgent(proxyUrl);
     this.store = createStore(rootReducer);
@@ -195,6 +207,7 @@ export class Session implements ISession {
     const cloneSession = new Session({
       logger: this.log,
       doiLimiter: this.doiLimiter,
+      executionSemaphore: this.executionSemaphore,
       configFiles: this.configFiles,
     });
     await cloneSession.reload();

packages/myst-cli/src/session/types.ts

@@ -11,13 +11,15 @@ import type { PreRendererData, RendererData, SingleCitationRenderer } from '../t
 import type { SessionManager } from '@jupyterlab/services';
 import type MystTemplate from 'myst-templates';
 import type { PluginInfo } from 'myst-config';
+import type { Semaphore } from 'async-mutex';
 
 export type ISession = {
   API_URL: string;
   configFiles: string[];
   store: Store<RootState>;
   log: Logger;
   doiLimiter: Limit;
+  executionSemaphore: Semaphore;
   reload(): Promise<ISession>;
   clone(): Promise<ISession>;
   sourcePath(): string;

packages/mystmd/package.json

@@ -41,6 +41,7 @@
     "build": "npm-run-all -l clean copy:version -p build:cli"
   },
   "devDependencies": {
+    "async-mutex": "^0.5.0",
     "chalk": "^5.2.0",
     "commander": "^10.0.1",
     "core-js": "^3.31.1",

packages/mystmd/src/clirun.ts

@@ -2,10 +2,13 @@ import type { Command } from 'commander';
 import type { ISession, Session } from 'myst-cli';
 import { checkNodeVersion, getNodeVersion, logVersions } from 'myst-cli';
 import { chalkLogger, LogLevel } from 'myst-cli-utils';
+import { Semaphore } from 'async-mutex';
+import { cpus } from 'node:os';
 
 type SessionOpts = {
   debug?: boolean;
   config?: string;
+  executeParallel?: number;
 };
 
 export function clirun(
@@ -25,13 +28,18 @@ export function clirun(
     keepAlive?: boolean | ((...args: any[]) => boolean);
   },
 ) {
-  return async (...args: any[]) => {
-    const opts = program.opts() as SessionOpts;
+  return async function (this: Command, ...args: any[]) {
+    // Use options from 'this' merged with parent program options
+    // Needed to pass options from e.g. the build command to the session
+    const opts = { ...program.opts(), ...this.opts() } as SessionOpts;
     const logger = chalkLogger(opts?.debug ? LogLevel.debug : LogLevel.info, process.cwd());
     // Override default myst.yml if --config option is given.
     const configFiles = opts?.config ? [opts.config] : null;
-    const session = new sessionClass({ logger, configFiles });
+    const parallelCount = opts?.executeParallel ?? Math.max(1, cpus().length - 1);
+    const executionSemaphore = new Semaphore(parallelCount);
+    const session = new sessionClass({ logger, configFiles, executionSemaphore });
     await session.reload();
+    session.log.info(`🍡 Execution parallelism set to: ${parallelCount}`);
     const versions = await getNodeVersion(session);
     logVersions(session, versions);
     const versionsInstalled = await checkNodeVersion(session);

packages/mystmd/tests/concurrent-execution/first.md

@@ -0,0 +1,23 @@
+---
+kernelspec:
+  name: python3
+  display_name: Python 3
+execute:
+  cache: false
+---
+
+```{code-cell} python3
+import time
+import os
+import os.path
+
+# 1. Hang until the file is created
+while True:
+    if os.path.exists("output-second.txt"):
+        break
+    time.sleep(10e-3)
+
+# 2. Write our own marker file
+with open("output-first.txt", "w") as f:
+    f.write("First completed!")
+```

packages/mystmd/tests/concurrent-execution/myst.yml

@@ -0,0 +1,10 @@
+# See docs at: https://mystmd.org/guide/frontmatter
+version: 1
+project:
+  toc:
+    - file: first.md
+    - file: second.md
+  execute:
+    cache: false
+site:
+  template: ../templates/site/myst/book-theme

packages/mystmd/tests/concurrent-execution/second.md

@@ -0,0 +1,16 @@
+---
+kernelspec:
+  name: python3
+  display_name: Python 3
+execute:
+  cache: false
+---
+
+```{code-cell} python3
+import time
+import os
+
+# Create marker for other process
+with open("output-second.txt", "w") as f:
+    f.write("Second completed!")
+```

packages/mystmd/tests/endToEnd.spec.ts

@@ -10,6 +10,8 @@ type TestFile = {
 type TestCase = {
   title: string;
   cwd: string;
+  env?: Record<string, any>;
+  timeout?: number;
   command: string;
   outputs: {
     path: string;
@@ -41,7 +43,7 @@ describe.concurrent('End-to-end cli export tests', { timeout: 15000 }, () => {
   const cases = loadCases('exports.yml');
   test.each(
     cases.filter((c) => !only || c.title === only).map((c): [string, TestCase] => [c.title, c]),
-  )('%s', async (_, { cwd, command, outputs }) => {
+  )('%s', async (_, { cwd, env, command, outputs, timeout }) => {
     // Clean expected outputs if they already exist
     await Promise.all(
       outputs.map(async (output) => {
@@ -51,7 +53,7 @@ describe.concurrent('End-to-end cli export tests', { timeout: 15000 }, () => {
       }),
     );
     // Run CLI command
-    await exec(command, { cwd: resolve(cwd) });
+    await exec(command, { cwd: resolve(cwd), env: { ...process.env, ...env }, timeout });
     // Expect correct output
     outputs.forEach((output) => {
       expect(fs.existsSync(resolve(output.path))).toBeTruthy();

packages/mystmd/tests/exports.yml

@@ -553,3 +553,19 @@ cases:
       - path: hidden-pages/_build/html/index.html
       - path: hidden-pages/_build/html/hidden/index.html
       - path: hidden-pages/_build/html/not-hidden/index.html
+  - title: Serial execution
+    cwd: serial-execution
+    command: myst build --site --execute --execute-parallel 1
+    env:
+      MYST_TEST_SLEEP_MS: 4000
+      MYST_TEST_DELAY_MS: 2000
+    outputs:
+      - path: serial-execution/stop-first
+      - path: serial-execution/stop-second
+  - title: Concurrent execution
+    cwd: concurrent-execution
+    command: myst build --site --execute --execute-parallel 2
+    timeout: 5000
+    outputs:
+      - path: concurrent-execution/output-first.txt
+      - path: concurrent-execution/output-second.txt