Reverted #231 (#252)

* Reverted basetenlabs/baseten#5830.

* Fix issue from merge.

* Bump truss version.

Author: squidarth

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.4.0"
+version = "0.4.1"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/tests/test_truss_handle.py

@@ -416,29 +416,6 @@ def test_update_requirements(custom_model_truss_dir_with_pre_and_post):
     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,7 +10,6 @@
 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
@@ -398,15 +397,11 @@ 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."""
 
-        # 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))
+        self._update_config(
+            lambda conf: replace(
+                conf, requirements=[*conf.requirements, python_requirement]
+            )
+        )
 
     def remove_python_requirement(self, python_requirement: str):
         """Remove a python requirement to truss model's config.
@@ -1111,8 +1106,3 @@ def _docker_image_from_labels(labels: Dict):
     images = get_images(labels)
     if images and isinstance(images, list):
         return images[0]
-
-
-def _python_req_name(python_requirement: str) -> str:
-    req = pkg_resources.Requirement.parse(python_requirement)
-    return req.name  # type: ignore