From eb942cfedd47588b733b8cbe3e174a767307c222 Mon Sep 17 00:00:00 2001 From: allegroai <> Date: Thu, 15 Oct 2020 23:23:46 +0300 Subject: [PATCH] Add agent.package_manager.conda_env_as_base_docker allowing "docker_cmd" to contain link to a full pre-packaged conda environment (conda-pack outputs a tar.gz). Use TRAINS_CONDA_ENV_PACKAGE to specify conda tar.gz file. --- docs/trains.conf | 1 + trains_agent/backend_api/session/defs.py | 1 + trains_agent/commands/worker.py | 14 ++++--- trains_agent/helper/base.py | 7 ++++ trains_agent/helper/package/conda_api.py | 44 ++++++++++++++++----- trains_agent/helper/package/pip_api/venv.py | 8 ++-- 6 files changed, 57 insertions(+), 18 deletions(-) diff --git a/docs/trains.conf b/docs/trains.conf index ffb6c54..f1931b6 100644 --- a/docs/trains.conf +++ b/docs/trains.conf @@ -63,6 +63,7 @@ agent { # additional conda channels to use when installing with conda package manager conda_channels: ["pytorch", "conda-forge", ] # conda_full_env_update: false + # conda_env_as_base_docker: false # set the priority packages to be installed before the rest of the required packages # priority_packages: ["cython", "numpy", "setuptools", ] diff --git a/trains_agent/backend_api/session/defs.py b/trains_agent/backend_api/session/defs.py index 92d0aba..4e374d8 100644 --- a/trains_agent/backend_api/session/defs.py +++ b/trains_agent/backend_api/session/defs.py @@ -8,3 +8,4 @@ ENV_ACCESS_KEY = EnvEntry("TRAINS_API_ACCESS_KEY", "TRAINS_API_ACCESS_KEY") ENV_SECRET_KEY = EnvEntry("TRAINS_API_SECRET_KEY", "TRAINS_API_SECRET_KEY") ENV_VERBOSE = EnvEntry("TRAINS_API_VERBOSE", "TRAINS_API_VERBOSE", type=bool, default=False) ENV_HOST_VERIFY_CERT = EnvEntry("TRAINS_API_HOST_VERIFY_CERT", "TRAINS_API_HOST_VERIFY_CERT", type=bool, default=True) +ENV_CONDA_ENV_PACKAGE = EnvEntry("TRAINS_CONDA_ENV_PACKAGE", "TRAINS_CONDA_ENV_PACKAGE") diff --git a/trains_agent/commands/worker.py b/trains_agent/commands/worker.py index c7c2a65..4210b58 100644 --- a/trains_agent/commands/worker.py +++ b/trains_agent/commands/worker.py @@ -1207,8 +1207,8 @@ class Worker(ServiceCommandSection): except: python_version = None - venv_folder, requirements_manager = self.install_virtualenv(venv_dir=target, - requested_python_version=python_version) + venv_folder, requirements_manager = self.install_virtualenv( + venv_dir=target, requested_python_version=python_version, execution_info=execution) if self._default_pip: if install_globally and self.global_package_api: @@ -1426,8 +1426,8 @@ class Worker(ServiceCommandSection): except: python_ver = None - venv_folder, requirements_manager = self.install_virtualenv(standalone_mode=standalone_mode, - requested_python_version=python_ver) + venv_folder, requirements_manager = self.install_virtualenv( + standalone_mode=standalone_mode, requested_python_version=python_ver, execution_info=execution) if not standalone_mode: if self._default_pip: @@ -2007,8 +2007,9 @@ class Worker(ServiceCommandSection): ) ) - def install_virtualenv(self, venv_dir=None, requested_python_version=None, standalone_mode=False): - # type: (str, str, bool) -> Tuple[Path, RequirementsManager] + def install_virtualenv( + self, venv_dir=None, requested_python_version=None, standalone_mode=False, execution_info=None): + # type: (str, str, bool, ExecutionInfo) -> Tuple[Path, RequirementsManager] """ Install a new python virtual environment, removing the old one if exists :return: virtualenv directory and requirements manager to use with task @@ -2055,6 +2056,7 @@ class Worker(ServiceCommandSection): python=executable_version_suffix if self.is_conda else executable_name, path=venv_dir, requirements_manager=requirements_manager, + execution_info=execution_info, ) global_package_manager_params = dict( diff --git a/trains_agent/helper/base.py b/trains_agent/helper/base.py index 79b1a10..4be2ff1 100644 --- a/trains_agent/helper/base.py +++ b/trains_agent/helper/base.py @@ -554,6 +554,7 @@ class ExecutionInfo(NonStrictAttrs): branch = nullable_string version_num = nullable_string tag = nullable_string + docker_cmd = nullable_string @classmethod def from_task(cls, task_info): @@ -571,6 +572,12 @@ class ExecutionInfo(NonStrictAttrs): execution.entry_point = entry_point execution.working_dir = working_dir or "" + # noinspection PyBroadException + try: + execution.docker_cmd = task_info.execution.docker_cmd + except Exception: + pass + return execution diff --git a/trains_agent/helper/package/conda_api.py b/trains_agent/helper/package/conda_api.py index 97c420b..8ea0615 100644 --- a/trains_agent/helper/package/conda_api.py +++ b/trains_agent/helper/package/conda_api.py @@ -19,13 +19,14 @@ from trains_agent.external.requirements_parser import parse from trains_agent.external.requirements_parser.requirement import Requirement from trains_agent.errors import CommandFailedError -from trains_agent.helper.base import rm_tree, NonStrictAttrs, select_for_platform, is_windows_platform +from trains_agent.helper.base import rm_tree, NonStrictAttrs, select_for_platform, is_windows_platform, ExecutionInfo from trains_agent.helper.process import Argv, Executable, DEVNULL, CommandSequence, PathLike from trains_agent.helper.package.requirements import SimpleVersion from trains_agent.session import Session from .base import PackageManager from .pip_api.venv import VirtualenvPip from .requirements import RequirementsManager, MarkerRequirement +from ...backend_api.session.defs import ENV_CONDA_ENV_PACKAGE package_normalize = partial(re.compile(r"""\[version=['"](.*)['"]\]""").sub, r"\1") @@ -41,8 +42,8 @@ def _package_diff(path, packages): class CondaPip(VirtualenvPip): def __init__(self, source=None, *args, **kwargs): - super(CondaPip, self).__init__(*args, interpreter=Path(kwargs.get('path'), "python.exe") \ - if is_windows_platform() and kwargs.get('path') else None, **kwargs) + super(CondaPip, self).__init__(*args, interpreter=Path(kwargs.get('path'), "python.exe") + if is_windows_platform() and kwargs.get('path') else None, **kwargs) self.source = source def run_with_env(self, command, output=False, **kwargs): @@ -62,8 +63,8 @@ class CondaAPI(PackageManager): MINIMUM_VERSION = "4.3.30" - def __init__(self, session, path, python, requirements_manager): - # type: (Session, PathLike, float, RequirementsManager) -> None + def __init__(self, session, path, python, requirements_manager, execution_info=None, **kwargs): + # type: (Session, PathLike, float, RequirementsManager, ExecutionInfo, Any) -> None """ :param python: base python version to use (e.g python3.6) :param path: path of env @@ -74,6 +75,13 @@ class CondaAPI(PackageManager): self.requirements_manager = requirements_manager self.path = path self.extra_channels = self.session.config.get('agent.package_manager.conda_channels', []) + self.conda_env_as_base_docker = \ + self.session.config.get('agent.package_manager.conda_env_as_base_docker', None) or \ + bool(ENV_CONDA_ENV_PACKAGE.get()) + if ENV_CONDA_ENV_PACKAGE.get(): + self.conda_pre_build_env_path = ENV_CONDA_ENV_PACKAGE.get() + else: + self.conda_pre_build_env_path = execution_info.docker_cmd if execution_info else None self.pip = CondaPip( session=self.session, source=self.source, @@ -81,10 +89,14 @@ class CondaAPI(PackageManager): requirements_manager=self.requirements_manager, path=self.path, ) - self.conda = ( - find_executable("conda") - or Argv(select_for_platform(windows="where", linux="which"), "conda").get_output(shell=True).strip() - ) + try: + self.conda = ( + find_executable("conda") + or Argv(select_for_platform(windows="where", linux="which"), "conda").get_output(shell=True).strip() + ) + except Exception: + raise ValueError("ERROR: package manager \"conda\" selected, " + "but \'conda\' executable could not be located") try: output = Argv(self.conda, "--version").get_output(stderr=subprocess.STDOUT) except subprocess.CalledProcessError as ex: @@ -136,11 +148,25 @@ class CondaAPI(PackageManager): if match else ("activate", self.path) ) + + if self.conda_env_as_base_docker and self.conda_pre_build_env_path: + print("Restoring Conda environment from {}".format(self.conda_pre_build_env_path)) + tar_path = find_executable("tar") + self.path.mkdir(parents=True, exist_ok=True) + output = Argv( + tar_path, + "-xzf", + self.conda_pre_build_env_path, + "-C", + self.path, + ).get_output() + conda_env = Path(self.conda).parent.parent / 'etc' / 'profile.d' / 'conda.sh' if conda_env.is_file() and not is_windows_platform(): self.source = self.pip.source = CommandSequence(('source', conda_env.as_posix()), self.source) # install cuda toolkit + # noinspection PyBroadException try: cuda_version = float(int(self.session.config['agent.cuda_version'])) / 10.0 if cuda_version > 0: diff --git a/trains_agent/helper/package/pip_api/venv.py b/trains_agent/helper/package/pip_api/venv.py index dc97b01..374cb62 100644 --- a/trains_agent/helper/package/pip_api/venv.py +++ b/trains_agent/helper/package/pip_api/venv.py @@ -1,6 +1,8 @@ +from typing import Any + from pathlib2 import Path -from trains_agent.helper.base import select_for_platform, rm_tree +from trains_agent.helper.base import select_for_platform, rm_tree, ExecutionInfo from trains_agent.helper.package.base import PackageManager from trains_agent.helper.process import Argv, PathLike from trains_agent.session import Session @@ -9,8 +11,8 @@ from ..requirements import RequirementsManager class VirtualenvPip(SystemPip, PackageManager): - def __init__(self, session, python, requirements_manager, path, interpreter=None): - # type: (Session, float, RequirementsManager, PathLike, PathLike) -> () + def __init__(self, session, python, requirements_manager, path, interpreter=None, execution_info=None, **kwargs): + # type: (Session, float, RequirementsManager, PathLike, PathLike, ExecutionInfo, Any) -> () """ Program interface to virtualenv pip. Must be given either path to virtualenv or source command.