Add train timing tests

src/utils/timer.py

@@ -30,6 +30,20 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+@timer
+def train_with_time(model, X_train: np.array, y_train: np.array) -> Tuple:
+    """Returns trained model with the time
+
+    Args:
+        model: Model to test latency on
+        X_test: Input data
+
+    Returns:
+        Predicted values and time taken to predict it
+    """
+    return model.fit(X_train, y_train)
+
+
 @timer
 def predict_with_time(model, X_test: np.array) -> Tuple[np.array]:
     """Returns model output with the time

tests/tree/test_decision_tree.py

@@ -5,7 +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
+from src.utils.timer import predict_with_time, train_with_time
 
 
 @pytest.fixture
@@ -349,7 +349,17 @@ def test_dt_evaluation(dummy_titanic_dt, dummy_titanic):
     assert auc_test > 0.84, 'AUC ROC on test should be > 0.84'
 
 
-def test_dt_latency(dummy_titanic):
+def test_dt_training_time(dummy_titanic):
+    X_train, y_train, X_test, y_test = dummy_titanic
+
+    # Standardize to use depth = 10
+    dt = DecisionTree(depth_limit=10)
+    latency_array = np.array([train_with_time(dt, X_train, y_train)[1] for i in range(100)])
+    time_p95 = np.quantile(latency_array, 0.95)
+    assert time_p95 < 1.0, 'Training time at 95th percentile should be < 1.0 sec'
+
+
+def test_dt_serving_latency(dummy_titanic):
     X_train, y_train, X_test, y_test = dummy_titanic
 
     # Standardize to use depth = 10
@@ -358,4 +368,4 @@ def test_dt_latency(dummy_titanic):
 
     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.004, 'Latency at 99th percentile should be < 0.004 sec'
+    assert latency_p99 < 0.004, 'Serving latency at 99th percentile should be < 0.004 sec'

tests/tree/test_random_forest.py

@@ -4,7 +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
+from src.utils.timer import predict_with_time, train_with_time
 
 
 @pytest.fixture
@@ -88,7 +88,17 @@ def test_dt_evaluation(dummy_titanic_rf, dummy_titanic):
     assert auc_test > 0.86, 'AUC ROC on test should be > 0.86'
 
 
-def test_dt_latency(dummy_titanic):
+def test_rf_training_time(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)
+    latency_array = np.array([train_with_time(rf, X_train, y_train)[1] for i in range(20)])
+    time_p95 = np.quantile(latency_array, 0.95)
+    assert time_p95 < 3, 'Training time at 95th percentile should be < 3.0 sec'
+
+
+def test_rf_serving_latency(dummy_titanic):
     X_train, y_train, X_test, y_test = dummy_titanic
 
     # Standardize to use depth = 10
@@ -97,4 +107,4 @@ def test_dt_latency(dummy_titanic):
 
     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.018, 'Latency at 99th percentile should be < 0.018 sec'
+    assert latency_p99 < 0.018, 'Serving latency at 99th percentile should be < 0.018 sec'