From 47b3ed4e44060e0b43c3d44b59d829efb29c7905 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 9 Aug 2026 13:29:30 +0800 Subject: [PATCH] feat: propagate attention backend across scheduler threads - InferenceEngine/Scheduler accept an explicit backend - capture request-level attn_backend context onto Task - split prefill/decode batches by backend instance - ASTR_BACKEND env overrides ContextVar as process-wide policy - report resolved backend and CUDA-graph state in benchmark --- astrai/extension/attention_backend.py | 105 ++++++++------- astrai/inference/engine.py | 13 ++ astrai/inference/runtime/executor.py | 12 +- astrai/inference/scheduler.py | 178 ++++++++++++++++++-------- astrai/inference/task.py | 9 +- scripts/tools/benchmark.py | 10 +- tests/extension/test_backend.py | 14 ++ tests/inference/test_engine.py | 38 ++++++ tests/inference/test_scheduler.py | 65 ++++++++++ 9 files changed, 334 insertions(+), 110 deletions(-) diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py index 2bc6510..e62ce27 100644 --- a/astrai/extension/attention_backend.py +++ b/astrai/extension/attention_backend.py @@ -56,8 +56,10 @@ if TYPE_CHECKING: _default_backend: Optional["AttentionBackend"] = None _default_backend_lock = threading.Lock() -_current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar( - "attn_backend" +_env_backend_name: Optional[str] = None +_env_backend: Optional["AttentionBackend"] = None +_current_backend: contextvars.ContextVar[Optional["AttentionBackend"]] = ( + contextvars.ContextVar("attn_backend", default=None) ) @@ -149,39 +151,70 @@ def _backend_supports( def _resolve_default_backend() -> "AttentionBackend": """Pick the highest-priority available backend (cuda -> flash -> torch). - Set ``ASTR_BACKEND`` to override: ``ASTR_BACKEND=cuda``, ``torch_native``, - or ``flash``. The value is the registered name (same as the - ``ATTN_BACKEND`` enum value). - Resolved lazily on first ``get_backend()`` and cached. Per-call capability fallback happens in ``attention()``, so the default is safe for training and fp32 models. """ - forced = os.environ.get("ASTR_BACKEND", "").strip().lower() - if forced: - try: - return AttentionBackendFactory.create(forced) - except (ValueError, RuntimeError): - pass return _priority_backends()[0] -def get_backend() -> "AttentionBackend": - """Return the active backend for the current thread/context. +def _environment_backend() -> Optional["AttentionBackend"]: + """Resolve the process-wide ``ASTR_BACKEND`` override, if configured.""" + global _env_backend, _env_backend_name + name = os.environ.get("ASTR_BACKEND", "").strip().lower() + if not name: + return None + if name != _env_backend_name: + with _default_backend_lock: + if name != _env_backend_name: + try: + _env_backend = AttentionBackendFactory.create(name) + except (ValueError, RuntimeError): + _env_backend = None + _env_backend_name = name + return _env_backend - Falls back to the highest-priority available backend (cuda -> flash -> - torch_native) when no backend has been activated via ``with``. Set - ``ASTR_BACKEND`` to override the default. + +def _resolve_backend( + backend: Optional[Union[str, ATTN_BACKEND, "AttentionBackend", type]] = None, +) -> "AttentionBackend": + """Resolve a backend configuration, defaulting to the process policy.""" + if backend is not None: + if isinstance(backend, ATTN_BACKEND): + return AttentionBackendFactory.create(backend.value) + if isinstance(backend, str): + return AttentionBackendFactory.create(backend) + if isinstance(backend, type) and issubclass(backend, AttentionBackend): + return backend() + if isinstance(backend, AttentionBackend): + return backend + raise TypeError( + f"expected a registered name, ATTN_BACKEND, AttentionBackend type, " + f"or instance, got {type(backend).__name__}" + ) + + global _default_backend + if _default_backend is None: + with _default_backend_lock: + if _default_backend is None: + _default_backend = _resolve_default_backend() + return _default_backend + + +def get_backend( + use_default: bool = True, +) -> Optional["AttentionBackend"]: + """Return the context override, optionally falling back to the process default. + + ``ASTR_BACKEND`` is a process-wide override and takes precedence over the + context value. Pass ``use_default=False`` at request submission to retain + only an environment override or the caller's :func:`attn_backend` value. """ - try: - return _current_backend.get() - except LookupError: - global _default_backend - if _default_backend is None: - with _default_backend_lock: - if _default_backend is None: - _default_backend = _resolve_default_backend() - return _default_backend + return ( + _environment_backend() + or _current_backend.get() + or (_resolve_backend() if use_default else None) + ) @contextmanager @@ -200,20 +233,7 @@ def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]): with attn_backend(TorchNativeBackend()): ... """ - if isinstance(backend, ATTN_BACKEND): - instance = AttentionBackendFactory.create(backend.value) - elif isinstance(backend, str): - instance = AttentionBackendFactory.create(backend) - elif isinstance(backend, type) and issubclass(backend, AttentionBackend): - instance = backend() - elif isinstance(backend, AttentionBackend): - instance = backend - else: - raise TypeError( - f"expected a registered name, ATTN_BACKEND, AttentionBackend type, " - f"or instance, " - f"got {type(backend).__name__}" - ) + instance = _resolve_backend(backend) token = _current_backend.set(instance) try: yield instance @@ -284,10 +304,7 @@ def attention( """ backend = get_backend() if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal): - try: - explicit = _current_backend.get() - except LookupError: - explicit = None + explicit = get_backend(use_default=False) if explicit is not None: raise RuntimeError( f"Explicitly-set backend {type(backend).__name__} cannot " diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index 18c2419..2f2b5a1 100644 --- a/astrai/inference/engine.py +++ b/astrai/inference/engine.py @@ -8,6 +8,7 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, import torch import torch.nn as nn +from astrai.extension import ATTN_BACKEND, AttentionBackend, get_backend from astrai.inference.cache import PagePool from astrai.inference.scheduler import InferenceScheduler from astrai.inference.task import STOP @@ -75,6 +76,7 @@ class InferenceEngine: max_seq_len: Optional[int] = None, cache: Optional[PagePool] = None, enable_cuda_graph: bool = True, + backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None, ): self.model = model self.tokenizer = tokenizer @@ -85,6 +87,7 @@ class InferenceEngine: max_seq_len=max_seq_len, cache=cache, enable_cuda_graph=enable_cuda_graph, + backend=backend, ) self.scheduler.start() @@ -174,6 +177,7 @@ class InferenceEngine: rep_window: int, ) -> Union[Generator, str, List[str]]: n = len(prompts) + request_backend = get_backend(use_default=False) result = GenerateResult(count=n) task_ids = [ self.scheduler.add_task( @@ -184,6 +188,7 @@ class InferenceEngine: top_k=top_k, frequency_penalty=frequency_penalty, rep_window=rep_window, + backend=request_backend, stream_callback=lambda token, idx=i: result.append(token, idx), ) for i, p in enumerate(prompts) @@ -223,6 +228,14 @@ class InferenceEngine: def get_stats(self) -> Dict[str, Any]: return self.scheduler.get_stats() + @property + def backend_name(self) -> str: + return self.scheduler.backend_name + + @property + def cuda_graph_enabled(self) -> bool: + return self.scheduler.cuda_graph_enabled + def shutdown(self): self.scheduler.stop() if torch.cuda.is_available(): diff --git a/astrai/inference/runtime/executor.py b/astrai/inference/runtime/executor.py index 7e76f52..810b965 100644 --- a/astrai/inference/runtime/executor.py +++ b/astrai/inference/runtime/executor.py @@ -8,9 +8,7 @@ import torch from torch import Tensor from astrai.extension.attention_backend import ( - ATTN_BACKEND, CudaBackend, - attn_backend, get_backend, ) from astrai.inference.cache import PagePool, TaskCacheManager @@ -153,7 +151,6 @@ def _warmup_cuda_graphs( with ( torch.inference_mode(), - attn_backend(ATTN_BACKEND.CUDA), timed(f"warmup decode b={b}", logger), ): for step in range(2): @@ -208,7 +205,10 @@ class Executor: config = model.config max_q_heads = config.num_attention_heads head_dim = config.hidden_size // config.num_attention_heads - self._graph_supported = CudaBackend.supports(head_dim=head_dim) + backend = get_backend() + self._graph_supported = backend.supports_graph() and CudaBackend.supports( + head_dim=head_dim + ) self._workspace = InferenceWorkspace( max_batch_size=kv_cache.max_batch_size, max_seq_len=kv_cache.max_seq_len, @@ -240,6 +240,10 @@ class Executor: device=self.device, ) + @property + def cuda_graph_enabled(self) -> bool: + return self._graph_ctx.enabled and self._graph_supported + def _sample_logits( self, logits: Tensor, diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index 35f0d21..e947b54 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -1,10 +1,17 @@ import logging import threading import uuid -from typing import Any, Dict, List, Optional, Tuple +from contextlib import nullcontext +from typing import Any, Dict, List, Optional, Tuple, Union import torch +from astrai.extension import ( + ATTN_BACKEND, + AttentionBackend, + attn_backend, + get_backend, +) from astrai.inference.cache import PagePool, TaskCacheManager from astrai.inference.metrics import MetricsCollector from astrai.inference.runtime.executor import Executor @@ -28,6 +35,7 @@ class InferenceScheduler: dtype: Optional[torch.dtype] = None, cache: Optional[PagePool] = None, enable_cuda_graph: bool = True, + backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None, ): config = model.config @@ -69,14 +77,31 @@ class InferenceScheduler: metrics=self._metrics, ) - self._executor = Executor( - model=model, - kv_cache=self._cache, - task_cache=self._task_cache, - device=self.device, - dtype=self.dtype, - enable_cuda_graph=enable_cuda_graph, - ) + if backend is None: + self._backend = None + default_backend = get_backend() + self._backend_name = type(default_backend).__name__ + with attn_backend(default_backend): + self._executor = Executor( + model=model, + kv_cache=self._cache, + task_cache=self._task_cache, + device=self.device, + dtype=self.dtype, + enable_cuda_graph=enable_cuda_graph, + ) + else: + with attn_backend(backend): + self._backend = get_backend() + self._backend_name = type(self._backend).__name__ + self._executor = Executor( + model=model, + kv_cache=self._cache, + task_cache=self._task_cache, + device=self.device, + dtype=self.dtype, + enable_cuda_graph=enable_cuda_graph, + ) self._stop_event = threading.Event() self._loop_thread: Optional[threading.Thread] = None @@ -91,6 +116,26 @@ class InferenceScheduler: def get_stats(self) -> Dict[str, Any]: return self._task_mgr.get_stats() + @property + def backend_name(self) -> str: + return self._backend_name + + @property + def cuda_graph_enabled(self) -> bool: + return self._executor.cuda_graph_enabled + + def _backend_context(self): + if self._backend is None: + return nullcontext() + return attn_backend(self._backend) + + @staticmethod + def _task_backend_groups(tasks: List[Task]): + groups = {} + for task in tasks: + groups.setdefault(task.backend, (task.backend, []))[1].append(task) + return groups.values() + def _step( self, tasks: List[Task], return_logprobs: bool = False ) -> Tuple[List[Task], List[Task]]: @@ -121,15 +166,24 @@ class InferenceScheduler: for t in to_prefill: t.input_tokens = len(t.prompt_ids) - groups: Dict[Tuple[int, int], List[Task]] = {} + groups: Dict[Tuple[int, int, Optional[AttentionBackend]], List[Task]] = {} for t in to_prefill: start_pos = min( self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1 ) - groups.setdefault((len(t.prompt_ids), start_pos), []).append(t) + groups.setdefault((len(t.prompt_ids), start_pos, t.backend), []).append( + t + ) - for (prompt_len, start_pos), group in groups.items(): - with self._metrics.record([t.task_id for t in group], "prefill"): + for (prompt_len, start_pos, _), group in groups.items(): + backend = group[0].backend + backend_context = ( + attn_backend(backend) if backend is not None else nullcontext() + ) + with ( + backend_context, + self._metrics.record([t.task_id for t in group], "prefill"), + ): prefilled, step_out = self._executor.execute_prefill( group, prompt_len, start_pos, return_logprobs=return_logprobs ) @@ -158,12 +212,18 @@ class InferenceScheduler: t.status = TaskStatus.ABORTED aborted.append(t) - if decoded: - with self._metrics.record([t.task_id for t in decoded], "decode"): + for backend, group in self._task_backend_groups(decoded): + backend_context = ( + attn_backend(backend) if backend is not None else nullcontext() + ) + with ( + backend_context, + self._metrics.record([t.task_id for t in group], "decode"), + ): step_out = self._executor.execute_decode( - decoded, return_logprobs=return_logprobs + group, return_logprobs=return_logprobs ) - for t, out in zip(decoded, step_out): + for t, out in zip(group, step_out): t.output_ids.append(out[0] if return_logprobs else out) t.output_tokens += 1 t.advance_kv() @@ -174,49 +234,52 @@ class InferenceScheduler: def _run_generation_loop(self): stop_ids = self._task_mgr.tokenizer.stop_ids try: - while not self._stop_event.is_set(): - finished = self._task_mgr.remove_finished_tasks(stop_ids) - for task in finished: - if task.status == TaskStatus.FINISHED: - self._task_cache.task_record_hashes( - task.task_id, - self._task_cache.task_cacheable_ids( - task.task_id, task.prompt_ids, task.output_ids - ), - ) - self._task_cache.task_free(task.task_id) + with self._backend_context(): + while not self._stop_event.is_set(): + finished = self._task_mgr.remove_finished_tasks(stop_ids) + for task in finished: + if task.status == TaskStatus.FINISHED: + self._task_cache.task_record_hashes( + task.task_id, + self._task_cache.task_cacheable_ids( + task.task_id, task.prompt_ids, task.output_ids + ), + ) + self._task_cache.task_free(task.task_id) - active = self._task_mgr.get_active_tasks() - available = self._task_mgr.max_batch_size - len(active) - if available > 0: - candidates = self._task_mgr.pull_candidates(available) - failed = [] - for task in candidates: - if self._task_cache.task_alloc(task.task_id, task.prompt_ids): - self._task_mgr.activate(task) - else: - failed.append(task) - if failed: - self._task_mgr.return_to_waiting(failed) + active = self._task_mgr.get_active_tasks() + available = self._task_mgr.max_batch_size - len(active) + if available > 0: + candidates = self._task_mgr.pull_candidates(available) + failed = [] + for task in candidates: + if self._task_cache.task_alloc( + task.task_id, task.prompt_ids + ): + self._task_mgr.activate(task) + else: + failed.append(task) + if failed: + self._task_mgr.return_to_waiting(failed) - if not self._task_mgr.has_work(): - self._task_mgr.wait_for_tasks(timeout=1.0) - continue + if not self._task_mgr.has_work(): + self._task_mgr.wait_for_tasks(timeout=1.0) + continue - active = self._task_mgr.get_active_tasks() + active = self._task_mgr.get_active_tasks() - decoded, aborted = self._step(active) + decoded, aborted = self._step(active) - for t in aborted: - self._task_mgr.invoke_callback(t.task_id, STOP) - - for t in decoded: - new_text = t.decode_new_token(self._task_mgr.tokenizer) - if new_text: - self._task_mgr.invoke_callback(t.task_id, new_text) - if t.is_finished(stop_ids): + for t in aborted: self._task_mgr.invoke_callback(t.task_id, STOP) + for t in decoded: + new_text = t.decode_new_token(self._task_mgr.tokenizer) + if new_text: + self._task_mgr.invoke_callback(t.task_id, new_text) + if t.is_finished(stop_ids): + self._task_mgr.invoke_callback(t.task_id, STOP) + except Exception as e: self._stop_event.set() logger.error(f"Scheduler loop crashed: {e}", exc_info=True) @@ -286,6 +349,7 @@ class InferenceScheduler: """ stop_ids = self._task_mgr.tokenizer.stop_ids seq_cap = self.max_seq_len + request_backend = get_backend(use_default=False) tasks: List[Task] = [] for ids in prompt_ids_list: @@ -309,6 +373,7 @@ class InferenceScheduler: top_k=top_k, frequency_penalty=frequency_penalty, rep_window=rep_window, + backend=request_backend, ) if not self._task_cache.task_alloc(task.task_id, task.prompt_ids): tasks.append(None) @@ -320,9 +385,10 @@ class InferenceScheduler: try: live = [t for t in tasks if t is not None] - while live: - decoded, _ = self._step(live, return_logprobs=return_logprobs) - live = [t for t in decoded if not t.is_finished(stop_ids)] + with self._backend_context(): + while live: + decoded, _ = self._step(live, return_logprobs=return_logprobs) + live = [t for t in decoded if not t.is_finished(stop_ids)] finally: for t in tasks: if t is not None: diff --git a/astrai/inference/task.py b/astrai/inference/task.py index 70a3f83..6b4bfc9 100644 --- a/astrai/inference/task.py +++ b/astrai/inference/task.py @@ -3,13 +3,16 @@ import time import uuid from collections import deque from enum import Enum -from typing import Any, Callable, Deque, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Optional from tokenizers.decoders import DecodeStream from astrai.inference.metrics import MetricsCollector from astrai.tokenize.tokenizer import AutoTokenizer +if TYPE_CHECKING: + from astrai.extension import AttentionBackend + STOP = object() @@ -62,6 +65,7 @@ class Task: top_k: int = 50, frequency_penalty: float = 0.0, rep_window: int = 64, + backend: Optional["AttentionBackend"] = None, ): self.task_id = task_id self.prompt_ids = prompt_ids @@ -71,6 +75,7 @@ class Task: self.top_k = top_k self.frequency_penalty = frequency_penalty self.rep_window = rep_window + self.backend = backend self.status = TaskStatus.PENDING self.output_ids: List[int] = [] @@ -152,6 +157,7 @@ class TaskManager: top_k: int = 50, frequency_penalty: float = 0.0, rep_window: int = 64, + backend: Optional["AttentionBackend"] = None, stream_callback: Optional[Callable[[str], None]] = None, ) -> str: task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}" @@ -173,6 +179,7 @@ class TaskManager: top_k=top_k, frequency_penalty=frequency_penalty, rep_window=rep_window, + backend=backend, ) with self._lock: diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py index 73e3051..2fc50c5 100644 --- a/scripts/tools/benchmark.py +++ b/scripts/tools/benchmark.py @@ -259,21 +259,20 @@ class GenerationBenchmark: max_seq_len=max_seq_len, cache=pool, enable_cuda_graph=self.cuda_graph, + backend=self.backend, ) prompts = [prompt] * batch_size try: # Capture graphs and populate the allocator before timing. The # first request also includes model/scheduler startup effects. - with attn_backend(self.backend): - engine.generate(prompts, max_tokens=gen_length, temperature=0.0) + engine.generate(prompts, max_tokens=gen_length, temperature=0.0) if self.device.startswith("cuda"): torch.cuda.synchronize() t0 = time.perf_counter() for _ in range(num_trials): - with attn_backend(self.backend): - engine.generate(prompts, max_tokens=gen_length, temperature=0.0) + engine.generate(prompts, max_tokens=gen_length, temperature=0.0) if self.device.startswith("cuda"): torch.cuda.synchronize() elapsed = time.perf_counter() - t0 @@ -291,7 +290,8 @@ class GenerationBenchmark: "benchmark_type": "engine_decode", "num_trials": num_trials, "prompt_length": prompt_tokens, - "engine": True, + "backend": engine.backend_name, + "cuda_graph": engine.cuda_graph_enabled, }, ) diff --git a/tests/extension/test_backend.py b/tests/extension/test_backend.py index 61b3417..cdbcb5f 100644 --- a/tests/extension/test_backend.py +++ b/tests/extension/test_backend.py @@ -42,6 +42,20 @@ def test_attn_backend_context_with_registered_name(): assert get_backend() is default +def test_backend_can_read_only_context_selection(): + assert get_backend(use_default=False) is None + with attn_backend("cuda") as backend: + assert get_backend(use_default=False) is backend + assert get_backend(use_default=False) is None + + +def test_environment_backend_overrides_context(monkeypatch): + monkeypatch.setenv("ASTR_BACKEND", "torch_native") + with attn_backend("cuda"): + assert type(get_backend()).__name__ == "TorchNativeBackend" + assert type(get_backend(use_default=False)).__name__ == "TorchNativeBackend" + + def test_attention_backend_factory_lists_builtin_backends(): assert AttentionBackendFactory.list_registered() == [ "cuda", diff --git a/tests/inference/test_engine.py b/tests/inference/test_engine.py index b240619..18a2026 100644 --- a/tests/inference/test_engine.py +++ b/tests/inference/test_engine.py @@ -3,6 +3,7 @@ import threading from unittest.mock import MagicMock, patch +from astrai.extension import TorchNativeBackend, attn_backend from astrai.inference import STOP from astrai.inference.engine import GenerateResult, InferenceEngine @@ -199,3 +200,40 @@ def test_engine_generate_zero_max_tokens_stream_is_empty(): eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1) assert list(eng.generate("hello", stream=True, max_tokens=0)) == [] instance.add_task.assert_not_called() + + +def test_engine_passes_backend_to_scheduler(): + mock_model = MagicMock() + mock_tokenizer = MagicMock() + + with patch("astrai.inference.engine.InferenceScheduler") as MockSched: + InferenceEngine( + mock_model, + mock_tokenizer, + max_batch_size=1, + backend="torch_native", + ) + + assert MockSched.call_args.kwargs["backend"] == "torch_native" + + +def test_generate_captures_calling_backend_context(): + mock_model = MagicMock() + mock_tokenizer = MagicMock() + captured = [] + + with patch("astrai.inference.engine.InferenceScheduler") as MockSched: + instance = MockSched.return_value + + def fake_add(prompt, **kwargs): + captured.append(kwargs["backend"]) + kwargs["stream_callback"](STOP) + return "task" + + instance.add_task.side_effect = fake_add + engine = InferenceEngine(mock_model, mock_tokenizer) + with attn_backend("torch_native"): + assert engine.generate("hello") == "" + + assert len(captured) == 1 + assert isinstance(captured[0], TorchNativeBackend) diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 487c48d..81c3626 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -7,7 +7,9 @@ from unittest.mock import MagicMock, patch import pytest import torch +from astrai.extension import CudaBackend, TorchNativeBackend, get_backend from astrai.inference import InferenceScheduler +from astrai.inference.metrics import MetricsCollector from astrai.inference.runtime.executor import DecodeSteadyState, Executor from astrai.inference.task import Task from astrai.model.transformer import AutoRegressiveLM @@ -76,6 +78,69 @@ def test_scheduler_concurrent_add_task(mock_model_and_tokenizer): assert len(results["task_ids"]) == 50 +def test_generation_loop_activates_backend_in_worker_thread(): + scheduler = object.__new__(InferenceScheduler) + scheduler._backend = TorchNativeBackend() + scheduler._stop_event = threading.Event() + scheduler._task_cache = MagicMock() + + observed = [] + task_mgr = MagicMock() + task_mgr.tokenizer.stop_ids = [0] + task_mgr.remove_finished_tasks.return_value = [] + task_mgr.get_active_tasks.return_value = [] + task_mgr.max_batch_size = 1 + task_mgr.pull_candidates.return_value = [] + task_mgr.has_work.return_value = False + + def observe_backend(*args, **kwargs): + observed.append(type(get_backend())) + scheduler._stop_event.set() + + task_mgr.wait_for_tasks.side_effect = observe_backend + scheduler._task_mgr = task_mgr + + thread = threading.Thread(target=scheduler._run_generation_loop) + thread.start() + thread.join(timeout=5) + + assert not thread.is_alive() + assert observed == [TorchNativeBackend] + + +def test_step_splits_decode_batch_by_request_backend(): + scheduler = object.__new__(InferenceScheduler) + scheduler._task_cache = MagicMock() + scheduler._task_cache.task_extend.return_value = True + scheduler._metrics = MetricsCollector() + scheduler._executor = MagicMock() + + observed = [] + + def execute(tasks, **kwargs): + observed.append((type(get_backend()), [task.task_id for task in tasks])) + return [1] * len(tasks) + + scheduler._executor.execute_decode.side_effect = execute + + torch_task = Task("torch", [1], backend=TorchNativeBackend()) + cuda_task = Task("cuda", [1], backend=CudaBackend()) + for task in (torch_task, cuda_task): + task.input_tokens = 1 + task.output_ids = [1] + task.mark_prefill_done() + scheduler._metrics.register(task.task_id) + + produced, aborted = scheduler._step([torch_task, cuda_task]) + + assert aborted == [] + assert produced == [torch_task, cuda_task] + assert observed == [ + (TorchNativeBackend, ["torch"]), + (CudaBackend, ["cuda"]), + ] + + def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer): """Test concurrent add and remove task operations.""" mock_model, mock_tokenizer = mock_model_and_tokenizer