From 9c3ef0c2a1f12049e129cf2686b98d8b5f9a4882 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 2 Sep 2026 10:15:48 +0800 Subject: [PATCH] fix: cancel abandoned generation tasks - 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 --- astrai/inference/cache/pool.py | 5 ++ astrai/inference/engine.py | 75 +++++++++++-------- astrai/inference/metrics.py | 74 +++++++++--------- astrai/inference/network/protocol.py | 50 +++++++------ astrai/inference/scheduler.py | 49 ++++++++++-- astrai/inference/task.py | 107 ++++++++++++++++++++------- tests/inference/test_engine.py | 47 ++++++++++++ tests/inference/test_scheduler.py | 99 +++++++++++++++++++++++++ tests/inference/test_server.py | 22 ++++-- tests/inference/test_task.py | 12 ++- 10 files changed, 410 insertions(+), 130 deletions(-) diff --git a/astrai/inference/cache/pool.py b/astrai/inference/cache/pool.py index 3f377b8..0550902 100644 --- a/astrai/inference/cache/pool.py +++ b/astrai/inference/cache/pool.py @@ -322,6 +322,11 @@ class TaskCacheManager: state.length = pos + 1 return True + @property + def task_count(self) -> int: + """Number of tasks currently holding KV request state.""" + return len(self._states) + def task_cached(self, task_id: str) -> int: state = self._states.get(task_id) return state.cached if state is not None else 0 diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index 8632dbc..e440cce 100644 --- a/astrai/inference/engine.py +++ b/astrai/inference/engine.py @@ -43,8 +43,7 @@ class GenerateResult: with self._cond: out = self.tokens.copy() self.tokens.clear() - if not out: - self._event.clear() + self._event.clear() return out def wait(self, timeout: Optional[float] = None) -> bool: @@ -141,25 +140,34 @@ class InferenceEngine: frequency_penalty: float = 0.0, rep_window: int = 64, ) -> AsyncGenerator[str, None]: - sync_gen = self._generate( - [prompt], - False, - True, - max_tokens, - temperature, - top_p, - top_k, - frequency_penalty, - rep_window, + request_backend = get_backend(use_default=False) + result = GenerateResult() + task_id = self.scheduler.add_task( + prompt=prompt, + max_tokens=max_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + frequency_penalty=frequency_penalty, + rep_window=rep_window, + backend=request_backend, + stream_callback=result.append, ) async def _agen(): - loop = asyncio.get_event_loop() - while True: - token = await loop.run_in_executor(None, next, sync_gen, None) - if token is None: - break - yield token + finished = False + try: + while not finished: + for _idx, token in result.pop_all(): + if token is STOP: + finished = True + break + yield token + if not finished: + await asyncio.to_thread(result.wait, 0.05) + finally: + if not finished: + self.scheduler.cancel_task(task_id) return _agen() @@ -198,10 +206,8 @@ class InferenceEngine: result.wait_completion() except TimeoutError: for tid in task_ids: - self.scheduler.remove_task(tid) + self.scheduler.cancel_task(tid) raise - for tid in task_ids: - self.scheduler.remove_task(tid) res = result.get_results() return res if is_batch else res[0] @@ -210,17 +216,22 @@ class InferenceEngine: def gen(): nonlocal remaining - while remaining > 0: - items = result.pop_all() - for idx, token in items: - if token is STOP: - if not finished[idx]: - finished[idx] = True - remaining -= 1 - else: - yield (idx, token) if is_batch else token - if remaining > 0: - result.wait(timeout=0.05) + try: + while remaining > 0: + items = result.pop_all() + for idx, token in items: + if token is STOP: + if not finished[idx]: + finished[idx] = True + remaining -= 1 + else: + yield (idx, token) if is_batch else token + if remaining > 0: + result.wait(timeout=0.05) + finally: + for idx, task_id in enumerate(task_ids): + if not finished[idx]: + self.scheduler.cancel_task(task_id) return gen() diff --git a/astrai/inference/metrics.py b/astrai/inference/metrics.py index b6abb6f..9742a7f 100644 --- a/astrai/inference/metrics.py +++ b/astrai/inference/metrics.py @@ -1,5 +1,6 @@ """Unified per-task perf/stats: timing records, context-manager scopes, aggregate reporting.""" +import threading import time from collections import deque from contextlib import contextmanager @@ -125,6 +126,7 @@ class MetricsCollector: def __init__(self, max_recent: int = 128): self._timings: Dict[str, TaskTiming] = {} self._completed: Deque[TaskTiming] = deque(maxlen=max_recent) + self._lock = threading.Lock() self._ttft_ms_sum = 0.0 self._ttft_ms_count = 0 @@ -135,18 +137,22 @@ class MetricsCollector: def register(self, task_id: str): """Create a timing record for a newly-created task.""" - self._timings[task_id] = TaskTiming(task_id=task_id, arrival_time=time.time()) + with self._lock: + self._timings[task_id] = TaskTiming( + task_id=task_id, arrival_time=time.time() + ) def mark_finished(self, task_id: str, input_tokens: int, output_tokens: int): """Close timing for a finished/aborted task and move it to completed.""" - timing = self._timings.pop(task_id, None) - if timing is None: - return - timing.finish_time = time.time() - timing.input_tokens = input_tokens - timing.output_tokens = output_tokens - self._completed.append(timing) - self._accumulate(timing) + with self._lock: + timing = self._timings.pop(task_id, None) + if timing is None: + return + timing.finish_time = time.time() + timing.input_tokens = input_tokens + timing.output_tokens = output_tokens + self._completed.append(timing) + self._accumulate(timing) # timing scopes @@ -158,34 +164,36 @@ class MetricsCollector: yield toc = time.time() dt = toc - tic - for tid in task_ids: - t = self._timings.get(tid) - if t is None: - continue - if phase == "prefill": - t.prefill_start_time = tic - t.first_token_time = toc - elif phase == "decode": - t._decode_steps += 1 - t._decode_total_s += dt + with self._lock: + for tid in task_ids: + t = self._timings.get(tid) + if t is None: + continue + if phase == "prefill": + t.prefill_start_time = tic + t.first_token_time = toc + elif phase == "decode": + t._decode_steps += 1 + t._decode_total_s += dt # aggregate stats def get_stats(self) -> Dict[str, Any]: - stats: Dict[str, Any] = {} - if self._ttft_ms_count > 0: - stats["avg_ttft_ms"] = round(self._ttft_ms_sum / self._ttft_ms_count, 2) - if self._decode_tps_count > 0: - stats["avg_decode_tps"] = round( - self._decode_tps_sum / self._decode_tps_count, 2 - ) - if self._e2e_ms_count > 0: - stats["avg_e2e_latency_ms"] = round( - self._e2e_ms_sum / self._e2e_ms_count, 2 - ) - if self._completed: - stats["recent_tasks"] = [t.to_dict() for t in self._completed] - return stats + with self._lock: + stats: Dict[str, Any] = {"in_flight_tasks": len(self._timings)} + if self._ttft_ms_count > 0: + stats["avg_ttft_ms"] = round(self._ttft_ms_sum / self._ttft_ms_count, 2) + if self._decode_tps_count > 0: + stats["avg_decode_tps"] = round( + self._decode_tps_sum / self._decode_tps_count, 2 + ) + if self._e2e_ms_count > 0: + stats["avg_e2e_latency_ms"] = round( + self._e2e_ms_sum / self._e2e_ms_count, 2 + ) + if self._completed: + stats["recent_tasks"] = [t.to_dict() for t in self._completed] + return stats # internal diff --git a/astrai/inference/network/protocol.py b/astrai/inference/network/protocol.py index 07c98e9..db6939c 100644 --- a/astrai/inference/network/protocol.py +++ b/astrai/inference/network/protocol.py @@ -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) diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index 7b56a59..f568fe2 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -107,12 +107,25 @@ class InferenceScheduler: def add_task(self, prompt: str, **kwargs) -> str: return self._task_mgr.add_task(prompt, **kwargs) - def remove_task(self, task_id: str): - for task in self._task_mgr.remove_task(task_id): - self._task_cache.task_free(task.task_id) + def cancel_task(self, task_id: str) -> bool: + """Cancel a waiting or active task without freeing in-use KV state.""" + immediate, cancelled = self._task_mgr.cancel_task(task_id) + for task in immediate: + self._metrics.mark_finished( + task.task_id, task.input_tokens, task.output_tokens + ) + if cancelled: + self._task_mgr.wake() + return cancelled + + def remove_task(self, task_id: str) -> bool: + """Backward-compatible alias for cancellation.""" + return self.cancel_task(task_id) def get_stats(self) -> Dict[str, Any]: - return self._task_mgr.get_stats() + stats = self._task_mgr.get_stats() + stats["kv_cache_tasks"] = self._task_cache.task_count + return stats @property def backend_name(self) -> str: @@ -254,7 +267,13 @@ class InferenceScheduler: if self._task_cache.task_alloc( task.task_id, task.prompt_ids ): - self._task_mgr.activate(task) + if not self._task_mgr.activate(task): + self._task_cache.task_free(task.task_id) + self._metrics.mark_finished( + task.task_id, + task.input_tokens, + task.output_tokens, + ) else: failed.append(task) if failed: @@ -264,7 +283,11 @@ class InferenceScheduler: self._task_mgr.wait_for_tasks(timeout=1.0) continue - active = self._task_mgr.get_active_tasks() + active = [ + task + for task in self._task_mgr.get_active_tasks() + if task.status != TaskStatus.ABORTED + ] decoded, aborted = self._step(active) @@ -272,6 +295,8 @@ class InferenceScheduler: self._task_mgr.invoke_callback(t.task_id, STOP) for t in decoded: + if t.status == TaskStatus.ABORTED: + continue new_text = t.decode_new_token(self._task_mgr.tokenizer) if new_text: self._task_mgr.invoke_callback(t.task_id, new_text) @@ -303,13 +328,21 @@ class InferenceScheduler: def _abort_and_clear(self, free_waiting: bool): """Invoke STOP callbacks, release cache slots, and clear task queues.""" - for task in self._task_mgr.get_active_tasks(): + active = self._task_mgr.get_active_tasks() + waiting = self._task_mgr.get_waiting_tasks() + for task in active: self._task_mgr.invoke_callback(task.task_id, STOP) self._task_cache.task_free(task.task_id) - for task in self._task_mgr.get_waiting_tasks(): + self._metrics.mark_finished( + task.task_id, task.input_tokens, task.output_tokens + ) + for task in waiting: self._task_mgr.invoke_callback(task.task_id, STOP) if free_waiting: self._task_cache.task_free(task.task_id) + self._metrics.mark_finished( + task.task_id, task.input_tokens, task.output_tokens + ) self._task_mgr.clear_queues() def run_batch( diff --git a/astrai/inference/task.py b/astrai/inference/task.py index d0b25f7..53eeff3 100644 --- a/astrai/inference/task.py +++ b/astrai/inference/task.py @@ -4,7 +4,17 @@ import uuid from collections import deque from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Literal, Optional +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Deque, + Dict, + List, + Literal, + Optional, + Tuple, +) from tokenizers.decoders import DecodeStream @@ -150,12 +160,14 @@ class TaskManager: self.waiting_queue: Deque[Task] = deque() self.active_tasks: List[Task] = [] self._callbacks: Dict[str, Callable[[str], None]] = {} + self._tasks: Dict[str, Task] = {} self._task_event = threading.Event() self._lock = threading.Lock() self._total_tasks = 0 self._total_tokens = 0 + self._cancelled_total = 0 self._metrics = metrics @@ -195,6 +207,7 @@ class TaskManager: with self._lock: self.waiting_queue.append(task) + self._tasks[task_id] = task self._total_tasks += 1 if stream_callback: self._callbacks[task_id] = stream_callback @@ -205,28 +218,49 @@ class TaskManager: self._task_event.set() return task_id - def remove_task(self, task_id: str) -> List[Task]: + def cancel_task(self, task_id: str) -> Tuple[List[Task], bool]: + """Mark a task cancelled and return tasks safe to clean immediately.""" with self._lock: - removed_active = [t for t in self.active_tasks if t.task_id == task_id] - self.waiting_queue = deque( - t for t in self.waiting_queue if t.task_id != task_id - ) - self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id] + task = self._tasks.get(task_id) self._callbacks.pop(task_id, None) - return removed_active + if task is None or task.status in ( + TaskStatus.FINISHED, + TaskStatus.ABORTED, + ): + return [], False + + task.status = TaskStatus.ABORTED + self._cancelled_total += 1 + if task in self.waiting_queue: + self.waiting_queue = deque( + waiting for waiting in self.waiting_queue if waiting is not task + ) + self._tasks.pop(task_id, None) + return [task], True + return [], True + + def remove_task(self, task_id: str) -> List[Task]: + """Backward-compatible alias for cancellation.""" + immediate, _ = self.cancel_task(task_id) + return immediate def invoke_callback(self, task_id: str, token: str): - cb = self._callbacks.get(task_id) + with self._lock: + cb = self._callbacks.get(task_id) if cb: cb(token) def get_stats(self) -> Dict[str, Any]: - stats: Dict[str, Any] = { - "total_tasks": self._total_tasks, - "total_tokens": self._total_tokens, - "active_tasks": len(self.active_tasks), - "waiting_queue": len(self.waiting_queue), - } + with self._lock: + waiting = len(self.waiting_queue) + stats: Dict[str, Any] = { + "total_tasks": self._total_tasks, + "total_tokens": self._total_tokens, + "active_tasks": len(self.active_tasks), + "waiting_tasks": waiting, + "waiting_queue": waiting, + "cancelled_total": self._cancelled_total, + } if self._metrics is not None: stats.update(self._metrics.get_stats()) return stats @@ -242,18 +276,21 @@ class TaskManager: finished.append(task) self._total_tokens += task.output_tokens - if self._metrics is not None: - for task in finished: - self._metrics.mark_finished( - task.task_id, task.input_tokens, task.output_tokens - ) - self.active_tasks = [ t for t in self.active_tasks if t.status not in (TaskStatus.FINISHED, TaskStatus.ABORTED) ] - return finished + for task in finished: + self._tasks.pop(task.task_id, None) + self._callbacks.pop(task.task_id, None) + + if self._metrics is not None: + for task in finished: + self._metrics.mark_finished( + task.task_id, task.input_tokens, task.output_tokens + ) + return finished def pull_candidates(self, n: int) -> List[Task]: to_add: List[Task] = [] @@ -263,18 +300,35 @@ class TaskManager: to_add.append(self.waiting_queue.popleft()) return to_add - def activate(self, task: Task): - task.status = TaskStatus.RUNNING + def activate(self, task: Task) -> bool: with self._lock: + if task.status == TaskStatus.ABORTED: + self._tasks.pop(task.task_id, None) + self._callbacks.pop(task.task_id, None) + return False + task.status = TaskStatus.RUNNING self.active_tasks.append(task) + return True def return_to_waiting(self, tasks: List[Task]): + cancelled = [] with self._lock: for task in reversed(tasks): - self.waiting_queue.appendleft(task) + if task.status == TaskStatus.ABORTED: + self._tasks.pop(task.task_id, None) + self._callbacks.pop(task.task_id, None) + cancelled.append(task) + else: + self.waiting_queue.appendleft(task) + if self._metrics is not None: + for task in cancelled: + self._metrics.mark_finished( + task.task_id, task.input_tokens, task.output_tokens + ) def has_work(self) -> bool: - return bool(self.active_tasks or self.waiting_queue) + with self._lock: + return bool(self.active_tasks or self.waiting_queue) def wait_for_tasks(self, timeout: float = 1.0): with self._lock: @@ -296,6 +350,7 @@ class TaskManager: self.waiting_queue.clear() self.active_tasks.clear() self._callbacks.clear() + self._tasks.clear() def wake(self): self._task_event.set() diff --git a/tests/inference/test_engine.py b/tests/inference/test_engine.py index 08ab819..049b728 100644 --- a/tests/inference/test_engine.py +++ b/tests/inference/test_engine.py @@ -157,6 +157,28 @@ def test_engine_generate_streaming_yields_tokens(): assert tokens == ["t1", "t2"] +def test_engine_stream_close_cancels_unfinished_task(): + mock_model, mock_tokenizer = _make_engine_mocks(decode="tok") + callbacks_saved = [] + + def capture_cb(prompt, **kwargs): + callbacks_saved.append(kwargs["stream_callback"]) + return "task-1" + + with patch("astrai.inference.engine.InferenceScheduler") as MockSched: + instance = MockSched.return_value + instance.add_task.side_effect = capture_cb + + engine = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1) + stream = engine.generate("hello", stream=True) + + callbacks_saved[0]("t1") + assert next(stream) == "t1" + stream.close() + + instance.cancel_task.assert_called_once_with("task-1") + + def test_engine_generate_async_yields_tokens_until_stop(): mock_model, mock_tokenizer = _make_engine_mocks(decode="tok") callbacks_saved = [] @@ -186,6 +208,31 @@ def test_engine_generate_async_yields_tokens_until_stop(): assert asyncio.run(collect()) == ["t1", "t2"] +def test_engine_async_close_cancels_unfinished_task(): + mock_model, mock_tokenizer = _make_engine_mocks(decode="tok") + callbacks_saved = [] + + def capture_cb(prompt, **kwargs): + callbacks_saved.append(kwargs["stream_callback"]) + return "task-1" + + with patch("astrai.inference.engine.InferenceScheduler") as MockSched: + instance = MockSched.return_value + instance.add_task.side_effect = capture_cb + + engine = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1) + stream = engine.generate_async("hello") + callbacks_saved[0]("t1") + + async def consume_then_close(): + assert await anext(stream) == "t1" + await stream.aclose() + + asyncio.run(consume_then_close()) + + instance.cancel_task.assert_called_once_with("task-1") + + def test_engine_generate_non_streaming_batch(): mock_model, mock_tokenizer = _make_engine_mocks(decode="r") diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index e864a87..4e65cbd 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -1,6 +1,7 @@ """Tests for scheduler concurrency.""" import threading +import time from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -259,6 +260,104 @@ def _make_real_scheduler(device): return scheduler, tokenizer, model +def test_cancel_waiting_task_storm_returns_to_baseline(device): + scheduler, _tok, _model = _make_real_scheduler(device) + try: + task_ids = [ + scheduler.add_task(f"waiting-{index}", max_tokens=32) for index in range(32) + ] + + assert all(scheduler.cancel_task(task_id) for task_id in task_ids) + stats = scheduler.get_stats() + assert stats["active_tasks"] == 0 + assert stats["waiting_tasks"] == 0 + assert stats["in_flight_tasks"] == 0 + assert stats["kv_cache_tasks"] == 0 + assert stats["cancelled_total"] == len(task_ids) + finally: + scheduler.stop() + + +def test_cancel_active_task_releases_metrics_and_kv(device): + scheduler, _tok, _model = _make_real_scheduler(device) + try: + task_id = scheduler.add_task("active", max_tokens=32) + task = scheduler._task_mgr.pull_candidates(1)[0] + assert scheduler._task_cache.task_alloc(task.task_id, task.prompt_ids) + assert scheduler._task_mgr.activate(task) + + before = scheduler.get_stats() + assert before["active_tasks"] == 1 + assert before["in_flight_tasks"] == 1 + assert before["kv_cache_tasks"] == 1 + + assert scheduler.cancel_task(task_id) + scheduler.start() + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + after = scheduler.get_stats() + if ( + after["active_tasks"] == 0 + and after["in_flight_tasks"] == 0 + and after["kv_cache_tasks"] == 0 + ): + break + time.sleep(0.01) + + assert after["active_tasks"] == 0 + assert after["waiting_tasks"] == 0 + assert after["in_flight_tasks"] == 0 + assert after["kv_cache_tasks"] == 0 + assert after["cancelled_total"] == 1 + finally: + scheduler.stop() + + +def test_cancel_during_kv_allocation_releases_metrics_and_kv(device): + scheduler, _tok, _model = _make_real_scheduler(device) + allocation_started = threading.Event() + continue_allocation = threading.Event() + original_alloc = scheduler._task_cache.task_alloc + + def blocking_alloc(*args, **kwargs): + allocation_started.set() + assert continue_allocation.wait(timeout=5) + return original_alloc(*args, **kwargs) + + try: + with patch.object( + scheduler._task_cache, + "task_alloc", + side_effect=blocking_alloc, + ): + scheduler.start() + task_id = scheduler.add_task("allocation-race", max_tokens=32) + assert allocation_started.wait(timeout=5) + assert scheduler.cancel_task(task_id) + continue_allocation.set() + + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + stats = scheduler.get_stats() + if ( + stats["active_tasks"] == 0 + and stats["waiting_tasks"] == 0 + and stats["in_flight_tasks"] == 0 + and stats["kv_cache_tasks"] == 0 + ): + break + time.sleep(0.01) + + assert stats["active_tasks"] == 0 + assert stats["waiting_tasks"] == 0 + assert stats["in_flight_tasks"] == 0 + assert stats["kv_cache_tasks"] == 0 + assert stats["cancelled_total"] == 1 + finally: + continue_allocation.set() + scheduler.stop() + + def test_run_batch_returns_token_sequences(device): scheduler, _tok, _model = _make_real_scheduler(device) try: diff --git a/tests/inference/test_server.py b/tests/inference/test_server.py index 19a5701..926e452 100644 --- a/tests/inference/test_server.py +++ b/tests/inference/test_server.py @@ -166,11 +166,15 @@ def test_messages_with_system(client, loaded_model): def test_chat_completions_stop_sequence(client, loaded_model): """POST /v1/chat/completions with stop parameter truncates at stop sequence.""" + closed = [] async def async_gen(): - yield "Hello" - yield "X" - yield "world" + try: + yield "Hello" + yield "X" + yield "world" + finally: + closed.append(True) get_app().state.engine = loaded_model loaded_model.generate_async.return_value = async_gen() @@ -188,15 +192,20 @@ def test_chat_completions_stop_sequence(client, loaded_model): content = data["choices"][0]["message"]["content"] assert "X" in content assert "world" not in content + assert closed == [True] def test_chat_completions_stop_sequence_stream(client, loaded_model): """POST /v1/chat/completions with stop parameter truncates SSE stream.""" + closed = [] async def async_gen(): - yield "Hello" - yield "X" - yield "world" + try: + yield "Hello" + yield "X" + yield "world" + finally: + closed.append(True) get_app().state.engine = loaded_model loaded_model.generate_async.return_value = async_gen() @@ -217,6 +226,7 @@ def test_chat_completions_stop_sequence_stream(client, loaded_model): assert any( "finish_reason" in line for line in content.split("\n") if "stop" in line ) + assert closed == [True] def test_chat_completions_real_engine(tmp_path, client): diff --git a/tests/inference/test_task.py b/tests/inference/test_task.py index a93b732..c65eaa1 100644 --- a/tests/inference/test_task.py +++ b/tests/inference/test_task.py @@ -73,15 +73,21 @@ def test_task_manager_remove_task(): assert len(tm.waiting_queue) == 0 -def test_task_manager_remove_active_task(): +def test_task_manager_cancel_active_task_defers_removal(): tm = TaskManager(tokenizer=_make_mock_tokenizer()) tid = tm.add_task("test") tasks = tm.pull_candidates(1) tm.activate(tasks[0]) assert len(tm.active_tasks) == 1 - removed = tm.remove_task(tid) - assert len(removed) == 1 + immediate, cancelled = tm.cancel_task(tid) + assert cancelled is True + assert immediate == [] + assert tm.active_tasks[0].status == TaskStatus.ABORTED + + removed = tm.remove_finished_tasks([]) + assert removed == tasks assert len(tm.active_tasks) == 0 + assert tm.get_stats()["cancelled_total"] == 1 def test_task_manager_pull_candidates_fifo():