host
78098c5594
每用户一台常驻 VM（暂停态，resume 到 SSH 5 秒内为待验证目标）， 池空/失败回落冷启动；CHV 配置加 virtio-balloon + free-page-reporting 与 vsock notify 最小实现（SSH 仍是命令通道与就绪判定）。 新模块 resident_pool.py / vsock.py + 客内 vsock_notify.py； _config 增加 CHV_RESIDENT_IDLE_MINUTES / CHV_BALLOON / CHV_VSOCK； worker_adapter 接常驻孤儿清理 facade；ADR-17 差距表见 docs/chv-muse-p1-gap.md。 恢复收尾（分离后回归，见 gap.md §4）：VmDeathTest 死亡模拟改杀 实际跑命令的 VM；两处冷路径旧测试显式关常驻；_balloon_vsock_payload 与 listen_notify 拆到 50 行内；resident_pool.py:208 mypy 修掉。 executor.py out 未注解与 marker_tools.py:142 是基线旧错，未动。 测试：全量 cloud 套件 1934 OK（skipped=8）；mypy 在改动文件上 干净（仅基线旧错 4 处）。真机验证待第三棒让出 WSL 测试机。
731 lines
29 KiB
Python
731 lines
29 KiB
Python
"""CHV 常驻池（feat-chv-muse-p1）测试。
|
||
|
||
覆盖 cloud/chv/resident_pool.py、executor 的常驻接线
|
||
（_run_with_vm 常驻/冷启动两路、finally 语义）与
|
||
worker_adapter.reconcile_resident_pool facade。全部用假
|
||
执行器/假 VMM，绝不起真实 VM、不连真实网络。
|
||
|
||
断言点（按任务书）：
|
||
- acquire/resume 失败回落冷启动，失败条目丢弃；
|
||
- release 只 pause 不销毁，pause 失败条目留 running；
|
||
- registry 损坏 quarantine、单坏条目不连坐；
|
||
- 重启后条目一律 orphan 不复用，只能 reconcile 清理；
|
||
- _run_with_vm 常驻命中跑完 pause、冷启动跑完 shutdown，
|
||
客内抛错时两路 finally 都执行；
|
||
- 带密钥 env / 未启用 / 建池失败时回落冷启动，任务不丢；
|
||
- 每用户隔离：release 只动本用户条目；
|
||
- 空闲回收只动超阈值 paused 条目，阈值 0 即关闭；
|
||
- WarmPool 实验模块不动（wired_to_dispatch 仍 False）。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import tempfile
|
||
import unittest
|
||
from types import SimpleNamespace
|
||
from unittest import mock
|
||
|
||
from cloud.chv import worker_adapter
|
||
from cloud.chv.errors import ChvFailure
|
||
from cloud.chv.executor import CloudHypervisorExecutor
|
||
from cloud.chv.experimental_pool import WarmPool
|
||
from cloud.chv.resident_pool import (
|
||
ResidentPool,
|
||
reconcile_resident_orphans_for_executor,
|
||
resident_task_for_user,
|
||
resident_user_key,
|
||
)
|
||
|
||
|
||
class _FakeExecutor:
|
||
"""ResidentPool 用的鸭子型执行器：记录调用，可注入失败。
|
||
|
||
wait 失败时模拟真实语义：先自行 shutdown 再抛错
|
||
（wait_guest_ready 的失败路径就是 stop_vm 后 raise）。
|
||
"""
|
||
|
||
def __init__(self, state_dir, env=None):
|
||
self.state_dir = state_dir
|
||
self._env = env or {}
|
||
self.resumes = []
|
||
self.pauses = []
|
||
self.wakes = []
|
||
self.waits = []
|
||
self.shutdowns = []
|
||
self.hibernates = []
|
||
self.emits = []
|
||
self.resume_fail = False
|
||
self.wait_fail = False
|
||
self.wake_fail = False
|
||
self.pause_fail = False
|
||
self.hibernate_fail = False
|
||
self.shutdown_fail = False
|
||
|
||
def vm_name(self, task_or_name):
|
||
raw = task_or_name.get("id") if isinstance(
|
||
task_or_name, dict) else str(task_or_name)
|
||
safe = re.sub(r"[^A-Za-z0-9_.-]", "-", raw)
|
||
if safe.startswith("chv-"):
|
||
return safe
|
||
return "chv-" + safe
|
||
|
||
def resume_vm(self, task_or_name):
|
||
self.resumes.append(self.vm_name(task_or_name))
|
||
if self.resume_fail:
|
||
raise RuntimeError("resume boom")
|
||
return True
|
||
|
||
def pause_vm(self, task_or_name):
|
||
self.pauses.append(self.vm_name(task_or_name))
|
||
if self.pause_fail:
|
||
raise RuntimeError("pause boom")
|
||
return True
|
||
|
||
def wake_vm(self, task_or_name):
|
||
self.wakes.append(self.vm_name(task_or_name))
|
||
if self.wake_fail:
|
||
raise RuntimeError("wake boom")
|
||
return {"vm": self.vm_name(task_or_name), "mode": "cold"}
|
||
|
||
def wait_guest_ready(self, task):
|
||
self.waits.append(self.vm_name(task))
|
||
if self.wait_fail:
|
||
self.shutdown_vm(task)
|
||
raise ChvFailure("guest_not_ready", "ssh 未就绪")
|
||
return True
|
||
|
||
def shutdown_vm(self, task_or_name):
|
||
self.shutdowns.append(self.vm_name(task_or_name))
|
||
if self.shutdown_fail:
|
||
raise RuntimeError("shutdown boom")
|
||
return True
|
||
|
||
def hibernate_vm(self, task_or_name, snapshot_url):
|
||
self.hibernates.append(
|
||
(self.vm_name(task_or_name), snapshot_url))
|
||
if self.hibernate_fail:
|
||
raise RuntimeError("hibernate boom")
|
||
return {"vm": self.vm_name(task_or_name),
|
||
"state": "hibernated", "snapshot": snapshot_url}
|
||
|
||
def _emit(self, event, task, detail="", kind=None):
|
||
self.emits.append(event)
|
||
|
||
|
||
def _write(path, content):
|
||
with open(path, "w") as fh:
|
||
fh.write(content)
|
||
return path
|
||
|
||
|
||
def _registry_path(state_dir):
|
||
return os.path.join(state_dir, "resident-pool.json")
|
||
|
||
|
||
class ResidentKeyTest(unittest.TestCase):
|
||
def test_user_key_sanitizes_and_defaults(self):
|
||
self.assertEqual(resident_user_key(None), "default")
|
||
self.assertEqual(resident_user_key(""), "default")
|
||
self.assertEqual(resident_user_key("alice"), "alice")
|
||
self.assertEqual(resident_user_key("a/b c"), "a-b-c")
|
||
self.assertEqual(resident_task_for_user("alice"),
|
||
{"id": "resident-alice"})
|
||
self.assertEqual(resident_task_for_user(None),
|
||
{"id": "resident-default"})
|
||
|
||
|
||
class ResidentPoolAcquireTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.fake = _FakeExecutor(self.tmp.name)
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _provisioned_pool(self, user="alice"):
|
||
pool = ResidentPool(self.fake)
|
||
got = pool.provision(user)
|
||
self.assertEqual(got["mode"], "resident")
|
||
return pool
|
||
|
||
def test_provision_then_acquire_then_release(self):
|
||
pool = self._provisioned_pool()
|
||
self.assertEqual(self.fake.wakes, ["chv-resident-alice"])
|
||
self.assertEqual(self.fake.pauses, ["chv-resident-alice"])
|
||
self.assertEqual(pool.paused_count(), 1)
|
||
got = pool.acquire("alice")
|
||
self.assertEqual(got["mode"], "resident")
|
||
self.assertEqual(got["vm"], "chv-resident-alice")
|
||
self.assertEqual(got["task"], {"id": "resident-alice"})
|
||
self.assertEqual(self.fake.resumes, ["chv-resident-alice"])
|
||
self.assertEqual(self.fake.waits, ["chv-resident-alice"])
|
||
self.assertEqual(pool.paused_count(), 0)
|
||
# running 中的条目不接受第二个 acquire（回落冷启动）。
|
||
self.assertEqual(pool.acquire("alice")["mode"], "cold")
|
||
self.assertEqual(len(self.fake.resumes), 1)
|
||
# release 只 pause，不销毁。
|
||
self.assertTrue(pool.release("alice"))
|
||
self.assertEqual(
|
||
self.fake.pauses,
|
||
["chv-resident-alice", "chv-resident-alice"])
|
||
self.assertEqual(self.fake.shutdowns, [])
|
||
self.assertEqual(pool.paused_count(), 1)
|
||
# 再取仍是同一台 VM。
|
||
got = pool.acquire("alice")
|
||
self.assertEqual(got["vm"], "chv-resident-alice")
|
||
|
||
def test_resume_failure_falls_back_to_cold_and_drops_entry(self):
|
||
pool = self._provisioned_pool()
|
||
self.fake.resume_fail = True
|
||
got = pool.acquire("alice")
|
||
self.assertEqual(got, {"mode": "cold", "vm": None,
|
||
"task": None})
|
||
self.assertEqual(pool.paused_count(), 0)
|
||
with open(_registry_path(self.tmp.name)) as fh:
|
||
self.assertEqual(json.load(fh), {})
|
||
|
||
def test_wait_failure_falls_back_to_cold_and_drops_entry(self):
|
||
pool = self._provisioned_pool()
|
||
self.fake.wait_fail = True
|
||
got = pool.acquire("alice")
|
||
self.assertEqual(got["mode"], "cold")
|
||
self.assertEqual(pool.paused_count(), 0)
|
||
# wait 的失败路径已自行 stop_vm（模拟真实语义）。
|
||
self.assertEqual(self.fake.shutdowns,
|
||
["chv-resident-alice"])
|
||
with open(_registry_path(self.tmp.name)) as fh:
|
||
self.assertEqual(json.load(fh), {})
|
||
|
||
def test_release_without_running_entry_is_false(self):
|
||
pool = self._provisioned_pool()
|
||
self.assertFalse(pool.release("alice"))
|
||
self.assertFalse(pool.release("nobody"))
|
||
self.assertEqual(self.fake.pauses, ["chv-resident-alice"])
|
||
|
||
def test_release_pause_failure_keeps_entry_running(self):
|
||
pool = self._provisioned_pool()
|
||
pool.acquire("alice")
|
||
self.fake.pause_fail = True
|
||
self.assertFalse(pool.release("alice"))
|
||
# 不删条目、不销毁：下次 acquire 仍是 running→cold，
|
||
# VM 本身没被 shutdown，留待 worker 关闭时清理。
|
||
self.assertEqual(self.fake.shutdowns, [])
|
||
self.assertEqual(pool.acquire("alice")["mode"], "cold")
|
||
|
||
def test_provision_is_idempotent_and_per_user_isolated(self):
|
||
pool = ResidentPool(self.fake)
|
||
self.assertEqual(pool.provision("alice")["mode"], "resident")
|
||
self.assertEqual(pool.provision("bob")["mode"], "resident")
|
||
self.assertEqual(pool.provision("alice")["mode"], "resident")
|
||
self.assertEqual(self.fake.wakes.count("chv-resident-alice"), 1)
|
||
self.assertEqual(self.fake.wakes.count("chv-resident-bob"), 1)
|
||
self.assertEqual(pool.paused_count(), 2)
|
||
# release 只动本用户条目。
|
||
self.assertFalse(pool.release("bob")) # bob 未 running
|
||
pool.acquire("alice")
|
||
self.assertTrue(pool.release("alice"))
|
||
self.assertEqual(pool.paused_count(), 2)
|
||
got = pool.acquire("bob")
|
||
self.assertEqual(got["vm"], "chv-resident-bob")
|
||
self.assertEqual(pool.paused_count(), 1)
|
||
self.assertEqual(self.fake.resumes, ["chv-resident-alice",
|
||
"chv-resident-bob"])
|
||
|
||
def test_provision_wake_failure_falls_back_and_cleans_up(self):
|
||
pool = ResidentPool(self.fake)
|
||
self.fake.wake_fail = True
|
||
got = pool.provision("alice")
|
||
self.assertEqual(got, {"mode": "cold", "vm": None,
|
||
"task": None})
|
||
# 建池失败尽力 shutdown，不留半醒 VM 与条目。
|
||
self.assertEqual(self.fake.shutdowns,
|
||
["chv-resident-alice"])
|
||
self.assertEqual(pool.paused_count(), 0)
|
||
with open(_registry_path(self.tmp.name)) as fh:
|
||
self.assertEqual(json.load(fh), {})
|
||
|
||
|
||
class ResidentPoolRegistryTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.fake = _FakeExecutor(self.tmp.name)
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def test_corrupt_registry_is_quarantined(self):
|
||
path = _write(_registry_path(self.tmp.name), "{not json")
|
||
pool = ResidentPool(self.fake)
|
||
self.assertTrue(pool.registry_corrupt)
|
||
self.assertEqual(pool.paused_count(), 0)
|
||
self.assertEqual(pool.orphan_count(), 0)
|
||
self.assertTrue(os.path.exists(path + ".corrupt"))
|
||
self.assertFalse(os.path.exists(path))
|
||
|
||
def test_top_level_non_dict_is_corrupt(self):
|
||
path = _write(_registry_path(self.tmp.name), "[1, 2]")
|
||
pool = ResidentPool(self.fake)
|
||
self.assertTrue(pool.registry_corrupt)
|
||
self.assertTrue(os.path.exists(path + ".corrupt"))
|
||
|
||
def test_bad_entries_dropped_without_tainting_good_ones(self):
|
||
_write(_registry_path(self.tmp.name), json.dumps({
|
||
"good": {"state": "paused", "vm": "chv-resident-good",
|
||
"user_id": "good"},
|
||
"no-vm": {"state": "paused"},
|
||
"not-a-dict": "x",
|
||
"wrong-state": {"state": "flying",
|
||
"vm": "chv-resident-x"},
|
||
}))
|
||
pool = ResidentPool(self.fake)
|
||
self.assertEqual(pool.registry_corrupt_entries, 3)
|
||
self.assertEqual(pool.orphan_count(), 1)
|
||
self.assertEqual(pool.paused_count(), 0)
|
||
|
||
def test_restart_never_reuses_entries(self):
|
||
pool1 = ResidentPool(self.fake)
|
||
pool1.provision("alice")
|
||
pool2 = ResidentPool(self.fake)
|
||
self.assertEqual(pool2.orphan_count(), 1)
|
||
self.assertEqual(pool2.paused_count(), 0)
|
||
# 新池绝不把上一世条目发给任务。
|
||
self.assertEqual(pool2.acquire("alice")["mode"], "cold")
|
||
self.assertEqual(self.fake.resumes, [])
|
||
cleaned = pool2.reconcile_orphans()
|
||
self.assertEqual(cleaned, ["chv-resident-alice"])
|
||
self.assertEqual(self.fake.shutdowns,
|
||
["chv-resident-alice"])
|
||
with open(_registry_path(self.tmp.name)) as fh:
|
||
self.assertEqual(json.load(fh), {})
|
||
|
||
def test_idle_minutes_env_parsing(self):
|
||
self.assertEqual(
|
||
ResidentPool(self.fake).idle_minutes, 0)
|
||
fake = _FakeExecutor(
|
||
self.tmp.name, env={"CHV_RESIDENT_IDLE_MINUTES": "45"})
|
||
self.assertEqual(ResidentPool(fake).idle_minutes, 45)
|
||
fake = _FakeExecutor(
|
||
self.tmp.name, env={"CHV_RESIDENT_IDLE_MINUTES": "junk"})
|
||
self.assertEqual(ResidentPool(fake).idle_minutes, 0)
|
||
fake = _FakeExecutor(
|
||
self.tmp.name, env={"CHV_RESIDENT_IDLE_MINUTES": "-5"})
|
||
self.assertEqual(ResidentPool(fake).idle_minutes, 0)
|
||
pool = ResidentPool(self.fake, idle_minutes=10)
|
||
self.assertEqual(pool.idle_minutes, 10)
|
||
|
||
|
||
class ResidentPoolIdleTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.fake = _FakeExecutor(self.tmp.name)
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def test_idle_disabled_by_default(self):
|
||
pool = ResidentPool(self.fake)
|
||
pool.provision("alice")
|
||
import time
|
||
self.assertEqual(
|
||
pool.reclaim_idle(now=time.time() + 10 ** 9), [])
|
||
self.assertEqual(self.fake.hibernates, [])
|
||
|
||
def test_overdue_paused_entry_is_hibernated(self):
|
||
import time
|
||
pool = ResidentPool(self.fake, idle_minutes=30)
|
||
pool.provision("alice")
|
||
pool.provision("bob")
|
||
now = time.time() + 31 * 60
|
||
reclaimed = pool.reclaim_idle(now=now)
|
||
self.assertEqual(reclaimed, ["chv-resident-alice",
|
||
"chv-resident-bob"])
|
||
self.assertEqual(len(self.fake.hibernates), 2)
|
||
for vm, url in self.fake.hibernates:
|
||
self.assertTrue(url.startswith("file://"))
|
||
self.assertIn(self.tmp.name, url)
|
||
self.assertTrue(url.endswith(
|
||
os.path.join(vm, "resident.snapshot")))
|
||
self.assertEqual(pool.paused_count(), 0)
|
||
with open(_registry_path(self.tmp.name)) as fh:
|
||
self.assertEqual(json.load(fh), {})
|
||
|
||
def test_not_yet_idle_entry_is_kept(self):
|
||
import time
|
||
pool = ResidentPool(self.fake, idle_minutes=30)
|
||
pool.provision("alice")
|
||
self.assertEqual(
|
||
pool.reclaim_idle(now=time.time() + 10 * 60), [])
|
||
self.assertEqual(pool.paused_count(), 1)
|
||
|
||
def test_hibernate_failure_keeps_entry_for_retry(self):
|
||
import time
|
||
pool = ResidentPool(self.fake, idle_minutes=30)
|
||
pool.provision("alice")
|
||
self.fake.hibernate_fail = True
|
||
self.assertEqual(
|
||
pool.reclaim_idle(now=time.time() + 31 * 60), [])
|
||
self.assertEqual(pool.paused_count(), 1)
|
||
|
||
|
||
class ResidentFacadeTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.fake = _FakeExecutor(self.tmp.name)
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def test_module_facade_cleans_and_empties_registry(self):
|
||
_write(_registry_path(self.tmp.name), json.dumps({
|
||
"a": {"state": "paused", "vm": "chv-resident-a"},
|
||
"b": {"state": "running", "vm": "chv-resident-b"},
|
||
"c": {"state": "provisioning", "vm": "chv-resident-c"},
|
||
"d": {"state": "warm", "vm": "chv-resident-d"},
|
||
"e": {"state": "paused"},
|
||
"f": "not-a-dict",
|
||
}))
|
||
cleaned = reconcile_resident_orphans_for_executor(self.fake)
|
||
self.assertEqual(cleaned, ["chv-resident-a", "chv-resident-b",
|
||
"chv-resident-c"])
|
||
self.assertEqual(self.fake.shutdowns, cleaned)
|
||
with open(_registry_path(self.tmp.name)) as fh:
|
||
self.assertEqual(json.load(fh), {})
|
||
self.assertIn("resident_pool_reconcile", self.fake.emits)
|
||
|
||
def test_module_facade_without_registry_is_empty(self):
|
||
self.assertEqual(
|
||
reconcile_resident_orphans_for_executor(self.fake), [])
|
||
self.assertEqual(self.fake.shutdowns, [])
|
||
|
||
def test_module_facade_without_state_dir_is_empty(self):
|
||
fake = _FakeExecutor(None)
|
||
fake.state_dir = None
|
||
self.assertEqual(
|
||
reconcile_resident_orphans_for_executor(fake), [])
|
||
|
||
class WorkerAdapterFacadeTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.fake = _FakeExecutor(self.tmp.name)
|
||
self.audits = []
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _audit(self, project_id, event, detail):
|
||
self.audits.append((project_id, event, detail))
|
||
|
||
def test_facade_cleans_resident_registry(self):
|
||
# worker_adapter 的 isinstance 门要求真执行器：用假
|
||
# 的 shutdown 替身包一层真执行器实例（只借类型）。
|
||
ex = _RealExecCase.make_executor(self.tmp.name, self.tmp.name)
|
||
ex.shutdown_vm = self.fake.shutdown_vm
|
||
ex._emit = self.fake._emit
|
||
_write(_registry_path(self.tmp.name), json.dumps({
|
||
"a": {"state": "paused", "vm": "chv-resident-a"},
|
||
"b": {"state": "running", "vm": "chv-resident-b"},
|
||
}))
|
||
cleaned = worker_adapter.reconcile_resident_pool(
|
||
ex, self._audit, "startup")
|
||
self.assertEqual(cleaned, ["chv-resident-a", "chv-resident-b"])
|
||
self.assertEqual(self.fake.shutdowns, cleaned)
|
||
self.assertTrue(any(
|
||
event == "chv_resident_pool_reconcile"
|
||
for _p, event, _d in self.audits))
|
||
with open(_registry_path(self.tmp.name)) as fh:
|
||
self.assertEqual(json.load(fh), {})
|
||
|
||
def test_facade_never_constructs_resident_pool(self):
|
||
ex = _RealExecCase.make_executor(self.tmp.name, self.tmp.name)
|
||
ex.shutdown_vm = self.fake.shutdown_vm
|
||
ex._emit = self.fake._emit
|
||
_write(_registry_path(self.tmp.name), json.dumps({
|
||
"a": {"state": "paused", "vm": "chv-resident-a"},
|
||
}))
|
||
with mock.patch(
|
||
"cloud.chv.resident_pool.ResidentPool",
|
||
side_effect=AssertionError("must not construct")):
|
||
cleaned = worker_adapter.reconcile_resident_pool(
|
||
ex, self._audit, "shutdown")
|
||
self.assertEqual(cleaned, ["chv-resident-a"])
|
||
|
||
def test_facade_ignores_non_chv_executor(self):
|
||
cleaned = worker_adapter.reconcile_resident_pool(
|
||
object(), self._audit, "startup")
|
||
self.assertEqual(cleaned, [])
|
||
self.assertEqual(self.audits, [])
|
||
|
||
def test_facade_does_not_touch_warm_pool_reconcile(self):
|
||
# reconcile_warm_pool 仍走 warm-pool.json，与常驻无关。
|
||
ex = _RealExecCase.make_executor(self.tmp.name, self.tmp.name)
|
||
ex.shutdown_vm = self.fake.shutdown_vm
|
||
ex._emit = self.fake._emit
|
||
warm_path = os.path.join(self.tmp.name, "warm-pool.json")
|
||
_write(warm_path, json.dumps({
|
||
"w": {"state": "warm", "vm": "chv-warm-1",
|
||
"task_id": "w"},
|
||
}))
|
||
cleaned = worker_adapter.reconcile_warm_pool(
|
||
ex, self._audit, "startup")
|
||
self.assertEqual(cleaned, ["chv-warm-1"])
|
||
self.assertFalse(os.path.exists(_registry_path(self.tmp.name)))
|
||
|
||
|
||
def _write_file(path, content=b"x"):
|
||
with open(path, "wb") as fh:
|
||
fh.write(content)
|
||
os.chmod(path, 0o755)
|
||
return path
|
||
|
||
|
||
class _RealExecCase(unittest.TestCase):
|
||
"""真执行器 + 全方法替身：测 _run_with_vm 的接线本身，
|
||
不测 VMM。"""
|
||
|
||
@staticmethod
|
||
def make_executor(state_dir, base_dir, env_extra=None):
|
||
env = {"CHV_EGRESS_PROXY": "http://10.0.0.1:8888",
|
||
"CHV_HOST_KEY_MODE": "keyscan",
|
||
"CHV_SSH_FORWARD": "1"}
|
||
if env_extra:
|
||
env.update(env_extra)
|
||
cpuinfo = os.path.join(base_dir, "cpuinfo")
|
||
with open(cpuinfo, "w") as fh:
|
||
fh.write("flags : fpu vmx sse\n")
|
||
ex = CloudHypervisorExecutor(
|
||
binary=_write_file(os.path.join(base_dir,
|
||
"cloud-hypervisor")),
|
||
ch_remote=_write_file(os.path.join(base_dir, "ch-remote")),
|
||
kernel=_write_file(os.path.join(base_dir, "vmlinux")),
|
||
rootfs=_write_file(os.path.join(base_dir, "disk.raw"),
|
||
b"R" * 100),
|
||
ssh_key=_write_file(os.path.join(base_dir, "id_chv")),
|
||
kvm_path=_write_file(os.path.join(base_dir, "kvm")),
|
||
cpuinfo_path=cpuinfo,
|
||
state_dir=state_dir,
|
||
sleeper=lambda s: None, env=env)
|
||
ex.audit_callback = lambda e, t, d, k=None: None
|
||
return ex
|
||
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.state_dir = os.path.join(self.tmp.name, "state")
|
||
self.calls = []
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _stubbed_exec(self, env_extra=None, wake_fail_resident=False):
|
||
ex = self.make_executor(self.state_dir, self.tmp.name, env_extra)
|
||
calls = self.calls
|
||
|
||
def wake(task, extra_env=None):
|
||
calls.append(("wake", ex.vm_name(task)))
|
||
if wake_fail_resident and ex.vm_name(task).startswith(
|
||
"chv-resident-"):
|
||
raise RuntimeError("provision boom")
|
||
return {"vm": ex.vm_name(task), "mode": "cold"}
|
||
|
||
def pause(task_or_name):
|
||
calls.append(("pause", ex.vm_name(task_or_name)))
|
||
return True
|
||
|
||
def resume(task_or_name):
|
||
calls.append(("resume", ex.vm_name(task_or_name)))
|
||
return True
|
||
|
||
def wait(task):
|
||
calls.append(("wait", ex.vm_name(task)))
|
||
return True
|
||
|
||
def shutdown(task_or_name):
|
||
calls.append(("shutdown", ex.vm_name(task_or_name)))
|
||
return True
|
||
|
||
def run_guest(task, inner_args, extra_env=None, **kwargs):
|
||
calls.append(("guest", task["id"], list(inner_args)))
|
||
return {"task_id": task["id"], "status": "completed",
|
||
"stdout": "ok", "vm": ex.vm_name(task)}
|
||
|
||
def run_guest_goal(task, inner_args, extra_env=None, **kwargs):
|
||
calls.append(("guest-goal", task["id"], list(inner_args)))
|
||
return {"task_id": task["id"], "status": "completed",
|
||
"goal_met": True, "vm": ex.vm_name(task)}
|
||
|
||
ex.wake_vm = wake
|
||
ex.pause_vm = pause
|
||
ex.resume_vm = resume
|
||
ex.wait_guest_ready = wait
|
||
ex.shutdown_vm = shutdown
|
||
ex._run_guest = run_guest
|
||
ex._run_guest_goal = run_guest_goal
|
||
return ex
|
||
|
||
|
||
class RunWithVmResidentTest(_RealExecCase):
|
||
def test_first_task_provisions_and_second_reuses(self):
|
||
ex = self._stubbed_exec()
|
||
task = {"id": "task-1", "prompt": "hi"}
|
||
result = ex._run_with_vm(task, ["dispatch", "--task", "hi"])
|
||
self.assertEqual(result["status"], "completed")
|
||
# 常驻伪任务跑命令，但 task_id 改回真实任务。
|
||
self.assertEqual(result["task_id"], "task-1")
|
||
self.assertEqual(result["vm"], "chv-resident-default")
|
||
kinds = [c[0] for c in self.calls]
|
||
self.assertEqual(kinds, ["wake", "pause", "resume", "wait",
|
||
"guest", "pause"])
|
||
self.assertEqual(self.calls[0][1], "chv-resident-default")
|
||
self.assertEqual(self.calls[4][1], "resident-default")
|
||
# 第二个任务复用同一台常驻 VM，不再 wake。
|
||
self.calls.clear()
|
||
task2 = {"id": "task-2", "prompt": "hi2"}
|
||
result2 = ex._run_with_vm(task2, ["dispatch", "--task", "hi2"])
|
||
self.assertEqual(result2["task_id"], "task-2")
|
||
self.assertEqual(result2["vm"], "chv-resident-default")
|
||
self.assertEqual([c[0] for c in self.calls],
|
||
["resume", "wait", "guest", "pause"])
|
||
# 落盘条目是 paused，可供下次 acquire。
|
||
with open(_registry_path(self.state_dir)) as fh:
|
||
data = json.load(fh)
|
||
self.assertEqual(data["default"]["state"], "paused")
|
||
self.assertEqual(data["default"]["vm"], "chv-resident-default")
|
||
|
||
def test_goal_path_runs_on_resident_and_rewrites_task_id(self):
|
||
ex = self._stubbed_exec()
|
||
task = {"id": "goal-1", "prompt": "g"}
|
||
result = ex._run_with_vm(task, ["goal", "--task", "g"],
|
||
goal=True)
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertEqual(result["task_id"], "goal-1")
|
||
self.assertTrue(result["goal_met"])
|
||
self.assertIn(("guest-goal", "resident-default",
|
||
["goal", "--task", "g"]), self.calls)
|
||
|
||
def test_dispatch_task_flows_through_resident(self):
|
||
ex = self._stubbed_exec()
|
||
result = ex.dispatch_task({"id": "d-1", "prompt": "hello"})
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertEqual(result["task_id"], "d-1")
|
||
self.assertEqual(result["vm"], "chv-resident-default")
|
||
|
||
def test_per_user_vms_are_isolated(self):
|
||
ex = self._stubbed_exec()
|
||
ex._run_with_vm({"id": "t-a", "user_id": "alice",
|
||
"prompt": "x"}, ["dispatch"])
|
||
ex._run_with_vm({"id": "t-b", "user_id": "bob",
|
||
"prompt": "x"}, ["dispatch"])
|
||
wakes = [c[1] for c in self.calls if c[0] == "wake"]
|
||
self.assertEqual(wakes, ["chv-resident-alice",
|
||
"chv-resident-bob"])
|
||
with open(_registry_path(self.state_dir)) as fh:
|
||
data = json.load(fh)
|
||
self.assertEqual(set(data), {"alice", "bob"})
|
||
|
||
def test_disabled_flag_goes_cold(self):
|
||
ex = self._stubbed_exec(
|
||
env_extra={"CHV_RESIDENT_ENABLED": "0"})
|
||
result = ex._run_with_vm({"id": "task-1", "prompt": "hi"},
|
||
["dispatch"])
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertEqual(result["task_id"], "task-1")
|
||
kinds = [c[0] for c in self.calls]
|
||
self.assertEqual(kinds, ["wake", "guest", "shutdown"])
|
||
self.assertEqual(self.calls[0][1], "chv-task-1")
|
||
self.assertIsNone(getattr(ex, "_resident_pool", None))
|
||
self.assertFalse(os.path.exists(
|
||
_registry_path(self.state_dir)))
|
||
|
||
def test_secret_env_goes_cold(self):
|
||
ex = self._stubbed_exec()
|
||
result = ex._run_with_vm(
|
||
{"id": "task-1", "prompt": "hi"}, ["dispatch"],
|
||
extra_env={"SUB2API_API_KEY": "sekret"})
|
||
self.assertEqual(result["status"], "completed")
|
||
kinds = [c[0] for c in self.calls]
|
||
self.assertEqual(kinds, ["wake", "guest", "shutdown"])
|
||
self.assertFalse(os.path.exists(
|
||
_registry_path(self.state_dir)))
|
||
|
||
def test_provision_failure_falls_back_to_cold(self):
|
||
ex = self._stubbed_exec(wake_fail_resident=True)
|
||
result = ex._run_with_vm({"id": "task-1", "prompt": "hi"},
|
||
["dispatch"])
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertEqual(result["task_id"], "task-1")
|
||
kinds = [c[0] for c in self.calls]
|
||
# 建池 wake 失败（尽力 shutdown）→ 冷启动照常跑任务。
|
||
self.assertEqual(kinds, ["wake", "shutdown", "wake",
|
||
"guest", "shutdown"])
|
||
self.assertEqual(self.calls[1][1], "chv-resident-default")
|
||
self.assertEqual(self.calls[2][1], "chv-task-1")
|
||
|
||
def test_guest_error_still_pauses_resident_in_finally(self):
|
||
ex = self._stubbed_exec()
|
||
|
||
def boom(task, inner_args, extra_env=None, **kwargs):
|
||
raise RuntimeError("guest boom")
|
||
|
||
ex._run_guest = boom
|
||
with self.assertRaises(RuntimeError):
|
||
ex._run_with_vm({"id": "task-1", "prompt": "hi"},
|
||
["dispatch"])
|
||
kinds = [c[0] for c in self.calls]
|
||
self.assertEqual(kinds, ["wake", "pause", "resume", "wait",
|
||
"pause"])
|
||
self.assertEqual(self.calls[-1][1], "chv-resident-default")
|
||
|
||
def test_guest_error_still_shuts_down_cold_in_finally(self):
|
||
ex = self._stubbed_exec(
|
||
env_extra={"CHV_RESIDENT_ENABLED": "0"})
|
||
|
||
def boom(task, inner_args, extra_env=None, **kwargs):
|
||
raise RuntimeError("guest boom")
|
||
|
||
ex._run_guest = boom
|
||
with self.assertRaises(RuntimeError):
|
||
ex._run_with_vm({"id": "task-1", "prompt": "hi"},
|
||
["dispatch"])
|
||
self.assertEqual([c[0] for c in self.calls],
|
||
["wake", "shutdown"])
|
||
|
||
def test_wake_failure_propagates_without_shutdown(self):
|
||
ex = self._stubbed_exec(
|
||
env_extra={"CHV_RESIDENT_ENABLED": "0"})
|
||
|
||
def wake_fail(task, extra_env=None):
|
||
raise ChvFailure("vm_create_failed", "no vm")
|
||
|
||
ex.wake_vm = wake_fail
|
||
with self.assertRaises(ChvFailure):
|
||
ex._run_with_vm({"id": "task-1", "prompt": "hi"},
|
||
["dispatch"])
|
||
self.assertEqual(self.calls, [])
|
||
|
||
|
||
class WarmPoolUntouchedTest(unittest.TestCase):
|
||
def test_wired_to_dispatch_still_false(self):
|
||
self.assertFalse(WarmPool.wired_to_dispatch)
|
||
|
||
def test_executor_source_never_imports_experimental(self):
|
||
import ast
|
||
from pathlib import Path
|
||
tree = ast.parse((Path(__file__).resolve().parents[2]
|
||
/ "cloud" / "chv" / "executor.py").read_text())
|
||
for node in ast.walk(tree):
|
||
if isinstance(node, ast.Import):
|
||
for alias in node.names:
|
||
self.assertNotIn(alias.name, (
|
||
"cloud.chv.pool",
|
||
"cloud.chv.experimental_pool",
|
||
"cloud.chv.orphans"))
|
||
if isinstance(node, ast.ImportFrom):
|
||
self.assertNotIn(node.module, (
|
||
"cloud.chv.pool",
|
||
"cloud.chv.experimental_pool",
|
||
"cloud.chv.orphans"))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|