send truss version on patch (basetenlabs#1268)

Distillative-AI · Dec 6, 2024 · 00e7e18 · 00e7e18
1 parent 3531513
commit 00e7e18
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 37 deletions.
diff --git a/truss/remote/baseten/api.py b/truss/remote/baseten/api.py
@@ -432,7 +432,7 @@ def patch_draft_truss_two_step(self, model_name, patch_request):
         query_string = f"""
         mutation {{
         stage_patch_for_draft_truss(name: "{model_name}",
-                    client_version: "TRUSS",
+                    client_version: "{truss.version()}",
                     patch: "{patch}",
     ) {{
             id,
@@ -457,7 +457,7 @@ def sync_draft_truss(self, model_name):
         query_string = f"""
         mutation {{
         sync_draft_truss(name: "{model_name}",
-                    client_version: "TRUSS",
+                    client_version: "{truss.version()}",
     ) {{
             id,
             name,
@@ -474,29 +474,6 @@ def sync_draft_truss(self, model_name):
             logging.debug(f"Failed to sync patch: {result}")
         return result
 
-    def patch_draft_truss(self, model_name, patch_request):
-        patch = base64_encoded_json_str(patch_request.to_dict())
-        query_string = f"""
-        mutation {{
-        patch_draft_truss(name: "{model_name}",
-                    client_version: "TRUSS",
-                    patch: "{patch}",
-    ) {{
-            id,
-            name,
-            version_id
-            succeeded
-            needs_full_deploy
-            error
-        }}
-        }}
-        """
-        resp = self._post_graphql_query(query_string)
-        result = resp["data"]["patch_draft_truss"]
-        if not result["succeeded"]:
-            logging.debug(f"Unsuccessful response: {result}")
-        return result
-
     def get_deployment(self, model_id: str, deployment_id: str) -> Any:
         headers = self._auth_token.header()
         resp = requests.get(

diff --git a/truss/remote/baseten/remote.py b/truss/remote/baseten/remote.py
@@ -519,22 +519,20 @@ def _patch(
             return PatchResult(
                 PatchStatus.SKIPPED, "No changes observed, skipping patching."
             )
+
+        def do_patch():
+            if should_create_patch:
+                resp = self._api.patch_draft_truss_two_step(model_name, patch_request)
+            else:
+                resp = self._api.sync_draft_truss(model_name)
+            return resp
+
         try:
             if console:
                 with console.status("Applying patch..."):
-                    if should_create_patch:
-                        resp = self._api.patch_draft_truss_two_step(
-                            model_name, patch_request
-                        )
-                    else:
-                        resp = self._api.sync_draft_truss(model_name)
+                    resp = do_patch()
             else:
-                if should_create_patch:
-                    resp = self._api.patch_draft_truss_two_step(
-                        model_name, patch_request
-                    )
-                else:
-                    resp = self._api.sync_draft_truss(model_name)
+                resp = do_patch()
 
         except ReadTimeout:
             return PatchResult(