Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

API for customizing the conda base environment. #38

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

API for customizing the conda base environment. #38

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a9d8597
rebasing branch
ssurbhi560 Oct 19, 2022
65a39b8
making changes as suggested in review
ssurbhi560 Nov 2, 2022
019a453
Use mamba if available
jaimergp Nov 7, 2022
a379faf
Add a print statement before starting updating environment
ssurbhi560 Nov 9, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
272 changes: 237 additions & 35 deletions condacolab.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,16 +15,23 @@
import sys
import shutil
from datetime import datetime, timedelta
from distutils.spawn import find_executable
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 distutils.spawn import find_executable
from IPython.display import display
from urllib.error import HTTPError

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
Expand All @@ -43,6 +50,7 @@
"Surbhi Sharma <[email protected]>"
)

yaml=YAML()

PREFIX = "/opt/conda"

Expand Down Expand Up @@ -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] = (),
channels: Iterable[str] = (),
pip_args: Iterable[str] = (),
extra_conda_args: Iterable[str] = (),
conda_exe: str = "conda",
):
"""
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.
"""
os.makedirs("/var/condacolab", exist_ok=True)
environment_file_path = "/var/condacolab/environment.yaml"

# 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)

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://")):
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:
shutil.copy(environment_file, environment_file_path)

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 isinstance(element, CommentedMap) and "pip" in element:
element["pip"].extend(pip_args)
break
# 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)

with open(environment_file_path, 'w') as f:
ssurbhi560 marked this conversation as resolved.
Show resolved Hide resolved
f.truncate(0)
yaml.dump(env_details, f)

_run_subprocess(
[conda_exe, "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] = (),
channels: Iterable[str] = (),
pip_args: Iterable[str] = (),
extra_conda_args: Iterable[str] = (),
):
"""
Download and run a constructor-like installer, patching
Expand Down Expand Up @@ -186,24 +297,38 @@ def install_from_url(
"pip_task.log"
)

print("📦 Updating enviornment using YAML file...")
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

print("📦 Updating environment using YAML file...")


_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,
conda_exe=f"{prefix}/bin/{conda_exe}",
)

env = env or {}
bin_path = f"{prefix}/bin"

os.rename(sys.executable, f"{sys.executable}.renamed_by_condacolab.bak")
with open(sys.executable, "w") as f:
f.write(
dedent(
f"""
#!/bin/bash
source {prefix}/etc/profile.d/conda.sh
conda activate
unset PYTHONPATH
mv /usr/bin/lsb_release /usr/bin/lsb_release.renamed_by_condacolab.bak
exec {bin_path}/python $@
"""
).lstrip()
)
run(["chmod", "+x", sys.executable])
if os.path.exists(sys.executable):
os.rename(sys.executable, f"{sys.executable}.renamed_by_condacolab.bak")
with open(sys.executable, "w") as f:
f.write(
dedent(
f"""
#!/bin/bash
source {prefix}/etc/profile.d/conda.sh
conda activate
unset PYTHONPATH
mv /usr/bin/lsb_release /usr/bin/lsb_release.renamed_by_condacolab.bak
exec {bin_path}/python $@
"""
).lstrip()
)
run(["chmod", "+x", sys.executable])

taken = timedelta(seconds=round((datetime.now() - t0).total_seconds(), 0))
print(f"⏲ Done in {taken}")
Expand All @@ -220,18 +345,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] = (),
python_version: str = None,
channels: Iterable[str] = (),
environment_file: str = None,
extra_conda_args: Iterable[str] = (),
pip_args: Iterable[str] = (),

):
"""
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
Expand All @@ -244,7 +377,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.
Expand All @@ -256,15 +388,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] = (),
python_version: str = None,
channels: Iterable[str] = (),
environment_file: str = None,
extra_conda_args: Iterable[str] = (),
pip_args: Iterable[str] = (),
):
"""
Install Mambaforge, built for Python 3.7.
Expand Down Expand Up @@ -298,11 +450,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] = (),
python_version: str = None,
channels: Iterable[str] = (),
environment_file: str = None,
extra_conda_args: Iterable[str] = (),
pip_args: Iterable[str] = (),
):
"""
Install Miniconda 4.12.0 for Python 3.7.
Expand Down Expand Up @@ -331,11 +504,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] = (),
python_version: str = None,
channels: Iterable[str] = (),
environment_file: str = None,
extra_conda_args: Iterable[str] = (),
pip_args: Iterable[str] = (),
):
"""
Install Anaconda 2022.05, the latest version built
Expand Down Expand Up @@ -365,7 +559,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):
Expand All @@ -382,10 +588,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')]}!"
Expand Down