add API reset optimzation engine

Summary: Pull Request resolved: https://github.com/facebookresearch/d2go/pull/640

Reviewed By: tglik

Differential Revision: D51908239

fbshipit-source-id: 7bcbad1fc7065b736cf4e38d155eed5d734758f7

d2go/runner/default_runner.py

@@ -198,6 +198,14 @@ class BaseRunner(object):
         """
         pass
 
+    def cleanup(self) -> None:
+        """
+        Override `cleanup` to add custom clean ups such as:
+            - de-register datasets.
+            - free up global variables.
+        """
+        pass
+
     @classmethod
     def create_shared_context(cls, cfg) -> D2GoSharedContext:
         """

tools/benchmark_data.py

@@ -136,6 +136,7 @@ def main(
     metrics = {"_name_": {dataset_name: results}}
     print_metrics_table(metrics)
 
+    runner.cleanup()
     return BenchmarkDataOutput(
         accuracy=metrics,
         metrics=metrics,

tools/evaluator.py

@@ -56,6 +56,7 @@ def main(
     predictor = create_predictor(predictor_path)
     metrics = runner.do_test(cfg, predictor)
     print_metrics_table(metrics)
+    runner.cleanup()
     return EvaluatorOutput(
         accuracy=metrics,
         metrics=metrics,

tools/exporter.py

@@ -95,6 +95,7 @@ def main(
             if not skip_if_fail:
                 raise e
 
+    runner.cleanup()
     return ExporterOutput(
         predictor_paths=predictor_paths,
         accuracy_comparison={},

tools/train_net.py

@@ -24,7 +24,6 @@ from d2go.setup import (
     setup_root_logger,
 )
 from d2go.trainer.api import TestNetOutput, TrainNetOutput
-from d2go.trainer.fsdp import is_fsdp_enabled
 from d2go.utils.mast import gather_mast_errors, mast_error_handler
 from d2go.utils.misc import (
     dump_trained_model_configs,
@@ -68,6 +67,7 @@ def main(
         model.eval()
         metrics = runner.do_test(cfg, model, train_iter=train_iter)
         print_metrics_table(metrics)
+        runner.cleanup()
         return TestNetOutput(
             accuracy=metrics,
             metrics=metrics,
@@ -98,6 +98,7 @@ def main(
 
     # dump config files for trained models
     trained_model_configs = dump_trained_model_configs(cfg.OUTPUT_DIR, trained_cfgs)
+    runner.cleanup()
     return TrainNetOutput(
         # for e2e_workflow
         accuracy=metrics,