Enable resume.

Author: Xreki

graph_net/test/agent_unittest_generator_test.sh

@@ -1,6 +1,7 @@
 #!/usr/bin/env bash
 
-ROOT_DIR="$(cd "$(dirname "$0")/../.." && pwd)"
+
+
 GRAPH_NET_ROOT=$(python -c "import graph_net, os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
 OUTPUT_DIR="/tmp/agent_unittests"
 
@@ -14,6 +15,7 @@ HANDLER_CONFIG=$(base64 -w 0 <<EOF
         "device": "auto",
         "generate_main": true,
         "try_run": true,
+        "resume": false,
         "data_input_predicator_filepath": "$GRAPH_NET_ROOT/graph_net/torch/constraint_util.py",                                                                                     
         "data_input_predicator_class_name": "NaiveDataInputPredicator"
     }

graph_net/torch/sample_passes/agent_unittest_generator.py

@@ -14,6 +14,7 @@
 
 from graph_net import imp_util
 from graph_net.sample_pass.sample_pass import SamplePass
+from graph_net.sample_pass.resumable_sample_pass_mixin import ResumableSamplePassMixin
 from graph_net.tensor_meta import TensorMeta
 
 
@@ -298,7 +299,7 @@ def _render_template(self, graph_module_desc):
         return template.render(graph_module_desc=graph_module_desc)
 
 
-class AgentUnittestGeneratorPass(SamplePass):
+class AgentUnittestGeneratorPass(SamplePass, ResumableSamplePassMixin):
     """SamplePass wrapper to generate Torch unittests via model_path_handler."""
 
     def __init__(self, config=None):
@@ -313,10 +314,15 @@ def declare_config(
         try_run: bool = False,
         data_input_predicator_filepath: str = None,
         data_input_predicator_class_name: str = None,
+        resume: bool = False,
+        limits_handled_models: int = None,
     ):
         pass
 
     def __call__(self, rel_model_path: str):
+        self.resumable_handle_sample(rel_model_path)
+
+    def resume(self, rel_model_path: str):
         model_path_prefix = Path(self.config["model_path_prefix"])
         output_dir = Path(self.config["output_dir"])
         generator = AgentUnittestGenerator(