Finish adding auto config generator

Author: iffiX

machin/__main__.py

@@ -1,7 +1,8 @@
+from . import envs
 from . import config
 from . import dataset
 from . import launcher
 from . import pl_logger
 from . import pl_plugin
 
-__all__ = ["config", "dataset", "launcher", "pl_logger", "pl_plugin"]
+__all__ = ["env", "config", "dataset", "launcher", "pl_logger", "pl_plugin"]

machin/auto/__init__.py

@@ -0,0 +1,82 @@
+import json
+import argparse
+from pprint import pprint
+from machin.auto.config import (
+    get_available_algorithms,
+    get_available_environments,
+    generate_algorithm_config,
+    generate_env_config,
+    generate_training_config,
+)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+
+    p_list = subparsers.add_parser(
+        "list", help="List available algorithms or environments."
+    )
+
+    p_list.add_argument(
+        "--algo", action="store_true", help="List available algorithms.",
+    )
+
+    p_list.add_argument(
+        "--env", action="store_true", help="List available environments."
+    )
+
+    p_generate = subparsers.add_parser("generate", help="Generate configuration.")
+
+    p_generate.add_argument(
+        "--algo", type=str, required=True, help="Algorithm name to use."
+    )
+    p_generate.add_argument(
+        "--env", type=str, required=True, help="Environment name to use."
+    )
+    p_generate.add_argument(
+        "--print", action="store_true", help="Direct config output to screen."
+    )
+    p_generate.add_argument(
+        "--output",
+        type=str,
+        default="config.json",
+        help="JSON config file output path.",
+    )
+
+    args = parser.parse_args()
+    if args.command == "list":
+        if args.env:
+            print("Available environments are:")
+            for env in get_available_environments():
+                print(env)
+        elif args.algo:
+            print("Available algorithms are:")
+            for algo in get_available_algorithms():
+                print(algo)
+        else:
+            print("You can list --algo or --env.")
+
+    elif args.command == "generate":
+        if args.algo not in get_available_algorithms():
+            print(
+                f"{args.algo} is not a valid algorithm, use list "
+                "--algo to get a list of available algorithms."
+            )
+            exit(0)
+        if args.env not in get_available_environments():
+            print(
+                f"{args.env} is not a valid environment, use list "
+                "--env to get a list of available environments."
+            )
+            exit(0)
+        config = {}
+        config = generate_env_config(args.env, config=config)
+        config = generate_algorithm_config(args.algo, config=config)
+        config = generate_training_config(config=config)
+
+        if args.print:
+            pprint(config)
+
+        with open(args.output, "w") as f:
+            json.dump(config, f, indent=4, sort_keys=True)
+        print(f"Config saved to {args.output}")

machin/auto/__main__.py

@@ -2,6 +2,7 @@
 from typing import Dict, Any, Union
 from machin.frame.algorithms import TorchFramework
 from machin.utils.conf import Config
+from . import envs
 import inspect
 import torch as t
 import machin.frame.algorithms as algorithms
@@ -16,7 +17,7 @@ def fill_default(
     return config
 
 
-def _get_available_algorithms():
+def get_available_algorithms():
     algos = []
     for algo in dir(algorithms):
         algo_cls = getattr(algorithms, algo)
@@ -29,8 +30,17 @@ def _get_available_algorithms():
     return algos
 
 
+def get_available_environments():
+    environments = []
+    for e in dir(envs):
+        e_module = getattr(envs, e)
+        if hasattr(e_module, "launch") and hasattr(e_module, "generate_env_config"):
+            environments.append(e)
+    return environments
+
+
 def generate_training_config(
-    root_dir: str = "./trial",
+    root_dir: str = "trial",
     episode_per_epoch: int = 10,
     max_episodes: int = 10000,
     config: Union[Dict[str, Any], Config] = None,
@@ -56,10 +66,23 @@ def generate_algorithm_config(
                 config["gpus"] = [0, 0, 0]
                 config["num_processes"] = 3
                 config["num_nodes"] = 1
-                config["batch_num"] = {"sampler": 10, "learner": 1}
+            else:
+                config["gpus"] = [0]
             return config
     raise ValueError(
-        f"Invalid algorithm: {algorithm}, valid ones are: {_get_available_algorithms()}"
+        f"Invalid algorithm: {algorithm}, valid ones are: {get_available_algorithms()}"
+    )
+
+
+def generate_env_config(environment: str, config: Union[Dict[str, Any], Config] = None):
+    config = deepcopy(config) or {}
+    if hasattr(envs, environment):
+        e_module = getattr(envs, environment)
+        if hasattr(e_module, "launch") and hasattr(e_module, "generate_env_config"):
+            return e_module.generate_env_config(config)
+    raise ValueError(
+        f"Invalid environment: {environment}, "
+        f"valid ones are: {get_available_algorithms()}"
     )
 
 
@@ -71,7 +94,7 @@ def init_algorithm_from_config(
     if not inspect.isclass(frame) or not issubclass(frame, TorchFramework):
         raise ValueError(
             f"Invalid algorithm: {config['frame']}, "
-            f"valid ones are: {_get_available_algorithms()}"
+            f"valid ones are: {get_available_algorithms()}"
         )
     return frame.init_from_config(config, model_device=model_device)
 
@@ -82,7 +105,7 @@ def is_algorithm_distributed(config: Union[Dict[str, Any], Config]):
     if not inspect.isclass(frame) or not issubclass(frame, TorchFramework):
         raise ValueError(
             f"Invalid algorithm: {config['frame']}, "
-            f"valid ones are: {_get_available_algorithms()}"
+            f"valid ones are: {get_available_algorithms()}"
         )
     return frame.is_distributed()
 

machin/auto/config.py

@@ -0,0 +1,3 @@
+from . import openai_gym
+
+__all__ = ["openai_gym"]

machin/auto/envs/__init__.py

@@ -244,34 +244,6 @@ def __next__(self):
         return result
 
 
-def generate_gym_env_config(
-    env_name: str = None, config: Union[Dict[str, Any], Config] = None
-):
-    """
-    Generate example OpenAI gym config.
-    """
-    config = deepcopy(config) or {}
-    return fill_default(
-        {
-            "trials_dir": "trials",
-            "gpus": 0,
-            "episode_per_epoch": 100,
-            "max_episodes": 1000000,
-            "train_env_config": {
-                "env_name": env_name or "CartPole-v1",
-                "render_every_episode": 100,
-                "act_kwargs": {},
-            },
-            "test_env_config": {
-                "env_name": env_name or "CartPole-v1",
-                "render_every_episode": 100,
-                "act_kwargs": {},
-            },
-        },
-        config,
-    )
-
-
 def gym_env_dataset_creator(frame, env_config):
     env = gym.make(env_config["env_name"])
     if _is_discrete_space(env.action_space):
@@ -295,9 +267,31 @@ def gym_env_dataset_creator(frame, env_config):
         )
 
 
-def launch_gym(
-    config: Union[Dict[str, Any], Config], pl_callbacks: List[Callback] = None
+def generate_env_config(
+    env_name: str = None, config: Union[Dict[str, Any], Config] = None
 ):
+    """
+    Generate example OpenAI gym config.
+    """
+    config = deepcopy(config) or {}
+    return fill_default(
+        {
+            "train_env_config": {
+                "env_name": env_name or "CartPole-v1",
+                "render_every_episode": 100,
+                "act_kwargs": {},
+            },
+            "test_env_config": {
+                "env_name": env_name or "CartPole-v1",
+                "render_every_episode": 100,
+                "act_kwargs": {},
+            },
+        },
+        config,
+    )
+
+
+def launch(config: Union[Dict[str, Any], Config], pl_callbacks: List[Callback] = None):
     """
     Args:
         config: All configs needed to launch a gym environment and initialize

machin/auto/env/openai_gym.py renamed to machin/auto/envs/openai_gym.py

@@ -1,6 +1,7 @@
 import os
 import logging
 import pytorch_lightning as pl
+from time import sleep
 from torch import distributed
 from pytorch_lightning.utilities.seed import seed_everything
 from pytorch_lightning.utilities.distributed import rank_zero_only
@@ -105,8 +106,12 @@ def pre_dispatch(self):
 
         # initialize framework in the launcher
         self._model.init_frame()
-        self._model.trainer.accelerator.optimizers = self._model.frame.optimizers
-        self._model.trainer.accelerator.lr_schedulers = self._model.frame.lr_schedulers
+        if self._model.frame.optimizers is not None:
+            self._model.trainer.accelerator.optimizers = self._model.frame.optimizers
+        if self._model.frame.lr_schedulers is not None:
+            self._model.trainer.accelerator.lr_schedulers = (
+                self._model.frame.lr_schedulers
+            )
 
         self.barrier()
 
@@ -199,8 +204,12 @@ def new_process(self, process_idx, trainer, mp_queue):
 
         # initialize framework in the launcher
         self._model.init_frame()
-        trainer.accelerator.optimizers = self._model.frame.optimizers
-        trainer.accelerator.lr_schedulers = self._model.frame.lr_schedulers
+        if self._model.frame.optimizers is not None:
+            self._model.trainer.accelerator.optimizers = self._model.frame.optimizers
+        if self._model.frame.lr_schedulers is not None:
+            self._model.trainer.accelerator.lr_schedulers = (
+                self._model.frame.lr_schedulers
+            )
 
         self.barrier()
 
@@ -231,7 +240,7 @@ def _spawn(self):
         ]
         for p in processes:
             p.start()
-        while all([p.is_alive() for p in processes]):
+        while True:
             should_exit = False
             for p in processes:
                 try:
@@ -240,9 +249,14 @@ def _spawn(self):
                     traceback.print_exc()
                     should_exit = True
             if should_exit:
+                for p in processes:
+                    p.terminate()
+                    p.join()
+                raise RuntimeError("One or more exceptions raised in sub-processes.")
+            elif not all([p.is_alive() for p in processes]):
                 break
+            sleep(0.1)
         for p in processes:
-            p.kill()
             p.join()
 
     def training_step(self, *args, **kwargs):
@@ -266,4 +280,3 @@ def post_training_step(self):
 # before the trainer is initialized.
 pl.trainer.connectors.accelerator_connector.DDPPlugin = DDPPlugin
 pl.trainer.connectors.accelerator_connector.DDPSpawnPlugin = DDPSpawnPlugin
-pl_logger.info("DDP plugin patched.")

machin/auto/pl_plugin.py

@@ -123,6 +123,10 @@ def optimizers(self, optimizers):
     def lr_schedulers(self):
         return []
 
+    @classmethod
+    def is_distributed(cls):
+        return True
+
     def set_sync(self, is_syncing):
         self.is_syncing = is_syncing
 

machin/frame/algorithms/a3c.py

@@ -187,6 +187,7 @@ def generate_config(cls, config: Dict[str, Any]):
         }
         config = deepcopy(config)
         config["frame"] = "DQNApex"
+        config["batch_num"] = {"sampler": 10, "learner": 1}
         if "frame_config" not in config:
             config["frame_config"] = default_values
         else:
@@ -461,6 +462,7 @@ def generate_config(cls, config: Union[Dict[str, Any], Config]):
         }
         config = deepcopy(config)
         config["frame"] = "DDPGApex"
+        config["batch_num"] = {"sampler": 10, "learner": 1}
         if "frame_config" not in config:
             config["frame_config"] = default_values
         else:

machin/frame/algorithms/apex.py

@@ -484,6 +484,7 @@ def generate_config(cls, config: Union[Dict[str, Any], Config]):
         }
         config = deepcopy(config)
         config["frame"] = "IMPALA"
+        config["batch_num"] = {"sampler": 10, "learner": 1}
         if "frame_config" not in config:
             config["frame_config"] = default_values
         else:

machin/frame/algorithms/impala.py

@@ -57,7 +57,7 @@
         "psutil",
         "numpy",
         "torch>=1.6.0",
-        "pytorch-lightning>=1.0",
+        "pytorch-lightning>=1.2.0",
         "torchviz",
         "moviepy",
         "matplotlib",

setup.py

@@ -1,4 +1,5 @@
 from machin.parallel.distributed import get_world, get_cur_rank
+from machin.utils.helper_classes import Object
 from torch.utils.data import DataLoader, TensorDataset
 import os
 import sys
@@ -24,6 +25,7 @@ class ParallelModule(pl.LightningModule):
     def __init__(self):
         super().__init__()
         self.nn_model = NNModule()
+        self.frame = Object({"optimizers": None, "lr_schedulers": None})
 
     def train_dataloader(self):
         return DataLoader(
@@ -41,6 +43,9 @@ def training_step(self, batch, _batch_idx):
             raise RuntimeError("World not initialized.")
         return None
 
+    def init_frame(self):
+        pass
+
     def configure_optimizers(self):
         return None