python-shell-cluster/src/index.ts

Utility to run python script in parallel from nodejs

import { PythonShell, Options } from "python-shell"

interface Task {
  task: string | Object
  resolve: Function
  reject: Function
}

class PythonShellCluster {
  private pyshells: PythonShell[]
  private queue: Task[]
  private scriptLocation: string

  constructor(numberOfinstances: number, path: string, options?: Options) {
    this.queue = []
    this.pyshells = []
    this.scriptLocation = path
    this.addInstances(numberOfinstances, options)
    console.log("Object inited")
  }

  addInstances(n: number, options?: Options) {
    for (let i = 0; i < n; i++) {
      this.pyshells.push(new PythonShell(this.scriptLocation, options))
    }
    console.log("Added " + n + " python-shells instances")
  }

  run(task: string | object) {
    console.log("Adding task to queue: " + task)
    return new Promise((resolve, reject) => {
      this.queue.push({
        task,
        resolve,
        reject,
      })
      console.log("Trying to execute task")
      const instance = this.pyshells.pop()
      if (instance) {
        console.log("Instance found")
        this.nextTask(instance)
      }
    })
  }
  private runner({ task, resolve, reject }: Task, instance: PythonShell) {
    console.log("Running task")
    let returnedInstance = instance
    instance.once("message", (...args) => {
      console.log("message recived: " + args)
      resolve(args)
      this.nextTask(returnedInstance)
    })

    instance.once("stderr", (...args) => {
      console.log("Error: ", args)
      reject(args)
      this.nextTask(returnedInstance)
    })

    instance.once("close", (...args) => {
      console.log("instance closed: " + args)
      reject(args)
      returnedInstance = new PythonShell(this.scriptLocation)
      this.nextTask(returnedInstance)
    })

    instance.send(task)
  }

  private nextTask(instance: PythonShell) {
    console.log("looking for next task")
    const task = this.queue.pop()
    if (task) {
      console.log("task found")
      this.runner(task, instance)
    } else {
      console.log("task not found")
      this.pyshells.push(instance)
    }
  }
}
export default PythonShellCluster