From 9517ab564bf213f1d52943d2c4a18bd8718fc216 Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Wed, 9 Oct 2024 14:51:09 -0400
Subject: [PATCH] Add memory benchmarks

---
 benchmarks/benchmarks.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/benchmarks/benchmarks.py b/benchmarks/benchmarks.py
index 18f1126d..6cb8bc50 100644
--- a/benchmarks/benchmarks.py
+++ b/benchmarks/benchmarks.py
@@ -26,6 +26,7 @@ def time_feature_creation():
             feature_extractor="malanchev"
         )
 
+
 def time_learn_loop(ml_model, strategy):
     """Benchmark how long it takes for a run of the learning loop."""
     # Use the precomputed features so we don't time their creation.
@@ -51,3 +52,29 @@ def time_learn_loop(ml_model, strategy):
     ["RandomForest", "KNN"],            # The different ML methods
     ["RandomSampling", "UncSampling"],  # The different strategies.
 ]
+
+
+def peakmem_learn_loop(ml_model):
+    """Benchmark how much memory it takes for a run of the learning loop."""
+    # Use the precomputed features so we don't time their creation.
+    features_file = str(_TEST_DATA_DIR / "test_features.csv")
+    with tempfile.TemporaryDirectory() as dir_name:
+        metrics_file = str(Path(dir_name) / "metrics.csv")
+        output_queried_file = str(Path(dir_name) / "queried.csv")
+        learn_loop(
+            nloops=25,
+            features_method="malanchev",
+            classifier=ml_model,
+            strategy="RandomSampling",
+            path_to_features=features_file,
+            output_metrics_file=metrics_file,
+            output_queried_file=output_queried_file,
+            training="original",
+            batch=1,
+        )  
+
+
+# Parameterize the ML models and strategies we benchmark.
+peakmem_learn_loop.params = [
+    ["RandomForest", "KNN"],            # The different ML methods
+]
\ No newline at end of file