Fix requirements handling and poetry support

trains_agent/commands/worker.py

@@ -64,6 +64,7 @@ from trains_agent.helper.console import ensure_text
 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.external_req import ExternalRequirements
 from trains_agent.helper.package.pip_api.system import SystemPip
 from trains_agent.helper.package.pip_api.venv import VirtualenvPip
 from trains_agent.helper.package.poetry_api import PoetryConfig, PoetryAPI
@@ -287,6 +288,7 @@ class Worker(ServiceCommandSection):
         PytorchRequirement,
         CythonRequirement,
         HorovodRequirement,
+        ExternalRequirements,
     )
 
     # poll queues every _polling_interval seconds
@@ -350,7 +352,6 @@ class Worker(ServiceCommandSection):
 
         self.is_venv_update = self._session.config.agent.venv_update.enabled
         self.poetry = PoetryConfig(self._session)
-        self.poetry.initialize()
         self.docker_image_func = None
         self._docker_image = None
         self._docker_arguments = None
@@ -934,7 +935,7 @@ class Worker(ServiceCommandSection):
             cached_requirements=requirements,
             cwd=vcs.location if vcs and vcs.location else directory,
         )
-        freeze = self.freeze_task_environment()
+        freeze = self.freeze_task_environment(requirements_manager=requirements_manager)
         script_dir = directory
 
         # Summary
@@ -1108,7 +1109,8 @@ class Worker(ServiceCommandSection):
 
         # do not update the task packages if we are using conda,
         # it will most likely make the task environment unreproducible
-        freeze = self.freeze_task_environment(current_task.id if not self.is_conda else None)
+        freeze = self.freeze_task_environment(current_task.id if not self.is_conda else None,
+                                              requirements_manager=requirements_manager)
         script_dir = (directory if isinstance(directory, Path) else Path(directory)).absolute().as_posix()
 
         # run code
@@ -1371,13 +1373,17 @@ class Worker(ServiceCommandSection):
                 status_message=self._task_status_change_message,
             )
 
-    def freeze_task_environment(self, task_id=None):
+    def freeze_task_environment(self, task_id=None, requirements_manager=None):
         try:
             freeze = self.package_api.freeze()
         except Exception as e:
             print("Could not freeze installed packages")
             self.log_traceback(e)
             return None
+
+        if requirements_manager:
+            freeze = requirements_manager.replace_back(freeze)
+
         if not task_id:
             return freeze
 
@@ -1402,6 +1408,9 @@ class Worker(ServiceCommandSection):
         if not repo_info:
             return None
         try:
+            if not self.poetry.enabled:
+                return None
+            self.poetry.initialize(cwd=repo_info.root)
             api = self.poetry.get_api(repo_info.root)
             if api.enabled:
                 api.install()
@@ -1447,7 +1456,7 @@ class Worker(ServiceCommandSection):
             except Exception as e:
                 self.log_traceback(e)
                 cached_requirements_failed = True
-                raise ValueError("Could not install task requirements!")
+                raise ValueError("Could not install task requirements!\n{}".format(e))
             else:
                 self.log("task requirements installation passed")
                 return

trains_agent/helper/package/base.py

@@ -107,10 +107,19 @@ class PackageManager(object):
         self._cwd = value
 
     @classmethod
-    def out_of_scope_install_package(cls, package_name):
+    def out_of_scope_install_package(cls, package_name, *args):
         if PackageManager._selected_manager is not None:
             try:
-                return PackageManager._selected_manager._install(package_name)
+                return PackageManager._selected_manager._install(package_name, *args)
             except Exception:
                 pass
-        return
+        return
+
+    @classmethod
+    def out_of_scope_freeze(cls):
+        if PackageManager._selected_manager is not None:
+            try:
+                return PackageManager._selected_manager.freeze()
+            except Exception:
+                pass
+        return []

trains_agent/helper/package/external_req.py

@@ -0,0 +1,60 @@
+from collections import OrderedDict
+from typing import Text
+
+from .base import PackageManager
+from .requirements import SimpleSubstitution
+
+
+class ExternalRequirements(SimpleSubstitution):
+
+    name = "external_link"
+
+    def __init__(self, *args, **kwargs):
+        super(ExternalRequirements, self).__init__(*args, **kwargs)
+        self.post_install_req = []
+        self.post_install_req_lookup = OrderedDict()
+
+    def match(self, req):
+        # match both editable or code or unparsed
+        if not (not req.name or req.req and (req.req.editable or req.req.vcs)):
+            return False
+        if not req.req or not req.req.line or not req.req.line.strip() or req.req.line.strip().startswith('#'):
+            return False
+        return True
+
+    def post_install(self):
+        post_install_req = self.post_install_req
+        self.post_install_req = []
+        for req in post_install_req:
+            try:
+                freeze_base = PackageManager.out_of_scope_freeze() or ''
+            except:
+                freeze_base = ''
+            PackageManager.out_of_scope_install_package(req.tostr(markers=False), "--no-deps")
+            try:
+                freeze_post = PackageManager.out_of_scope_freeze() or ''
+                package_name = list(set(freeze_post['pip']) - set(freeze_base['pip']))
+                if package_name and package_name[0] not in self.post_install_req_lookup:
+                    self.post_install_req_lookup[package_name[0]] = req.req.line
+            except:
+                pass
+            PackageManager.out_of_scope_install_package(req.tostr(markers=False), "--ignore-installed")
+
+    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('')
+
+    def replace_back(self, list_of_requirements):
+        if 'pip' in list_of_requirements:
+            original_requirements = list_of_requirements['pip']
+            list_of_requirements['pip'] = [r for r in original_requirements
+                                           if r not in self.post_install_req_lookup]
+            list_of_requirements['pip'] += [self.post_install_req_lookup.get(r, '')
+                                            for r in self.post_install_req_lookup.keys() if r in original_requirements]
+        return list_of_requirements

trains_agent/helper/package/pip_api/system.py

@@ -35,7 +35,7 @@ class SystemPip(PackageManager):
         self._install(*(packages + self.install_flags()))
 
     def _install(self, *args):
-        self.run_with_env(('install',) + args)
+        self.run_with_env(('install',) + args, cwd=self.cwd)
 
     def uninstall_packages(self, *packages):
         self.run_with_env(('uninstall', '-y') + packages)

trains_agent/helper/package/poetry_api.py

@@ -1,6 +1,7 @@
 from functools import wraps
 
 import attr
+import sys
 from pathlib2 import Path
 from trains_agent.helper.process import Argv, DEVNULL
 from trains_agent.session import Session, POETRY
@@ -35,10 +36,12 @@ def prop_guard(prop, log_prop=None):
 
 class PoetryConfig:
 
-    def __init__(self, session):
+    def __init__(self, session, interpreter=None):
-        # type: (Session) -> ()
+        # type: (Session, str) -> ()
         self.session = session
         self._log = session.get_logger(__name__)
+        self._python = interpreter or sys.executable
+        self._initialized = False
 
     @property
     def log(self):
@@ -53,7 +56,7 @@ class PoetryConfig:
     def run(self, *args, **kwargs):
         func = kwargs.pop("func", Argv.get_output)
         kwargs.setdefault("stdin", DEVNULL)
-        argv = Argv("poetry", "-n", *args)
+        argv = Argv(self._python, "-m", "poetry", *args)
         self.log.debug("running: %s", argv)
         return func(argv, **kwargs)
 
@@ -61,10 +64,12 @@ class PoetryConfig:
         return self.run("config", *args, **kwargs)
 
     @_guard_enabled
-    def initialize(self):
+    def initialize(self, cwd=None):
-        self._config("settings.virtualenvs.in-project",  "true")
+        if not self._initialized:
-        # self._config("repositories.{}".format(self.REPO_NAME), PYTHON_INDEX)
+            self._initialized = True
-        # self._config("http-basic.{}".format(self.REPO_NAME), *PYTHON_INDEX_CREDENTIALS)
+            self._config("--local", "virtualenvs.in-project",  "true", cwd=cwd)
+            # self._config("repositories.{}".format(self.REPO_NAME), PYTHON_INDEX)
+            # self._config("http-basic.{}".format(self.REPO_NAME), *PYTHON_INDEX_CREDENTIALS)
 
     def get_api(self, path):
         # type: (Path) -> PoetryAPI
@@ -81,7 +86,7 @@ class PoetryAPI(object):
     def install(self):
         # type: () -> bool
         if self.enabled:
-            self.config.run("install", cwd=str(self.path), func=Argv.check_call)
+            self.config.run("install", "-n", cwd=str(self.path), func=Argv.check_call)
             return True
         return False
 
@@ -92,7 +97,9 @@ class PoetryAPI(object):
         )
 
     def freeze(self):
-        return {"poetry": self.config.run("show", cwd=str(self.path)).splitlines()}
+        lines = self.config.run("show", cwd=str(self.path)).splitlines()
+        lines = [[p for p in line.split(' ') if p] for line in lines]
+        return {"pip": [parts[0]+'=='+parts[1]+' # '+' '.join(parts[2:]) for parts in lines]}
 
     def get_python_command(self, extra):
         return Argv("poetry", "run", "python", *extra)

trains_agent/helper/package/requirements.py

@@ -8,7 +8,7 @@ from copy import deepcopy
 from itertools import chain, starmap
 from operator import itemgetter
 from os import path
-from typing import Text, List, Type, Optional, Tuple
+from typing import Text, List, Type, Optional, Tuple, Dict
 
 from packaging import version as packaging_version
 from pathlib2 import Path
@@ -184,6 +184,13 @@ class SimpleSubstitution(RequirementSubstitution):
         req.specs = [('==', version_number + self.suffix)]
         return Text(req)
 
+    def replace_back(self, list_of_requirements):  # type: (Dict) -> Dict
+        """
+        :param list_of_requirements: {'pip': ['a==1.0', ]}
+        :return: {'pip': ['a==1.0', ]}
+        """
+        return list_of_requirements
+
 
 @six.add_metaclass(ABCMeta)
 class CudaSensitiveSubstitution(SimpleSubstitution):
@@ -235,15 +242,17 @@ class RequirementsManager(object):
         return None
 
     def replace(self, requirements):  # type: (Text) -> Text
+        def safe_parse(req_str):
+            try:
+                return next(parse(req_str))
+            except Exception as ex:
+                return Requirement(req_str)
+
         parsed_requirements = tuple(
             map(
                 MarkerRequirement,
-                filter(
+                [safe_parse(line) for line in (requirements.splitlines()
-                    None,
+                                               if isinstance(requirements, six.text_type) else requirements)]
-                    parse(requirements)
-                    if isinstance(requirements, six.text_type)
-                    else (next(parse(line), None) for line in requirements)
-                )
             )
         )
         if not parsed_requirements:
@@ -280,6 +289,14 @@ class RequirementsManager(object):
             except Exception as ex:
                 print('RequirementsManager handler {} raised exception: {}'.format(h, ex))
 
+    def replace_back(self, requirements):
+        for h in self.handlers:
+            try:
+                requirements = h.replace_back(requirements)
+            except Exception:
+                pass
+        return requirements
+
     @staticmethod
     def get_cuda_version(config):  # type: (ConfigTree) -> (Text, Text)
         # we assume os.environ already updated the config['agent.cuda_version'] & config['agent.cudnn_version']