归档:Kata / microVM runner 实验(2026-09-17 工作区快照)
从旧工作区 chore/recover-old-workspace 清理时保存,内容与 2026-09-17
stash@{0} 快照中的版本一致;未合并、未在 main 上使用,仅作参考,不开 PR。
被 gitignore 的 tfstate 与凭据文件不在此分支,仍留在本地工作区。
Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
@@ -0,0 +1,138 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Capacity-bounded JetStream consumer that launches one ephemeral VM per job."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import secrets
|
||||
import ssl
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import nats
|
||||
from aiohttp import web
|
||||
from nats.errors import TimeoutError
|
||||
from nats.js.api import AckPolicy, ConsumerConfig
|
||||
|
||||
CAPACITY = int(os.environ.get("RUNNER_CAPACITY", "1"))
|
||||
SUBJECT = os.environ.get("NATS_SUBJECT", "ci.runner.kind-microvm")
|
||||
STREAM = os.environ.get("NATS_STREAM", "CI_RUNNER")
|
||||
DURABLE = os.environ.get("NATS_DURABLE", "kind-microvm")
|
||||
MAX_INFLIGHT = int(os.environ.get("RUNNER_MAX_INFLIGHT", "64"))
|
||||
NATS_URL = os.environ.get("NATS_URL", "tls://nats.ad.ddupan.top:4222")
|
||||
NATS_USER = os.environ.get("NATS_USER", "ci-worker")
|
||||
NATS_PASSWORD_FILE = Path(os.environ.get("NATS_PASSWORD_FILE", "/etc/microvm-runner/nats-password"))
|
||||
NATS_CA_FILE = os.environ.get("NATS_CA_FILE", "/etc/ssl/certs/ca-certificates.crt")
|
||||
REGISTRATION_TOKEN_FILE = Path(os.environ.get("REGISTRATION_TOKEN_FILE", "/etc/microvm-runner/registration-token"))
|
||||
LAUNCHER = os.environ.get("LAUNCHER", "/usr/local/libexec/microvm-runner-launch")
|
||||
TOKEN_LISTEN = os.environ.get("TOKEN_LISTEN", "172.30.0.1")
|
||||
TOKEN_PORT = int(os.environ.get("TOKEN_PORT", "8787"))
|
||||
|
||||
tokens: dict[str, bytes] = {}
|
||||
token_lock = asyncio.Lock()
|
||||
|
||||
|
||||
async def token(request: web.Request) -> web.Response:
|
||||
nonce = request.match_info["nonce"]
|
||||
async with token_lock:
|
||||
value = tokens.pop(nonce, None)
|
||||
if value is None:
|
||||
raise web.HTTPNotFound()
|
||||
return web.Response(body=value, headers={"Cache-Control": "no-store"})
|
||||
|
||||
|
||||
async def heartbeat(message: object, stop: asyncio.Event) -> None:
|
||||
while True:
|
||||
try:
|
||||
await asyncio.wait_for(stop.wait(), timeout=60)
|
||||
return
|
||||
except asyncio.TimeoutError:
|
||||
await message.in_progress()
|
||||
|
||||
|
||||
async def run_one(message: object) -> None:
|
||||
instance_id = str(uuid.uuid4())
|
||||
nonce = secrets.token_urlsafe(32)
|
||||
async with token_lock:
|
||||
tokens[nonce] = REGISTRATION_TOKEN_FILE.read_bytes().strip()
|
||||
stop = asyncio.Event()
|
||||
pulse = asyncio.create_task(heartbeat(message, stop))
|
||||
try:
|
||||
process = await asyncio.create_subprocess_exec(LAUNCHER, instance_id, nonce)
|
||||
return_code = await process.wait()
|
||||
if return_code == 0:
|
||||
await message.ack()
|
||||
else:
|
||||
await message.nak(delay=30)
|
||||
except Exception:
|
||||
await message.nak(delay=30)
|
||||
raise
|
||||
finally:
|
||||
stop.set()
|
||||
await pulse
|
||||
async with token_lock:
|
||||
tokens.pop(nonce, None)
|
||||
|
||||
|
||||
async def consume() -> None:
|
||||
if CAPACITY < 1:
|
||||
raise ValueError("RUNNER_CAPACITY must be at least 1")
|
||||
tls = ssl.create_default_context(cafile=NATS_CA_FILE)
|
||||
nc = await nats.connect(
|
||||
NATS_URL,
|
||||
user=NATS_USER,
|
||||
password=NATS_PASSWORD_FILE.read_text().strip(),
|
||||
tls=tls,
|
||||
name=DURABLE,
|
||||
)
|
||||
js = nc.jetstream()
|
||||
subscription = await js.pull_subscribe(
|
||||
SUBJECT,
|
||||
durable=DURABLE,
|
||||
stream=STREAM,
|
||||
config=ConsumerConfig(
|
||||
durable_name=DURABLE,
|
||||
filter_subject=SUBJECT,
|
||||
ack_policy=AckPolicy.EXPLICIT,
|
||||
ack_wait=5 * 60,
|
||||
# This durable consumer is shared by every host of this runner type.
|
||||
# Local CAPACITY controls each host; this is only a global safety cap.
|
||||
max_ack_pending=MAX_INFLIGHT,
|
||||
max_deliver=5,
|
||||
),
|
||||
)
|
||||
active: set[asyncio.Task[None]] = set()
|
||||
try:
|
||||
while True:
|
||||
active = {task for task in active if not task.done()}
|
||||
free = CAPACITY - len(active)
|
||||
if free == 0:
|
||||
await asyncio.wait(active, return_when=asyncio.FIRST_COMPLETED)
|
||||
continue
|
||||
try:
|
||||
messages = await subscription.fetch(batch=free, timeout=5)
|
||||
except TimeoutError:
|
||||
continue
|
||||
for message in messages:
|
||||
task = asyncio.create_task(run_one(message))
|
||||
task.add_done_callback(lambda done: done.exception())
|
||||
active.add(task)
|
||||
finally:
|
||||
if active:
|
||||
await asyncio.gather(*active, return_exceptions=True)
|
||||
await nc.drain()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
app = web.Application()
|
||||
app.router.add_get("/token/{nonce}", token)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
await web.TCPSite(runner, TOKEN_LISTEN, TOKEN_PORT).start()
|
||||
try:
|
||||
await consume()
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user