From 62ecc35551d123b3c76ce3011adb101039f9f09a Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Sat, 18 Sep 2021 22:18:02 -0700 Subject: [PATCH 1/3] Use contextvars to maintain a call stack during the usage calls Signed-off-by: Achal Shah --- sdk/python/feast/usage.py | 30 ++++++++++++++++++++++-------- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/sdk/python/feast/usage.py b/sdk/python/feast/usage.py index b70b4f2c72..1172c2ee27 100644 --- a/sdk/python/feast/usage.py +++ b/sdk/python/feast/usage.py @@ -12,16 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. import concurrent.futures +import contextvars import enum import logging import os import sys import uuid +from collections import defaultdict from datetime import datetime from functools import wraps from os.path import expanduser, join from pathlib import Path -from typing import List, Optional, Tuple +from typing import List, Optional, Tuple, Union import requests @@ -31,6 +33,7 @@ _logger = logging.getLogger(__name__) executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) +call_stack: contextvars.ContextVar = contextvars.ContextVar("call_stack", default=[]) @enum.unique @@ -69,7 +72,7 @@ def check_env_and_configure(self): usage_filepath = join(feast_home_dir, "usage") self._is_test = os.getenv("FEAST_IS_USAGE_TEST", "False") == "True" - self._usage_counter = {} + self._usage_counter = defaultdict(lambda: 0) if os.path.exists(usage_filepath): with open(usage_filepath, "r") as f: @@ -106,9 +109,10 @@ def _send_usage_request(self, json): def log_function(self, function_name: str): self.check_env_and_configure() if self._usage_enabled and self.usage_id: - if ( - function_name == "get_online_features" - and not self.should_log_for_get_online_features_event(function_name) + if "get_online_features" in call_stack.get( + [] + ) and not self.should_log_for_get_online_features_event( + "get_online_features" ): return json = { @@ -121,10 +125,10 @@ def log_function(self, function_name: str): } self._send_usage_request(json) - def should_log_for_get_online_features_event(self, event_name: str): - if event_name not in self._usage_counter: - self._usage_counter[event_name] = 0 + def increment_event_count(self, event_name: Union[UsageEvent, str]): self._usage_counter[event_name] += 1 + + def should_log_for_get_online_features_event(self, event_name: str): if self._usage_counter[event_name] % 10000 != 2: return False self._usage_counter[event_name] = 2 # avoid overflow @@ -174,6 +178,7 @@ def log_exceptions(func): @wraps(func) def exception_logging_wrapper(*args, **kwargs): try: + call_stack.set(call_stack.get() + [func.__name__]) result = func(*args, **kwargs) except Exception as e: error_type = type(e).__name__ @@ -190,6 +195,9 @@ def exception_logging_wrapper(*args, **kwargs): tb = tb.tb_next usage.log_exception(error_type, trace_to_log) raise + finally: + if len(call_stack.get([])) > 0: + call_stack.set(call_stack.get([])[:-1]) return result return exception_logging_wrapper @@ -199,6 +207,8 @@ def log_exceptions_and_usage(func): @wraps(func) def exception_logging_wrapper(*args, **kwargs): try: + call_stack.set(call_stack.get() + [func.__name__]) + usage.increment_event_count(func.__name__) result = func(*args, **kwargs) usage.log_function(func.__name__) except Exception as e: @@ -216,12 +226,16 @@ def exception_logging_wrapper(*args, **kwargs): tb = tb.tb_next usage.log_exception(error_type, trace_to_log) raise + finally: + if len(call_stack.get([])) > 0: + call_stack.set(call_stack.get([])[:-1]) return result return exception_logging_wrapper def log_event(event: UsageEvent): + usage.increment_event_count(event) usage.log_event(event) From c0ed4ae75086914de1a0cc3ae0ed32aaeda3c746 Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Sat, 18 Sep 2021 22:27:08 -0700 Subject: [PATCH 2/3] refactor Signed-off-by: Achal Shah --- sdk/python/feast/usage.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/sdk/python/feast/usage.py b/sdk/python/feast/usage.py index 1172c2ee27..224dfade63 100644 --- a/sdk/python/feast/usage.py +++ b/sdk/python/feast/usage.py @@ -54,6 +54,8 @@ def __str__(self): class Usage: def __init__(self): self._usage_enabled: bool = False + self._is_test = os.getenv("FEAST_IS_USAGE_TEST", "False") == "True" + self._usage_counter = defaultdict(lambda: 0) self.check_env_and_configure() def check_env_and_configure(self): @@ -71,9 +73,6 @@ def check_env_and_configure(self): Path(feast_home_dir).mkdir(exist_ok=True) usage_filepath = join(feast_home_dir, "usage") - self._is_test = os.getenv("FEAST_IS_USAGE_TEST", "False") == "True" - self._usage_counter = defaultdict(lambda: 0) - if os.path.exists(usage_filepath): with open(usage_filepath, "r") as f: self._usage_id = f.read() From 6a4f8761b6ef24de608ec13da01b9c420d732a7f Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Sat, 18 Sep 2021 22:37:43 -0700 Subject: [PATCH 3/3] Remove default values Signed-off-by: Achal Shah --- sdk/python/feast/usage.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/sdk/python/feast/usage.py b/sdk/python/feast/usage.py index 224dfade63..dd225bfe04 100644 --- a/sdk/python/feast/usage.py +++ b/sdk/python/feast/usage.py @@ -108,9 +108,7 @@ def _send_usage_request(self, json): def log_function(self, function_name: str): self.check_env_and_configure() if self._usage_enabled and self.usage_id: - if "get_online_features" in call_stack.get( - [] - ) and not self.should_log_for_get_online_features_event( + if "get_online_features" in call_stack.get() and not self.should_log_for_get_online_features_event( "get_online_features" ): return @@ -195,8 +193,8 @@ def exception_logging_wrapper(*args, **kwargs): usage.log_exception(error_type, trace_to_log) raise finally: - if len(call_stack.get([])) > 0: - call_stack.set(call_stack.get([])[:-1]) + if len(call_stack.get()) > 0: + call_stack.set(call_stack.get()[:-1]) return result return exception_logging_wrapper @@ -226,8 +224,8 @@ def exception_logging_wrapper(*args, **kwargs): usage.log_exception(error_type, trace_to_log) raise finally: - if len(call_stack.get([])) > 0: - call_stack.set(call_stack.get([])[:-1]) + if len(call_stack.get()) > 0: + call_stack.set(call_stack.get()[:-1]) return result return exception_logging_wrapper