From 04db603b7e21435b65743d26fb0fc609e154da81 Mon Sep 17 00:00:00 2001
From: bxdd <bxddream@gmail.com>
Date: Sun, 15 Oct 2023 17:11:54 +0800
Subject: [PATCH] [MNT] update client.load_learnware to fix now code

---
 learnware/client/container.py        | 9 +++++++--
 learnware/client/learnware_client.py | 6 +++---
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/learnware/client/container.py b/learnware/client/container.py
index 8b03aee..4a74dc5 100644
--- a/learnware/client/container.py
+++ b/learnware/client/container.py
@@ -1,5 +1,6 @@
 import os
 import pickle
+import atexit
 import tempfile
 import shortuuid
 from concurrent.futures import ProcessPoolExecutor
@@ -56,10 +57,13 @@ class ModelEnvContainer(BaseModel):
         input_shape = output_results["metadata"]["input_shape"]
         output_shape = output_results["metadata"]["output_shape"]
         super(ModelEnvContainer, self).__init__(input_shape, output_shape)
+        atexit.register(self.remove_env)
 
     def remove_env(self):
-        remove_enviroment(self.conda_env)
-
+        if self.conda_env is not None:
+            self.conda_env = None
+            remove_enviroment(self.conda_env)
+        
     def run_model_with_script(self, method, **kargs):
         with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir:
             input_path = os.path.join(tempdir, "input.pkl")
@@ -137,6 +141,7 @@ class LearnwaresContainer:
         model_list = [_learnware.get_model() for _learnware in self.learnware_list]
         with ProcessPoolExecutor(max_workers=max(os.cpu_count() // 2, 1)) as executor:
             executor.map(self._initialize_model_container, model_list)
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if not self.cleanup:
diff --git a/learnware/client/learnware_client.py b/learnware/client/learnware_client.py
index 88abd3e..715bfe1 100644
--- a/learnware/client/learnware_client.py
+++ b/learnware/client/learnware_client.py
@@ -401,9 +401,9 @@ class LearnwareClient:
                 for i in range(len(learnware_list)):
                     learnware_list[i].instantiate_model()
             elif runnable_option == "conda_env":
-                env_container = LearnwaresContainer(learnware_list, zip_paths)
-                learnware_list = env_container.get_learnwares_with_container()
-
+                with LearnwaresContainer(learnware_list, zip_paths, cleanup=False) as env_container:
+                    learnware_list = env_container.get_learnwares_with_container()
+                    
         if len(learnware_list) == 1:
             return learnware_list[0]
         else: