extend patch support for truss dir patch applier (#553)

* extend patch support for truss dir patch applier

* bump version

* add missing continue

* correct version

* always overwrite config when applicable

* add tests, remove redundant update

Author: joostinyi

pyproject.toml

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

truss/patch/truss_dir_patch_applier.py

@@ -1,5 +1,4 @@
 import logging
-from dataclasses import replace
 from pathlib import Path
 from typing import List
 
@@ -16,6 +15,9 @@
 )
 from truss.templates.control.control.helpers.types import (
     Action,
+    ConfigPatch,
+    EnvVarPatch,
+    ExternalDataPatch,
     ModelCodePatch,
     Patch,
     PythonRequirementPatch,
@@ -40,6 +42,7 @@ def __call__(self, patches: List[Patch]):
         # Aggregate config patches and apply at end
         reqs = reqs_by_name(self._truss_config.requirements)
         pkgs = system_packages_set(self._truss_config.system_packages)
+        new_config = self._truss_config
         for patch in patches:
             self._logger.debug(f"Applying patch {patch.to_dict()}")
             action = patch.body.action
@@ -67,11 +70,15 @@ def __call__(self, patches: List[Patch]):
                 if action == Action.ADD or Action.UPDATE:
                     pkgs.add(pkg)
                     continue
+            # Each of EnvVarPatch and ExternalDataPatch can be expressed through an overwrite of the config,
+            # handled below
+            if isinstance(patch.body, EnvVarPatch):
+                continue
+            if isinstance(patch.body, ExternalDataPatch):
+                continue
+            if isinstance(patch.body, ConfigPatch):
+                new_config = TrussConfig.from_dict(patch.body.config)
+                continue
             raise UnsupportedPatch(f"Unknown patch type {patch.type}")
 
-        self._truss_config = replace(
-            self._truss_config,
-            requirements=list(reqs.values()),
-            system_packages=list(pkgs),
-        )
-        self._truss_config.write_to_yaml_file(self._truss_config_path)
+        new_config.write_to_yaml_file(self._truss_config_path)

truss/tests/patch/test_truss_dir_patch_applier.py

@@ -1,9 +1,11 @@
 import logging
 from pathlib import Path
 
+import yaml
 from truss.patch.truss_dir_patch_applier import TrussDirPatchApplier
 from truss.templates.control.control.helpers.types import (
     Action,
+    ConfigPatch,
     ModelCodePatch,
     Patch,
     PatchType,
@@ -33,6 +35,8 @@ def test_model_code_patch(custom_model_truss_dir: Path):
 def test_python_requirement_patch(custom_model_truss_dir: Path):
     req = "git+https://github.com/huggingface/transformers.git"
     applier = TrussDirPatchApplier(custom_model_truss_dir, TEST_LOGGER)
+    config = yaml.safe_load((custom_model_truss_dir / "config.yaml").read_text())
+    config["requirements"] = [req]
     applier(
         [
             Patch(
@@ -41,7 +45,13 @@ def test_python_requirement_patch(custom_model_truss_dir: Path):
                     action=Action.ADD,
                     requirement=req,
                 ),
-            )
+            ),
+            Patch(
+                type=PatchType.CONFIG,
+                body=ConfigPatch(
+                    action=Action.UPDATE, config=config, path="config.yaml"
+                ),
+            ),
         ]
     )
     assert TrussConfig.from_yaml(
@@ -51,15 +61,23 @@ def test_python_requirement_patch(custom_model_truss_dir: Path):
 
 def test_system_requirement_patch(custom_model_truss_dir: Path):
     applier = TrussDirPatchApplier(custom_model_truss_dir, TEST_LOGGER)
+    config = yaml.safe_load((custom_model_truss_dir / "config.yaml").read_text())
+    config["system_packages"] = ["curl"]
     applier(
         [
+            Patch(
+                type=PatchType.CONFIG,
+                body=ConfigPatch(
+                    action=Action.UPDATE, config=config, path="config.yaml"
+                ),
+            ),
             Patch(
                 type=PatchType.SYSTEM_PACKAGE,
                 body=SystemPackagePatch(
                     action=Action.ADD,
                     package="curl",
                 ),
-            )
+            ),
         ]
     )
     assert TrussConfig.from_yaml(