Slow memory leak in high load tfjs and discord.js application

Viewed 74

Context

  • This application's event functions (listener.js) execute roughly 14 times per second to meet traffic
  • listener.js is the file containing the event functions.
  • handler.js is used to handle the listener.js event.
  • sharder.js is the file to shard the application
  • index.js is file executed each time by sharder.js
  • gc.js is the file to manually call the garbage collector (this does not work, but was recommended)

System/Dependencies

  • node.js v16.13.1
  • discord.js v13.6.0
  • @tensorflow/tfjs v3.14.0
  • @tensorflow/tfjs-node v3.14.0

Problem

A memory leak is detectable after all shards of the bot are online. It is slight but noticeable and would require me to restart the node process every day (64GB of RAM on the host). All tensors are properly disposed of (tensors remain at 263 because the model is loaded outside of the event listener and is not disposed of). I have a listener to manually call the garbage collector, but it does not work. In the listener.js file I even null-ed all possible variables (unsure if this has any effect or not).

Question

Are there any aspects that I am overlooking which would cause the memory leak/are there any solutions to this?

(all code below)

listener.js

const { Readable } = require('stream');
const PImage = require('pureimage');
const tf = require(`@tensorflow/tfjs`)
const tfnode = require('@tensorflow/tfjs-node');

let nameArr = [
  // array of names here
]

let bufferToStream = (binary) => {
  let readableInstanceStream = new Readable({
    read() {
      this.push(binary);
      this.push(null);
    }
  });
  return readableInstanceStream;
}

const predict = async (imageUrl, model) => {

  let data = await fetch(imageUrl);
  let fileType = data.headers.get("Content-Type");
  let buffer = await data.buffer();

  let stream = bufferToStream(buffer);
  let image;
  if ((/png/).test(fileType)) {
    image = await PImage.decodePNGFromStream(stream);
  }
  else if ((/jpe?g/).test(fileType)) {
    image = await PImage.decodeJPEGFromStream(stream);
  }
  else {
    return `Error. Invalid file type.`
  }

  let rawTensor;
  rawTensor = tf.tidy(() => {
    let tensorImage;
    tensorImage = tf.browser.fromPixels(image).toFloat();
    tensorImage = tf.image.resizeNearestNeighbor(tensorImage, [model.inputs[0].shape[1], model.inputs[0].shape[2]]);
    let offset = tf.scalar(127.5);
    tensorImage = tensorImage.sub(offset).div(offset);
    offset = null;
    tensorImage = tensorImage.reshape([1, model.inputs[0].shape[1], model.inputs[0].shape[2], model.inputs[0].shape[3]]);

    return model.predict(tensorImage);
  });

  let classes = []
  for (let i = 1; i < 181; i++) {
    classes.push(`${i}`)
  }

  let sorted = tf.topk(rawTensor, classes.length);
  let predictions = [ sorted.values.arraySync(), sorted.indices.arraySync() ];

  let rawArray;
  rawArray = await rawTensor.data();
  rawArray = Array.from(rawArray);

  tf.dispose([rawTensor, sorted])

  let predInd = predictions[1][0][0];
  let predVal = (predictions[0][0][0]*100).toFixed(2);
  let msg = `${classes[predInd]} (${predVal}%) -`;

  data = null;
  fileType = null;
  buffer = null;
  image = null;
  rawTensor = null;
  classes = null;
  sorted = null;
  predictions = null;
  rawArray = null;
  predInd = null;
  predVal = null;
  i = null;

  return msg
};

module.exports = {
  event: 'messageCreate',
  run: async (message, client, Discord, model) => {

    let mb = message.embeds[0];
    if (!mb) return;
    if (mb.title) {

      var link = mb.image[`proxyURL`];
      let first = Date.now()
      let prediction = await predict(`${link}`, model)
      let second = Date.now()

      let pred1 = prediction.split(` `)
      let pred2 = nameArr[((pred1[0]*1)-1)]
      let logPred = `${pred2} ${pred1[1]} ${pred1[2]} ${second-first}ms`
      console.log(logPred)
      message.channel.send(logPred)

      mb = null;
      link = null;
      first = null;
      prediction = null;
      second = null;
      pred1 = null;
      pred2 = null;
      x = null;
      logPred = null;
    }

  },
};

handler.js

    if (err) return console.error(err);
    files.forEach(async (file) => {
        const eventFunction = require(`./../events/${folder}${file}`);
        if (eventFunction.disabled) return;

        const event = eventFunction.event || file.split('.')[0];
        const emitter =
            (typeof eventFunction.emitter === 'string'
                ? client[eventFunction.emitter]
                : eventFunction.emitter) || client;
        const once = eventFunction.once;

        try {
            emitter[once ? 'once' : 'on'](event, (...args) =>
                eventFunction.run(...args, client, Discord, model),
            );
        }
        catch (error) {
            console.error(error.stack);
        }
    });

};

sharder.js

const { token } = require('./config.json');

const manager = new ShardingManager('./index.js', { token: `${token}` });

manager.on('shardCreate', async shard => {
  console.log(`Launched shard ${shard.id}`)
});

manager.spawn({ amount: 90 , delay: 10000, timeout: 1 * 1000 * 60 })

index.js

const Discord = require('discord.js');
const { token } = require('./config.json');
const client = new Discord.Client({ intents: [ Discord.Intents.FLAGS.GUILDS, Discord.Intents.FLAGS.GUILD_MESSAGES ] });
const db = require("quick.db");
const eco = { bot: new db.table("bot") };
module.exports = { eco };

const folders = [ "interactionCreate/" ]
for (let i = 0; i < folders.length; i++) {
    const folder = folders[i]
    fs.readdir(`./events/${folder}`, async (err, files) => {
        const eventHandler = require("./data/eventHandler.js");
        const tf = require(`@tensorflow/tfjs-node`);
        let model = await tf.loadLayersModel(`file://./models/model.json`);

        eventHandler(err, files, client, Discord, folder, model);
    });
}

client.login(token);

gc.js

module.exports = {
  event: 'messageCreate',
  run: async (message, client, Discord) => {

    if (!message.content.startsWith(`clear`)) return

    const col = async (client) => {
      try {
        if (global.gc) {global.gc();}
        console.log(`Garbage Collected`)
      } catch (e) {
        console.log(`Unable to collect`)
      }
    }

    const exec = async () => {
      await client.shard.broadcastEval(col)
    }

    await exec();

  },
};
2 Answers

first, use either @tensorflow/tfjs OR @tensorflow/tfjs-node - never both.

second, if there is no tensor leak, then its something on heap and you want to read up on nodejs profiling (its not trivial and too long for how-to here).

As mentioned in my previous comment, I have experienced a similar problem when batch processing many images, also the allocated was never freed and remained in memory. I found that tf.browser.fromPixels(image) is responsible for the huge memory usage but only after processing larger images.

So I scaled down all images to the preferred size of the model (320 x 320 for me) before sending them into fromPixels it seems that this solved my memory issues for now.

Related