Add latency tests

.gitignore

@@ -3,3 +3,6 @@
 .ipynb_checkpoints/
 notebooks/
 readme.txt
+*.pyc
+.coverage
+coverage.xml

src/utils/timer.py

@@ -0,0 +1,29 @@
+from functools import wraps
+from time import perf_counter
+from typing import Callable
+
+
+def timer(func: Callable) -> Callable:
+    """Decorator to time a function.
+
+    Args:
+        func: Function to time
+
+    Returns:
+        Function results and time (in seconds)
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        start = perf_counter()
+        results = func(*args, **kwargs)
+        end = perf_counter()
+        run_time = end - start
+        return results, run_time
+
+    return wrapper
+
+
+@timer
+def predict_with_time(model, X_test):
+    return model.predict(X_test)

tests/tree/test_decision_tree.py

@@ -5,6 +5,7 @@
 
 from src.data_prep.prep_titanic import load_df, prep_df, split_df, get_feats_and_labels
 from src.tree.decision_tree import gini_gain, gini_impurity, DecisionTree
+from src.utils.timer import predict_with_time
 
 
 @pytest.fixture
@@ -341,3 +342,15 @@ def test_dt_evaluation(dummy_titanic_dt, dummy_titanic):
 
     assert acc_test > 0.83, 'Accuracy on test should be > 0.83'
     assert auc_test > 0.84, 'AUC ROC on test should be > 0.84'
+
+
+def test_dt_latency(dummy_titanic):
+    X_train, y_train, X_test, y_test = dummy_titanic
+
+    # Standardize to use depth = 10
+    dt = DecisionTree(depth_limit=10)
+    dt.fit(X_train, y_train)
+
+    latency_array = np.array([predict_with_time(dt, X_test)[1] for i in range(500)])
+    latency_p99 = np.quantile(latency_array, 0.99)
+    assert latency_p99 < 0.002, 'Latency at 99th percentile should be < 0.002 sec'

tests/tree/test_random_forest.py

@@ -4,6 +4,7 @@
 
 from src.tree.random_forest import RandomForest, DecisionTree
 from tests.tree.test_decision_tree import dummy_titanic, dummy_feats_and_labels
+from src.utils.timer import predict_with_time
 
 
 @pytest.fixture
@@ -85,3 +86,15 @@ def test_dt_evaluation(dummy_titanic_rf, dummy_titanic):
 
     assert acc_test > 0.85, 'Accuracy on test should be > 0.83'
     assert auc_test > 0.86, 'AUC ROC on test should be > 0.86'
+
+
+def test_dt_latency(dummy_titanic):
+    X_train, y_train, X_test, y_test = dummy_titanic
+
+    # Standardize to use depth = 10
+    rf = RandomForest(depth_limit=10, num_trees=5, col_subsampling=0.8, row_subsampling=0.8)
+    rf.fit(X_train, y_train)
+
+    latency_array = np.array([predict_with_time(rf, X_test)[1] for i in range(500)])
+    latency_p99 = np.quantile(latency_array, 0.99)
+    assert latency_p99 < 0.01, 'Latency at 99th percentile should be < 0.01 sec'