add checkpoint function

energon/cli/service.py

@@ -7,7 +7,6 @@
 def service():
     pass
 
-
 @service.command()
 @click.option("--model_name", default="bert_small", type=str)
 @click.option("--model_type", default="bert", type=str)

energon/engine/engine.py

@@ -30,22 +30,19 @@ def __init__(self,
                 pp_init_size: int = -1,
                 host: str = 'localhost',
                 port: int = 29500,
-                dtype=None,
-                checkpoint=None
+                dtype=None
                 ):
         """
         Args:
             model: torch.nn.Module
             dtype: data-type by which inference is executed
-            checkpoint: load parameter.
         """
         super().__init__()
 
         self.model_class = model_class
         self.model_config = model_config
         self.model_type = model_type
         self.dtype = dtype
-        self.checkpoint = checkpoint
         self.max_batch_size = max_batch_size
 
         # for gpc
@@ -78,7 +75,7 @@ def _init_model(self):
         for i in range(self.global_world_size):
             print(f'[INFO] rank{self.rank} calls rank{i} to init.')
             ob_info = rpc.get_worker_info(self.WORKER_NAME.format(i))
-            self.rrefs.append(rpc.remote(ob_info, RPCWorker, args=(self.model_class, self.model_config, self.model_type, self.dtype, self.checkpoint, self.max_batch_size)))
+            self.rrefs.append(rpc.remote(ob_info, RPCWorker, args=(self.model_class, self.model_config, self.model_type, self.dtype, self.max_batch_size)))
 
     def run(self, inputs): 
         res_rref = 0

energon/engine/rpc_worker.py

@@ -20,12 +20,10 @@ def __init__(self,
                  model_config,
                  model_type,
                  dtype,
-                 checkpoint,
                  max_batch_size:int  = 1) -> None:
         self.model_class = model_class
         self.model_config = model_config
         self.dtype = dtype
-        self.checkpoint = checkpoint
         self.max_batch_size = max_batch_size
         self.pipe_wrapper = WRAPPER_TYPES[model_type]
 

energon/server/engine_server.py

@@ -56,13 +56,13 @@ def launch_engine(model_name,
                 host: str = "localhost",
                 port: int = 29500,
                 dtype = torch.float,
-                checkpoint = None,
+                checkpoint: str = None,
                 server_host = "localhost",
                 server_port = 8005,
                 log_level = "info"
                 ):
 
-    model_config = {'dtype': dtype}
+    model_config = {'dtype': dtype, 'checkpoint': True, 'checkpoint_path': checkpoint}
     global engine
     engine = InferenceEngine(MODEL_CLASSES[model_name], 
                             model_config,
@@ -72,8 +72,7 @@ def launch_engine(model_name,
                             pp_init_size = pp_init_size, 
                             host = host,
                             port = port,
-                            dtype = dtype,
-                            checkpoint = checkpoint)
+                            dtype = dtype)
 
     global server
     config = uvicorn.Config(app, host=server_host, port=server_port, log_level=log_level)