"googlemock/include/vscode:/vscode.git/clone" did not exist on "adc8a88ff0c5c85413d37af20343faa6c197abad"
hello_world.py 3.99 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

16
17
import logging

18
19
from pydantic import BaseModel

20
from dynamo.sdk import DYNAMO_IMAGE, api, depends, dynamo_endpoint, service
21
22
from dynamo.sdk.lib.config import ServiceConfig
from dynamo.sdk.lib.logging import configure_server_logging
23

24
25
logger = logging.getLogger(__name__)

26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
"""
Pipeline Architecture:

Users/Clients (HTTP)


┌─────────────┐
│  Frontend   │  HTTP API endpoint (/generate)
└─────────────┘
      │ dynamo/runtime

┌─────────────┐
│   Middle    │
└─────────────┘
      │ dynamo/runtime

┌─────────────┐
│  Backend    │
└─────────────┘
"""


class RequestType(BaseModel):
    text: str


class ResponseType(BaseModel):
    text: str


@service(
    dynamo={
        "enabled": True,
        "namespace": "inference",
    },
61
    image=DYNAMO_IMAGE,
62
63
64
)
class Backend:
    def __init__(self) -> None:
65
        logger.info("Starting backend")
66
67
68
        config = ServiceConfig.get_instance()
        self.message = config.get("Backend", {}).get("message", "back")
        logger.info(f"Backend config message: {self.message}")
69
70
71
72
73

    @dynamo_endpoint()
    async def generate(self, req: RequestType):
        """Generate tokens."""
        req_text = req.text
74
        logger.info(f"Backend received: {req_text}")
75
        text = f"{req_text}-{self.message}"
76
77
78
79
80
81
        for token in text.split():
            yield f"Backend: {token}"


@service(
    dynamo={"enabled": True, "namespace": "inference"},
82
    image=DYNAMO_IMAGE,
83
84
85
86
87
)
class Middle:
    backend = depends(Backend)

    def __init__(self) -> None:
88
        logger.info("Starting middle")
89
90
91
        config = ServiceConfig.get_instance()
        self.message = config.get("Middle", {}).get("message", "mid")
        logger.info(f"Middle config message: {self.message}")
92
93
94
95
96

    @dynamo_endpoint()
    async def generate(self, req: RequestType):
        """Forward requests to backend."""
        req_text = req.text
97
        logger.info(f"Middle received: {req_text}")
98
        text = f"{req_text}-{self.message}"
99
100
        next_request = RequestType(text=text).model_dump_json()
        async for response in self.backend.generate(next_request):
101
            logger.info(f"Middle received response: {response}")
102
103
104
            yield f"Middle: {response}"


105
106
107
@service(
    image=DYNAMO_IMAGE,
)  # Regular HTTP API
108
109
110
111
class Frontend:
    middle = depends(Middle)

    def __init__(self) -> None:
112
113
114
        # Configure logging
        configure_server_logging(service_name="Frontend")

115
116
117
118
119
120
        logger.info("Starting frontend")
        config = ServiceConfig.get_instance()
        self.message = config.get("Frontend", {}).get("message", "front")
        self.port = config.get("Frontend", {}).get("port", 8000)
        logger.info(f"Frontend config message: {self.message}")
        logger.info(f"Frontend config port: {self.port}")
121
122
123
124

    @api
    async def generate(self, text):
        """Stream results from the pipeline."""
125
126
        logger.info(f"Frontend received: {text}")
        logger.info(f"Frontend received type: {type(text)}")
127
        txt = RequestType(text=text)
128
        logger.info(f"Frontend sending: {type(txt)}")
129
130
        async for response in self.middle.generate(txt.model_dump_json()):
            yield f"Frontend: {response}"