From 284271c65487891cfec21d4f5ae46127802ede34 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Wed, 22 Jan 2020 12:02:12 +0200
Subject: [PATCH] Support limiting pip version, limit to <20 by default

---
 docs/trains.conf                                |  4 ++++
 .../backend_api/config/default/agent.conf       |  5 ++++-
 trains_agent/commands/worker.py                 |  6 ++++--
 trains_agent/helper/package/base.py             | 17 ++++++++++++++++-
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/docs/trains.conf b/docs/trains.conf
index 13e215e..59c6c26 100644
--- a/docs/trains.conf
+++ b/docs/trains.conf
@@ -40,6 +40,10 @@ agent {
     package_manager: {
         # supported options: pip, conda
         type: pip,
+
+        # specify pip version to use (examples "<20", "==19.3.1", "", empty string will install the latest version)
+        # pip_version: "<20"
+
         # virtual environment inheres packages from system
         system_site_packages: false,
         # install with --upgrade
diff --git a/trains_agent/backend_api/config/default/agent.conf b/trains_agent/backend_api/config/default/agent.conf
index 6a4a43f..d8cb404 100644
--- a/trains_agent/backend_api/config/default/agent.conf
+++ b/trains_agent/backend_api/config/default/agent.conf
@@ -22,9 +22,12 @@
     # currently supported pip and conda
     # poetry is used if pip selected and repository contains poetry.lock file
     package_manager: {
-        # supported options: pip, conda
+        # supported options: pip, conda, poetry
         type: pip,
 
+        # specify pip version to use (examples "<20", "==19.3.1", "", empty string will install the latest version)
+        pip_version: "<20",
+
         # virtual environment inheres packages from system
         system_site_packages: false,
 
diff --git a/trains_agent/commands/worker.py b/trains_agent/commands/worker.py
index 8550153..f419e49 100644
--- a/trains_agent/commands/worker.py
+++ b/trains_agent/commands/worker.py
@@ -355,6 +355,7 @@ class Worker(ServiceCommandSection):
         self.docker_image_func = None
         self._docker_image = None
         self._docker_arguments = None
+        PackageManager.set_pip_version(self._session.config.get("agent.package_manager.pip_version", None))
         self._extra_docker_arguments = self._session.config.get("agent.extra_docker_arguments", None)
         self._extra_shell_script = self._session.config.get("agent.extra_docker_shell_script", None)
         self._docker_force_pull = self._session.config.get("agent.docker_force_pull", False)
@@ -1867,10 +1868,11 @@ class Worker(ServiceCommandSection):
                 "chown -R root /root/.cache/pip ; " \
                 "apt-get update ; " \
                 "apt-get install -y git libsm6 libxext6 libxrender-dev libglib2.0-0 {python_single_digit}-pip ; " \
-                "{python} -m pip install -U pip ; " \
+                "{python} -m pip install -U \"pip{pip_version}\" ; " \
                 "{python} -m pip install -U trains-agent{specify_version} ; ".format(
                     python_single_digit=python_version.split('.')[0],
-                    python=python_version, specify_version=specify_version)
+                    python=python_version, pip_version=PackageManager.get_pip_version(),
+                    specify_version=specify_version)
 
         base_cmd += [
             '-v', conf_file+':/root/trains.conf',
diff --git a/trains_agent/helper/package/base.py b/trains_agent/helper/package/base.py
index 5296722..af6d383 100644
--- a/trains_agent/helper/package/base.py
+++ b/trains_agent/helper/package/base.py
@@ -17,6 +17,7 @@ class PackageManager(object):
 
     _selected_manager = None
     _cwd = None
+    _pip_version = None
 
     @abc.abstractproperty
     def bin(self):
@@ -65,7 +66,7 @@ class PackageManager(object):
         pass
 
     def upgrade_pip(self):
-        return self._install("pip", "--upgrade")
+        return self._install("pip"+self.get_pip_version(), "--upgrade")
 
     def get_python_command(self, extra=()):
         # type: (...) -> Executable
@@ -123,3 +124,17 @@ class PackageManager(object):
             except Exception:
                 pass
         return []
+
+    @classmethod
+    def set_pip_version(cls, version):
+        if not version:
+            return
+        version = version.replace(' ', '')
+        if ('=' in version) or ('~' in version) or ('<' in version) or ('>' in version):
+            cls._pip_version = version
+        else:
+            cls._pip_version = "=="+version
+
+    @classmethod
+    def get_pip_version(cls):
+        return cls._pip_version or ''