@himalaya-quant/synapse
Version:
A lightweight TypeScript utility to spawn and interact with Python modules from Node.js with a native, message-based protocol over stdin/stdout.
253 lines • 10.2 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.InstanceManger = void 0;
const path_1 = require("path");
const fs_1 = require("fs");
const msgpack_1 = require("@msgpack/msgpack");
const rxjs_1 = require("rxjs");
const child_process_1 = require("child_process");
class InstanceManger {
inputsQueue = [];
instanceOutputs$ = new rxjs_1.Subject();
instanceLogs$ = new rxjs_1.Subject();
instanceInputs$ = new rxjs_1.Subject();
messageBuffer = Buffer.alloc(0);
instance;
currentInputResolver;
instanceInputStreamSubscription;
instanceOutputStreamSubscription;
/**
* Subscription to the instance logs.
* If you want your script logs to be seen by Synapse, in your python script
* always print to the stderr like:
*
* ```py
* import sys
*
* print("my log", file=sys.stderr)
* ```
*
* every log written this way will pass through this stream.
*
* @returns The observable to which you can subscribe to access
* the logs stream.
*/
get instanceLogs() {
return this.instanceLogs$.asObservable();
}
/**
* Calls the spawned python instance with the given input.
* Throws if the instance has not been spawned first, or if the script sends
* an error response.
*
* To send an error response, just send the usual message from the script,
* but passing a a dictionary with an "error" key and the error message
* as value. Eg: {"error": "my error message"}
*
* @param input Any simple JSON structure will be accepted.
* For more details see: https://msgpack.org/
* @param forceJSONParse Forcefully tries to parse the result. If it
* fails, will return the payload as it is.
*
* @returns The result returned from your python script.
* @throws {Error} If the instance has not been spawned or an error response
* is sent back from the python script
*/
call(input, forceJSONParse = false) {
if (!this.instance)
throw new Error(`Cannot send inputs to instance before spawning it.`);
return new Promise((resolver, rejector) => {
this.inputsQueue.push({
input,
resolver,
rejector,
parse: forceJSONParse,
});
if (this.inputsQueue.length === 1)
this.instanceInputs$.next(this.inputsQueue[0]);
});
}
/**
* Spawns a new python script instance and keeps it alive until dispose is
* called. After the spawning you can safely start sending messages to it.
* Throws if there's an error during the spawning process.
*
* What it does:
* - Postfixes the extension .py to the entrypoint if missing
* - Ensures that the paths are correct and the requirements.txt exists
* - Creates a dedicated Python virtual environment
* - Installs dependencies via requirements.txt
* - Spawns Python script as subprocess
* - Reuses instance until explicit disposal avoiding spawn overhead
*
* @param directory The path pointing to the python module directory.
* @param entrypoint The entrypoint file name.
*
* @returns A promise that resolves once the spawn completes.
* @throws {Error} If there's an error during the spawning process.
*/
async spawn(directory, entrypoint) {
if (this.instance)
return;
entrypoint = this.postfixExtension(entrypoint);
this.ensureExistsOrThrow(directory, entrypoint);
if (!(0, fs_1.existsSync)((0, path_1.join)(directory, '.venv'))) {
this.createVirtualEnv(directory);
await this.installDependencies(directory);
}
await this.spawnInstance(directory, entrypoint);
this.openSubscriptions();
}
/**
* Disposes the instance, closing the stdin stream, all the subscriptions
* and tries to kill the instance.
* Manages graceful and forceful termination, first tries with a SIGTERM, if
* after 500ms is not killed, will force a SIGKILL.
* Throws if after the SIGKILL the instance is still alive.
*
* After dispose has been called, you will have to call spawn again. Trying
* to send any messages after dispose, will result in an error.
*
* @returns Resolves once the dispose has been done.
* @throws {Error} If after forceful termination the instance is still alive
*/
async dispose() {
this.instanceInputStreamSubscription?.unsubscribe();
this.instanceOutputStreamSubscription?.unsubscribe();
if (!this.instance)
return;
this.instance.stdin.end();
this.instance.kill('SIGTERM');
const isTerminated = await this.waitForTermination(500);
if (!isTerminated) {
console.log('Instance did not close gracefully in 500ms. Forcing SIGKILL.');
this.instance.kill('SIGKILL');
await this.waitForTermination(500);
}
if (!this.instance.killed)
throw new Error(`Cannot kill instance with PID: ${this.instance.pid}`);
this.instance = null;
}
waitForTermination(timeout) {
return new Promise((resolve) => {
setTimeout(() => {
resolve(this.instance.killed);
}, timeout);
});
}
ensureExistsOrThrow(directory, entrypoint) {
if (!(0, fs_1.existsSync)(directory))
throw new Error(`Directory "${directory}" does not exist.`);
if (!(0, fs_1.existsSync)((0, path_1.join)(directory, entrypoint)))
throw new Error(`Entrypoint "${entrypoint}" does not exist.`);
if (!(0, fs_1.existsSync)((0, path_1.join)(directory, 'requirements.txt')))
throw new Error(`"requirements.txt" file is missing.`);
}
createVirtualEnv(directory) {
(0, child_process_1.spawnSync)('python', ['-m', 'venv', (0, path_1.join)(directory, '.venv')]);
}
installDependencies(directory) {
const python = this.getVirtualEnvPythonInterpreter(directory);
const requirementsPath = (0, path_1.join)(directory, 'requirements.txt');
const installDeps = (0, child_process_1.spawn)(python, [
'-m',
'pip',
'install',
'-r',
requirementsPath,
]);
return new Promise((resolve) => {
installDeps.stdout.on('data', (data) => {
console.log(`stdout: ${data}`);
});
installDeps.stderr.on('data', (data) => {
console.log(`stderr: ${data}`);
});
installDeps.on('close', async (code) => {
if (code === 0) {
resolve();
}
else {
throw new Error(`Dependencies install failure: ${code}`);
}
});
});
}
getVirtualEnvPythonInterpreter(directory) {
return (0, path_1.join)(directory, '.venv', 'bin', 'python');
}
async spawnInstance(directory, entrypoint) {
this.instance = (0, child_process_1.spawn)(this.getVirtualEnvPythonInterpreter(directory), [
(0, path_1.join)(directory, entrypoint),
]);
this.instance.stdout.on('data', (chunk) => {
this.handleChunk(chunk);
});
// we expect all messages (logs, errors, etc) that is not the actual
// response on the stderr stream so we don't touch the stdin encoding
this.instance.stderr.on('data', (chunk) => {
const msg = `🐍 Python: ${chunk.toString()}`;
this.instanceLogs$.next(msg);
});
}
packAndSend(message) {
const payload = (0, msgpack_1.encode)(message);
const lengthBuffer = Buffer.alloc(4); // 4 bytes for the payload size
lengthBuffer.writeUint32LE(payload.length, 0);
this.instance.stdin.write(Buffer.concat([lengthBuffer, payload]));
}
postfixExtension(entrypoint) {
return entrypoint.endsWith('.py') ? entrypoint : `${entrypoint}.py`;
}
handleChunk = (chunk) => {
this.messageBuffer = Buffer.concat([this.messageBuffer, chunk]);
while (this.messageBuffer.length >= 4) {
const messageLength = this.messageBuffer.readUInt32LE(0);
if (this.messageBuffer.length >= 4 + messageLength) {
const messagePayload = this.messageBuffer.subarray(4, 4 + messageLength);
const decoded = (0, msgpack_1.decode)(messagePayload);
this.instanceOutputs$.next(decoded);
this.messageBuffer = this.messageBuffer.subarray(4 + messageLength);
}
else {
// Not enough data yet
break;
}
}
};
openSubscriptions() {
this.instanceInputStreamSubscription = this.instanceInputs$
.pipe((0, rxjs_1.tap)(({ input }) => this.packAndSend(input)))
.subscribe();
this.instanceOutputStreamSubscription = this.instanceOutputs$
.pipe((0, rxjs_1.tap)((output) => {
const { parse, resolver, rejector } = this.inputsQueue.shift();
let result = this.extractResult(parse, output);
if (result.error)
rejector(result.error);
else
resolver(result);
if (this.inputsQueue.length)
this.instanceInputs$.next(this.inputsQueue[0]);
}))
.subscribe();
}
extractResult(parse, output) {
if (!parse)
return output;
let result;
try {
// console.log(output);
result = JSON.parse(output);
}
catch (e) {
let msg = `[ERROR] forced parsing failed. Expected parsable str from py: `;
msg += e instanceof Error ? e.message : JSON.stringify(e);
this.instanceLogs$.next(msg);
result = output;
}
return result;
}
}
exports.InstanceManger = InstanceManger;
//# sourceMappingURL=instance-manager.service.js.map