|
1 | 1 | # pylint: disable=import-outside-toplevel
|
| 2 | +from os import environ |
2 | 3 | import asyncio
|
3 | 4 | import threading
|
4 | 5 | from concurrent.futures import Future
|
@@ -87,7 +88,8 @@ def generate_text(
|
87 | 88 | messages: ModelInput,
|
88 | 89 | parameters: dict[str, Any],
|
89 | 90 | ) -> tuple[LazyMessage, PdlLazy[Any]]:
|
90 |
| - print(f"Asynchronous model call started to {model_id}", file=stderr) |
| 91 | + if "PDL_VERBOSE_ASYNC" in environ: |
| 92 | + print(f"Asynchronous model call started to {model_id}", file=stderr) |
91 | 93 | # global _BACKGROUND_TASKS
|
92 | 94 | future = asyncio.run_coroutine_threadsafe(
|
93 | 95 | LitellmModel.async_generate_text(
|
@@ -130,22 +132,23 @@ def update_end_nanos(future):
|
130 | 132 | else 0
|
131 | 133 | )
|
132 | 134 | exec_nanos = block.pdl__timing.end_nanos - start
|
133 |
| - print( |
134 |
| - f"Asynchronous model call to {model_id} completed in {(exec_nanos)/1000000}ms", |
135 |
| - file=stderr, |
136 |
| - ) |
137 |
| - msg = future.result()[0] |
138 |
| - if msg["content"] is not None: |
139 |
| - from termcolor import colored |
140 |
| - |
141 |
| - from .pdl_ast import BlockKind |
142 |
| - from .pdl_scheduler import color_of |
143 |
| - |
| 135 | + if "PDL_VERBOSE_ASYNC" in environ: |
144 | 136 | print(
|
145 |
| - colored(msg["content"], color=color_of(BlockKind.MODEL)), |
| 137 | + f"Asynchronous model call to {model_id} completed in {(exec_nanos)/1000000}ms", |
146 | 138 | file=stderr,
|
147 | 139 | )
|
148 |
| - print("\n", file=stderr) |
| 140 | + msg = future.result()[0] |
| 141 | + if msg["content"] is not None: |
| 142 | + from termcolor import colored |
| 143 | + |
| 144 | + from .pdl_ast import BlockKind |
| 145 | + from .pdl_scheduler import color_of |
| 146 | + |
| 147 | + print( |
| 148 | + colored(msg["content"], color=color_of(BlockKind.MODEL)), |
| 149 | + file=stderr, |
| 150 | + ) |
| 151 | + print("\n", file=stderr) |
149 | 152 |
|
150 | 153 | future.add_done_callback(update_end_nanos)
|
151 | 154 |
|
|
0 commit comments