From caa47a2f47f6c6ce936a0762fee5bbaa39fc492d Mon Sep 17 00:00:00 2001
From: Timothy Pogue <me@pogue.dev>
Date: Wed, 28 Sep 2022 03:06:05 +0000
Subject: [PATCH] close subproc env on shutdown

---
 freqtrade/freqai/freqai_interface.py               |  9 +++++++++
 .../ReinforcementLearner_multiproc.py              | 14 +++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/freqtrade/freqai/freqai_interface.py b/freqtrade/freqai/freqai_interface.py
index 1a847a25e..f8ca34ddb 100644
--- a/freqtrade/freqai/freqai_interface.py
+++ b/freqtrade/freqai/freqai_interface.py
@@ -158,6 +158,13 @@ class IFreqaiModel(ABC):
         self.model = None
         self.dk = None
 
+    def _on_stop(self):
+        """
+        Callback for Subclasses to override to include logic for shutting down resources
+        when SIGINT is sent.
+        """
+        return
+
     def shutdown(self):
         """
         Cleans up threads on Shutdown, set stop event. Join threads to wait
@@ -166,6 +173,8 @@ class IFreqaiModel(ABC):
         logger.info("Stopping FreqAI")
         self._stop_event.set()
 
+        self._on_stop()
+
         logger.info("Waiting on Training iteration")
         for _thread in self._threads:
             _thread.join()
diff --git a/freqtrade/freqai/prediction_models/ReinforcementLearner_multiproc.py b/freqtrade/freqai/prediction_models/ReinforcementLearner_multiproc.py
index 0e6449dcd..efdd4883c 100644
--- a/freqtrade/freqai/prediction_models/ReinforcementLearner_multiproc.py
+++ b/freqtrade/freqai/prediction_models/ReinforcementLearner_multiproc.py
@@ -73,7 +73,7 @@ class ReinforcementLearner_multiproc(BaseReinforcementLearningModel):
         test_df = data_dictionary["test_features"]
 
         env_id = "train_env"
-        num_cpu = int(self.freqai_info["rl_config"]["thread_count"])
+        num_cpu = int(self.freqai_info["rl_config"].get("cpu_count", 2))
         self.train_env = SubprocVecEnv([make_env(self.MyRLEnv, env_id, i, 1, train_df, prices_train,
                                         self.reward_params, self.CONV_WIDTH, monitor=True,
                                         config=self.config) for i
@@ -88,3 +88,15 @@ class ReinforcementLearner_multiproc(BaseReinforcementLearningModel):
         self.eval_callback = EvalCallback(self.eval_env, deterministic=True,
                                           render=False, eval_freq=len(train_df),
                                           best_model_save_path=str(dk.data_path))
+
+
+    def _on_stop(self):
+        """
+        Hook called on bot shutdown. Close SubprocVecEnv subprocesses for clean shutdown.
+        """
+
+        if hasattr(self, "train_env") and self.train_env:
+            self.train_env.close()
+
+        if hasattr(self, "eval_env") and self.eval_env:
+            self.eval_env.close()
\ No newline at end of file