fix: resolve audited training and inference bugs
- reject prompts that encode to zero tokens in add_task instead of admitting a task whose prefill can never run, and surface empty-id run_batch calls as prompt_empty errors - deliver the STOP stream callback when cancelling a live task so clients observe termination instead of hanging until socket timeout - strip the torch.compile _orig_mod. prefix at every unwrap_model site and when loading checkpoints so FSDP state dicts and saved weights no longer leak the wrapper name into downstream keys - reject online_* train strategies with nprocs > 1 at config validation time, explaining the NCCL all-gather deadlock they would otherwise hit mid-run - apply the frequency penalty before temperature scaling (OpenAI semantics) so the penalty survives temperature=0 instead of being annihilated by the 1e8 logit blowup, and exclude penalty pipelines from the greedy fast path - return logprobs from the raw pre-strategy distribution so they match training-side policy logprobs for PPO/GRPO importance ratios
This commit is contained in:
@@ -2,7 +2,9 @@
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from astrai.inference import Task, TaskManager, TaskStatus
|
||||
import pytest
|
||||
|
||||
from astrai.inference import STOP, Task, TaskManager, TaskStatus
|
||||
|
||||
|
||||
def _make_mock_tokenizer():
|
||||
@@ -178,3 +180,40 @@ def test_task_manager_get_stats():
|
||||
assert stats["total_tasks"] == 1
|
||||
assert stats["waiting_queue"] == 1
|
||||
assert stats["active_tasks"] == 0
|
||||
|
||||
|
||||
def test_task_manager_add_task_rejects_empty_prompt():
|
||||
tm = TaskManager(tokenizer=_make_mock_tokenizer())
|
||||
tm.tokenizer.encode.return_value = []
|
||||
|
||||
with pytest.raises(ValueError, match="zero tokens"):
|
||||
tm.add_task("")
|
||||
|
||||
|
||||
def test_task_manager_cancel_delivers_stop_callback():
|
||||
tm = TaskManager(tokenizer=_make_mock_tokenizer())
|
||||
received = []
|
||||
tm.add_task("test", stream_callback=received.append)
|
||||
|
||||
immediate, cancelled = tm.cancel_task("does-not-exist")
|
||||
assert not cancelled and immediate == [] and received == []
|
||||
|
||||
task_id = next(iter(tm._tasks))
|
||||
immediate, cancelled = tm.cancel_task(task_id)
|
||||
assert cancelled
|
||||
assert len(immediate) == 1
|
||||
assert received == [STOP]
|
||||
|
||||
|
||||
def test_task_manager_cancel_active_task_delivers_stop_callback():
|
||||
tm = TaskManager(tokenizer=_make_mock_tokenizer())
|
||||
received = []
|
||||
task_id = tm.add_task("test", stream_callback=received.append)
|
||||
task = tm._tasks[task_id]
|
||||
tm.waiting_queue.clear()
|
||||
tm.active_tasks.append(task)
|
||||
task.status = TaskStatus.RUNNING
|
||||
|
||||
immediate, cancelled = tm.cancel_task(task_id)
|
||||
assert cancelled and immediate == []
|
||||
assert received == [STOP]
|
||||
|
||||
Reference in New Issue
Block a user