[Core] *Prompt* logprobs support in Multi-step (#8199)

2024-09-18 11:38:43 -04:00
parent 7c7714d856
commit a8c1d161a7
5 changed files with 299 additions and 58 deletions
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -493,6 +493,7 @@ async def completions_with_server_args(
    '''

    outputs = None
+    max_wait_seconds = 240 * 3  # 240 is default
    with RemoteOpenAIServer(model_name,
                            server_cli_args,
                            max_wait_seconds=max_wait_seconds) as server:
@@ -503,7 +504,7 @@ async def completions_with_server_args(
                                                  stream=False,
                                                  max_tokens=5,
                                                  logprobs=num_logprobs)
-    assert outputs is not None
+    assert outputs is not None, "Completion API call failed."

    return outputs