Add num dpus parameter to query offline API

fennel/CHANGELOG.md

@@ -1,4 +1,6 @@
 # Changelog
+## [1.5.78] - 2024-11-12
+- Add support for num_dpus parameter in offline query
 
 ## [1.5.77] - 2025-02-13
 - Revert join on optional fields in LHS

fennel/client/client.py

@@ -459,6 +459,7 @@ def query_offline(
         feature_to_column_map: Optional[Dict[Feature, str]] = None,
         workflow: Optional[str] = None,
         use_v2: bool = True,
+        num_dpus: Optional[int] = None,
     ) -> Dict[str, Any]:
         """Extract point in time correct values of output features.
 
@@ -498,6 +499,9 @@ def query_offline(
             workflow param is equivalent to tagging (ex: fraud, finance etc.).
              If not provided, the "default" value will be used.
 
+        num_dpus (Optional[int]): Number of glue DPUs to use while executing the offline query.
+            We do not retry the job if it fails with the given DPUs
+
         Returns: Dict[str, Any]:
         ----------
         A dictionary containing the request_id, the output s3 bucket and prefix,
@@ -615,7 +619,7 @@ def query_offline(
                     ]
                 )
 
-        req = {
+        req: dict[str, Any] = {
             "input_features": input_feature_names,
             "output_features": output_feature_names,
             "input": extract_historical_input,
@@ -625,6 +629,8 @@ def query_offline(
         }
         if workflow is not None:
             req["workflow"] = workflow
+        if num_dpus is not None:
+            req["num_dpus"] = num_dpus
         response = self._post_json(
             "{}/query_offline".format(
                 V1_API,

fennel/testing/mock_client.py

@@ -312,6 +312,7 @@ def query_offline(
         feature_to_column_map: Optional[Dict[Feature, str]] = None,
         workflow: Optional[str] = "default",
         use_v2: bool = False,
+        num_dpus: Optional[int] = None,
     ) -> Union[pd.DataFrame, pd.Series]:
         if input_dataframe is None:
             raise ValueError(

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fennel-ai"
-version = "1.5.77"
+version = "1.5.78"
 description = "The modern realtime feature engineering platform"
 authors = ["Fennel AI <[email protected]>"]
 packages = [{ include = "fennel" }]