Patch retrieve requests for remote models

ecmwf-lab · Mar 18, 2024 · 5213785 · 5213785
1 parent 89138c1
commit 5213785
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/ai_models/remote.py b/ai_models/remote.py
@@ -50,6 +50,10 @@ def run(self):
     def parse_model_args(self, args):
         return None
 
+    def patch_retrieve_request(self, request):
+        patched = self.api.patch_retrieve_request(self.cfg, request)
+        request.update(patched)
+
     def load_parameters(self):
         params = self.api.metadata(
             self.model,
@@ -221,6 +225,14 @@ def models(self):
 
         return results
 
+    def patch_retrieve_request(self, cfg, request):
+        cfg["patchrequest"] = request
+        result = self._request(requests.post, "patch", json=cfg)
+        if status := result.get("status"):
+            LOG.error(status)
+            sys.exit(1)
+        return result
+
     def _request(self, type, href, data=None, json=None, auth=None):
         response = robust(type, retry_after=30)(
             urljoin(self.url, href),