Skip to content

Commit ae61346

Browse files
author
Dmytro Parfeniuk
committed
Backend OpenAI tests
* base unit tests are added * OpenAI unit tests are added * `Backend.submit` with OpenaAI backend integration test is added * `.env.example` is added as a source of a project configuration * `README.md` now has project configuring and testing small guide
1 parent 4af96de commit ae61346

24 files changed

+562
-235
lines changed

‎.env.example

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
# OpenAI compatible server address.
2+
3+
# If you are going to run custom OpenAI API compatible service change this configuration.
4+
# Could be specified by --openai-base-url CLI parameter
5+
OPENAI_BASE_URL=http://127.0.0.1:8080
6+
7+
# The OpenAI API Key.
8+
# Could be specified by --openai-api-key CLI parameter
9+
OPENAI_API_KEY=invalid

‎README.md

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,20 @@
1-
# guidellm
1+
# guidellm
2+
3+
# Project configuration
4+
5+
The project is configured with environment variables. Check the example in `.env.example`.
6+
7+
```sh
8+
# Create .env file and update the configuration
9+
cp .env.example .env
10+
11+
# Export all variables
12+
set -o allexport; source .env; set +o allexport
13+
```
14+
15+
## Environment Variables
16+
17+
| Variable | Default Value | Description |
18+
| --------------- | --------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- |
19+
| OPENAI_BASE_URL | http://127.0.0.1:8080 | The host where the `openai` library will make requests to. For running integration tests it is required to have the external OpenAI compatible server running. |
20+
| OPENAI_API_KEY | invalid | [OpenAI Platform](https://platform.openai.com/api-keys) to create a new API key. This value is not used for tests. |

‎src/__init__.py

Whitespace-only changes.

‎src/guidellm/backend/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1-
from .base import Backend, BackendTypes, GenerativeResponse
1+
from .base import Backend, BackendEngine, GenerativeResponse
22
from .openai import OpenAIBackend
33

44
__all__ = [
55
"Backend",
6-
"BackendTypes",
6+
"BackendEngine",
77
"GenerativeResponse",
88
"OpenAIBackend",
99
]

‎src/guidellm/backend/base.py

Lines changed: 37 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,23 @@
1-
import uuid
21
from abc import ABC, abstractmethod
32
from dataclasses import dataclass
43
from enum import Enum
54
from typing import Iterator, List, Optional, Type, Union
65

76
from loguru import logger
87

9-
from guidellm.core.request import TextGenerationRequest
10-
from guidellm.core.result import TextGenerationResult
8+
from guidellm.core import TextGenerationRequest, TextGenerationResult
119

12-
__all__ = ["Backend", "BackendTypes", "GenerativeResponse"]
10+
__all__ = ["Backend", "BackendEngine", "GenerativeResponse"]
1311

1412

15-
class BackendTypes(Enum):
13+
class BackendEngine(str, Enum):
14+
"""
15+
Determines the Engine of the LLM Backend.
16+
All the implemented backends in the project have the engine.
17+
18+
NOTE: the `TEST` engine has to be used only for testing purposes.
19+
"""
20+
1621
TEST = "test"
1722
OPENAI_SERVER = "openai_server"
1823

@@ -33,43 +38,46 @@ class GenerativeResponse:
3338

3439
class Backend(ABC):
3540
"""
36-
An abstract base class for generative AI backends.
41+
An abstract base class with template methods for generative AI backends.
3742
"""
3843

3944
_registry = {}
4045

41-
@staticmethod
42-
def register_backend(backend_type: BackendTypes):
46+
@classmethod
47+
def register(cls, backend_type: BackendEngine):
4348
"""
4449
A decorator to register a backend class in the backend registry.
4550
4651
:param backend_type: The type of backend to register.
47-
:type backend_type: BackendTypes
52+
:type backend_type: BackendType
4853
"""
4954

5055
def inner_wrapper(wrapped_class: Type["Backend"]):
51-
Backend._registry[backend_type] = wrapped_class
56+
cls._registry[backend_type] = wrapped_class
5257
return wrapped_class
5358

5459
return inner_wrapper
5560

56-
@staticmethod
57-
def create_backend(backend_type: Union[str, BackendTypes], **kwargs) -> "Backend":
61+
@classmethod
62+
def create(cls, backend_type: Union[str, BackendEngine], **kwargs) -> "Backend":
5863
"""
5964
Factory method to create a backend based on the backend type.
6065
6166
:param backend_type: The type of backend to create.
62-
:type backend_type: BackendTypes
67+
:type backend_type: BackendType
6368
:param kwargs: Additional arguments for backend initialization.
6469
:type kwargs: dict
6570
:return: An instance of a subclass of Backend.
6671
:rtype: Backend
6772
"""
73+
6874
logger.info(f"Creating backend of type {backend_type}")
69-
if backend_type not in Backend._registry:
75+
76+
if backend_type not in cls._registry:
7077
logger.error(f"Unsupported backend type: {backend_type}")
7178
raise ValueError(f"Unsupported backend type: {backend_type}")
72-
return Backend._registry[backend_type](**kwargs)
79+
80+
return cls._registry[backend_type](**kwargs)
7381

7482
def submit(self, request: TextGenerationRequest) -> TextGenerationResult:
7583
"""
@@ -80,23 +88,23 @@ def submit(self, request: TextGenerationRequest) -> TextGenerationResult:
8088
:return: The populated result result.
8189
:rtype: TextGenerationResult
8290
"""
91+
8392
logger.info(f"Submitting request with prompt: {request.prompt}")
84-
result_id = str(uuid.uuid4())
85-
result = TextGenerationResult(result_id)
93+
94+
result = TextGenerationResult(request=request)
8695
result.start(request.prompt)
8796

88-
for response in self.make_request(request):
97+
for response in self.make_request(request): # GenerativeResponse
8998
if response.type_ == "token_iter" and response.add_token:
9099
result.output_token(response.add_token)
91100
elif response.type_ == "final":
92101
result.end(
93-
response.output,
94102
response.prompt_token_count,
95103
response.output_token_count,
96104
)
97-
break
98105

99106
logger.info(f"Request completed with output: {result.output}")
107+
100108
return result
101109

102110
@abstractmethod
@@ -111,7 +119,8 @@ def make_request(
111119
:return: An iterator over the generative responses.
112120
:rtype: Iterator[GenerativeResponse]
113121
"""
114-
raise NotImplementedError()
122+
123+
pass
115124

116125
@abstractmethod
117126
def available_models(self) -> List[str]:
@@ -121,8 +130,10 @@ def available_models(self) -> List[str]:
121130
:return: A list of available models.
122131
:rtype: List[str]
123132
"""
124-
raise NotImplementedError()
125133

134+
pass
135+
136+
@property
126137
@abstractmethod
127138
def default_model(self) -> str:
128139
"""
@@ -131,7 +142,8 @@ def default_model(self) -> str:
131142
:return: The default model.
132143
:rtype: str
133144
"""
134-
raise NotImplementedError()
145+
146+
pass
135147

136148
@abstractmethod
137149
def model_tokenizer(self, model: str) -> Optional[str]:
@@ -143,4 +155,5 @@ def model_tokenizer(self, model: str) -> Optional[str]:
143155
:return: The tokenizer for the model, or None if it cannot be created.
144156
:rtype: Optional[str]
145157
"""
146-
raise NotImplementedError()
158+
159+
pass

‎src/guidellm/backend/openai.py

Lines changed: 71 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,19 @@
1-
from typing import Any, Iterator, List, Optional
1+
import functools
2+
import os
3+
from typing import Any, Dict, Iterator, List, Optional
24

3-
import openai
45
from loguru import logger
6+
from openai import OpenAI, Stream
7+
from openai.types import Completion
58
from transformers import AutoTokenizer
69

7-
from guidellm.backend import Backend, BackendTypes, GenerativeResponse
8-
from guidellm.core.request import TextGenerationRequest
10+
from guidellm.backend import Backend, BackendEngine, GenerativeResponse
11+
from guidellm.core import TextGenerationRequest
912

1013
__all__ = ["OpenAIBackend"]
1114

1215

13-
@Backend.register_backend(BackendTypes.OPENAI_SERVER)
16+
@Backend.register(BackendEngine.OPENAI_SERVER)
1417
class OpenAIBackend(Backend):
1518
"""
1619
An OpenAI backend implementation for the generative AI result.
@@ -33,34 +36,37 @@ class OpenAIBackend(Backend):
3336

3437
def __init__(
3538
self,
36-
target: Optional[str] = None,
37-
host: Optional[str] = None,
38-
port: Optional[int] = None,
39-
path: Optional[str] = None,
39+
openai_api_key: Optional[str] = None,
40+
internal_callback_url: Optional[str] = None,
4041
model: Optional[str] = None,
41-
api_key: Optional[str] = None,
42-
**request_args,
42+
**request_args: Any,
4343
):
44-
self.target = target
45-
self.model = model
46-
self.request_args = request_args
47-
48-
if not self.target:
49-
if not host:
50-
raise ValueError("Host is required if target is not provided.")
51-
52-
port_incl = f":{port}" if port else ""
53-
path_incl = path if path else ""
54-
self.target = f"http://{host}{port_incl}{path_incl}"
44+
"""
45+
Initialize an OpenAI Client
46+
"""
5547

56-
openai.api_base = self.target
57-
openai.api_key = api_key
48+
self.request_args = request_args
5849

59-
if not model:
60-
self.model = self.default_model()
50+
if not (_api_key := (openai_api_key or os.getenv("OPENAI_API_KEY", None))):
51+
raise ValueError(
52+
"`OPENAI_API_KEY` environment variable "
53+
"or --openai-api-key CLI parameter "
54+
"must be specify for the OpenAI backend"
55+
)
56+
57+
if not (
58+
_base_url := (internal_callback_url or os.getenv("OPENAI_BASE_URL", None))
59+
):
60+
raise ValueError(
61+
"`OPENAI_BASE_URL` environment variable "
62+
"or --openai-base-url CLI parameter "
63+
"must be specify for the OpenAI backend"
64+
)
65+
self.openai_client = OpenAI(api_key=_api_key, base_url=_base_url)
66+
self.model = model or self.default_model
6167

6268
logger.info(
63-
f"Initialized OpenAIBackend with target: {self.target} "
69+
f"Initialized OpenAIBackend with callback url: {internal_callback_url} "
6470
f"and model: {self.model}"
6571
)
6672

@@ -75,52 +81,46 @@ def make_request(
7581
:return: An iterator over the generative responses.
7682
:rtype: Iterator[GenerativeResponse]
7783
"""
84+
7885
logger.debug(f"Making request to OpenAI backend with prompt: {request.prompt}")
79-
num_gen_tokens = request.params.get("generated_tokens", None)
80-
request_args = {
81-
"n": 1,
82-
}
8386

84-
if num_gen_tokens:
85-
request_args["max_tokens"] = num_gen_tokens
86-
request_args["stop"] = None
87+
# How many completions to generate for each prompt
88+
request_args: Dict = {"n": 1}
89+
90+
if (num_gen_tokens := request.params.get("generated_tokens", None)) is not None:
91+
request_args.update(max_tokens=num_gen_tokens, stop=None)
8792

8893
if self.request_args:
8994
request_args.update(self.request_args)
9095

91-
response = openai.Completion.create(
92-
engine=self.model,
96+
response: Stream[Completion] = self.openai_client.completions.create(
97+
model=self.model,
9398
prompt=request.prompt,
9499
stream=True,
95100
**request_args,
96101
)
97102

98103
for chunk in response:
99-
if chunk.get("choices"):
100-
choice = chunk["choices"][0]
101-
if choice.get("finish_reason") == "stop":
102-
logger.debug("Received final response from OpenAI backend")
103-
yield GenerativeResponse(
104-
type_="final",
105-
output=choice["text"],
106-
prompt=request.prompt,
107-
prompt_token_count=(
108-
request.token_count
109-
if request.token_count
110-
else self._token_count(request.prompt)
111-
),
112-
output_token_count=(
113-
num_gen_tokens
114-
if num_gen_tokens
115-
else self._token_count(choice["text"])
116-
),
117-
)
118-
break
119-
else:
120-
logger.debug("Received token from OpenAI backend")
121-
yield GenerativeResponse(
122-
type_="token_iter", add_token=choice["text"]
123-
)
104+
chunk_content: str = getattr(chunk, "content", "")
105+
106+
if getattr(chunk, "stop", True) is True:
107+
logger.debug("Received final response from OpenAI backend")
108+
109+
yield GenerativeResponse(
110+
type_="final",
111+
prompt=getattr(chunk, "prompt", request.prompt),
112+
prompt_token_count=(
113+
request.prompt_token_count or self._token_count(request.prompt)
114+
),
115+
output_token_count=(
116+
num_gen_tokens
117+
if num_gen_tokens
118+
else self._token_count(chunk_content)
119+
),
120+
)
121+
else:
122+
logger.debug("Received token from OpenAI backend")
123+
yield GenerativeResponse(type_="token_iter", add_token=chunk_content)
124124

125125
def available_models(self) -> List[str]:
126126
"""
@@ -129,21 +129,28 @@ def available_models(self) -> List[str]:
129129
:return: A list of available models.
130130
:rtype: List[str]
131131
"""
132-
models = [model["id"] for model in openai.Engine.list()["data"]]
132+
133+
models: list[str] = [
134+
model.id for model in self.openai_client.models.list().data
135+
]
133136
logger.info(f"Available models: {models}")
137+
134138
return models
135139

140+
@property
141+
@functools.lru_cache(maxsize=1)
136142
def default_model(self) -> str:
137143
"""
138144
Get the default model for the backend.
139145
140146
:return: The default model.
141147
:rtype: str
142148
"""
143-
models = self.available_models()
144-
if models:
149+
150+
if models := self.available_models():
145151
logger.info(f"Default model: {models[0]}")
146152
return models[0]
153+
147154
logger.error("No models available.")
148155
raise ValueError("No models available.")
149156

0 commit comments

Comments
 (0)