ray-llm

1
import logging
2

3
from ray import serve
4

5
from rayllm.backend.llm.trtllm.trtllm_engine import TRTLLMEngine
6
from rayllm.backend.llm.trtllm.trtllm_models import (
7
    TRTLLMApp,
8
    TRTLLMGenerationRequest,
9
    TRTLLMSamplingParams,
10
)
11
from rayllm.common.models import Prompt
12

13
logger = logging.getLogger(__name__)
14

15

16
@serve.deployment
17
class TRTLLMDeployment:
18
    _generation_request_cls = TRTLLMGenerationRequest
19

20
    def __init__(self, base_config: TRTLLMApp, generation_request_cls=None):
21
        self.engine = TRTLLMEngine(base_config)
22
        self.base_config = base_config
23
        self.generation_request_cls = (
24
            generation_request_cls or self._generation_request_cls
25
        )
26

27
    async def stream(
28
        self,
29
        request_id: str,
30
        prompt: Prompt,
31
        priority=None,
32
    ):
33
        sample_params = TRTLLMSamplingParams.parse_obj(prompt.parameters)
34
        logger.info(f"Received request {request_id}")
35

36
        prompt_text = (
37
            self.base_config.engine_config.generation.prompt_format.generate_prompt(
38
                prompt
39
            )
40
        )
41

42
        request = self.generation_request_cls(
43
            prompt=prompt_text,
44
            request_id=request_id,
45
            sampling_params=sample_params,
46
            stream=prompt.parameters.get("stream", True),
47
        )
48

49
        async for aviary_model_response in self.engine.generate(request):
50
            yield aviary_model_response
51