@livekit/agents 0.1.0 → 0.3.0

package/src/ipc/job_main.ts

@@ -1,83 +1,163 @@
 // SPDX-FileCopyrightText: 2024 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
-import { type JobAssignment, ServerMessage } from '@livekit/protocol';
-import { Room } from '@livekit/rtc-node';
-import { type ChildProcess, fork } from 'child_process';
+import { Room, RoomEvent } from '@livekit/rtc-node';
+import type { ChildProcess } from 'child_process';
+import { fork } from 'child_process';
 import { EventEmitter, once } from 'events';
-import { JobContext } from '../job_context.js';
-import { log } from '../log.js';
-import { IPC_MESSAGE, type JobMainArgs, type Message, type Ping } from './protocol.js';
+import type { Logger } from 'pino';
+import { fileURLToPath } from 'url';
+import { type Agent, isAgent } from '../generator.js';
+import type { RunningJobInfo } from '../job.js';
+import { JobContext } from '../job.js';
+import { JobProcess } from '../job.js';
+import { initializeLogger, log } from '../log.js';
+import { defaultInitializeProcessFunc } from '../worker.js';
+import type { IPCMessage } from './message.js';
 
-export const runJob = (args: JobMainArgs): ChildProcess => {
-  return fork(import.meta.filename, [args.raw, args.entry, args.fallbackURL]);
+const ORPHANED_TIMEOUT = 15 * 1000;
+
+type StartArgs = {
+  agentFile: string;
+  // userArguments: unknown;
+};
+
+type JobTask = {
+  ctx: JobContext;
+  task: Promise<void>;
+};
+
+export const runProcess = (args: StartArgs): ChildProcess => {
+  return fork(fileURLToPath(import.meta.url), [args.agentFile]);
+};
+
+const startJob = (
+  proc: JobProcess,
+  func: (ctx: JobContext) => Promise<void>,
+  info: RunningJobInfo,
+  closeEvent: EventEmitter,
+  logger: Logger,
+): JobTask => {
+  let connect = false;
+  let shutdown = false;
+
+  const room = new Room();
+  room.on(RoomEvent.Disconnected, () => {
+    closeEvent.emit('close', false);
+  });
+
+  const onConnect = () => {
+    connect = true;
+  };
+  const onShutdown = (reason: string) => {
+    shutdown = true;
+    closeEvent.emit('close', true, reason);
+  };
+
+  const ctx = new JobContext(proc, info, room, onConnect, onShutdown);
+
+  const task = new Promise<void>(async () => {
+    const unconnectedTimeout = setTimeout(() => {
+      if (!(connect || shutdown)) {
+        logger.warn(
+          'room not connect after job_entry was called after 10 seconds, ',
+          'did you forget to call ctx.connect()?',
+        );
+      }
+    }, 10000);
+    func(ctx).finally(() => clearTimeout(unconnectedTimeout));
+
+    await once(closeEvent, 'close').then((close) => {
+      logger.debug('shutting down');
+      process.send!({ case: 'exiting', value: { reason: close[1] } });
+    });
+
+    await room.disconnect();
+    logger.debug('disconnected from room');
+
+    const shutdownTasks = [];
+    for (const callback of ctx.shutdownCallbacks) {
+      shutdownTasks.push(callback());
+    }
+    await Promise.all(shutdownTasks).catch(() => logger.error('error while shutting down the job'));
+
+    process.send!({ case: 'done' });
+    process.exit();
+  });
+
+  return { ctx, task };
 };
 
 if (process.send) {
   // process.argv:
   //   [0] `node'
   //   [1] import.meta.filename
-  //   [2] proto.JobAssignment, serialized to JSON string
-  //   [3] import.meta.filename of function containing entry file
-  //   [4] fallback URL in case JobAssignment.url is empty
-
-  const msg = new ServerMessage();
-  msg.fromJsonString(process.argv[2]);
-  const args = msg.message.value as JobAssignment;
-
-  const room = new Room();
-  const closeEvent = new EventEmitter();
-  let shuttingDown = false;
-  let closed = false;
-
-  process.on('message', (msg: Message) => {
-    if (msg.type === IPC_MESSAGE.ShutdownRequest) {
-      shuttingDown = true;
-      closed = true;
-      closeEvent.emit('close');
-    } else if (msg.type === IPC_MESSAGE.Ping) {
-      process.send!({
-        type: IPC_MESSAGE.Pong,
-        lastTimestamp: (msg as Ping).timestamp,
-        timestamp: Date.now(),
-      });
+  //   [2] import.meta.filename of function containing entry file
+  const moduleFile = process.argv[2];
+  const agent: Agent = await import(moduleFile).then((module) => {
+    const agent = module.default;
+    if (agent === undefined || !isAgent(agent)) {
+      throw new Error(`Unable to load agent: Missing or invalid default export in ${moduleFile}`);
     }
+    return agent;
   });
+  if (!agent.prewarm) {
+    agent.prewarm = defaultInitializeProcessFunc;
+  }
 
   // don't do anything on C-c
+  // this is handled in cli, triggering a termination of all child processes at once.
   process.on('SIGINT', () => {});
 
-  const conn = room.connect(args.url || process.argv[4], args.token);
+  await once(process, 'message').then(([msg]: IPCMessage[]) => {
+    if (msg.case !== 'initializeRequest') {
+      throw new Error('first message must be InitializeRequest');
+    }
+    initializeLogger(msg.value.loggerOptions);
+  });
+  const proc = new JobProcess();
+  let logger = log().child({ pid: proc.pid });
 
-  const start = () => {
-    if (room.isConnected && !closed) {
-      process.send!({ type: IPC_MESSAGE.StartJobResponse });
+  logger.debug('initializing job runner');
+  agent.prewarm(proc);
+  logger.debug('job runner initialized');
+  process.send({ case: 'initializeResponse' });
 
-      // here we import the file containing the exported entry function, and call it.
-      // the file must export default an Agent, usually using defineAgent().
-      import(process.argv[3]).then((agent) => {
-        agent.default.entry(new JobContext(closeEvent, args.job!, room));
-      });
-    }
-  };
+  let job: JobTask | undefined = undefined;
+  const closeEvent = new EventEmitter();
 
-  new Promise(() => {
-    conn
-      .then(() => {
-        if (!closed) start();
-      })
-      .catch((err) => {
-        if (!closed) process.send!({ type: IPC_MESSAGE.StartJobResponse, err });
-      });
-  });
+  const orphanedTimeout = setTimeout(() => {
+    logger.warn('process orphaned, shutting down');
+    process.exit();
+  }, ORPHANED_TIMEOUT);
 
-  await once(closeEvent, 'close');
-  log.debug('disconnecting from room');
-  await room.disconnect();
-  if (shuttingDown) {
-    process.send({ type: IPC_MESSAGE.ShutdownResponse });
-  } else {
-    process.send({ type: IPC_MESSAGE.UserExit });
-  }
-  process.exit();
+  process.on('message', (msg: IPCMessage) => {
+    switch (msg.case) {
+      case 'pingRequest': {
+        orphanedTimeout.refresh();
+        process.send!({
+          case: 'pongResponse',
+          value: { lastTimestamp: msg.value.timestamp, timestamp: Date.now() },
+        });
+        break;
+      }
+      case 'startJobRequest': {
+        if (job) {
+          throw new Error('job task already running');
+        }
+
+        logger = logger.child({ jobID: msg.value.runningJob.job.id });
+
+        job = startJob(proc, agent.entry, msg.value.runningJob, closeEvent, logger);
+        logger.debug('job started');
+        break;
+      }
+      case 'shutdownRequest': {
+        if (!job) {
+          break;
+        }
+        closeEvent.emit('close', '');
+      }
+    }
+  });
 }

package/src/ipc/message.ts

@@ -0,0 +1,39 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { RunningJobInfo } from '../job.js';
+import type { LoggerOptions } from '../log.js';
+
+export type IPCMessage =
+  | {
+      case: 'initializeRequest';
+      value: { loggerOptions: LoggerOptions };
+    }
+  | {
+      case: 'initializeResponse';
+      value: undefined;
+    }
+  | {
+      case: 'pingRequest';
+      value: { timestamp: number };
+    }
+  | {
+      case: 'pongResponse';
+      value: { lastTimestamp: number; timestamp: number };
+    }
+  | {
+      case: 'startJobRequest';
+      value: { runningJob: RunningJobInfo };
+    }
+  | {
+      case: 'shutdownRequest';
+      value: { reason?: string };
+    }
+  | {
+      case: 'exiting';
+      value: { reason?: string };
+    }
+  | {
+      case: 'done';
+      value: undefined;
+    };

package/src/ipc/proc_job_executor.ts

@@ -0,0 +1,162 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { ChildProcess } from 'child_process';
+import { once } from 'events';
+import type { RunningJobInfo } from '../job.js';
+import { log, loggerOptions } from '../log.js';
+import { Future } from '../utils.js';
+import type { ProcOpts } from './job_executor.js';
+import { JobExecutor } from './job_executor.js';
+import type { IPCMessage } from './message.js';
+
+export class ProcJobExecutor extends JobExecutor {
+  #opts: ProcOpts;
+  #started = false;
+  #closing = false;
+  #runningJob?: RunningJobInfo = undefined;
+  #proc?: ChildProcess;
+  #pingInterval?: ReturnType<typeof setInterval>;
+  #pongTimeout?: ReturnType<typeof setTimeout>;
+  #init = new Future();
+  #join = new Future();
+  #logger = log().child({ runningJob: this.#runningJob });
+
+  constructor(agent: string, initializeTimeout: number, closeTimeout: number) {
+    super();
+    this.#opts = {
+      agent,
+      initializeTimeout,
+      closeTimeout,
+    };
+  }
+
+  get started(): boolean {
+    return this.#started;
+  }
+
+  get runningJob(): RunningJobInfo | undefined {
+    return this.#runningJob;
+  }
+
+  async start() {
+    if (this.#started) {
+      throw new Error('runner already started');
+    } else if (this.#closing) {
+      throw new Error('runner is closed');
+    }
+
+    this.#proc = await import('./job_main.js').then((m) =>
+      m.runProcess({
+        agentFile: this.#opts.agent,
+      }),
+    );
+
+    this.#started = true;
+    this.run();
+  }
+
+  async run() {
+    await this.#init.await;
+
+    this.#pingInterval = setInterval(() => {
+      this.#proc!.send({ case: 'pingRequest', value: { timestamp: Date.now() } });
+    }, this.PING_INTERVAL);
+
+    this.#pongTimeout = setTimeout(() => {
+      this.#logger.warn('job is unresponsive');
+      clearTimeout(this.#pongTimeout);
+      clearInterval(this.#pingInterval);
+      this.#proc!.kill();
+      this.#join.resolve();
+    }, this.PING_TIMEOUT);
+
+    const listener = (msg: IPCMessage) => {
+      switch (msg.case) {
+        case 'pongResponse': {
+          const delay = Date.now() - msg.value.timestamp;
+          if (delay > this.HIGH_PING_THRESHOLD) {
+            this.#logger.child({ delay }).warn('job executor is unresponsive');
+          }
+          this.#pongTimeout?.refresh();
+          break;
+        }
+        case 'exiting': {
+          this.#logger.child({ reason: msg.value.reason }).debug('job exiting');
+          break;
+        }
+        case 'done': {
+          this.#closing = true;
+          this.#proc!.off('message', listener);
+          this.#join.resolve();
+          break;
+        }
+      }
+    };
+    this.#proc!.on('message', listener);
+    this.#proc!.on('error', (err) => {
+      if (this.#closing) return;
+      this.#logger.child({ err }).warn('job process exited unexpectedly');
+      clearTimeout(this.#pongTimeout);
+      clearInterval(this.#pingInterval);
+      this.#join.resolve();
+    });
+
+    await this.#join.await;
+  }
+
+  async join() {
+    if (!this.#started) {
+      throw new Error('runner not started');
+    }
+
+    await this.#join.await;
+  }
+
+  async initialize() {
+    const timer = setTimeout(() => {
+      const err = new Error('runner initialization timed out');
+      this.#init.reject(err);
+      throw err;
+    }, this.#opts.initializeTimeout);
+    this.#proc!.send({ case: 'initializeRequest', value: { loggerOptions } });
+    await once(this.#proc!, 'message').then(([msg]: IPCMessage[]) => {
+      clearTimeout(timer);
+      if (msg.case !== 'initializeResponse') {
+        throw new Error('first message must be InitializeResponse');
+      }
+    });
+    this.#init.resolve();
+  }
+
+  async close() {
+    if (!this.#started) {
+      return;
+    }
+    this.#closing = true;
+
+    if (!this.#runningJob) {
+      this.#proc!.kill();
+      this.#join.resolve();
+    }
+
+    this.#proc!.send({ case: 'shutdownRequest' });
+
+    const timer = setTimeout(() => {
+      this.#logger.error('job shutdown is taking too much time');
+    }, this.#opts.closeTimeout);
+    await this.#join.await.then(() => {
+      clearTimeout(timer);
+      clearTimeout(this.#pongTimeout);
+      clearInterval(this.#pingInterval);
+    });
+  }
+
+  async launchJob(info: RunningJobInfo) {
+    if (this.#runningJob) {
+      throw new Error('executor already has a running job');
+    }
+    this.#runningJob = info;
+    this.#proc!.send({ case: 'startJobRequest', value: { runningJob: info } });
+  }
+}

package/src/ipc/proc_pool.ts

@@ -0,0 +1,109 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { RunningJobInfo } from '../job.js';
+import { Mutex, Queue } from '../utils.js';
+import type { JobExecutor } from './job_executor.js';
+import { ProcJobExecutor } from './proc_job_executor.js';
+
+export class ProcPool {
+  agent: string;
+  initializeTimeout: number;
+  closeTimeout: number;
+  executors: JobExecutor[] = [];
+  tasks: Promise<void>[] = [];
+  started = false;
+  closed = false;
+  controller = new AbortController();
+  initMutex = new Mutex();
+  procMutex: Mutex;
+  procUnlock?: () => void;
+  warmedProcQueue = new Queue<JobExecutor>();
+
+  constructor(
+    agent: string,
+    numIdleProcesses: number,
+    initializeTimeout: number,
+    closeTimeout: number,
+  ) {
+    this.agent = agent;
+    this.procMutex = new Mutex(numIdleProcesses);
+    this.initializeTimeout = initializeTimeout;
+    this.closeTimeout = closeTimeout;
+  }
+
+  get processes(): JobExecutor[] {
+    return this.executors;
+  }
+
+  getByJobId(id: string): JobExecutor | null {
+    return this.executors.find((x) => x.runningJob && x.runningJob.job.id === id) || null;
+  }
+
+  async launchJob(info: RunningJobInfo) {
+    const proc = await this.warmedProcQueue.get();
+    if (this.procUnlock) {
+      this.procUnlock();
+      this.procUnlock = undefined;
+    }
+    await proc.launchJob(info);
+  }
+
+  async procWatchTask() {
+    const proc = new ProcJobExecutor(this.agent, this.initializeTimeout, this.closeTimeout);
+
+    try {
+      this.executors.push(proc);
+
+      const unlock = await this.initMutex.lock();
+      if (this.closed) {
+        return;
+      }
+
+      await proc.start();
+      try {
+        await proc.initialize();
+        await this.warmedProcQueue.put(proc);
+      } catch {
+        if (this.procUnlock) {
+          this.procUnlock();
+          this.procUnlock = undefined;
+        }
+      }
+
+      unlock();
+      await proc.join();
+    } finally {
+      this.executors.splice(this.executors.indexOf(proc));
+    }
+  }
+
+  start() {
+    if (this.started) {
+      return;
+    }
+
+    this.started = true;
+    this.run(this.controller.signal);
+  }
+
+  async run(signal: AbortSignal) {
+    while (!signal.aborted) {
+      this.procUnlock = await this.procMutex.lock();
+      const task = this.procWatchTask();
+      this.tasks.push(task);
+      task.finally(() => this.tasks.splice(this.tasks.indexOf(task)));
+    }
+  }
+
+  async close() {
+    if (!this.started) {
+      return;
+    }
+    this.closed = true;
+    this.controller.abort();
+    this.warmedProcQueue.items.forEach((e) => e.close());
+    this.executors.forEach((e) => e.close());
+    await Promise.allSettled(this.tasks);
+  }
+}