wehub-resource-sync
113 lines
3.8 KiB
Python
113 lines
3.8 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
|
|
from ray.llm._internal.common.callbacks.base import (
|
|
CallbackBase,
|
|
)
|
|
from ray.llm._internal.common.utils.download_utils import NodeModelDownloadable
|
|
from ray.llm._internal.serve.core.configs.llm_config import (
|
|
LLMConfig,
|
|
ModelLoadingConfig,
|
|
)
|
|
|
|
|
|
class TestingCallback(CallbackBase):
|
|
def __init__(self, llm_config, raise_error_on_callback: bool = True, **kwargs):
|
|
super().__init__(llm_config, raise_error_on_callback, **kwargs)
|
|
self.before_init_called = False
|
|
self.after_init_called = False
|
|
self.before_init_ctx = None
|
|
self.after_init_ctx = None
|
|
assert kwargs["kwargs_test_key"] == "kwargs_test_value"
|
|
|
|
async def on_before_node_init(self) -> None:
|
|
assert (
|
|
self.ctx.worker_node_download_model
|
|
== NodeModelDownloadable.MODEL_AND_TOKENIZER
|
|
)
|
|
self.ctx.worker_node_download_model = NodeModelDownloadable.NONE
|
|
|
|
self.ctx.custom_data["ctx_test_key"] = "ctx_test_value"
|
|
self.before_init_called = True
|
|
self.ctx.run_init_node = False
|
|
|
|
async def on_after_node_init(self) -> None:
|
|
assert self.ctx.worker_node_download_model == NodeModelDownloadable.NONE
|
|
|
|
self.after_init_called = True
|
|
assert self.ctx.custom_data["ctx_test_key"] == "ctx_test_value"
|
|
|
|
|
|
class TestCallbackBase:
|
|
@pytest.fixture
|
|
def llm_config(self):
|
|
config = LLMConfig(
|
|
model_loading_config=ModelLoadingConfig(model_id="test-model"),
|
|
llm_engine="vLLM",
|
|
callback_config={
|
|
"callback_class": TestingCallback,
|
|
"callback_kwargs": {"kwargs_test_key": "kwargs_test_value"},
|
|
},
|
|
)
|
|
return config
|
|
|
|
def test_callback_methods_called(self, llm_config):
|
|
"""Test that callback methods are called during initialization."""
|
|
|
|
# Run initialization
|
|
async def run_initialization():
|
|
callback = llm_config.get_or_create_callback()
|
|
await callback.run_callback("on_before_node_init")
|
|
if callback.ctx.run_init_node:
|
|
raise Exception("run_init_node is True")
|
|
await callback.run_callback("on_after_node_init")
|
|
|
|
asyncio.run(run_initialization())
|
|
# Verify callback was created and methods were called
|
|
|
|
callback = llm_config.get_or_create_callback()
|
|
assert callback is not None
|
|
assert isinstance(callback, TestingCallback)
|
|
assert callback.before_init_called is True
|
|
assert callback.after_init_called is True
|
|
|
|
def test_callback_singleton_behavior(self, llm_config):
|
|
"""Test that callback instance is cached (singleton pattern)."""
|
|
# Get callback multiple times
|
|
callback1 = llm_config.get_or_create_callback()
|
|
callback2 = llm_config.get_or_create_callback()
|
|
|
|
# Should be the same instance
|
|
assert callback1 is callback2
|
|
|
|
def test_callback_must_inherit_from_callback_class(self):
|
|
"""Test that callback_class must be a subclass of Callback, not just implement the same methods."""
|
|
|
|
class FakeCallback:
|
|
"""A class that implements the same methods as Callback but doesn't inherit from it."""
|
|
|
|
def __init__(self, **kwargs):
|
|
pass
|
|
|
|
async def on_before_node_init(self):
|
|
pass
|
|
|
|
async def on_after_node_init(self):
|
|
pass
|
|
|
|
# Should raise an error when trying to create callback
|
|
with pytest.raises(Exception, match="is-subclass"):
|
|
LLMConfig(
|
|
model_loading_config=ModelLoadingConfig(model_id="test-model"),
|
|
llm_engine="vLLM",
|
|
callback_config={
|
|
"callback_class": FakeCallback,
|
|
"callback_kwargs": {},
|
|
},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main(["-v", __file__])
|