156 lines
4.8 KiB
Python
156 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Capacity-bounded JetStream consumer that launches one ephemeral VM per job."""
|
|
|
|
import asyncio
|
|
import logging
|
|
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
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
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:
|
|
LOG.error("launcher for %s exited with %d", instance_id, return_code)
|
|
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)
|
|
|
|
|
|
def _report_task(task: asyncio.Task[None]) -> None:
|
|
if task.cancelled():
|
|
return
|
|
error = task.exception()
|
|
if error is not None:
|
|
LOG.error(
|
|
"runner task failed",
|
|
exc_info=(type(error), error, error.__traceback__),
|
|
)
|
|
|
|
|
|
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,
|
|
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(_report_task)
|
|
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()
|
|
|
|
|
|
def cli() -> None:
|
|
logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))
|
|
asyncio.run(main())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|