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')]}!"