Skip to content

Commit

Permalink
rebasing branch
Browse files Browse the repository at this point in the history
  • Loading branch information
@ssurbhi560
ssurbhi560 committed Oct 30, 2022
1 parent 7f720b5 commit a9d8597
Showing 1 changed file with 236 additions and 18 deletions.
254 changes: 236 additions & 18 deletions condacolab.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
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] = 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
Expand Down Expand Up @@ -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"

Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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')]}!"
Expand Down

0 comments on commit a9d8597

Please sign in to comment.