add_python_req modified to handle duplicates, change the version if exists (#231)

Author: tugrulguner

truss/tests/test_truss_handle.py

@@ -427,6 +427,21 @@ def test_update_requirements(custom_model_truss_dir_with_pre_and_post):
     sc_requirements = th.spec.requirements
     assert sc_requirements == requirements
 
+def test_add_python_requirements_new(custom_model_truss_dir_with_pre_and_post):
+    th = TrussHandle(custom_model_truss_dir_with_pre_and_post)
+    python_package_to_add = "tensorflow==2.3.1"
+    th.add_python_requirement(python_package_to_add)
+    sc_requirements = th.spec.requirements
+    assert python_package_to_add in sc_requirements
+
+def test_add_python_requirements_version_change(custom_model_truss_dir_with_pre_and_post):
+    th = TrussHandle(custom_model_truss_dir_with_pre_and_post)
+    th.add_python_requirement("tensorflow==2.3.1")
+    
+    python_package_to_change_version = "tensorflow==2.3.5"
+    th.add_python_requirement(python_package_to_change_version)
+    sc_requirements = th.spec.requirements
+    assert python_package_to_change_version.split('==')[1]==sc_requirements[0].split('==')[1]
 
 def test_update_requirements_from_file(
     custom_model_truss_dir_with_pre_and_post, tmp_path

truss/truss_handle.py

@@ -10,6 +10,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from urllib.error import HTTPError
 
+import pkg_resources
 import requests
 import yaml
 from requests import exceptions
@@ -79,6 +80,10 @@
     logger.addHandler(logging.StreamHandler(sys.stdout))
 
 
+def _python_req_name(python_requirement: str) -> str:
+    return pkg_resources.Requirement.parse(python_requirement).name
+
+
 class TrussHandle:
     def __init__(self, truss_dir: Path, validate: bool = True) -> None:
         self._truss_dir = truss_dir
@@ -106,7 +111,6 @@ def spec(self) -> TrussSpec:
     def _wait_for_predict(
         model_base_url: str, request: Dict, binary: bool = False
     ) -> Response:
-
         url = f"{model_base_url}/v1/models/model:predict"
 
         if binary:
@@ -397,11 +401,15 @@ def training_docker_build_setup(self, build_dir: Optional[Path] = None):
 
     def add_python_requirement(self, python_requirement: str):
         """Add a python requirement to truss model's config."""
-        self._update_config(
-            lambda conf: replace(
-                conf, requirements=[*conf.requirements, python_requirement]
-            )
-        )
+
+        # Parse the added python requirements
+        input_python_req_name = _python_req_name(python_requirement)
+        new_reqs = [
+            req for req in self._spec.config.requirements
+                if _python_req_name(req) != input_python_req_name
+        ]
+        new_reqs.append(python_requirement)
+        self._update_config(lambda conf: replace(conf, requirements=new_reqs))
 
     def remove_python_requirement(self, python_requirement: str):
         """Remove a python requirement to truss model's config.