130 lines
3.5 KiB
Python
Executable File
130 lines
3.5 KiB
Python
Executable File
#!/usr/local/sbin/charm-env python3
|
|
|
|
import os
|
|
import sys
|
|
|
|
from charmhelpers.core.hookenv import (
|
|
action_set,
|
|
action_get,
|
|
action_fail,
|
|
config,
|
|
)
|
|
|
|
from charmhelpers.fetch import (
|
|
apt_hold,
|
|
apt_install,
|
|
apt_update,
|
|
apt_unhold,
|
|
)
|
|
|
|
from charmhelpers.fetch.ubuntu_apt_pkg import PkgVersion
|
|
|
|
from charmhelpers.core.host import service_restart
|
|
|
|
from charms.reactive import is_state, remove_state
|
|
|
|
from reactive.containerd import CONTAINERD_PACKAGE, apt_packages, configure_nvidia
|
|
|
|
|
|
class ActionError(Exception):
|
|
pass
|
|
|
|
|
|
def _gpu_packages():
|
|
"""Returns list of packages required for specific gpu support"""
|
|
if is_state("containerd.nvidia.ready"):
|
|
return config("nvidia_apt_packages").split()
|
|
return []
|
|
|
|
|
|
def _package_list(containerd, gpu):
|
|
package_list = []
|
|
if not containerd and not gpu:
|
|
raise ActionError("Must select at-least one of container and gpu")
|
|
|
|
if containerd:
|
|
package_list += [CONTAINERD_PACKAGE]
|
|
|
|
if gpu and _gpu_packages():
|
|
package_list += _gpu_packages()
|
|
|
|
return set(package_list)
|
|
|
|
|
|
def _dry_run(containerd, gpu):
|
|
"""Determine if a new package is available."""
|
|
apt_update(fatal=True)
|
|
package_list = _package_list(containerd, gpu)
|
|
search = apt_packages(package_list)
|
|
for name in package_list:
|
|
if name not in search:
|
|
raise ActionError(f"Package '{name}' not found in apt.")
|
|
|
|
result = {}
|
|
for name, pkg in search.items():
|
|
available, installed = map(PkgVersion, (pkg.version, pkg.current_ver.ver_str))
|
|
result[f"{name}.available"] = available.version
|
|
result[f"{name}.installed"] = installed.version
|
|
result[f"{name}.upgrade-available"] = available > installed
|
|
|
|
return result
|
|
|
|
|
|
def _upgrade(containerd, gpu):
|
|
"""Do actual upgrade."""
|
|
|
|
if not containerd and not gpu:
|
|
raise ActionError("Must select at-least one of container and gpu")
|
|
|
|
upgrade_list = _dry_run(containerd, gpu)
|
|
try:
|
|
pkg = CONTAINERD_PACKAGE
|
|
if upgrade_list.get(f"{pkg}.upgrade-available"):
|
|
apt_update(fatal=True)
|
|
apt_unhold(pkg)
|
|
apt_install(pkg, fatal=True)
|
|
apt_hold(pkg)
|
|
upgrade_list[f"{pkg}.upgrade-complete"] = True
|
|
|
|
if any(upgrade_list.get(f"{pkg}.upgrade-available") for pkg in _gpu_packages()):
|
|
configure_nvidia(reconfigure=False)
|
|
for pkg in _gpu_packages():
|
|
upgrade_list[f"{pkg}.upgrade-complete"] = True
|
|
|
|
if any(upgrade_list.get(f"{pkg}.upgrade-complete") for pkg in upgrade_list):
|
|
service_restart(CONTAINERD_PACKAGE)
|
|
remove_state("containerd.version-published")
|
|
|
|
return upgrade_list
|
|
|
|
except Exception as e:
|
|
raise ActionError("Failed to complete upgrades") from e
|
|
|
|
|
|
def upgrade_main(containerd, gpu):
|
|
"""Upgrade containerd to the latest in apt."""
|
|
dry_run = action_get().get("dry-run")
|
|
|
|
try:
|
|
if dry_run:
|
|
result = _dry_run(containerd, gpu)
|
|
else:
|
|
result = _upgrade(containerd, gpu)
|
|
action_set(result)
|
|
except ActionError as ae:
|
|
action_fail(str(ae))
|
|
|
|
|
|
def main(args):
|
|
action_name = os.path.basename(args[0])
|
|
if action_name == "upgrade-containerd":
|
|
upgrade_main(True, False)
|
|
elif action_name == "upgrade-packages":
|
|
containerd = action_get().get("containerd")
|
|
gpu = action_get().get("gpu")
|
|
upgrade_main(containerd, gpu)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(sys.argv))
|