From 1a3c48a50091878ed9dca1b5db7aa6d051d30fa8 Mon Sep 17 00:00:00 2001 From: Min RK Date: Tue, 21 Mar 2023 14:21:00 +0100 Subject: [PATCH] Update base user environment to mambaforge 22.11.1-4 shift some duplicated code into utility functions and constants --- tests/test_conda.py | 26 +++++------- tljh/conda.py | 41 +++++++++++++++---- tljh/installer.py | 98 +++++++++++++++++++++++++++++++-------------- tljh/utils.py | 9 +++++ 4 files changed, 120 insertions(+), 54 deletions(-) diff --git a/tests/test_conda.py b/tests/test_conda.py index a13ab39..d38a85a 100644 --- a/tests/test_conda.py +++ b/tests/test_conda.py @@ -2,6 +2,7 @@ Test conda commandline wrappers """ from tljh import conda +from tljh import installer import os import pytest import subprocess @@ -13,25 +14,20 @@ def prefix(): """ Provide a temporary directory with a mambaforge conda environment """ - # see https://github.com/conda-forge/miniforge/releases - mambaforge_version = "4.10.3-7" - if os.uname().machine == "aarch64": - installer_sha256 = ( - "ac95f137b287b3408e4f67f07a284357b1119ee157373b788b34e770ef2392b2" - ) - elif os.uname().machine == "x86_64": - installer_sha256 = ( - "fc872522ec427fcab10167a93e802efaf251024b58cc27b084b915a9a73c4474" - ) - installer_url = "https://github.com/conda-forge/miniforge/releases/download/{v}/Mambaforge-{v}-Linux-{arch}.sh".format( - v=mambaforge_version, arch=os.uname().machine - ) + machine = os.uname().machine + installer_url, checksum = installer._mambaforge_url() with tempfile.TemporaryDirectory() as tmpdir: with conda.download_miniconda_installer( - installer_url, installer_sha256 + installer_url, checksum ) as installer_path: conda.install_miniconda(installer_path, tmpdir) - conda.ensure_conda_packages(tmpdir, ["conda==4.10.3"]) + conda.ensure_conda_packages( + tmpdir, + [ + f"conda=={installer.MAMBAFORGE_CONDA_VERSION}", + f"mamba=={installer.MAMBAFORGE_MAMBA_VERSION}", + ], + ) yield tmpdir diff --git a/tljh/conda.py b/tljh/conda.py index 88923f6..7aa864a 100644 --- a/tljh/conda.py +++ b/tljh/conda.py @@ -8,8 +8,8 @@ import hashlib import contextlib import tempfile import requests -from distutils.version import LooseVersion as V from tljh import utils +from tljh.utils import parse_version as V def sha256_file(fname): @@ -29,19 +29,44 @@ def check_miniconda_version(prefix, version): """ Return true if a miniconda install with version exists at prefix """ + versions = get_mamba_versions(prefix) + if "conda" not in versions: + return False + + return V(versions["conda"]) >= V(version) + + +def get_mamba_versions(prefix): + """Parse `mamba --version` output into a dict + + which looks like: + + mamba 1.1.0 + conda 22.11.1 + + into: + + { + "mamba": "1.1.0", + "conda": "22.11.1", + } + """ + versions = {} try: - installed_version = ( + out = ( subprocess.check_output( - [os.path.join(prefix, "bin", "conda"), "-V"], stderr=subprocess.STDOUT + [os.path.join(prefix, "bin", "mamba"), "--version"], + stderr=subprocess.STDOUT, ) .decode() .strip() - .split()[1] ) - return V(installed_version) >= V(version) except (subprocess.CalledProcessError, FileNotFoundError): - # Conda doesn't exist - return False + return versions + for line in out.strip().splitlines(): + pkg, version = line.split() + versions[pkg] = version + return versions @contextlib.contextmanager @@ -53,7 +78,7 @@ def download_miniconda_installer(installer_url, sha256sum): of given version, verifies the sha256sum & provides path to it to the `with` block to run. """ - with tempfile.NamedTemporaryFile("wb") as f: + with tempfile.NamedTemporaryFile("wb", suffix=".sh") as f: f.write(requests.get(installer_url).content) # Remain in the NamedTemporaryFile context, but flush changes, see: # https://docs.python.org/3/library/os.html#os.fsync diff --git a/tljh/installer.py b/tljh/installer.py index 7e9948d..c11eb3e 100644 --- a/tljh/installer.py +++ b/tljh/installer.py @@ -26,6 +26,7 @@ from tljh import ( traefik, user, ) + from .config import ( CONFIG_DIR, CONFIG_FILE, @@ -34,6 +35,7 @@ from .config import ( STATE_DIR, USER_ENV_PREFIX, ) +from .utils import parse_version as V from .yaml import yaml HERE = os.path.abspath(os.path.dirname(__file__)) @@ -154,58 +156,92 @@ def ensure_usergroups(): f.write("Defaults exempt_group = jupyterhub-admins\n") +# Install mambaforge using an installer from +# https://github.com/conda-forge/miniforge/releases +MAMBAFORGE_VERSION = "22.11.1-4" +# sha256 checksums +MAMBAFORGE_CHECKSUMS = { + "aarch64": "96191001f27e0cc76612d4498d34f9f656d8a7dddee44617159e42558651479c", + "x86_64": "16c7d256de783ceeb39970e675efa4a8eb830dcbb83187f1197abfea0bf07d30", +} +# run `mamba --version` to get the conda and mamba versions +# conda/mamba will be _upgraded_ to these versions, if they differ from what's in +# the mambaforge distribution +MAMBAFORGE_MAMBA_VERSION = "1.1.0" +MAMBAFORGE_CONDA_VERSION = "22.11.1" + + +def _mambaforge_url(version=MAMBAFORGE_VERSION, arch=None): + """Return (URL, checksum) for mambaforge download for a given version and arch + + Default values provided for both version and arch + """ + if arch is None: + arch = os.uname().machine + installer_url = "https://github.com/conda-forge/miniforge/releases/download/{v}/Mambaforge-{v}-Linux-{arch}.sh".format( + v=version, + arch=arch, + ) + # Check system architecture, set appropriate installer checksum + checksum = MAMBAFORGE_CHECKSUMS.get(arch) + if not checksum: + raise ValueError( + f"Unsupported architecture: {arch}. TLJH only supports {','.join(MAMBAFORGE_CHECKSUMS.keys())}" + ) + return installer_url, checksum + + def ensure_user_environment(user_requirements_txt_file): """ Set up user conda environment with required packages """ logger.info("Setting up user environment...") + # note: these must be in descending order + conda_upgrade_versions = { + # format: "conda version": (conda_version, mamba_version), + # mambaforge 4.10.3-7 (2023-03-21) + "22.11.1": (MAMBAFORGE_CONDA_VERSION, MAMBAFORGE_MAMBA_VERSION), + # tljh up to 0.2.0 (2021-10-18) + "4.10.3": ("4.10.3", "0.16.0"), + # very old versions, do these still work? + "4.7.10": ("4.8.1", "0.16.0"), + "4.5.4": ("4.5.8", "0.16.0"), + } - miniconda_old_version = "4.5.4" - miniconda_new_version = "4.7.10" - # Install mambaforge using an installer from - # https://github.com/conda-forge/miniforge/releases - mambaforge_new_version = "4.10.3-7" - # Check system architecture, set appropriate installer checksum - if os.uname().machine == "aarch64": - installer_sha256 = ( - "ac95f137b287b3408e4f67f07a284357b1119ee157373b788b34e770ef2392b2" - ) - elif os.uname().machine == "x86_64": - installer_sha256 = ( - "fc872522ec427fcab10167a93e802efaf251024b58cc27b084b915a9a73c4474" - ) # Check OS, set appropriate string for conda installer path if os.uname().sysname != "Linux": raise OSError("TLJH is only supported on Linux platforms.") - # Then run `mamba --version` to get the conda and mamba versions - # Keep these in sync with tests/test_conda.py::prefix - mambaforge_conda_new_version = "4.10.3" - mambaforge_mamba_version = "0.16.0" + found_conda = False + have_versions = conda.get_mamba_versions(USER_ENV_PREFIX) + have_conda_version = have_versions.get("conda") + if have_conda_version: + for check_version, conda_mamba_version in conda_upgrade_versions.items(): + if V(have_conda_version) >= V(check_version): + found_conda = True + conda_version, mamba_version = conda_mamba_version + break - if conda.check_miniconda_version(USER_ENV_PREFIX, mambaforge_conda_new_version): - conda_version = "4.10.3" - elif conda.check_miniconda_version(USER_ENV_PREFIX, miniconda_new_version): - conda_version = "4.8.1" - elif conda.check_miniconda_version(USER_ENV_PREFIX, miniconda_old_version): - conda_version = "4.5.8" - # If no prior miniconda installation is found, we can install a newer version - else: + if not found_conda: + if os.path.exists(USER_ENV_PREFIX): + logger.warning( + f"Found prefix at {USER_ENV_PREFIX}, but too old or missing conda/mamba ({have_versions}). Rebuilding env from scratch!!" + ) + # FIXME: should this fail? I'm not sure how destructive it is logger.info("Downloading & setting up user environment...") - installer_url = "https://github.com/conda-forge/miniforge/releases/download/{v}/Mambaforge-{v}-Linux-{arch}.sh".format( - v=mambaforge_new_version, arch=os.uname().machine - ) + installer_url, installer_sha256 = _mambaforge_url() with conda.download_miniconda_installer( installer_url, installer_sha256 ) as installer_path: conda.install_miniconda(installer_path, USER_ENV_PREFIX) - conda_version = "4.10.3" + conda_version = MAMBAFORGE_CONDA_VERSION + mamba_version = MAMBAFORGE_MAMBA_VERSION conda.ensure_conda_packages( USER_ENV_PREFIX, [ # Conda's latest version is on conda much more so than on PyPI. "conda==" + conda_version, - "mamba==" + mambaforge_mamba_version, + "mamba==" + MAMBAFORGE_MAMBA_VERSION, ], ) diff --git a/tljh/utils.py b/tljh/utils.py index 0b61da6..a78fb80 100644 --- a/tljh/utils.py +++ b/tljh/utils.py @@ -59,3 +59,12 @@ def get_plugin_manager(): pm.load_setuptools_entrypoints("tljh") return pm + + +def parse_version(version_string): + """Parse version string to tuple + + Finds all numbers and returns a tuple of ints + _very_ loose version parsing, like the old distutils.version.LooseVersion + """ + return tuple(int(part) for part in version_string.split("."))