From 566427d5507387b2aea1b2adac49791e6193f655 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sun, 31 May 2020 14:02:42 +0300
Subject: [PATCH] Fix build failing due to missing session

---
 trains_agent/commands/worker.py               |  1 +
 trains_agent/helper/package/pip_api/system.py |  8 +++++---
 trains_agent/helper/package/pip_api/venv.py   | 10 ++++------
 trains_agent/helper/package/translator.py     |  2 +-
 4 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/trains_agent/commands/worker.py b/trains_agent/commands/worker.py
index 8e74e91..2012414 100644
--- a/trains_agent/commands/worker.py
+++ b/trains_agent/commands/worker.py
@@ -1887,6 +1887,7 @@ class Worker(ServiceCommandSection):
 
         global_package_manager_params = dict(
             interpreter=executable_name,
+            session=self._session,
         )
 
         if not self.is_conda and standalone_mode:
diff --git a/trains_agent/helper/package/pip_api/system.py b/trains_agent/helper/package/pip_api/system.py
index b5f47b4..856abff 100644
--- a/trains_agent/helper/package/pip_api/system.py
+++ b/trains_agent/helper/package/pip_api/system.py
@@ -1,22 +1,24 @@
 import sys
 from itertools import chain
-from typing import Text
+from typing import Text, Optional
 
 from trains_agent.definitions import PIP_EXTRA_INDICES, PROGRAM_NAME
 from trains_agent.helper.package.base import PackageManager
 from trains_agent.helper.process import Argv, DEVNULL
+from trains_agent.session import Session
 
 
 class SystemPip(PackageManager):
 
     indices_args = None
 
-    def __init__(self, interpreter=None):
-        # type: (Text) -> ()
+    def __init__(self, interpreter=None, session=None):
+        # type: (Optional[Text], Optional[Session]) -> ()
         """
         Program interface to the system pip.
         """
         self._bin = interpreter or sys.executable
+        self.session = session
 
     @property
     def bin(self):
diff --git a/trains_agent/helper/package/pip_api/venv.py b/trains_agent/helper/package/pip_api/venv.py
index ecb600f..f9c14c2 100644
--- a/trains_agent/helper/package/pip_api/venv.py
+++ b/trains_agent/helper/package/pip_api/venv.py
@@ -15,19 +15,17 @@ class VirtualenvPip(SystemPip, PackageManager):
         Program interface to virtualenv pip.
         Must be given either path to virtualenv or source command.
         Either way, ``self.source`` is exposed.
+        :param session: a Session object for communication
         :param python: interpreter path
         :param path: path of virtual environment to create/manipulate
         :param python: python version
         :param interpreter: path of python interpreter
         """
         super(VirtualenvPip, self).__init__(
-            interpreter
-            or Path(
-                path,
-                select_for_platform(linux="bin/python", windows="scripts/python.exe"),
-            )
+            session=session,
+            interpreter=interpreter or Path(
+                path, select_for_platform(linux="bin/python", windows="scripts/python.exe"))
         )
-        self.session = session
         self.path = path
         self.requirements_manager = requirements_manager
         self.python = python
diff --git a/trains_agent/helper/package/translator.py b/trains_agent/helper/package/translator.py
index e80a56e..4d594e2 100644
--- a/trains_agent/helper/package/translator.py
+++ b/trains_agent/helper/package/translator.py
@@ -22,7 +22,7 @@ class RequirementsTranslator(object):
         self.enabled = config["agent.pip_download_cache.enabled"]
         Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
         self.config = Config()
-        self.pip = SystemPip(interpreter=interpreter)
+        self.pip = SystemPip(interpreter=interpreter, session=self._session)
 
     def download(self, url):
         self.pip.download_package(url, cache_dir=self.cache_dir)