vm-bench/backend/provisioner.py

465 lines
18 KiB
Python

"""
VM provisioning on the Proxmox host — real qm integration.
Runs qemu-img convert, virt-resize (optional), qm create / importdisk / set.
Jobs run in background threads with status tracked in memory.
"""
from __future__ import annotations
import json
import logging
import os
import re
import shutil
import subprocess
import threading
import time
from pathlib import Path
from typing import Optional
from models import (
JobSubmissionRequest,
JobStatusResponse,
CloneRequest,
JobStatus,
DiskType,
BootType,
)
from converter import extract_if_needed, discover_disk, detect_efi
logger = logging.getLogger("backend.provisioner")
STAGING_ROOT = Path(__file__).resolve().parent.parent / "tmp"
STAGING_IN = STAGING_ROOT / "in"
STAGING_OUT = STAGING_ROOT / "out"
def _staging_in(session_id: str = "") -> Path:
base = STAGING_ROOT / session_id if session_id else STAGING_ROOT
return base / "in"
def _staging_out(session_id: str = "") -> Path:
base = STAGING_ROOT / session_id if session_id else STAGING_ROOT
return base / "out"
# In-memory job tracking — survives as long as the process runs
_jobs: dict[str, dict] = {}
_jobs_lock = threading.Lock()
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
def submit_job(req: JobSubmissionRequest) -> JobStatusResponse:
"""Queue a conversion job and start it in a background thread."""
job_id = f"job_{req.vmid}_{int(time.time())}"
with _jobs_lock:
_jobs[job_id] = {
"vmid": req.vmid,
"status": JobStatus.QUEUED,
"progress": 0,
"message": "Queued",
"error": None,
"request": req,
"source_files": [],
"disk_paths": [],
}
thread = threading.Thread(target=_process_job, args=(job_id, req), daemon=True)
thread.start()
logger.info("Job %s queued for VM %d — thread started", job_id, req.vmid)
return JobStatusResponse(
job_id=job_id,
vmid=req.vmid,
status=JobStatus.QUEUED,
progress_percentage=0,
message="Job accepted and queued.",
)
def get_job_status(job_id: str) -> JobStatusResponse:
"""Return current status of a job from the in-memory store."""
with _jobs_lock:
job = _jobs.get(job_id)
if not job:
raise KeyError(job_id)
return JobStatusResponse(
job_id=job_id,
vmid=job["vmid"],
status=job["status"],
progress_percentage=job["progress"],
message=job.get("message", ""),
error_details=job.get("error"),
)
def clone_vm(req: CloneRequest) -> JobStatusResponse:
"""qm clone — instant full clone of an existing VM."""
safe_name = re.sub(r'[^a-zA-Z0-9-]', '-', req.target_name).strip('-').lower()
if not safe_name:
safe_name = f"vm-{req.target_vmid}"
job_id = f"clone_{req.source_vmid}_{req.target_vmid}_{int(time.time())}"
logger.info("Cloning VM %d%d (%s)", req.source_vmid, req.target_vmid, safe_name)
try:
_run(["qm", "clone", str(req.source_vmid), str(req.target_vmid),
"--name", safe_name, "--full"], timeout=600)
except Exception as exc:
logger.error("Clone failed: %s", exc)
return JobStatusResponse(
job_id=job_id, vmid=req.target_vmid,
status=JobStatus.FAILED, progress_percentage=0,
message="Clone failed", error_details=str(exc),
)
logger.info("Clone complete: %d%d", req.source_vmid, req.target_vmid)
return JobStatusResponse(
job_id=job_id, vmid=req.target_vmid,
status=JobStatus.COMPLETED, progress_percentage=100,
message=f"VM {req.target_vmid} cloned from {req.source_vmid}",
)
def cleanup_staging(vmid: int, delete: bool, session_id: str = "") -> dict:
"""Remove staging files. If delete=True and session_id is set,
removes the entire session directory (both in/ and out/)."""
if not delete:
return {
"job_id": f"job_{vmid}",
"action_taken": "retained",
"message": f"Staging files preserved for VM {vmid}",
}
# Delete the entire session tmp directory if we have a session ID
if session_id:
session_dir = STAGING_ROOT / session_id
if session_dir.exists():
shutil.rmtree(session_dir, ignore_errors=True)
logger.info("Cleaned up session dir: %s", session_dir)
return {
"job_id": f"job_{vmid}",
"action_taken": "purged",
"message": "Session files deleted.",
}
# Fallback: delete just the VM output dir
out_dir = _staging_out(session_id) / str(vmid)
if out_dir.exists():
shutil.rmtree(out_dir, ignore_errors=True)
logger.info("Cleaned up output dir: %s", out_dir)
return {
"job_id": f"job_{vmid}",
"action_taken": "purged",
"message": f"Staging files deleted for VM {vmid}",
}
# ---------------------------------------------------------------------------
# Background job processor
# ---------------------------------------------------------------------------
def _update_job(job_id: str, *, status: Optional[JobStatus] = None,
progress: Optional[int] = None, message: Optional[str] = None,
error: Optional[str] = None):
with _jobs_lock:
j = _jobs.get(job_id)
if not j:
return
if status is not None:
j["status"] = status
if progress is not None:
j["progress"] = progress
if message is not None:
j["message"] = message
if error is not None:
j["error"] = error
def _run(cmd: list[str], check: bool = True, timeout: int = 3600) -> subprocess.CompletedProcess:
logger.debug("CMD: %s", " ".join(cmd))
result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
if result.stdout:
logger.debug(" stdout: %s", result.stdout.strip()[:500])
if result.stderr:
logger.debug(" stderr: %s", result.stderr.strip()[:500])
if check and result.returncode != 0:
raise RuntimeError(
f"Command failed (rc={result.returncode}): {' '.join(cmd)}\n"
f"stderr: {result.stderr.strip()[:500]}"
)
return result
def _convert_with_progress(
job_id: str, source_path: Path, output_path: Path, format: str,
start_pct: int, end_pct: int,
):
"""Run qemu-img convert in background and update job progress
by polling the output file size against the source virtual size."""
# Get source virtual size for percentage estimation
info = _run(["qemu-img", "info", "--output=json", str(source_path)])
source_size = json.loads(info.stdout).get("virtual-size", 0)
if source_size == 0:
source_size = source_path.stat().st_size
proc = subprocess.Popen(
["qemu-img", "convert", "-f", format, "-O", "qcow2",
str(source_path), str(output_path)],
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
)
# Poll output file size and update progress
while proc.poll() is None:
time.sleep(2)
if output_path.exists():
current = output_path.stat().st_size
if source_size > 0:
ratio = min(current / source_size, 1.0)
pct = start_pct + int(ratio * (end_pct - start_pct))
gb_done = current / (1024**3)
gb_total = source_size / (1024**3)
_update_job(job_id, progress=pct,
message=f"Converting disk... {gb_done:.1f} / {gb_total:.1f} GiB")
if proc.returncode != 0:
raise RuntimeError(f"qemu-img convert failed (rc={proc.returncode})")
_update_job(job_id, progress=end_pct, message="Disk converted.")
def _import_ovf(job_id: str, req: JobSubmissionRequest):
"""Import an OVA/OVF appliance using qm importovf."""
vmid = req.vmid
ovf_path = req.ovf_path
logger.info("OVF import: vmid=%d ovf=%s storage=%s", vmid, ovf_path, req.target_storage)
_update_job(job_id, status=JobStatus.PROCESSING_CONVERSION, progress=20,
message="Importing OVF appliance...")
safe_name = re.sub(r'[^a-zA-Z0-9-]', '-', req.vm_name).strip('-').lower()
if not safe_name:
safe_name = f"vm-{vmid}"
existing = subprocess.run(["qm", "status", str(vmid)], capture_output=True, text=True)
if "does not exist" not in existing.stderr and "does not exist" not in existing.stdout:
logger.warning("VM %d already exists — destroying", vmid)
subprocess.run(["qm", "stop", str(vmid)], capture_output=True)
subprocess.run(["qm", "destroy", str(vmid), "--purge"], capture_output=True)
time.sleep(2)
_run(["qm", "importovf", str(vmid), ovf_path, req.target_storage, "--format", "qcow2"],
timeout=3600)
_update_job(job_id, status=JobStatus.IMPORTING_STORAGE, progress=80,
message="Applying configuration...")
if req.cpu_cores and req.cpu_cores != 4:
_run(["qm", "set", str(vmid), "--cores", str(req.cpu_cores)])
if req.ram_mb and req.ram_mb != 8192:
_run(["qm", "set", str(vmid), "--memory", str(req.ram_mb)])
if safe_name:
_run(["qm", "set", str(vmid), "--name", safe_name])
final_boot = req.boot_type
if req.auto_detect_boot:
ovf_dir = Path(ovf_path).parent
efi_disk = None
for f in ovf_dir.iterdir():
if f.is_file() and f.suffix.lower() in {".vmdk", ".qcow2", ".img", ".raw", ".vhd", ".vhdx", ".vdi"}:
efi_disk = f
break
if efi_disk:
efi = detect_efi(efi_disk)
if efi is True:
final_boot = BootType.UEFI
elif efi is False:
final_boot = BootType.LEGACY
bios = "ovmf" if final_boot.value == "uefi" else "seabios"
_run(["qm", "set", str(vmid), "--bios", bios])
has_net = subprocess.run(["qm", "config", str(vmid)], capture_output=True, text=True)
if "net0" not in has_net.stdout:
_run(["qm", "set", str(vmid), "--net0", "virtio,bridge=vmbr0"])
_run(["qm", "set", str(vmid), "--serial0", "socket"])
config = _run(["qm", "config", str(vmid)], check=False)
logger.info("OVF import complete:\n%s", config.stdout.strip())
_update_job(job_id, status=JobStatus.COMPLETED, progress=100,
message=f"VM {vmid} imported and ready.")
def _process_job(job_id: str, req: JobSubmissionRequest):
try:
vmid = req.vmid
logger.info("=== Job %s started for VM %d ===", job_id, vmid)
if req.ovf_path:
_import_ovf(job_id, req)
return
_update_job(job_id, status=JobStatus.PROCESSING_CONVERSION, progress=10,
message="Converting disk image...")
_update_job(job_id, status=JobStatus.PROCESSING_CONVERSION, progress=5,
message="Preparing source files...")
out_dir = _staging_out(req.session_id) / str(vmid)
out_dir.mkdir(parents=True, exist_ok=True)
# Collect all disks (boot disk + additional)
disk_specs = [req.boot_disk] + (req.additional_disks or [])
disk_paths = []
for idx, spec in enumerate(disk_specs):
output_path = out_dir / f"vm-{vmid}-disk-{idx}.qcow2"
if spec.disk_type == DiskType.IMAGE_FILE:
# Extract archive if needed, then locate the actual disk image
src_path = _staging_in(req.session_id) / spec.source_filename
source = extract_if_needed(str(src_path))
disk = discover_disk(source)
source_path = disk.path
logger.info("Disk %d: source=%s format=%s size=%.1f GiB",
idx, source_path, disk.format, disk.size_gb)
# Delete archive after extraction to free space
if src_path.exists() and src_path != source and src_path.is_file():
try:
src_path.unlink()
logger.info("Deleted archive: %s", src_path)
except Exception:
pass
# Import directly — qm importdisk handles conversion natively
# Only convert if shrink was explicitly requested for this disk
need_shrink = getattr(spec, 'target_disk_size_gb', None)
if need_shrink and need_shrink < disk.size_gb:
_convert_with_progress(
job_id, source_path, output_path, disk.format,
start_pct=10 + idx * 10, end_pct=35 + idx * 10,
)
else:
output_path = source_path # skip convert, point to original file
else:
size_gb = spec.size_gb or 32
logger.info("Creating empty disk %d: %d GiB", idx, size_gb)
_run(["qemu-img", "create", "-f", "qcow2",
str(output_path), f"{size_gb}G"])
_update_job(job_id, progress=20 + (idx + 1) * 15,
message=f"Created empty disk {idx + 1}")
disk_paths.append(output_path)
# ── Detach any existing VM with this ID ─────────────────────
existing = subprocess.run(["qm", "status", str(vmid)],
capture_output=True, text=True)
if "does not exist" not in existing.stderr and "does not exist" not in existing.stdout:
logger.warning("VM %d already exists — destroying it first", vmid)
subprocess.run(["qm", "stop", str(vmid)], capture_output=True)
subprocess.run(["qm", "destroy", str(vmid), "--purge"], capture_output=True)
time.sleep(2)
# ── Create VM ──────────────────────────────────────────────
_update_job(job_id, status=JobStatus.IMPORTING_STORAGE, progress=70,
message="Creating VM...")
# Sanitize VM name: Proxmox requires DNS-safe names (a-z, 0-9, hyphens)
safe_name = re.sub(r'[^a-zA-Z0-9-]', '-', req.vm_name).strip('-').lower()
if not safe_name:
safe_name = f"vm-{vmid}"
# ── Auto-detect boot type if requested ─────────────────────
final_boot = req.boot_type
if req.auto_detect_boot and disk_paths:
logger.info("Auto-detecting boot type...")
efi = detect_efi(disk_paths[0])
if efi is True:
final_boot = BootType.UEFI
logger.info(" EFI partition detected → UEFI (OVMF)")
elif efi is False:
final_boot = BootType.LEGACY
logger.info(" No EFI partition → Legacy BIOS (SeaBIOS)")
else:
logger.info(" EFI detection inconclusive — using %s", final_boot.value)
bios = "ovmf" if final_boot.value == "uefi" else "seabios"
_run(["qm", "create", str(vmid),
"--name", safe_name,
"--bios", bios,
"--cores", str(req.cpu_cores),
"--memory", str(req.ram_mb),
"--net0", "virtio,bridge=vmbr0",
"--scsihw", "virtio-scsi-single"])
# ── Import and attach each disk ────────────────────────────
for idx, disk_path in enumerate(disk_paths):
logger.info("Importing disk %d: %s", idx, disk_path)
result = _run(["qm", "disk", "import", str(vmid), str(disk_path),
req.target_storage])
output = result.stdout + result.stderr
# Parse volume ID from qm importdisk output
match = re.search(r"imported disk '([^']+)'", output)
if not match:
match = re.search(r"Formatting '([^']+)'", output)
if not match:
match = re.search(r"Successfully imported disk as '([^']+)'", output)
if not match:
raise RuntimeError(
f"Could not parse volume ID from qm importdisk output:\n{output}"
)
vol_id = match.group(1)
logger.info(" Volume ID: %s", vol_id)
scsi_id = f"scsi{idx}"
_run(["qm", "set", str(vmid), f"--{scsi_id}", vol_id])
_update_job(job_id, progress=75 + idx * 10,
message=f"Imported disk {idx + 1}/{len(disk_paths)}")
# Delete local QCOW2 after import to free space
if disk_path.exists():
try:
disk_path.unlink()
logger.info("Deleted local QCOW2: %s", disk_path)
except Exception:
pass
# ── Add EFI disk if UEFI ──────────────────────────────────
if final_boot.value == "uefi":
logger.info("Adding EFI disk")
_run(["qm", "set", str(vmid), "--efidisk0",
f"{req.target_storage}:0,format=raw,size=4M"])
# ── Finalise ───────────────────────────────────────────────
_update_job(job_id, progress=95, message="Configuring VM...")
_run(["qm", "set", str(vmid), "--boot", "order=scsi0"])
_run(["qm", "set", str(vmid), "--serial0", "socket"])
# Log final config for debugging
config = _run(["qm", "config", str(vmid)], check=False)
logger.info("Final VM config:\n%s", config.stdout.strip())
_update_job(job_id, status=JobStatus.COMPLETED, progress=100,
message=f"VM {vmid} created and ready.")
with _jobs_lock:
j = _jobs.get(job_id)
if j:
j["disk_paths"] = disk_paths
j["source_files"] = [spec.source_filename for spec in disk_specs
if spec.source_filename]
logger.info("=== Job %s completed: VM %d ===", job_id, vmid)
except Exception as exc:
logger.exception("Job %s failed", job_id)
_update_job(job_id, status=JobStatus.FAILED, progress=0,
message="Job failed", error=str(exc))