ProxMenux/oci/remote/oci_nvidia_refresh.py

#!/usr/bin/env python3
"""Explicit stopped-CT NVIDIA refresh. Never changes the desired deployment."""
from __future__ import annotations

import argparse
import copy
import os
from pathlib import Path
import subprocess

import oci_instances as instances
import oci_nvidia_runtime as nv
from oci_ui import translate, msg_info, msg_ok, msg_error


def destination(root, name):
    if not name.startswith('/') or '..' in Path(name).parts:
        raise ValueError(translate('Invalid NVIDIA destination'))
    parent = (root / name.lstrip('/')).parent.resolve()
    if parent != root and root not in parent.parents:
        raise ValueError(translate('The NVIDIA destination escapes the rootfs'))
    return parent / Path(name).name


def prepare(root, plan):
    root = root.resolve()
    # Validate every destination before making any rootfs change.
    files = {name: destination(root, name) for name in plan['inventory']['files']}
    links = {name: destination(root, name) for name in plan['links']}
    for path in list(files.values()) + list(links.values()):
        if path.exists() and not path.is_file() and not path.is_symlink():
            raise ValueError(translate('The NVIDIA destination cannot be replaced'))
    for path in files.values():
        create_parent(path.parent, root)
        if path.is_symlink():
            path.unlink()
        if not path.exists():
            fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o644)
            os.close(fd)
            owner = path.parent.stat()
            os.chown(path, owner.st_uid, owner.st_gid)
    for name, path in links.items():
        create_parent(path.parent, root)
        if path.exists() or path.is_symlink():
            path.unlink()
        path.symlink_to(plan['links'][name])
    # Native LXC cannot mount onto an alias in a usr-merged image.
    lines = []
    for line in plan['config'].decode().splitlines():
        if line.startswith('lxc.mount.entry: '):
            fields = line.split(': ', 1)[1].split()
            fields[1] = str(files['/' + fields[1]].relative_to(root))
            line = 'lxc.mount.entry: ' + ' '.join(fields)
        lines.append(line)
    return ('\n'.join(lines) + '\n').encode()


def create_parent(path, root):
    if path.exists():
        if not path.is_dir():
            raise ValueError(translate('The parent of an NVIDIA destination is not a directory'))
        return
    if path == root:
        raise ValueError(translate('The rootfs is not mounted'))
    create_parent(path.parent, root)
    path.mkdir(mode=0o755)
    owner = path.parent.stat()
    os.chown(path, owner.st_uid, owner.st_gid)


def rebuild(root, vmid, same_gpu=False):
    """Point a stopped container at the NVIDIA driver of this host. The caller
    holds the registry and the container is not started: it is the step a
    restore on a host with another driver needs before the first start.
    Returns whether anything had to change."""
    record = instances.read(root, vmid)
    config = instances.command('pct', 'config', str(vmid))
    if not instances.same_config_except_notes(record, config):
        raise ValueError(translate('The container identity or configuration changed'))
    plan = nv.refresh_plan(config, record['observed']['gpu_devices'][nv.KEY], same_gpu=same_gpu)
    if not plan['changed']:
        return False
    if instances.command('pct', 'status', str(vmid)).strip() != b'status: stopped':
        raise ValueError(translate('Stop the container before the NVIDIA refresh'))
    instances.command('pct', 'mount', str(vmid))
    try:
        candidate = prepare(Path(f'/var/lib/lxc/{vmid}/rootfs'), plan)
    finally:
        instances.command('pct', 'unmount', str(vmid))
    Path(f'/etc/pve/lxc/{vmid}.conf').write_bytes(candidate)
    updated = copy.deepcopy(record)
    updated['observed'] = instances.observe(vmid, record['installation_id'],
        record['observed']['archive_path'], record['observed']['resolved_registry_digest'],
        record['observed']['image'])
    nv.check_mounts(updated['observed']['config'].encode(), plan['inventory'])
    nv.check_devices(updated['observed']['config'].encode(), plan['inventory'])
    instances.write(instances.location(root, vmid), updated)
    return True


def refresh(root, vmid, apply=False):
    with instances.locked(root):
        record = instances.read(root, vmid)
        if record['status'] != 'installed' or record.get('pending_transaction'):
            raise ValueError(translate('The instance has a pending operation'))
        if not nv.enabled(record['deployment']):
            raise ValueError(translate('The instance does not use NVIDIA'))
        config = instances.command('pct', 'config', str(vmid))
        if not instances.same_config_except_notes(record, config):
            raise ValueError(translate('The container identity or configuration changed'))
        previous = record['observed']['gpu_devices'][nv.KEY]
        plan = nv.refresh_plan(config, previous)
        journal = instances.location(root, vmid).parent / 'nvidia-refresh.json'
        if journal.exists() or journal.is_symlink():
            raise ValueError(translate('A previous NVIDIA refresh is pending review'))
        if not apply:
            msg_ok(translate('Current NVIDIA inventory resolved: a refresh is required') if plan['changed']
                   else translate('Current NVIDIA inventory resolved: no refresh is required'))
            return
        if not plan['changed']:
            msg_ok(translate('The NVIDIA runtime is up to date; the container is not modified or started'))
            return
        if instances.command('pct', 'status', str(vmid)).strip() != b'status: stopped':
            raise ValueError(translate('Stop the container before the NVIDIA refresh'))
        instances.write(journal, {'phase': 'preparing', 'record': record,
                                  'inventory': plan['inventory']})
        msg_info(translate('Refreshing the NVIDIA runtime...'))
        instances.command('pct', 'mount', str(vmid))
        try:
            candidate = prepare(Path(f'/var/lib/lxc/{vmid}/rootfs'), plan)
        finally:
            instances.command('pct', 'unmount', str(vmid))
        nv.verify(plan['inventory'])
        if instances.command('pct', 'config', str(vmid)) != config:
            raise ValueError(translate('The configuration changed during the NVIDIA refresh'))
        conf = Path(f'/etc/pve/lxc/{vmid}.conf')
        # Same native configuration file used by the common installer.
        conf.write_bytes(candidate)
        instances.write(journal, {'phase': 'validating', 'record': record,
                                  'inventory': plan['inventory']})
        instances.command('pct', 'start', str(vmid))
        try:
            nv.validate_runtime(vmid, plan['inventory'])
        finally:
            instances.command('pct', 'shutdown', str(vmid), '--timeout', '30')
        updated = copy.deepcopy(record)
        updated['observed'] = instances.observe(vmid, record['installation_id'],
            record['observed']['archive_path'], record['observed']['resolved_registry_digest'],
            record['observed']['image'])
        nv.check_mounts(updated['observed']['config'].encode(), plan['inventory'])
        nv.check_devices(updated['observed']['config'].encode(), plan['inventory'])
        if updated['observed']['gpu_devices'][nv.KEY] != plan['inventory']:
            raise ValueError(translate('The observed inventory differs from the validated runtime'))
        nv.verify(plan['inventory'])
        instances.write(instances.location(root, vmid), updated)
        journal.unlink()
        msg_ok(translate('NVIDIA refresh validated; the container is stopped and its settings are kept'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('vmid', type=int)
    parser.add_argument('--root', type=Path, default=instances.ROOT)
    parser.add_argument('--apply', action='store_true')
    args = parser.parse_args()
    try:
        refresh(args.root, args.vmid, args.apply)
    except BlockingIOError:
        msg_error(translate('Another OCI operation is using the registry. This operation was not started.'))
        raise SystemExit(1)
    except (OSError, ValueError, KeyError, RuntimeError, subprocess.SubprocessError) as error:
        msg_error(str(error) if not isinstance(error, KeyError) else
                  translate('The saved OCI record is incomplete or has an unexpected format.'))
        raise SystemExit(1)