add backoff, use direct requests instead of inference client

hf_calculator.py

@@ -6,17 +6,19 @@ LICENSE file in the root directory of this source tree.
 """
 
 import hashlib
-import json
 import os
 from pathlib import Path
 
 import ase
+import backoff
 import gradio as gr
 import huggingface_hub as hf_hub
+import requests
 from ase.calculators.calculator import Calculator
 from ase.db.core import now
 from ase.db.row import AtomsRow
 from ase.io.jsonio import decode, encode
+from requests.exceptions import HTTPError
 
 
 def hash_save_file(atoms: ase.Atoms, task_name, path: Path | str):
@@ -65,37 +67,56 @@ class HFEndpointCalculator(Calculator):
                     "You need to log in to HF and have gated model access to UMA before running your own simulations!"
                 )
 
-        self.client = hf_hub.InferenceClient(
-            model=endpoint_url, token=os.environ["HF_TOKEN"]
-        )
+        self.url = endpoint_url
+        self.token = os.environ["HF_TOKEN"]
         self.atoms = atoms
         self.task_name = task_name
 
         super().__init__(*args, **kwargs)
 
+    @staticmethod
+    @backoff.on_exception(
+        backoff.expo,
+        (requests.exceptions.RequestException,),
+        max_tries=10,
+        jitter=backoff.full_jitter,
+    )
+    def _post_with_backoff(url, headers, payload):
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        return response
+
     def calculate(self, atoms, properties, system_changes):
         Calculator.calculate(self, atoms, properties, system_changes)
 
         task_name = self.task_name.lower()
 
-        # Run inference via a post request to the endpoint
+        payload = {
+            "inputs": atoms_to_json(atoms, data=atoms.info),
+            "properties": properties,
+            "system_changes": system_changes,
+            "task_name": task_name,
+        }
+
+        headers = {
+            "Accept": "application/json",
+            "Authorization": f"Bearer {self.token}",
+            "Content-Type": "application/json",
+        }
+
+        print(payload)
+
         try:
-            response = self.client.post(
-                json={
-                    "inputs": atoms_to_json(atoms, data=atoms.info),
-                    "properties": properties,
-                    "system_changes": system_changes,
-                    "task_name": task_name,
-                }
-            )
-        except hf_hub.errors.BadRequestError:
+            response = self._post_with_backoff(self.url, headers, payload)
+            response_dict = response.json()
+        except HTTPError as error:
             hash_save_file(atoms, task_name, "/data/custom_inputs/errors/")
             raise gr.Error(
-                "Backend failure during your calculation; if you have continued issues please file an issue in the main FAIR chemistry repo (https://github.com/facebookresearch/fairchem)."
+                f"Backend failure during your calculation; if you have continued issues please file an issue in the main FAIR chemistry repo (https://github.com/facebookresearch/fairchem).\n{error}"
             )
 
         # Load the response and store the results in the calc and atoms object
-        response_dict = decode(json.loads(response))
+        response_dict = decode(response_dict)
         self.results = response_dict["results"]
         atoms.info = response_dict["info"]
 
@@ -118,6 +139,8 @@ def atoms_to_json(atoms, data=None):
 
     if data:
         dct["data"] = data
+    else:
+        dct["data"] = {}
 
     constraints = row.get("constraints")
     if constraints:

requirements.txt

@@ -1,6 +1,5 @@
 gradio<5.30.0
 numpy
 ase
-
-# Pinned since huggingface_hub>=0.31 does not support custom inference endpoint posts
-huggingface_hub<0.31
+huggingface_hub
+backoff