Merge pull request #774 from roboflow/send-count-inference-to-usage-c…

…ollector Pass countinference to usage collector
roboflow · Nov 5, 2024 · 4656d70 · 4656d70
2 parents 9ffd56c + ffcefa7
commit 4656d70
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py
@@ -2239,6 +2239,10 @@ async def legacy_infer_from_request(
                             f"Invalid Content-Type: {request.headers['Content-Type']}"
                         )
 
+                if not countinference and service_secret != ROBOFLOW_SERVICE_SECRET:
+                    raise MissingServiceSecretError(
+                        "Service secret is required to disable inference usage tracking"
+                    )
                 if LAMBDA:
                     request_model_id = (
                         request.scope["aws.event"]["requestContext"]["authorizer"][
@@ -2301,6 +2305,7 @@ async def legacy_infer_from_request(
                     active_learning_target_dataset=active_learning_target_dataset,
                     source=source,
                     source_info=source_info,
+                    usage_billable=countinference,
                     **args,
                 )
 

diff --git a/inference/usage_tracking/collector.py b/inference/usage_tracking/collector.py
@@ -504,12 +504,15 @@ def _extract_usage_params_from_func_kwargs(
         usage_workflow_id: str,
         usage_workflow_preview: bool,
         usage_inference_test_run: bool,
+        usage_billable: bool,
         func: Callable[[Any], Any],
         args: List[Any],
         kwargs: Dict[str, Any],
     ) -> Dict[str, Any]:
         func_kwargs = collect_func_params(func, args, kwargs)
-        resource_details = {}
+        resource_details = {
+            "billable": usage_billable,
+        }
         resource_id = ""
         category = None
         # TODO: add requires_api_key, True if workflow definition comes from platform or model comes from workspace
@@ -608,6 +611,7 @@ def sync_wrapper(
             usage_workflow_id: str = "",
             usage_workflow_preview: bool = False,
             usage_inference_test_run: bool = False,
+            usage_billable: bool = True,
             **kwargs: P.kwargs,
         ) -> T:
             self.record_usage(
@@ -617,6 +621,7 @@ def sync_wrapper(
                     usage_workflow_id=usage_workflow_id,
                     usage_workflow_preview=usage_workflow_preview,
                     usage_inference_test_run=usage_inference_test_run,
+                    usage_billable=usage_billable,
                     func=func,
                     args=args,
                     kwargs=kwargs,
@@ -632,6 +637,7 @@ async def async_wrapper(
             usage_workflow_id: str = "",
             usage_workflow_preview: bool = False,
             usage_inference_test_run: bool = False,
+            usage_billable: bool = True,
             **kwargs: P.kwargs,
         ) -> T:
             await self.async_record_usage(
@@ -641,6 +647,7 @@ async def async_wrapper(
                     usage_workflow_id=usage_workflow_id,
                     usage_workflow_preview=usage_workflow_preview,
                     usage_inference_test_run=usage_inference_test_run,
+                    usage_billable=usage_billable,
                     func=func,
                     args=args,
                     kwargs=kwargs,