API for customizing the conda base environment.

condacolab.py

@@ -18,13 +18,20 @@
 from pathlib import Path
 from subprocess import check_output, run, PIPE, STDOUT
 from textwrap import dedent
-from typing import Dict, AnyStr
+from typing import Dict, AnyStr, Iterable
 from urllib.request import urlopen
+from urllib.error import HTTPError
 from distutils.spawn import find_executable
 from IPython.display import display
 
 from IPython import get_ipython
 
+try:
+    from ruamel.yaml import YAML
+    from ruamel.yaml.comments import CommentedMap
+except ImportError as e:
+    raise RuntimeError("Could not find ruamel.yaml, plese install using `!pip install ruamel.yaml`!") from e
+
 try:
     import ipywidgets as widgets
     HAS_IPYWIDGETS = True
@@ -43,6 +50,7 @@
     "Surbhi Sharma <[email protected]>"
 )
 
+yaml=YAML()
 
 PREFIX = "/opt/conda"
 
@@ -87,12 +95,115 @@ def _run_subprocess(command, logs_filename):
     assert (task.returncode == 0), f"💥💔💥 The installation failed! Logs are available at `{logs_file_path}/{logs_filename}`."
 
 
+def _update_environment(
+    prefix:os.PathLike = PREFIX,
+    environment_file: str = None,
+    python_version: str = None,
+    specs: Iterable[str] = None,
+    channels: Iterable[str] = None,
+    pip_args: Iterable[str] = None,
+    extra_conda_args: Iterable[str] = None,
+    ):
+    """
+    Install the dependencies in conda base environment during
+    the condacolab installion.
+
+    Parameters
+    ----------
+    prefix
+        Target location for the installation.
+    environment_file
+        Path or URL of the environment.yaml file to use for
+        updating the conda base enviornment.
+    python_version
+        Python version to use in the conda base environment, eg. "3.9".
+    specs
+        List of additional specifications (packages) to install.
+    channels
+        Comma separated list of channels to use in the conda
+        base environment.
+    pip_args
+        List of additional packages to be installed using pip.
+    extra_conda_args
+        Any extra conda arguments to be used during the installation.
+    """
+
+    # When environment.yaml file is not provided.
+    if environment_file is None:
+        env_details = {}
+        if channels:
+            env_details["channels"] = channels
+        if specs:
+            env_details["dependencies"] = specs
+        if python_version:
+            env_details["dependencies"] += [f"python={python_version}"]
+        if pip_args:
+            pip_args_dict = {"pip": pip_args}
+            env_details["dependencies"].append(pip_args_dict)
+        environment_file_path = "/environment.yaml"
+
+        with open(environment_file_path, 'w') as f:
+            yaml.indent(mapping=2, sequence=4, offset=2)
+            yaml.dump(env_details, f)
+    else:
+        # If URL is given for environment.yaml file
+        if environment_file.startswith(("http://", "https://")):
+            environment_file_path = "/environment.yaml"
+            try:
+                with urlopen(environment_file) as response, open(environment_file_path, "wb") as out:
+                    shutil.copyfileobj(response, out)
+            except HTTPError as e:
+                raise HTTPError("The URL you entered is not working, please check it again.") from e
+
+        # If path is given for environment.yaml file
+        else:
+            environment_file_path = environment_file
+
+        with open(environment_file_path, 'r') as f:
+            env_details = yaml.load(f.read())
+
+        for key in env_details:
+            if channels and key == "channels":
+                env_details["channels"].extend(channels)
+            if key == "dependencies":
+                if specs:
+                    env_details["dependencies"].extend(specs)
+                if python_version:
+                    env_details["dependencies"].extend([f"python={python_version}"])
+                if pip_args:
+                    for element in env_details["dependencies"]:
+                        # if pip dependencies are already specified.
+                        if type(element) is CommentedMap and "pip" in element:
+                            element["pip"].extend(pip_args)
+                        # if there are no pip dependencies specified in the yaml file.
+                        else:
+                            pip_args_dict = CommentedMap([("pip", [*pip_args])])
+                            env_details["dependencies"].append(pip_args_dict)
+                        break
+        with open(environment_file_path, 'w') as f:
+            f.truncate(0)
+            yaml.dump(env_details, f)
+
+    extra_conda_args = extra_conda_args or ()
+
+    _run_subprocess(
+        [f"{prefix}/bin/python", "-m", "conda_env", "update", "-n", "base", "-f", environment_file_path, *extra_conda_args],
+        "environment_file_update.log",
+    )
+
+
 def install_from_url(
     installer_url: AnyStr,
     prefix: os.PathLike = PREFIX,
     env: Dict[AnyStr, AnyStr] = None,
     run_checks: bool = True,
     restart_kernel: bool = True,
+    environment_file: str = None,
+    python_version: str = None,
+    specs: Iterable[str] = None,
+    channels: Iterable[str] = None,
+    pip_args: Iterable[str] = None,
+    extra_conda_args: Iterable[str] = None,
 ):
     """
     Download and run a constructor-like installer, patching
@@ -186,6 +297,36 @@ def install_from_url(
         "pip_task.log"
         )
 
+    print("📌 Adjusting configuration...")
+    cuda_version = ".".join(os.environ.get("CUDA_VERSION", "*.*.*").split(".")[:2])
+    prefix = Path(prefix)
+    condameta = prefix / "conda-meta"
+    condameta.mkdir(parents=True, exist_ok=True)
+
+
+    with open(condameta / "pinned", "a") as f:
+        f.write(f"cudatoolkit {cuda_version}.*\n")
+
+    with open(prefix / ".condarc", "a") as f:
+        f.write("always_yes: true\n")
+
+    if environment_file and not specs and not channels and not pip_args and not python_version:
+        extra_conda_args = extra_conda_args or ()
+        _run_subprocess(
+            [f"{prefix}/bin/python", "-m", "conda_env", "update", "-n", "base", "-f", environment_file],
+            "environment_file_update.log",
+        )
+    else:
+        _update_environment(
+            prefix=prefix,
+            environment_file=environment_file,
+            specs=specs,
+            channels=channels,
+            python_version=python_version,
+            pip_args=pip_args,
+            extra_conda_args=extra_conda_args,
+            )
+
     env = env or {}
     bin_path = f"{prefix}/bin"
 
@@ -220,18 +361,26 @@ def install_from_url(
     else:
         print("🔁 Please restart kernel by clicking on Runtime > Restart runtime.")
 
+
 def install_mambaforge(
-    prefix: os.PathLike = PREFIX, env: Dict[AnyStr, AnyStr] = None, run_checks: bool = True, restart_kernel: bool = True,
+    prefix: os.PathLike = PREFIX,
+    env: Dict[AnyStr, AnyStr] = None,
+    run_checks: bool = True,
+    restart_kernel: bool = True,
+    specs: Iterable[str] = None,
+    python_version: str = None,
+    channels: Iterable[str] = None,
+    environment_file: str = None,
+    extra_conda_args: Iterable[str] = None,
+    pip_args: Iterable[str] = None,
+
 ):
     """
     Install Mambaforge, built for Python 3.7.
-
     Mambaforge consists of a Miniconda-like distribution optimized
     and preconfigured for conda-forge packages, and includes ``mamba``,
     a faster ``conda`` implementation.
-
     Unlike the official Miniconda, this is built with the latest ``conda``.
-
     Parameters
     ----------
     prefix
@@ -244,7 +393,6 @@ def install_mambaforge(
         to add those yourself in the raw string. They will
         end up added to a line like ``exec env VAR=VALUE python3...``.
         For example, a value with spaces should be passed as::
-
             env={"VAR": '"a value with spaces"'}
     run_checks
         Run checks to see if installation was run previously.
@@ -256,15 +404,35 @@ def install_mambaforge(
         automatically and get a button instead to do it.
     """
     installer_url = r"https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh"
-    install_from_url(installer_url, prefix=prefix, env=env, run_checks=run_checks, restart_kernel=restart_kernel)
-
+    install_from_url(
+        installer_url,
+        prefix=prefix, 
+        env=env,
+        run_checks=run_checks,
+        restart_kernel=restart_kernel,
+        specs=specs,
+        python_version=python_version,
+        channels=channels,
+        environment_file=environment_file,
+        extra_conda_args=extra_conda_args,
+        pip_args=pip_args,
+        )
 
 # Make mambaforge the default
 install = install_mambaforge
 
 
 def install_miniforge(
-    prefix: os.PathLike = PREFIX, env: Dict[AnyStr, AnyStr] = None, run_checks: bool = True, restart_kernel: bool = True,
+    prefix: os.PathLike = PREFIX,
+    env: Dict[AnyStr, AnyStr] = None,
+    run_checks: bool = True,
+    restart_kernel: bool = True,
+    specs: Iterable[str] = None,
+    python_version: str = None,
+    channels: Iterable[str] = None,
+    environment_file: str = None,
+    extra_conda_args: Iterable[str] = None,
+    pip_args: Iterable[str] = None,
 ):
     """
     Install Mambaforge, built for Python 3.7.
@@ -298,11 +466,32 @@ def install_miniforge(
         automatically and get a button instead to do it.
     """
     installer_url = r"https://github.com/jaimergp/miniforge/releases/latest/download/Miniforge-colab-Linux-x86_64.sh"
-    install_from_url(installer_url, prefix=prefix, env=env, run_checks=run_checks, restart_kernel=restart_kernel)
+    install_from_url(
+        installer_url,
+        prefix=prefix,
+        env=env,
+        run_checks=run_checks,
+        restart_kernel=restart_kernel,
+        specs=specs,
+        python_version=python_version,
+        channels=channels,
+        environment_file=environment_file,
+        extra_conda_args=extra_conda_args,
+        pip_args=pip_args,
+        )
 
 
 def install_miniconda(
-    prefix: os.PathLike = PREFIX, env: Dict[AnyStr, AnyStr] = None, run_checks: bool = True, restart_kernel: bool = True,
+    prefix: os.PathLike = PREFIX,
+    env: Dict[AnyStr, AnyStr] = None,
+    run_checks: bool = True,
+    restart_kernel: bool = True,
+    specs: Iterable[str] = None,
+    python_version: str = None,
+    channels: Iterable[str] = None,
+    environment_file: str = None,
+    extra_conda_args: Iterable[str] = None,
+    pip_args: Iterable[str] = None,
 ):
     """
     Install Miniconda 4.12.0 for Python 3.7.
@@ -331,11 +520,32 @@ def install_miniconda(
         automatically and get a button instead to do it.
     """
     installer_url = r"https://repo.anaconda.com/miniconda/Miniconda3-py37_4.12.0-Linux-x86_64.sh"
-    install_from_url(installer_url, prefix=prefix, env=env, run_checks=run_checks, restart_kernel=restart_kernel)
+    install_from_url(
+        installer_url,
+        prefix=prefix,
+        env=env,
+        run_checks=run_checks,
+        restart_kernel=restart_kernel,
+        specs=specs,
+        python_version=python_version,
+        channels=channels,
+        environment_file=environment_file,
+        extra_conda_args=extra_conda_args,
+        pip_args=pip_args,
+    )
 
 
 def install_anaconda(
-    prefix: os.PathLike = PREFIX, env: Dict[AnyStr, AnyStr] = None, run_checks: bool = True, restart_kernel: bool = True,
+    prefix: os.PathLike = PREFIX,
+    env: Dict[AnyStr, AnyStr] = None,
+    run_checks: bool = True,
+    restart_kernel: bool = True,
+    specs: Iterable[str] = None,
+    python_version: str = None,
+    channels: Iterable[str] = None,
+    environment_file: str = None,
+    extra_conda_args: Iterable[str] = None,
+    pip_args: Iterable[str] = None,
 ):
     """
     Install Anaconda 2022.05, the latest version built
@@ -365,7 +575,19 @@ def install_anaconda(
         automatically and get a button instead to do it.
     """
     installer_url = r"https://repo.anaconda.com/archive/Anaconda3-2022.05-Linux-x86_64.sh"
-    install_from_url(installer_url, prefix=prefix, env=env, run_checks=run_checks, restart_kernel=restart_kernel)
+    install_from_url(        
+        installer_url,
+        prefix=prefix,
+        env=env,
+        run_checks=run_checks,
+        restart_kernel=restart_kernel,
+        specs=specs,
+        python_version=python_version,
+        channels=channels,
+        environment_file=environment_file,
+        extra_conda_args=extra_conda_args,
+        pip_args=pip_args,
+    )
 
 
 def check(prefix: os.PathLike = PREFIX, verbose: bool = True):
@@ -382,10 +604,6 @@ def check(prefix: os.PathLike = PREFIX, verbose: bool = True):
         Print success message if True
     """
     assert find_executable("conda"), "💥💔💥 Conda not found!"
-
-    pymaj, pymin = sys.version_info[:2]
-    sitepackages = f"{prefix}/lib/python{pymaj}.{pymin}/site-packages"
-    assert sitepackages in sys.path, f"💥💔💥 PYTHONPATH was not patched! Value: {sys.path}"
     assert all(
         not path.startswith("/usr/local/") for path in sys.path
     ), f"💥💔💥 PYTHONPATH include system locations: {[path for path in sys.path if path.startswith('/usr/local')]}!"