308 lines
11 KiB
Python
308 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""JetStream worker that creates one disposable Kubernetes Pod per job."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
import ssl
|
|
from urllib.parse import quote
|
|
import uuid
|
|
|
|
import nats
|
|
from aiohttp import ClientResponseError
|
|
from nats.errors import TimeoutError
|
|
from nats.js.api import AckPolicy, ConsumerConfig
|
|
|
|
from .kubernetes import KubernetesClient
|
|
from .models import IdentityBinding, RunnerRequest
|
|
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
STREAM = os.environ.get("NATS_STREAM", "CI_RUNNER")
|
|
REQUEST_SUBJECT = os.environ.get("NATS_SUBJECT", "ci.runner.pod")
|
|
BINDING_SUBJECT = os.environ.get("NATS_BINDING_SUBJECT", "ci.runner.pod.binding")
|
|
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", "/run/secrets/nats/password"))
|
|
NATS_CA_FILE = os.environ.get("NATS_CA_FILE", "/etc/ssl/certs/ca-certificates.crt")
|
|
KUBERNETES_API = os.environ.get("KUBERNETES_API", "https://kubernetes.default.svc")
|
|
KUBERNETES_TOKEN_FILE = Path(
|
|
os.environ.get(
|
|
"KUBERNETES_TOKEN_FILE",
|
|
"/var/run/secrets/kubernetes.io/serviceaccount/token",
|
|
)
|
|
)
|
|
KUBERNETES_CA_FILE = Path(
|
|
os.environ.get(
|
|
"KUBERNETES_CA_FILE",
|
|
"/var/run/secrets/kubernetes.io/serviceaccount/ca.crt",
|
|
)
|
|
)
|
|
NAMESPACE = os.environ.get("RUNNER_NAMESPACE", "gitea-actions")
|
|
RUNNER_IMAGE = os.environ.get("RUNNER_IMAGE", "docker.io/gitea/runner:2")
|
|
RUNNER_SERVICE_ACCOUNT = os.environ.get("RUNNER_SERVICE_ACCOUNT", "gitea-dynamic-runner")
|
|
RUNNER_TOKEN_SECRET = os.environ.get("RUNNER_TOKEN_SECRET", "gitea-dynamic-runner")
|
|
GITEA_INSTANCE = os.environ.get("GITEA_INSTANCE", "https://git.ddupan.top")
|
|
CAPACITY = int(os.environ.get("RUNNER_CAPACITY", "4"))
|
|
POD_TIMEOUT = int(os.environ.get("RUNNER_POD_TIMEOUT", str(4 * 60 * 60)))
|
|
|
|
|
|
def identity_path(repository: str, job_name: str) -> str:
|
|
parts = repository.split("/")
|
|
if len(parts) != 2 or not all(parts):
|
|
raise ValueError("repository must be owner/name")
|
|
encoded = [quote(part, safe="-._~") for part in (*parts, job_name)]
|
|
if not job_name.strip() or any(not part for part in encoded):
|
|
raise ValueError("identity components must not be empty")
|
|
return "/".join(encoded)
|
|
|
|
|
|
def pod_manifest(request: RunnerRequest, pod_name: str) -> dict[str, object]:
|
|
if request.backend != "pod":
|
|
raise ValueError("Pod backend only accepts pod requests")
|
|
return {
|
|
"apiVersion": "v1",
|
|
"kind": "Pod",
|
|
"metadata": {
|
|
"name": pod_name,
|
|
"namespace": NAMESPACE,
|
|
"labels": {
|
|
"app.kubernetes.io/name": "gitea-dynamic-runner",
|
|
"app.kubernetes.io/component": "runner",
|
|
"ci.ddupan.top/backend": "pod",
|
|
},
|
|
"annotations": {
|
|
"ci.ddupan.top/queued-job-id": str(request.job_id),
|
|
"ci.ddupan.top/queued-run-id": str(request.run_id),
|
|
},
|
|
},
|
|
"spec": {
|
|
"serviceAccountName": RUNNER_SERVICE_ACCOUNT,
|
|
"restartPolicy": "Never",
|
|
"terminationGracePeriodSeconds": 30,
|
|
"containers": [
|
|
{
|
|
"name": "runner",
|
|
"image": RUNNER_IMAGE,
|
|
"imagePullPolicy": "IfNotPresent",
|
|
"securityContext": {"privileged": True},
|
|
"env": [
|
|
{"name": "GITEA_INSTANCE_URL", "value": GITEA_INSTANCE},
|
|
{
|
|
"name": "GITEA_RUNNER_NAME",
|
|
"valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}},
|
|
},
|
|
{
|
|
"name": "GITEA_RUNNER_REGISTRATION_TOKEN_FILE",
|
|
"value": "/run/secrets/gitea/token",
|
|
},
|
|
{"name": "GITEA_RUNNER_LABELS", "value": "self-hosted:host,pod:host"},
|
|
{"name": "GITEA_RUNNER_EPHEMERAL", "value": "1"},
|
|
{"name": "GITEA_RUNNER_ONCE", "value": "1"},
|
|
{"name": "CONFIG_FILE", "value": "/etc/gitea-runner/config.yaml"},
|
|
],
|
|
"volumeMounts": [
|
|
{
|
|
"name": "registration-token",
|
|
"mountPath": "/run/secrets/gitea",
|
|
"readOnly": True,
|
|
},
|
|
{
|
|
"name": "spire-agent-socket",
|
|
"mountPath": "/run/spire/agent-sockets",
|
|
"readOnly": True,
|
|
},
|
|
],
|
|
}
|
|
],
|
|
"volumes": [
|
|
{
|
|
"name": "registration-token",
|
|
"secret": {
|
|
"secretName": RUNNER_TOKEN_SECRET,
|
|
"items": [{"key": "token", "path": "token"}],
|
|
},
|
|
},
|
|
{
|
|
"name": "spire-agent-socket",
|
|
"csi": {"driver": "csi.spiffe.io", "readOnly": True},
|
|
},
|
|
],
|
|
},
|
|
}
|
|
|
|
|
|
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 wait_for_pod(client: KubernetesClient, name: str) -> bool:
|
|
deadline = asyncio.get_running_loop().time() + POD_TIMEOUT
|
|
while asyncio.get_running_loop().time() < deadline:
|
|
pod = await client.get_pod(name)
|
|
if pod is None:
|
|
raise RuntimeError(f"runner Pod {name} disappeared")
|
|
status = pod.get("status")
|
|
phase = status.get("phase") if isinstance(status, dict) else None
|
|
if phase == "Succeeded":
|
|
return True
|
|
if phase == "Failed":
|
|
return False
|
|
await asyncio.sleep(2)
|
|
raise asyncio.TimeoutError(f"runner Pod {name} timed out")
|
|
|
|
|
|
async def run_request(message: object, client: KubernetesClient) -> None:
|
|
try:
|
|
request = RunnerRequest.from_json(message.data)
|
|
except (json.JSONDecodeError, UnicodeDecodeError, ValueError) as error:
|
|
LOG.error("discarding invalid Pod request: %s", error)
|
|
await message.ack()
|
|
return
|
|
if request.backend != "pod":
|
|
LOG.error("discarding %s request received by Pod worker", request.backend)
|
|
await message.ack()
|
|
return
|
|
|
|
pod_name = f"gitea-pod-{uuid.uuid4().hex[:12]}"
|
|
stop = asyncio.Event()
|
|
pulse = asyncio.create_task(heartbeat(message, stop))
|
|
try:
|
|
await client.create_pod(pod_manifest(request, pod_name))
|
|
if await wait_for_pod(client, pod_name):
|
|
await message.ack()
|
|
else:
|
|
LOG.error("runner Pod %s failed", pod_name)
|
|
await message.nak(delay=30)
|
|
except Exception:
|
|
await message.nak(delay=30)
|
|
raise
|
|
finally:
|
|
stop.set()
|
|
await pulse
|
|
await client.delete_pod(pod_name)
|
|
|
|
|
|
async def bind_request(message: object, client: KubernetesClient) -> None:
|
|
try:
|
|
document = json.loads(message.data)
|
|
binding = IdentityBinding(**document)
|
|
if binding.backend != "pod" or not binding.runner_name.startswith("gitea-pod-"):
|
|
raise ValueError("invalid Pod identity binding")
|
|
path = identity_path(binding.repository, binding.job_name)
|
|
await client.bind_identity(binding.runner_name, path)
|
|
except ClientResponseError as error:
|
|
if error.status == 404:
|
|
await message.nak(delay=2)
|
|
return
|
|
await message.nak(delay=30)
|
|
raise
|
|
except (json.JSONDecodeError, TypeError, ValueError) as error:
|
|
LOG.error("discarding invalid identity binding: %s", error)
|
|
await message.ack()
|
|
return
|
|
await message.ack()
|
|
|
|
|
|
async def consume_requests(subscription: object, client: KubernetesClient) -> None:
|
|
active: set[asyncio.Task[None]] = set()
|
|
while True:
|
|
active = {task for task in active if not task.done()}
|
|
free = CAPACITY - len(active)
|
|
if free < 1:
|
|
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_request(message, client))
|
|
task.add_done_callback(_report_task)
|
|
active.add(task)
|
|
|
|
|
|
async def consume_bindings(subscription: object, client: KubernetesClient) -> None:
|
|
while True:
|
|
try:
|
|
messages = await subscription.fetch(batch=16, timeout=5)
|
|
except TimeoutError:
|
|
continue
|
|
await asyncio.gather(*(bind_request(message, client) for message in messages))
|
|
|
|
|
|
def _report_task(task: asyncio.Task[None]) -> None:
|
|
if not task.cancelled() and (error := task.exception()) is not None:
|
|
LOG.error("Pod runner task failed", exc_info=(type(error), error, error.__traceback__))
|
|
|
|
|
|
async def main() -> None:
|
|
if CAPACITY < 1:
|
|
raise ValueError("RUNNER_CAPACITY must be at least 1")
|
|
context = 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=context,
|
|
name="pod-runner-worker",
|
|
)
|
|
js = nc.jetstream()
|
|
requests = await js.pull_subscribe(
|
|
REQUEST_SUBJECT,
|
|
durable="pod",
|
|
stream=STREAM,
|
|
config=ConsumerConfig(
|
|
durable_name="pod",
|
|
filter_subject=REQUEST_SUBJECT,
|
|
ack_policy=AckPolicy.EXPLICIT,
|
|
ack_wait=5 * 60,
|
|
max_ack_pending=max(CAPACITY, 1),
|
|
max_deliver=5,
|
|
),
|
|
)
|
|
bindings = await js.pull_subscribe(
|
|
BINDING_SUBJECT,
|
|
durable="pod-binding",
|
|
stream=STREAM,
|
|
config=ConsumerConfig(
|
|
durable_name="pod-binding",
|
|
filter_subject=BINDING_SUBJECT,
|
|
ack_policy=AckPolicy.EXPLICIT,
|
|
ack_wait=30,
|
|
max_ack_pending=64,
|
|
max_deliver=10,
|
|
),
|
|
)
|
|
async with KubernetesClient(
|
|
api_url=KUBERNETES_API,
|
|
token_file=KUBERNETES_TOKEN_FILE,
|
|
ca_file=KUBERNETES_CA_FILE,
|
|
namespace=NAMESPACE,
|
|
) as client:
|
|
try:
|
|
await asyncio.gather(
|
|
consume_requests(requests, client),
|
|
consume_bindings(bindings, client),
|
|
)
|
|
finally:
|
|
await nc.drain()
|
|
|
|
|
|
def cli() -> None:
|
|
logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))
|
|
asyncio.run(main())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|