ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: cancel abandoned generation tasks

Browse Source 
- Propagate stream closure and stop-sequence termination into scheduler cancellation
- Defer active KV release to the scheduler owner and close metrics safely
- Expose lifecycle counters and cover waiting, active, and allocation-race cleanup
This commit is contained in:
0z5a
2026-09-02 12:55:05 +08:00
parent 90de5bc1bd
commit 9c3ef0c2a1
10 changed files with 410 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/network/protocol.py
+28 -22
View File
@@ -146,25 +146,28 @@ class ProtocolHandler:
yielded = ""
matched = None
token_ids: List[int] = []
async for token in agen:
body += token
try:
async for token in agen:
body += token
new_ids = self.engine.tokenizer.encode(token)
token_ids.extend(new_ids)
new_ids = self.engine.tokenizer.encode(token)
token_ids.extend(new_ids)
matched = checker.check(body)
if matched:
break
matched = checker.check(body)
if matched:
break
ctx.completion_tokens += 1
for event in self.builder.format_chunk(
token,
body=body,
current_token_ids=token_ids,
delta_token_ids=new_ids,
):
yield event
yielded += token
ctx.completion_tokens += 1
for event in self.builder.format_chunk(
token,
body=body,
current_token_ids=token_ids,
delta_token_ids=new_ids,
):
yield event
yielded += token
finally:
await agen.aclose()
stop = StopInfo(matched=matched, body=body, yielded=yielded)
for event in self.builder.format_stream_end(ctx, stop):
@@ -184,14 +187,17 @@ class ProtocolHandler:
body = ""
matched = None
async for token in agen:
body += token
try:
async for token in agen:
body += token
matched = checker.check(body)
if matched:
break
matched = checker.check(body)
if matched:
break
ctx.completion_tokens += 1
ctx.completion_tokens += 1
finally:
await agen.aclose()
stop = StopInfo(matched=matched, body=body)
return self.builder.format_response(ctx, body, stop)