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

@@ -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", ]

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

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(

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
 
 

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:

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.