develop

day5/requirements.txt

@@ -4,4 +4,5 @@ mlflow
 pandas
 pytest
 great_expectations
-black
+black
+psutil

day5/演習3/tests/test_model.py

@@ -11,6 +11,7 @@
 from sklearn.preprocessing import OneHotEncoder, StandardScaler
 from sklearn.compose import ColumnTransformer
 from sklearn.pipeline import Pipeline
+import psutil
 
 # テスト用データとモデルパスを定義
 DATA_PATH = os.path.join(os.path.dirname(__file__), "../data/Titanic.csv")
@@ -171,3 +172,23 @@ def test_model_reproducibility(sample_data, preprocessor):
     assert np.array_equal(
         predictions1, predictions2
     ), "モデルの予測結果に再現性がありません"
+
+
+def memory_usage_amount(train_model):
+    """メモリ使用量を検証"""
+    model, X_test, y_test = train_model
+    # プロセス取得
+    process = psutil.Process(os.getpid())
+
+    # predict前のメモリ
+    mem_before = process.memory_info().rss / 1024**2  # MB単位
+
+    # 実行
+    model.predict(X_test)
+
+    # predict後のメモリ
+    mem_after = process.memory_info().rss / 1024**2
+
+    assert (
+        mem_after - mem_before < 10
+    ), f"推論にメモリを使いすぎです. {mem_after - mem_before}MB"