the process crashes when stream() in ChatLlamaCpp is called in parallel.

[Motiveko]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import { ChatLlamaCpp } from "@langchain/community/chat_models/llama_cpp";
import {ChatPromptTemplate} from '@langchain/core/prompts';

const modelPath = '/Users/donggi/Library/Application Support/amigo/models/_Meta-Llama-3.1-8B-Instruct-Q4_K_S.gguf'
const model = new ChatLlamaCpp({modelPath, temperature: 0.5});
const template = ChatPromptTemplate.fromMessages([
  ['system', 'answer the question.'],
  ['human', '{input}'],
]);
const run = async () => {
  const prompt = await template.invoke({input: 'hi. what is your name?'})
  const stream = await model.stream(prompt);
  for await (const chunk of stream) {
    console.log(chunk.content);
  }

  console.log('end');
}
run();
run();

Error Message and Stack Trace (if applicable)

...
llama_new_context_with_model:        CPU  output buffer size =     0.98 MiB
llama_new_context_with_model:      Metal compute buffer size =   507.00 MiB
llama_new_context_with_model:        CPU compute buffer size =    23.01 MiB
llama_new_context_with_model: graph nodes  = 1690
llama_new_context_with_model: graph splits = 2
zsh: segmentation fault  node src/node-llama-cpp.js

Description

I want to run the Gemma model using node-llama-cpp through langchain.js (ChatLlamaCpp) and provide the model's stream() method as an API to clients. It should handle multiple concurrent requests.

I expect that even when multiple stream() calls are made in parallel on a single ChatLlamaCpp instance, the stream will work correctly.

A single stream() call works fine, but as soon as two stream() calls are made simultaneously, the process crashes.

System Info

dependncis

langchain: 0.2.17
@langchain/community: 0.2.31

platform : mac, 14.3.1, m1 max(32gb)
node: v18.18.0
yarn : v1.22.19

[dosubot[bot]]

@Motiveko

Hey there! I'm here to help you out with any bugs, questions, or contributions you have in mind. Let's tackle this issue together while we wait for a human maintainer to join in.

The process crashes when stream() in ChatLlamaCpp is called in parallel because the ChatLlamaCpp class is likely not designed to handle concurrent streaming requests. This can lead to race conditions, resource conflicts, or other concurrency issues.

To handle multiple concurrent requests correctly, you should ensure that each request is isolated and does not interfere with others. One way to achieve this is by creating a new instance of ChatLlamaCpp for each request. This ensures that each instance maintains its own state and resources.

Here is an example of how you can handle multiple concurrent requests by creating separate instances of ChatLlamaCpp:

import { ChatLlamaCpp } from "@langchain/community/chat_models/llama_cpp";
import { ChatPromptTemplate } from '@langchain/core/prompts';

const modelPath = '/Users/donggi/Library/Application Support/amigo/models/_Meta-Llama-3.1-8B-Instruct-Q4_K_S.gguf';

async function handleConcurrentRequests(requests) {
  const results = await Promise.all(
    requests.map(async (request) => {
      const model = new ChatLlamaCpp({ modelPath, temperature: 0.5 });
      const template = ChatPromptTemplate.fromMessages([
        ['system', 'answer the question.'],
        ['human', '{input}'],
      ]);
      const prompt = await template.invoke({ input: request });
      const stream = await model.stream(prompt);
      const chunks = [];
      for await (const chunk of stream) {
        chunks.push(chunk.content);
      }
      return chunks;
    })
  );
  return results;
}

// Example usage
const requests = [
  'hi. what is your name?',
  'Tell me a short story about a happy Llama.'
];

handleConcurrentRequests(requests).then((results) => {
  console.log(results);
});

This code creates a new ChatLlamaCpp instance for each request, ensuring that the requests are handled independently and concurrently without causing crashes ^[1].

To continue talking to Dosu, mention @dosu.

[Motiveko]

Creating a new ChatLlamaCpp instance for each stream() call means reloading the model into memory each time. This will result in memory usage skyrocketing to tens of gigabytes with just a few parallel requests, eventually causing the process to stop.

When using node-llama-cpp directly, as shown in the code below, multiple calls can be performed in parallel, and the model internally responds to each stream sequentially.

import { ChatLlamaCpp } from "@langchain/community/chat_models/llama_cpp";
import {ChatPromptTemplate} from '@langchain/core/prompts';

const modelPath = '/Users/donggi/Library/Application Support/amigo/models/Gemma-2-Ataraxy-9B-Q4_K_S.gguf'
const model = new ChatLlamaCpp({modelPath, temperature: 0.5});
const template = ChatPromptTemplate.fromMessages([
  ['system', 'answer the question.'],
  ['human', '{input}'],
]);

const run = async () => {
  const prompt = await template.invoke({input: 'hi. what is your name?'})
  const stream = await model.stream(prompt);
  for await (const chunk of stream) {
    console.log(chunk.content);
  }

  console.log('end');
}

// works fine, no crash
run();
run();
run();

I expect the same behavior when running LangChain.js. Why does the process crash just from calling stream() twice?

[dosubot[bot]]

The process crashes when calling stream() twice in ChatLlamaCpp using LangChain.js because ChatLlamaCpp does not support multiple streaming calls concurrently. This limitation is not present when using node-llama-cpp directly, which allows for multiple parallel calls.

The implementation of ChatLlamaCpp in LangChain.js shows that it is designed to handle single streaming calls. The _streamResponseChunks method in ChatLlamaCpp uses a single session (_session) and context (_context) for streaming responses. This design does not support concurrent streaming calls, as the session and context are shared resources that can lead to conflicts or resource management issues when accessed simultaneously by multiple streams.

async *_streamResponseChunks(
  input: BaseMessage[],
  _options: this["ParsedCallOptions"],
  runManager?: CallbackManagerForLLMRun
): AsyncGenerator<ChatGenerationChunk> {
  const promptOptions = {
    temperature: this?.temperature,
    topK: this?.topK,
    topP: this?.topP,
  };

  const prompt = this._buildPrompt(input);

  const stream = await this.caller.call(async () =>
    this._context.evaluate(this._context.encode(prompt), promptOptions)
  );

  for await (const chunk of stream) {
    yield new ChatGenerationChunk({
      text: this._context.decode([chunk]),
      message: new AIMessageChunk({
        content: this._context.decode([chunk]),
      }),
      generationInfo: {},
    });
    await runManager?.handleLLMNewToken(this._context.decode([chunk]) ?? "");
  }
}

In this method, the context and session are used to evaluate and decode the prompt, which means that initiating another stream would interfere with the ongoing process, leading to crashes. This is why multiple concurrent streaming calls are not supported in ChatLlamaCpp ^[1][2].

To handle multiple concurrent requests without reloading the model each time, you might need to manage the concurrency at a higher level, such as by queuing the requests or using a pool of ChatLlamaCpp instances. However, this would require additional implementation to ensure that the model is not reloaded for each request.

To continue talking to Dosu, mention @dosu.

[Motiveko]

Isn't there a design issue? For example, ChatOpenAI provided by LangChain supports parallel streaming just fine. Here's an example code.

import {ChatPromptTemplate} from '@langchain/core/prompts';
import {ChatOpenAI} from '@langchain/openai';
process.env.OPENAI_API_KEY="APIKEY"
const model = new ChatOpenAI({
  model: 'gpt-4o-mini',
});

const template = ChatPromptTemplate.fromMessages([
  ['system', 'answer the question.'],
  ['human', '{input}'],
]);

const run = async () => {
  const prompt = await template.invoke({input: 'hi. what is your name?'})
  const stream = await model.stream(prompt);
  for await (const chunk of stream) {
    
    console.log(chunk.content);
  }

  console.log('end');
}

// works fine
run();
run();

Isn't the fact that the process crashes just because parallel calls aren't supported a problem in itself?

[dosubot[bot]]

Yes, the fact that the process crashes because parallel calls aren't supported in ChatLlamaCpp can be considered a design issue, especially since ChatOpenAI provided by LangChain supports parallel streaming. The provided tests for ChatLlamaCpp do not include any handling for parallel calls, which indicates a limitation in its design. Additionally, the ChatLlamaCpp implementation does not include mechanisms to manage concurrent sessions or parallel processing, further highlighting this limitation. In contrast, ChatOpenAI has extensive tests demonstrating its ability to handle parallel calls, streaming, and various other scenarios, showcasing its robust design for concurrent operations ^[1][2][3].

To continue talking to Dosu, mention @dosu.