Avoid downgrading user-env conda/mamba

- Only check lower bound on conda/mamba, upgrade unbounded if not matched (let conda apply upper bound according to existing pins, such as Python)
- handle missing mamba
- avoid upgrading Python by aborting the install, instead of keeping old envs
- minimum supported Python for user env is 3.9
- Fix output reporting of conda install step (no need for json capture when we don't parse the output - exit codes will do)

dev-requirements.txt

@@ -1,3 +1,4 @@
+packaging
 pytest
 pytest-cov
 pytest-mock

tests/test_conda.py

@@ -21,13 +21,6 @@ def prefix():
             installer_url, checksum
         ) as installer_path:
             conda.install_miniconda(installer_path, tmpdir)
-        conda.ensure_conda_packages(
-            tmpdir,
-            [
-                f"conda=={installer.MAMBAFORGE_CONDA_VERSION}",
-                f"mamba=={installer.MAMBAFORGE_MAMBA_VERSION}",
-            ],
-        )
         yield tmpdir
 
 

tests/test_installer.py

@@ -6,11 +6,12 @@ import os
 from unittest import mock
 from subprocess import run, PIPE
 
+from packaging.version import parse as V
+from packaging.specifiers import SpecifierSet
 import pytest
 
 from tljh import conda
 from tljh import installer
-from tljh.utils import parse_version as V
 from tljh.yaml import yaml
 
 
@@ -48,13 +49,18 @@ def setup_conda(distro, version, prefix):
     """Install mambaforge or miniconda in a prefix"""
     if distro == "mambaforge":
         installer_url, _ = installer._mambaforge_url(version)
+    elif distro == "miniforge":
+        installer_url, _ = installer._mambaforge_url(version)
+        installer_url = installer_url.replace("Mambaforge", "Miniforge3")
     elif distro == "miniconda":
         arch = os.uname().machine
         installer_url = (
             f"https://repo.anaconda.com/miniconda/Miniconda3-{version}-Linux-{arch}.sh"
         )
     else:
-        raise ValueError(f"{distro=} must be 'miniconda' or 'mambaforge'")
+        raise ValueError(
+            f"{distro=} must be 'miniconda' or 'mambaforge' or 'miniforge'"
+        )
     with conda.download_miniconda_installer(installer_url, None) as installer_path:
         conda.install_miniconda(installer_path, str(prefix))
     # avoid auto-updating conda when we install other packages
@@ -79,68 +85,111 @@ def user_env_prefix(tmp_path):
         yield user_env_prefix
 
 
+def _specifier(version):
+    """Convert version string to SpecifierSet
+
+    If just a version number, add == to make it a specifier
+
+    Any missing fields are replaced with .*
+
+    If it's already a specifier string, pass it directly to SpecifierSet
+
+    e.g.
+
+    - 3.7 -> ==3.7.*
+    - 1.2.3 -> ==1.2.3
+    """
+    if version[0].isdigit():
+        # it's a version number, not a specifier
+        if version.count(".") < 2:
+            # pad missing fields
+            version += ".*"
+        version = f"=={version}"
+    return SpecifierSet(version)
+
+
 @pytest.mark.parametrize(
-    "distro, version, conda_version, mamba_version",
+    "distro, distro_version, expected_versions",
     [
+        # No previous install, start fresh
         (
             None,
             None,
-            installer.MAMBAFORGE_CONDA_VERSION,
+            {
-            installer.MAMBAFORGE_MAMBA_VERSION,
+                "python": "3.10.*",
-        ),
+                "conda": "22.11.1",
-        (
+                "mamba": "1.1.0",
-            "exists",
+            },
-            None,
-            installer.MAMBAFORGE_CONDA_VERSION,
-            installer.MAMBAFORGE_MAMBA_VERSION,
         ),
+        # previous install, 1.0
         (
             "mambaforge",
             "22.11.1-4",
-            installer.MAMBAFORGE_CONDA_VERSION,
+            {
-            installer.MAMBAFORGE_MAMBA_VERSION,
+                "python": "3.10.*",
+                "conda": "22.11.1",
+                "mamba": "1.1.0",
+            },
         ),
-        ("mambaforge", "4.10.3-7", "4.10.3", "0.16.0"),
+        # 0.2 install, no upgrade needed
+        (
+            "mambaforge",
+            "4.10.3-7",
+            {
+                "conda": "4.10.3",
+                "mamba": "0.16.0",
+                "python": "3.9.*",
+            },
+        ),
+        # simulate missing mamba
+        # will be installed but not pinned
+        # to avoid conflicts
+        (
+            "miniforge",
+            "4.10.3-7",
+            {
+                "conda": "4.10.3",
+                "mamba": ">=1.1.0",
+                "python": "3.9.*",
+            },
+        ),
+        # too-old Python (3.7), abort
         (
             "miniconda",
             "4.7.10",
-            installer.MAMBAFORGE_CONDA_VERSION,
+            ValueError,
-            installer.MAMBAFORGE_MAMBA_VERSION,
-        ),
-        (
-            "miniconda",
-            "4.5.1",
-            installer.MAMBAFORGE_CONDA_VERSION,
-            installer.MAMBAFORGE_MAMBA_VERSION,
         ),
     ],
 )
 def test_ensure_user_environment(
     user_env_prefix,
     distro,
-    version,
+    distro_version,
-    conda_version,
+    expected_versions,
-    mamba_version,
 ):
-    if version and V(version) < V("4.10.1") and os.uname().machine == "aarch64":
+    if (
-        pytest.skip(f"Miniconda {version} not available for aarch64")
+        distro_version
+        and V(distro_version) < V("4.10.1")
+        and os.uname().machine == "aarch64"
+    ):
+        pytest.skip(f"{distro} {distro_version} not available for aarch64")
     canary_file = user_env_prefix / "test-file.txt"
     canary_package = "types-backports_abc"
     if distro:
-        if distro == "exists":
+        setup_conda(distro, distro_version, user_env_prefix)
-            user_env_prefix.mkdir()
-        else:
-            setup_conda(distro, version, user_env_prefix)
         # install a noarch: python package that won't be used otherwise
         # should depend on Python, so it will interact with possible upgrades
+        pkgs = [canary_package]
         run(
             [
                 str(user_env_prefix / "bin/conda"),
                 "install",
+                "-S",
                 "-y",
-                    "-c" "conda-forge",
+                "-c",
-                    canary_package,
+                "conda-forge",
-                ],
+            ]
+            + pkgs,
             input="",
             check=True,
         )
@@ -149,7 +198,14 @@ def test_ensure_user_environment(
         with canary_file.open("w") as f:
             f.write("I'm here\n")
 
+    if isinstance(expected_versions, type) and issubclass(expected_versions, Exception):
+        exc_class = expected_versions
+        with pytest.raises(exc_class):
             installer.ensure_user_environment("")
+        return
+    else:
+        installer.ensure_user_environment("")
+
     p = run(
         [str(user_env_prefix / "bin/conda"), "list", "--json"],
         stdout=PIPE,
@@ -158,14 +214,23 @@ def test_ensure_user_environment(
     )
     package_list = json.loads(p.stdout)
     packages = {package["name"]: package for package in package_list}
+
     if distro:
         # make sure we didn't wipe out files
         assert canary_file.exists()
-        if distro != "exists":
         # make sure we didn't delete the installed package
         assert canary_package in packages
 
-    assert "conda" in packages
+    for pkg, version in expected_versions.items():
-    assert packages["conda"]["version"] == conda_version
+        assert pkg in packages
-    assert "mamba" in packages
+        assert V(packages[pkg]["version"]) in _specifier(version)
-    assert packages["mamba"]["version"] == mamba_version
+
+
+def test_ensure_user_environment_no_clobber(user_env_prefix):
+    # don't clobber existing user-env dir if it's non-empty and not a conda install
+    user_env_prefix.mkdir()
+    canary_file = user_env_prefix / "test-file.txt"
+    with canary_file.open("w") as f:
+        pass
+    with pytest.raises(OSError):
+        installer.ensure_user_environment("")

tljh/conda.py

@@ -13,7 +13,6 @@ import time
 import requests
 
 from tljh import utils
-from tljh.utils import parse_version as V
 
 
 def sha256_file(fname):
@@ -29,47 +28,20 @@ def sha256_file(fname):
     return hash_sha256.hexdigest()
 
 
-def check_miniconda_version(prefix, version):
+def get_conda_package_versions(prefix):
-    """
+    """Get conda package versions, via `conda list --json`"""
-    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:
-        out = (
+        out = subprocess.check_output(
-            subprocess.check_output(
+            [os.path.join(prefix, "bin", "conda"), "list", "--json"],
-                [os.path.join(prefix, "bin", "mamba"), "--version"],
+            text=True,
-                stderr=subprocess.STDOUT,
-            )
-            .decode()
-            .strip()
         )
     except (subprocess.CalledProcessError, FileNotFoundError):
         return versions
-    for line in out.strip().splitlines():
+
-        pkg, version = line.split()
+    packages = json.loads(out)
-        versions[pkg] = version
+    for package in packages:
+        versions[package["name"]] = package["version"]
     return versions
 
 
@@ -133,39 +105,24 @@ def ensure_conda_packages(prefix, packages):
     Note that conda seem to update dependencies by default, so there is probably
     no need to have a update parameter exposed for this function.
     """
-    conda_executable = [os.path.join(prefix, "bin", "mamba")]
+    conda_executable = os.path.join(prefix, "bin", "mamba")
+    if not os.path.isfile(conda_executable):
+        # fallback on conda if mamba is not present (e.g. for mamba to install itself)
+        conda_executable = os.path.join(prefix, "bin", "conda")
+
     abspath = os.path.abspath(prefix)
-    # Let subprocess errors propagate
+
-    # Explicitly do *not* capture stderr, since that's not always JSON!
+    utils.run_subprocess(
-    # Scripting conda is a PITA!
+        [
-    # FIXME: raise different exception when using
+            conda_executable,
-    raw_output = subprocess.check_output(
-        conda_executable
-        + [
             "install",
             "-c",
             "conda-forge",  # Make customizable if we ever need to
-            "--json",
             "--prefix",
             abspath,
         ]
-        + packages
+        + packages,
-    ).decode()
-    # `conda install` outputs JSON lines for fetch updates,
-    # and a undelimited output at the end. There is no reasonable way to
-    # parse this outside of this kludge.
-    filtered_output = "\n".join(
-        [
-            l
-            for l in raw_output.split("\n")
-            # Sometimes the JSON messages start with a \x00. The lstrip removes these.
-            # conda messages seem to randomly throw \x00 in places for no reason
-            if not l.lstrip("\x00").startswith('{"fetch"')
-        ]
     )
-    output = json.loads(filtered_output.lstrip("\x00"))
-    if "success" in output and output["success"] == True:
-        return
     fix_permissions(prefix)
 
 

tljh/installer.py

@@ -164,11 +164,15 @@ 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
+# minimum versions of packages
-# the mambaforge distribution
+MINIMUM_VERSIONS = {
-MAMBAFORGE_MAMBA_VERSION = "1.1.0"
+    # if conda/mamba are lower than this, upgrade them before installing the user packages
-MAMBAFORGE_CONDA_VERSION = "22.11.1"
+    "mamba": "0.16.0",
+    "conda": "4.10",
+    # minimum Python version (if not matched, abort to avoid big disruptive updates)
+    "python": "3.9",
+}
 
 
 def _mambaforge_url(version=MAMBAFORGE_VERSION, arch=None):
@@ -196,53 +200,70 @@ 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 (since 2021-10-18)
-        "4.10.3": ("4.10.3", "0.16.0"),
-    }
-
     # Check OS, set appropriate string for conda installer path
     if os.uname().sysname != "Linux":
         raise OSError("TLJH is only supported on Linux platforms.")
-    found_conda = False
-    have_versions = conda.get_mamba_versions(USER_ENV_PREFIX)
-    have_conda_version = have_versions.get("conda")
-    if have_conda_version:
-        logger.info(
-            f"Found prefix at {USER_ENV_PREFIX}, with conda/mamba({have_versions})"
-        )
-        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 not found_conda:
+    # Check the existing environment for what to do
-        if os.path.exists(USER_ENV_PREFIX):
+    package_versions = conda.get_conda_package_versions(USER_ENV_PREFIX)
-            logger.warning(
+
-                f"Found prefix at {USER_ENV_PREFIX}, but too old or missing conda/mamba ({have_versions}). Upgrading from mambaforge."
+    # Case 1: no existing environment
-            )
+    if not package_versions:
-            # FIXME: should this fail? I'm not sure how destructive it is
+        # 1a. no environment, but prefix exists.
+        # Abort to avoid clobbering something we don't recognize
+        if os.path.exists(USER_ENV_PREFIX) and os.listdir(USER_ENV_PREFIX):
+            msg = f"Found non-empty directory that is not a conda install in {USER_ENV_PREFIX}. Please remove it (or rename it to preserve files) and run tljh again."
+            logger.error(msg)
+            raise OSError(msg)
+
+        # 1b. No environment, directory empty or doesn't exist
+        # start fresh install
         logger.info("Downloading & setting up user environment...")
         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 = MAMBAFORGE_CONDA_VERSION
+        package_versions = conda.get_conda_package_versions(USER_ENV_PREFIX)
-        mamba_version = MAMBAFORGE_MAMBA_VERSION
+        # quick sanity check: we should have conda and mamba!
+        assert "conda" in package_versions
+        assert "mamba" in package_versions
 
+    # next, check Python
+    python_version = package_versions["python"]
+    logger.debug(f"Found python={python_version} in {USER_ENV_PREFIX}")
+    if V(python_version) < V(MINIMUM_VERSIONS["python"]):
+        msg = (
+            f"TLJH requires Python >={MINIMUM_VERSIONS['python']}, found python={python_version} in {USER_ENV_PREFIX}."
+            f"\nPlease upgrade Python (may be highly disruptive!), or remove/rename {USER_ENV_PREFIX} to allow TLJH to make a fresh install."
+            f"\nYou can use `{USER_ENV_PREFIX}/bin/conda list` to save your current list of packages."
+        )
+        logger.error(msg)
+        raise ValueError(msg)
+
+    # at this point, we know we have an env ready with conda and are going to start installing
+    # first, check if we should upgrade/install conda and/or mamba
+    to_upgrade = []
+    for pkg in ("conda", "mamba"):
+        version = package_versions.get(pkg)
+        min_version = MINIMUM_VERSIONS[pkg]
+        if not version:
+            logger.warning(f"{USER_ENV_PREFIX} is missing {pkg}, installing it...")
+            to_upgrade.append(pkg)
+        else:
+            logger.debug(f"Found {pkg}=={version} in {USER_ENV_PREFIX}")
+            if V(version) < V(min_version):
+                logger.info(
+                    f"{USER_ENV_PREFIX} has {pkg}=={version}, it will be upgraded to {pkg}>={min_version}"
+                )
+                to_upgrade.append(pkg)
+
+    if to_upgrade:
         conda.ensure_conda_packages(
             USER_ENV_PREFIX,
-        [
+            # we _could_ explicitly pin Python here,
-            # Conda's latest version is on conda much more so than on PyPI.
+            # but conda already does this by default
-            "conda==" + conda_version,
+            to_upgrade,
-            "mamba==" + mamba_version,
-        ],
         )
 
     conda.ensure_pip_requirements(