Merge pull request #1621 from roboflow/feat/include-errors-in-usage

Feat/include errors in usage

Author: PawelPeczek-Roboflow

inference/usage_tracking/collector.py

@@ -728,25 +728,31 @@ def sync_wrapper(
                         )
                     )
                 except Exception as exc:
+                    t2 = time.time()
                     if GCP_SERVERLESS is True:
-                        t2 = time.time()
                         execution_duration = max(t2 - t1, 0.1)
-                        self.record_usage(
-                            **self._extract_usage_params_from_func_kwargs(
-                                usage_fps=usage_fps,
-                                usage_api_key=usage_api_key,
-                                usage_workflow_id=usage_workflow_id,
-                                usage_workflow_preview=usage_workflow_preview,
-                                usage_inference_test_run=usage_inference_test_run,
-                                usage_billable=usage_billable,
-                                execution_duration=execution_duration,
-                                func=func,
-                                category=category,
-                                exc=str(exc),
-                                args=args,
-                                kwargs=kwargs,
-                            )
+                    else:
+                        execution_duration = t2 - t1
+                    exc_type = type(exc).__name__
+                    if hasattr(exc, "inner_error_type"):
+                        exc_type = exc.inner_error_type
+                    exc_str = f"{exc_type}: {str(exc)}"
+                    self.record_usage(
+                        **self._extract_usage_params_from_func_kwargs(
+                            usage_fps=usage_fps,
+                            usage_api_key=usage_api_key,
+                            usage_workflow_id=usage_workflow_id,
+                            usage_workflow_preview=usage_workflow_preview,
+                            usage_inference_test_run=usage_inference_test_run,
+                            usage_billable=usage_billable,
+                            execution_duration=execution_duration,
+                            func=func,
+                            category=category,
+                            exc=exc_str,
+                            args=args,
+                            kwargs=kwargs,
                         )
+                    )
                     raise
                 return res
 
@@ -786,25 +792,31 @@ async def async_wrapper(
                         )
                     )
                 except Exception as exc:
+                    t2 = time.time()
                     if GCP_SERVERLESS is True:
-                        t2 = time.time()
                         execution_duration = max(t2 - t1, 0.1)
-                        await self.async_record_usage(
-                            **self._extract_usage_params_from_func_kwargs(
-                                usage_fps=usage_fps,
-                                usage_api_key=usage_api_key,
-                                usage_workflow_id=usage_workflow_id,
-                                usage_workflow_preview=usage_workflow_preview,
-                                usage_inference_test_run=usage_inference_test_run,
-                                usage_billable=usage_billable,
-                                execution_duration=execution_duration,
-                                func=func,
-                                category=category,
-                                exc=str(exc),
-                                args=args,
-                                kwargs=kwargs,
-                            )
+                    else:
+                        execution_duration = t2 - t1
+                    exc_type = type(exc).__name__
+                    if hasattr(exc, "inner_error_type"):
+                        exc_type = exc.inner_error_type
+                    exc_str = f"{exc_type}: {str(exc)}"
+                    await self.async_record_usage(
+                        **self._extract_usage_params_from_func_kwargs(
+                            usage_fps=usage_fps,
+                            usage_api_key=usage_api_key,
+                            usage_workflow_id=usage_workflow_id,
+                            usage_workflow_preview=usage_workflow_preview,
+                            usage_inference_test_run=usage_inference_test_run,
+                            usage_billable=usage_billable,
+                            execution_duration=execution_duration,
+                            func=func,
+                            category=category,
+                            exc=exc_str,
+                            args=args,
+                            kwargs=kwargs,
                         )
+                    )
                     raise
                 return res
 

tests/inference/unit_tests/usage_tracking/test_collector.py

@@ -955,7 +955,7 @@ def test_func(api_key="test_key"):
         test_func()
 
     # then
-    assert len(usage_collector._usage) == 0
+    assert len(usage_collector._usage) == 1
 
 
 def test_record_usage_with_exception_on_GCP(usage_collector_with_mocked_threads):
@@ -979,5 +979,5 @@ def test_func(api_key="test_key"):
         json.loads(
             usage_collector._usage["test_key"]["model:unknown"]["resource_details"]
         ).get("error")
-        == "test exception"
+        == "Exception: test exception"
     )