try to fix streaming responses
This commit is contained in:
@@ -286,7 +286,7 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._send_json(200, completion)
|
self._send_json(200, completion)
|
||||||
|
|
||||||
def _proxy_to_head_worker(self, url: str, body_bytes: bytes) -> None:
|
def _proxy_to_head_worker(self, url: str, body_bytes: bytes) -> None:
|
||||||
"""Forward a raw request body to a head worker and stream the response back."""
|
"""Forward a raw request body to a head worker and relay SSE without buffering."""
|
||||||
target_url = f"{url}/v1/chat/completions"
|
target_url = f"{url}/v1/chat/completions"
|
||||||
req = urllib.request.Request(
|
req = urllib.request.Request(
|
||||||
target_url,
|
target_url,
|
||||||
@@ -297,6 +297,19 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
try:
|
try:
|
||||||
with urllib.request.urlopen(req, timeout=30.0) as r:
|
with urllib.request.urlopen(req, timeout=30.0) as r:
|
||||||
content_type = r.headers.get("Content-Type", "application/json")
|
content_type = r.headers.get("Content-Type", "application/json")
|
||||||
|
if "text/event-stream" in content_type:
|
||||||
|
self.send_response(r.status)
|
||||||
|
self.send_header("Content-Type", content_type)
|
||||||
|
self.send_header("Cache-Control", "no-cache")
|
||||||
|
self.send_header("X-Accel-Buffering", "no")
|
||||||
|
self.end_headers()
|
||||||
|
while True:
|
||||||
|
line = r.readline()
|
||||||
|
if not line:
|
||||||
|
break
|
||||||
|
self.wfile.write(line)
|
||||||
|
self.wfile.flush()
|
||||||
|
return
|
||||||
resp_body = r.read()
|
resp_body = r.read()
|
||||||
status = r.status
|
status = r.status
|
||||||
except urllib.error.HTTPError as exc:
|
except urllib.error.HTTPError as exc:
|
||||||
|
|||||||
@@ -1241,7 +1241,7 @@ async function sendChat() {
|
|||||||
.map(msg => ({ role: msg.role, content: msg.content })),
|
.map(msg => ({ role: msg.role, content: msg.content })),
|
||||||
{ role: "user", content: prompt },
|
{ role: "user", content: prompt },
|
||||||
],
|
],
|
||||||
stream: false,
|
stream: true,
|
||||||
max_tokens: 15120,
|
max_tokens: 15120,
|
||||||
};
|
};
|
||||||
chatBusy = true;
|
chatBusy = true;
|
||||||
@@ -1249,36 +1249,83 @@ async function sendChat() {
|
|||||||
promptEl.value = "";
|
promptEl.value = "";
|
||||||
promptEl.style.height = "auto";
|
promptEl.style.height = "auto";
|
||||||
chatHistory.push({ role: "user", content: prompt, model: selectedChatModel });
|
chatHistory.push({ role: "user", content: prompt, model: selectedChatModel });
|
||||||
|
const assistantMessage = { role: "assistant", content: "", model: selectedChatModel };
|
||||||
|
chatHistory.push(assistantMessage);
|
||||||
renderChatHistory();
|
renderChatHistory();
|
||||||
persistActiveChatSession();
|
persistActiveChatSession();
|
||||||
renderChatStatus("sending request…");
|
renderChatStatus("streaming response…");
|
||||||
const r = await apiCall("/v1/chat/completions", "POST", body, bearerToken);
|
|
||||||
chatBusy = false;
|
try {
|
||||||
$("chat-send").disabled = false;
|
const headers = { "Content-Type": "application/json" };
|
||||||
if (!r.ok) {
|
if (bearerToken) headers["Authorization"] = "Bearer " + bearerToken;
|
||||||
const error = r.data && r.data.error
|
const response = await fetch("/v1/chat/completions", {
|
||||||
? (typeof r.data.error === "string" ? r.data.error : r.data.error.message || "request failed")
|
method: "POST",
|
||||||
: "request failed";
|
headers,
|
||||||
chatHistory.push({ role: "error", content: error, model: selectedChatModel });
|
credentials: "same-origin",
|
||||||
|
body: JSON.stringify(body),
|
||||||
|
});
|
||||||
|
if (!response.ok) {
|
||||||
|
const data = await response.json().catch(() => ({}));
|
||||||
|
const error = data && data.error
|
||||||
|
? (typeof data.error === "string" ? data.error : data.error.message || "request failed")
|
||||||
|
: "request failed";
|
||||||
|
assistantMessage.role = "error";
|
||||||
|
assistantMessage.content = error;
|
||||||
|
renderChatHistory();
|
||||||
|
persistActiveChatSession();
|
||||||
|
renderChatStatus(error);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if (!response.body) throw new Error("stream response body is missing");
|
||||||
|
|
||||||
|
const reader = response.body.getReader();
|
||||||
|
const decoder = new TextDecoder();
|
||||||
|
let buffer = "";
|
||||||
|
let receivedAny = false;
|
||||||
|
while (true) {
|
||||||
|
const { value, done } = await reader.read();
|
||||||
|
if (done) break;
|
||||||
|
buffer += decoder.decode(value, { stream: true });
|
||||||
|
const events = buffer.split("\n\n");
|
||||||
|
buffer = events.pop() || "";
|
||||||
|
for (const eventText of events) {
|
||||||
|
for (const line of eventText.split("\n")) {
|
||||||
|
if (!line.startsWith("data: ")) continue;
|
||||||
|
const data = line.slice(6).trim();
|
||||||
|
if (data === "[DONE]") {
|
||||||
|
buffer = "";
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
try {
|
||||||
|
const chunk = JSON.parse(data);
|
||||||
|
const delta = chunk.choices && chunk.choices[0] && chunk.choices[0].delta;
|
||||||
|
if (delta && delta.content) {
|
||||||
|
assistantMessage.content += delta.content;
|
||||||
|
receivedAny = true;
|
||||||
|
renderChatHistory();
|
||||||
|
persistActiveChatSession();
|
||||||
|
}
|
||||||
|
} catch {
|
||||||
|
// Ignore malformed SSE keepalive/event lines.
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (!receivedAny) assistantMessage.content = "(empty response)";
|
||||||
renderChatHistory();
|
renderChatHistory();
|
||||||
persistActiveChatSession();
|
persistActiveChatSession();
|
||||||
renderChatStatus(error);
|
renderChatStatus("done");
|
||||||
|
} catch (err) {
|
||||||
|
assistantMessage.role = "error";
|
||||||
|
assistantMessage.content = err && err.message ? err.message : "request failed";
|
||||||
|
renderChatHistory();
|
||||||
|
persistActiveChatSession();
|
||||||
|
renderChatStatus(assistantMessage.content);
|
||||||
|
} finally {
|
||||||
|
chatBusy = false;
|
||||||
|
$("chat-send").disabled = false;
|
||||||
promptEl.focus();
|
promptEl.focus();
|
||||||
return;
|
|
||||||
}
|
}
|
||||||
const reply = (r.data && r.data.choices && r.data.choices[0] && r.data.choices[0].message && r.data.choices[0].message.content) || "";
|
|
||||||
const usage = r.data && r.data.usage;
|
|
||||||
chatHistory.push({
|
|
||||||
role: "assistant",
|
|
||||||
content: reply || "(empty response)",
|
|
||||||
model: selectedChatModel,
|
|
||||||
});
|
|
||||||
renderChatHistory();
|
|
||||||
persistActiveChatSession();
|
|
||||||
renderChatStatus(usage
|
|
||||||
? `done: ${usage.total_tokens ?? "?"} tokens`
|
|
||||||
: "done");
|
|
||||||
promptEl.focus();
|
|
||||||
}
|
}
|
||||||
|
|
||||||
function bindChatPromptShortcuts() {
|
function bindChatPromptShortcuts() {
|
||||||
|
|||||||
@@ -33,6 +33,21 @@ def test_dashboard_served_with_all_panels():
|
|||||||
tracker.stop()
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_dashboard_chat_uses_streaming_fetch():
|
||||||
|
tracker = TrackerServer(billing=BillingLedger())
|
||||||
|
port = tracker.start()
|
||||||
|
try:
|
||||||
|
html = urllib.request.urlopen(
|
||||||
|
f"http://127.0.0.1:{port}/dashboard"
|
||||||
|
).read().decode()
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
assert "stream: true" in html
|
||||||
|
assert ".body.getReader()" in html
|
||||||
|
assert 'data === "[DONE]"' in html
|
||||||
|
|
||||||
|
|
||||||
def test_dashboard_served_by_follower():
|
def test_dashboard_served_by_follower():
|
||||||
"""A tracker that is not the leader (unreachable peers → never elected)
|
"""A tracker that is not the leader (unreachable peers → never elected)
|
||||||
still serves the dashboard from its own replicated state."""
|
still serves the dashboard from its own replicated state."""
|
||||||
|
|||||||
@@ -142,6 +142,21 @@ def _send_chat_request(gateway_url: str, prompt: str) -> dict:
|
|||||||
return json.loads(r.read())
|
return json.loads(r.read())
|
||||||
|
|
||||||
|
|
||||||
|
def _send_streaming_chat_request(gateway_url: str, prompt: str):
|
||||||
|
data = json.dumps({
|
||||||
|
"model": GPT2_MODEL,
|
||||||
|
"messages": [{"role": "user", "content": prompt}],
|
||||||
|
"stream": True,
|
||||||
|
}).encode()
|
||||||
|
req = urllib.request.Request(
|
||||||
|
f"{gateway_url}/v1/chat/completions",
|
||||||
|
data=data,
|
||||||
|
headers={"Content-Type": "application/json"},
|
||||||
|
method="POST",
|
||||||
|
)
|
||||||
|
return urllib.request.urlopen(req)
|
||||||
|
|
||||||
|
|
||||||
def test_all_responses_valid_openai_format(tracker_node_setup):
|
def test_all_responses_valid_openai_format(tracker_node_setup):
|
||||||
"""Ten requests via gateway all return valid OpenAI chat completion format."""
|
"""Ten requests via gateway all return valid OpenAI chat completion format."""
|
||||||
gateway_url, _, _ = tracker_node_setup
|
gateway_url, _, _ = tracker_node_setup
|
||||||
@@ -155,6 +170,30 @@ def test_all_responses_valid_openai_format(tracker_node_setup):
|
|||||||
assert isinstance(message.get("content"), str), f"request {i}: content must be a string"
|
assert isinstance(message.get("content"), str), f"request {i}: content must be a string"
|
||||||
|
|
||||||
|
|
||||||
|
def test_streaming_head_worker_response_is_not_buffered_with_content_length(tracker_node_setup):
|
||||||
|
"""Gateway must relay head-worker SSE as a live stream, not a buffered JSON-sized body."""
|
||||||
|
gateway_url, _, _ = tracker_node_setup
|
||||||
|
|
||||||
|
with _send_streaming_chat_request(gateway_url, "stream through head worker") as resp:
|
||||||
|
assert resp.status == 200
|
||||||
|
assert "text/event-stream" in resp.headers["Content-Type"]
|
||||||
|
assert "Content-Length" not in resp.headers
|
||||||
|
data_lines = []
|
||||||
|
while len(data_lines) < 4:
|
||||||
|
line = resp.readline().decode().strip()
|
||||||
|
if line.startswith("data: "):
|
||||||
|
data_lines.append(line)
|
||||||
|
if line == "data: [DONE]":
|
||||||
|
break
|
||||||
|
|
||||||
|
assert data_lines[-1] == "data: [DONE]"
|
||||||
|
content = "".join(
|
||||||
|
json.loads(line[6:])["choices"][0].get("delta", {}).get("content", "")
|
||||||
|
for line in data_lines[:-1]
|
||||||
|
)
|
||||||
|
assert "head-worker" in content
|
||||||
|
|
||||||
|
|
||||||
def test_both_tracker_nodes_receive_load(tracker_node_setup):
|
def test_both_tracker_nodes_receive_load(tracker_node_setup):
|
||||||
"""Both head workers handle at least one request each out of ten."""
|
"""Both head workers handle at least one request each out of ten."""
|
||||||
gateway_url, tracker_node_a, tracker_node_b = tracker_node_setup
|
gateway_url, tracker_node_a, tracker_node_b = tracker_node_setup
|
||||||
|
|||||||
Reference in New Issue
Block a user