Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions evals/benchmark/stresscli/locust/aistress.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

import gevent
import sseclient
import transformers
from locust import HttpUser, between, events, task
from locust.runners import STATE_CLEANUP, STATE_STOPPED, STATE_STOPPING, MasterRunner, WorkerRunner

Expand Down Expand Up @@ -85,13 +86,17 @@ def _(parser):
bench_package = ""
console_logger = logging.getLogger("locust.stats_logger")

tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")


class AiStressUser(HttpUser):
request = 0
_lock = threading.Lock()

def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
global tokenizer
self.environment.tokenizer = tokenizer

@task
def bench_main(self):
Expand Down
6 changes: 5 additions & 1 deletion evals/benchmark/stresscli/locust/tokenresponse.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,11 @@ def testFunc():


def respStatics(environment, req, resp):
tokenizer = transformers.AutoTokenizer.from_pretrained(environment.parsed_options.llm_model)
if not hasattr(environment, "tokenizer"):
tokenizer = transformers.AutoTokenizer.from_pretrained(environment.parsed_options.llm_model)
else:
tokenizer = environment.tokenizer

if environment.parsed_options.bench_target in [
"chatqnafixed",
"chatqnabench",
Expand Down