Add post_packages post_optional_packages to control packages installed after all the rest (e.g. horovod)

Rename CythonReq to PriorityPackageRequirement and HorovodReq to PostRequirement

docs/trains.conf

@@ -58,6 +58,22 @@ agent {
         # additional conda channels to use when installing with conda package manager
         conda_channels: ["pytorch", "conda-forge", ]
 
+        # set the priority packages to be installed before the rest of the required packages
+        # priority_packages: ["cython", "numpy", "setuptools", ]
+
+        # set the optional priority packages to be installed before the rest of the required packages,
+        # In case a package installation fails, the package will be ignored,
+        # and the virtual environment process will continue
+        # priority_optional_packages: ["pygobject", ]
+
+        # set the post packages to be installed after all the rest of the required packages
+        # post_packages: ["horovod", ]
+
+        # set the optional post packages to be installed after all the rest of the required packages,
+        # In case a package installation fails, the package will be ignored,
+        # and the virtual environment process will continue
+        # post_optional_packages: []
+
         # set to True to support torch nightly build installation,
         # notice: torch nightly builds are ephemeral and are deleted from time to time
         torch_nightly: false,

trains_agent/backend_api/config/default/agent.conf

@@ -44,6 +44,22 @@
         # additional conda channels to use when installing with conda package manager
         conda_channels: ["defaults", "conda-forge", "pytorch", ]
 
+        # set the priority packages to be installed before the rest of the required packages
+        # priority_packages: ["cython", "numpy", "setuptools", ]
+
+        # set the optional priority packages to be installed before the rest of the required packages,
+        # In case a package installation fails, the package will be ignored,
+        # and the virtual environment process will continue
+        # priority_optional_packages: ["pygobject", ]
+
+        # set the post packages to be installed after all the rest of the required packages
+        # post_packages: ["horovod", ]
+
+        # set the optional post packages to be installed after all the rest of the required packages,
+        # In case a package installation fails, the package will be ignored,
+        # and the virtual environment process will continue
+        # post_optional_packages: []
+
         # set to True to support torch nightly build installation,
         # notice: torch nightly builds are ephemeral and are deleted from time to time
         torch_nightly: false,

trains_agent/backend_api/session/session.py

@@ -434,16 +434,15 @@ class Session(TokenManager):
     @classmethod
     def get_api_server_host(cls, config=None):
         if not config:
-            from ...config import config_obj
+            return None
-            config = config_obj
+
         return ENV_HOST.get(default=(config.get("api.api_server", None) or
                                      config.get("api.host", None) or cls.default_host))
 
     @classmethod
     def get_app_server_host(cls, config=None):
         if not config:
-            from ...config import config_obj
+            return None
-            config = config_obj
 
         # get from config/environment
         web_host = ENV_WEB_HOST.get(default=config.get("api.web_server", None))
@@ -470,8 +469,8 @@ class Session(TokenManager):
     @classmethod
     def get_files_server_host(cls, config=None):
         if not config:
-            from ...config import config_obj
+            return None
-            config = config_obj
+
         # get from config/environment
         files_host = ENV_FILES_HOST.get(default=(config.get("api.files_server", None)))
         if files_host:

trains_agent/commands/worker.py

@@ -71,7 +71,7 @@ from trains_agent.helper.console import ensure_text, print_text, decode_binary_l
 from trains_agent.helper.os.daemonize import daemonize_process
 from trains_agent.helper.package.base import PackageManager
 from trains_agent.helper.package.conda_api import CondaAPI
-from trains_agent.helper.package.horovod_req import HorovodRequirement
+from trains_agent.helper.package.post_req import PostRequirement
 from trains_agent.helper.package.external_req import ExternalRequirements
 from trains_agent.helper.package.pip_api.system import SystemPip
 from trains_agent.helper.package.pip_api.venv import VirtualenvPip
@@ -91,7 +91,7 @@ from trains_agent.helper.process import (
     get_docker_id,
     commit_docker, terminate_process,
 )
-from trains_agent.helper.package.cython_req import CythonRequirement
+from trains_agent.helper.package.priority_req import PriorityPackageRequirement
 from trains_agent.helper.repo import clone_repository_cached, RepoInfo, VCS
 from trains_agent.helper.resource_monitor import ResourceMonitor
 from trains_agent.session import Session
@@ -303,8 +303,8 @@ class Worker(ServiceCommandSection):
 
     _requirement_substitutions = (
         PytorchRequirement,
-        CythonRequirement,
+        PriorityPackageRequirement,
-        HorovodRequirement,
+        PostRequirement,
         ExternalRequirements,
     )
 
@@ -1717,7 +1717,7 @@ class Worker(ServiceCommandSection):
         package_api.set_selected_package_manager()
         # always install cython,
         # if we have a specific version in the requirements,
-        # the CythonRequirement(SimpleSubstitution) will reinstall cython with the specific version
+        # the PriorityPackageRequirement(SimpleSubstitution) will reinstall cython with the specific version
         if not self.is_conda:
             package_api.out_of_scope_install_package('Cython')
 

trains_agent/helper/package/cython_req.py

@@ -1,25 +0,0 @@
-from typing import Text
-
-from .base import PackageManager
-from .requirements import SimpleSubstitution
-
-
-class CythonRequirement(SimpleSubstitution):
-
-    name = ("cython", "numpy", )
-
-    def __init__(self, *args, **kwargs):
-        super(CythonRequirement, self).__init__(*args, **kwargs)
-
-    def match(self, req):
-        # match both Cython & cython
-        return req.name and req.name.lower() in self.name
-
-    def replace(self, req):
-        """
-        Replace a requirement
-        :raises: ValueError if version is pre-release
-        """
-        # install Cython before
-        PackageManager.out_of_scope_install_package(str(req))
-        return Text(req)

trains_agent/helper/package/horovod_req.py

@@ -1,32 +0,0 @@
-from typing import Text
-
-from .base import PackageManager
-from .requirements import SimpleSubstitution
-
-
-class HorovodRequirement(SimpleSubstitution):
-
-    name = "horovod"
-
-    def __init__(self, *args, **kwargs):
-        super(HorovodRequirement, self).__init__(*args, **kwargs)
-        self.post_install_req = None
-
-    def match(self, req):
-        # match both horovod
-        return req.name and self.name == req.name.lower()
-
-    def post_install(self, session):
-        if self.post_install_req:
-            PackageManager.out_of_scope_install_package(self.post_install_req.tostr(markers=False))
-            self.post_install_req = None
-
-    def replace(self, req):
-        """
-        Replace a requirement
-        :raises: ValueError if version is pre-release
-        """
-        # Store in post req install, and return nothing
-        self.post_install_req = req
-        # mark skip package, we will install it in post install hook
-        return Text('')

trains_agent/helper/package/post_req.py

@@ -0,0 +1,48 @@
+from typing import Text
+
+from .base import PackageManager
+from .requirements import SimpleSubstitution
+
+
+class PostRequirement(SimpleSubstitution):
+
+    name = ("horovod", )
+    optional_package_names = tuple()
+
+    def __init__(self, *args, **kwargs):
+        super(PostRequirement, self).__init__(*args, **kwargs)
+        self.post_install_req = []
+        # check if we need to replace the packages:
+        post_packages = self.config.get('agent.package_manager.post_packages', None)
+        if post_packages:
+            self.__class__.name = post_packages
+        post_optional_packages = self.config.get('agent.package_manager.post_optional_packages', None)
+        if post_optional_packages:
+            self.__class__.optional_package_names = post_optional_packages
+
+    def match(self, req):
+        # match both horovod
+        return req.name and (req.name.lower() in self.name or req.name.lower() in self.optional_package_names)
+
+    def post_install(self, session):
+        for req in self.post_install_req:
+            if req.name in self.optional_package_names:
+                # noinspection PyBroadException
+                try:
+                    PackageManager.out_of_scope_install_package(req.tostr(markers=False))
+                except Exception:
+                    pass
+            else:
+                PackageManager.out_of_scope_install_package(req.tostr(markers=False))
+
+        self.post_install_req = []
+
+    def replace(self, req):
+        """
+        Replace a requirement
+        :raises: ValueError if version is pre-release
+        """
+        # Store in post req install, and return nothing
+        self.post_install_req.append(req)
+        # mark skip package, we will install it in post install hook
+        return Text('')

trains_agent/helper/package/priority_req.py

@@ -0,0 +1,40 @@
+from typing import Text
+
+from .base import PackageManager
+from .requirements import SimpleSubstitution
+
+
+class PriorityPackageRequirement(SimpleSubstitution):
+
+    name = ("cython", "numpy", "setuptools", )
+    optional_package_names = tuple()
+
+    def __init__(self, *args, **kwargs):
+        super(PriorityPackageRequirement, self).__init__(*args, **kwargs)
+        # check if we need to replace the packages:
+        priority_packages = self.config.get('agent.package_manager.priority_packages', None)
+        if priority_packages:
+            self.__class__.name = priority_packages
+        priority_optional_packages = self.config.get('agent.package_manager.priority_optional_packages', None)
+        if priority_optional_packages:
+            self.__class__.optional_package_names = priority_optional_packages
+
+    def match(self, req):
+        # match both Cython & cython
+        return req.name and (req.name.lower() in self.name or req.name.lower() in self.optional_package_names)
+
+    def replace(self, req):
+        """
+        Replace a requirement
+        :raises: ValueError if version is pre-release
+        """
+        if req.name in self.optional_package_names:
+            # noinspection PyBroadException
+            try:
+                if PackageManager.out_of_scope_install_package(str(req)):
+                    return Text(req)
+            except Exception:
+                pass
+            return Text('')
+        PackageManager.out_of_scope_install_package(str(req))
+        return Text(req)