stop spark context to prevent memory leaks

MaastrichtU-IDS · Sep 19, 2023 · a495b4c · a495b4c
1 parent 38b8a56
commit a495b4c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/openpredict_model/train.py b/src/openpredict_model/train.py
@@ -462,6 +462,7 @@ def calculateCombinedSimilarity(pairs_train, pairs_test, classes_train, classes_
         test_df = sparkBuildFeatures(spark_context, pairs_test, classes_test,
                                      knownDrugDis_bc.value,  drug_df_bc.value, disease_df_bc.value)
         log.info("Finishing Spark jobs 🏁")
+        spark_context.stop()
     else:
         log.info("Spark cluster not found, using pandas 🐼")
         train_df = createFeatureDF(
@@ -569,6 +570,7 @@ def createFeaturesSparkOrDF(pairs, classes, drug_df, disease_df):
         feature_df = sparkBuildFeatures(
             spark_context, pairs, classes, knownDrugDis_bc.value,  drug_df_bc.value, disease_df_bc.value)
         log.info("Finishing Spark jobs 🏁")
+        spark_context.stop()
     else:
         log.info("Spark cluster not found, using pandas 🐼")
         feature_df = createFeatureDF(