[ENH] add docker container interface

learnware/client/container.py

@@ -16,16 +16,69 @@ from ..logger import get_module_logger
 
 logger = get_module_logger(module_name="client_container")
 
-
-class ModelEnvContainer(BaseModel):
-    def __init__(self, model_config: dict, learnware_zippath: str):
+class ModelContainer(BaseModel):
+    def __init__(self, model_config: dict, learnware_zippath: str, build: bool=True):
         self.model_script = os.path.join(C.package_path, "client", "scripts", "run_model.py")
         self.model_config = model_config
-        self.conda_env = f"learnware_{shortuuid.uuid()}"
         self.learnware_zippath = learnware_zippath
+        self.build = build
+        self.cleanup_flag = False
+    
+    def reset(self, **kwargs):
+        for _k, _v in kwargs.items():
+            if hasattr(self, _k):
+                setattr(_k, _v)
+                
+    def init_env_and_set_metadata(self):
+        """We must set `input_shape` and `output_shape` 
+        """
+        if self.build:
+            self.cleanup_flag = True
+            self._init_env()
+            atexit.register(self.remove_env)
+        
+        self._set_metadata()
+        
+    def remove_env(self):
+        if self.cleanup_flag is True:
+            self.cleanup_flag = False
+            try:
+                self._remove_env()
+            except Exception as err:
+                self.cleanup_flag = True
+                raise err
+    
+    def _set_metadata(self):
+        raise NotImplementedError('_set_metadata method is not implemented!')
+        
+    def _init_env(self):
+        raise NotImplementedError('_init_env method is not implemented!')
+        
+    def _remove_env(self):
+        raise NotImplementedError('_remove_env method is not implemented!')
+    
+    def fit(self, X, y):
+        raise NotImplementedError('fit method is not implemented!')
 
-    def init_env_and_metadata(self):
+    def predict(self, X):
+        raise NotImplementedError('predict method is not implemented!')
+
+    def finetune(self, X, y) -> None:
+        raise NotImplementedError('finetune method is not implemented!')
+        
+class ModelCondaContainer(ModelContainer):
+    
+    def __init__(self, model_config: dict, learnware_zippath: str, conda_env: str=None, build: bool=True):
+        self.conda_env = f"learnware_{shortuuid.uuid()}" if conda_env is None else conda_env
+        super(ModelCondaContainer, self).__init__(model_config, learnware_zippath, build)
+        
+    def _init_env(self):
         install_environment(self.learnware_zippath, self.conda_env)
+    
+    def _remove_env(self):
+        remove_enviroment(self.conda_env)
+        
+    def _set_metadata(self):
         with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir:
             output_path = os.path.join(tempdir, "output.pkl")
             model_path = os.path.join(tempdir, "model.pkl")
@@ -56,20 +109,9 @@ class ModelEnvContainer(BaseModel):
             raise output_results["error_info"]
         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)
+        self.reset(input_shape=input_shape, output_shape=output_shape)
 
-    def remove_env(self):
-        if self.conda_env is not None:
-            conda_env = self.conda_env
-            self.conda_env = None
-            try:
-                remove_enviroment(self.conda_env)
-            except Exception as err:
-                self.conda_env = conda_env
-                raise err
-
-    def run_model_with_script(self, method, **kargs):
+    def _run_model_with_script(self, method, **kargs):
         with tempfile.TemporaryDirectory(prefix="learnware_") as tempdir:
             input_path = os.path.join(tempdir, "input.pkl")
             output_path = os.path.join(tempdir, "output.pkl")
@@ -108,18 +150,50 @@ class ModelEnvContainer(BaseModel):
         return output_results[method]
 
     def fit(self, X, y):
-        self.run_model_with_script("fit", X=X, y=y)
+        self._run_model_with_script("fit", X=X, y=y)
 
     def predict(self, X):
-        return self.run_model_with_script("predict", X=X)
+        return self._run_model_with_script("predict", X=X)
 
     def finetune(self, X, y) -> None:
-        self.run_model_with_script("finetune", X=X, y=y)
+        self._run_model_with_script("finetune", X=X, y=y)
+
+class ModelDockerContainer(ModelCondaContainer):
+    def __init__(self, model_config: dict, learnware_zippath: str, docker_img=None, conda_env: str=None, build: bool=True):
+        """_summary_
+
+        Parameters
+        ----------
+        build : bool, optional
+            Whether to build the docker env, by default True
+        """
+        if docker_img is None:
+            self.docker_img = None # create docker img
+        self.conda_env = f"learnware_{shortuuid.uuid()}" if conda_env is None else conda_env
+        # call init method of parent of parent class
+        super(ModelCondaContainer, self).__init__(model_config, learnware_zippath, True if docker_img is None else build)
+        
+    def _set_metadata(self):
+        raise NotImplementedError('_set_metadata method is not implemented!')
+        
+    def _init_env(self):
+        raise NotImplementedError('_init_env method is not implemented!')
+        
+    def _remove_env(self):
+        raise NotImplementedError('_remove_env method is not implemented!')
+    
+    def fit(self, X, y):
+        raise NotImplementedError('fit method is not implemented!')
 
+    def predict(self, X):
+        raise NotImplementedError('predict method is not implemented!')
 
+    def finetune(self, X, y) -> None:
+        raise NotImplementedError('finetune method is not implemented!')
+        
 class LearnwaresContainer:
     def __init__(
-        self, learnwares: Union[List[Learnware], Learnware], learnware_zippaths: Union[List[str], str], cleanup=True
+        self, learnwares: Union[List[Learnware], Learnware], learnware_zippaths: Union[List[str], str], cleanup=True, mode='conda'
     ):
         """The initializaiton method for base reuser
 
@@ -136,41 +210,69 @@ class LearnwaresContainer:
         assert all(
             [isinstance(_learnware.get_model(), dict) for _learnware in learnwares]
         ), "the learnwre model should not be instantiated for reuser with containter"
-        self.learnware_list = [
-            Learnware(
-                _learnware.id, ModelEnvContainer(_learnware.get_model(), _zippath), _learnware.get_specification()
-            )
-            for _learnware, _zippath in zip(learnwares, learnware_zippaths)
-        ]
+        
+        if mode == 'conda':
+            self.learnware_list = [
+                Learnware(
+                    _learnware.id, ModelCondaContainer(_learnware.get_model(), _zippath), _learnware.get_specification()
+                )
+                for _learnware, _zippath in zip(learnwares, learnware_zippaths)
+            ]
+        elif mode == 'docker':
+            docker_img = self._generate_docker_img()
+            self.learnware_list = [
+                Learnware(
+                    _learnware.id, ModelDockerContainer(_learnware.get_model(), _zippath, docker_img), _learnware.get_specification()
+                )
+                for _learnware, _zippath in zip(learnwares, learnware_zippaths)
+            ]
+        
+        else:
+            raise ValueError(f"mode must be in ['conda', 'docker'], should not be {mode}")
+        self.results = [True] * len(learnwares)
         self.cleanup = cleanup
+        print('234', self.learnware_list)
 
+    def _generate_docker_img():
+        return None
+        
     def __enter__(self):
         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)
+            results = executor.map(self._initialize_model_container, model_list)
+        self.results = list(results)
+        print('234', self.results, sum(self.results), len(self.learnware_list))
+        if sum(self.results) < len(self.learnware_list):
+            logger.warning(f'{len(self.learnware_list) - sum(results)} of {len(self.learnware_list)} learnwares init failed! This learnware will be ignored')
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if not self.cleanup:
-            logger.warning(f"Notice, the learnware container env is not clean up!")
+            logger.warning(f"Notice, the learnware container env is not cleaned up!")
             return
         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._destroy_model_container, model_list)
 
     @staticmethod
-    def _initialize_model_container(model: ModelEnvContainer):
+    def _initialize_model_container(model: ModelCondaContainer):
         try:
-            model.init_env_and_metadata()
+            model.init_env_and_set_metadata()
         except Exception as err:
             logger.error(f"build env {model.conda_env} failed due to {err}")
-
+            return False
+        return True
+    
     @staticmethod
-    def _destroy_model_container(model: ModelEnvContainer):
+    def _destroy_model_container(model: ModelCondaContainer):
         try:
             model.remove_env()
         except Exception as err:
             logger.error(f"remove env {model.conda_env} failed due to {err}")
+            return False
+        return True
 
     def get_learnwares_with_container(self):
-        return self.learnware_list
+        learnwares = [_learnware for _learnware, _result in zip(self.learnware_list, self.results) if _result is True]
+        print('233', learnwares, list(self.results))
+        return learnwares

learnware/client/utils.py

@@ -14,7 +14,7 @@ def system_execute(args, timeout=None):
     try:
         com_process.check_returncode()
     except subprocess.CalledProcessError as err:
-        print("System Execute Error:", str(com_process.stderr))
+        print("System Execute Error:", com_process.stderr.decode())
         raise err
 
 

tests/test_learnware_client/test_learnware.py

@@ -18,6 +18,6 @@ if __name__ == "__main__":
             id="test", semantic_spec=semantic_specification, learnware_dirpath=learnware_dirpath
         )
 
-    with LearnwaresContainer(learnware, zip_path) as env_container:
-        learnware = env_container.get_learnwares_with_container()[0]
-        EasyMarket.check_learnware(learnware)
+        with LearnwaresContainer(learnware, zip_path) as env_container:
+            learnware = env_container.get_learnwares_with_container()[0]
+            EasyMarket.check_learnware(learnware)

tests/test_learnware_client/test_load.py

@@ -6,7 +6,7 @@ import numpy as np
 import learnware
 from learnware.learnware import get_learnware_from_dirpath
 from learnware.client import LearnwareClient
-from learnware.client.container import ModelEnvContainer, LearnwaresContainer
+from learnware.client.container import ModelCondaContainer, LearnwaresContainer
 from learnware.learnware.reuse import AveragingReuser