trade

examples/trade/main.py

@@ -0,0 +1,158 @@
+import re
+import os
+import argparse
+import yaml
+from executor import Executor
+import warnings
+import redis
+import subprocess
+
+warnings.filterwarnings("ignore")
+
+from util import merge_dicts
+
+loader = yaml.FullLoader
+loader.add_implicit_resolver(
+    "tag:yaml.org,2002:float",
+    re.compile(
+        """^(?:
+     [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)?
+    |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
+    |\\.[0-9_]+(?:[eE][-+][0-9]+)?
+    |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]*
+    |[-+]?\\.(?:inf|Inf|INF)
+    |\\.(?:nan|NaN|NAN))$""",
+        re.X,
+    ),
+    list("-+0123456789."),
+)
+
+
+def get_full_config(config, dir_name):
+    while "base" in config:
+        base_config = os.path.normpath(os.path.join(dir_name, config.pop("base")))
+        dir_name = os.path.dirname(base_config)
+        with open(base_config, "r") as f:
+            base_config = yaml.load(base_config, Loader=yaml.FullLoader)
+        config = merge_dicts(base_config, config)
+    return config
+
+
+def run(config):
+    log_dir = config["log_dir"]
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+    with open(log_dir + "/config.yml", "w") as f:
+        yaml.dump(config, f)
+    executor = Executor(**config)
+    if config["task"] == "train":
+        return executor.train(**config["optim"])
+    elif config["task"] == "eval":
+        return executor.eval(
+            config["test_paths"]["order_dir"],
+            save_res=True,
+            logdir=config["log_dir"] + "/test/",
+        )
+    else:
+        raise NotImplementedError
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-c", "--config", type=str)
+    parser.add_argument("-n", "--index", type=int, default=None)
+    args = parser.parse_args()
+
+    print(os.cpu_count())
+
+    EXP_PATH = os.environ["EXP_PATH"]
+    config_path = os.path.normpath(os.path.join(EXP_PATH, args.config))
+    EXP_NAME = os.path.relpath(config_path, EXP_PATH)
+    if os.path.isdir(config_path):
+        if not args.index is None:
+            with open(config_path + "/configs.yml") as f:
+                config_list = list(yaml.load_all(f, Loader=loader))
+            config = config_list[args.index]
+            if "PT_OUTPUT_DIR" in os.environ:
+                config["log_dir"] = os.environ["PT_OUTPUT_DIR"]
+            else:
+                log_prefix = (
+                    os.environ["OUTPUT_DIR"] if "OUTPUT_DIR" in os.environ else "../log"
+                )
+                config["log_dir"] = os.path.join(log_prefix, config["log_dir"])
+            config = get_full_config(config, config_path)
+            run(config)
+        else:
+            redis_server = redis.Redis(
+                host=os.environ["REDIS_SERVER"],
+                port=os.environ["REDIS_PORT"],
+                db=0,
+                charset="utf-8",
+                decode_responses=True,
+            )
+            with open(config_path + "/configs.yml") as f:
+                config_list = list(yaml.load_all(f, Loader=loader))
+            config_num = len(config_list)
+            if not redis_server.exists(EXP_NAME):
+                for i in range(config_num):
+                    redis_server.rpush(EXP_NAME, i)
+                    redis_server.set(f"{EXP_NAME}_{i}", "Pending")
+            else:
+                if redis_server.llen(EXP_NAME) == 0:
+                    for i in range(config_num):
+                        if (
+                            not redis_server.exists(f"{EXP_NAME}_{i}")
+                            or redis_server.get(f"{EXP_NAME}_{i}") == "Failed"
+                        ):
+                            redis_server.rpush(EXP_NAME, i)
+                            redis_server.set(f"{EXP_NAME}_{i}", "Pending")
+            print(f"Starting..., {redis_server.llen(EXP_NAME)} trails to run")
+            while True:
+                index = redis_server.lpop(EXP_NAME)
+                if index is None:
+                    print("All done")
+                    break
+                index = int(index)
+                redis_server.set(f"{EXP_NAME}_{index}", "Running")
+                print(f"Trail_{index} is running")
+                try:
+                    res = subprocess.run(
+                        [
+                            "python",
+                            "main.py",
+                            "--config",
+                            args.config,
+                            "--index",
+                            str(index),
+                        ],
+                    )
+                except KeyboardInterrupt:
+                    redis_server.set(f"{EXP_NAME}_{index}", "Failed")
+                    print(
+                        f"Trail_{index} has failed, {redis_server.llen(EXP_NAME)} trails to run"
+                    )
+                    break
+                if res.returncode == 0:
+                    redis_server.set(f"{EXP_NAME}_{index}", "Finished")
+                    print(
+                        f"Finish running one trail, {redis_server.llen(EXP_NAME)} trails to run"
+                    )
+                else:
+                    redis_server.set(f"{EXP_NAME}_{index}", "Failed")
+                    print(
+                        f"Trail_{index} has failed, {redis_server.llen(EXP_NAME)} trails to run"
+                    )
+
+    elif os.path.isfile(config_path):
+        assert config_path.endswith(".yml"), "Config file should be an yaml file"
+        EXP_NAME = EXP_NAME[:-4]
+        with open(config_path, "r") as f:
+            config = yaml.load(f, Loader=loader)
+        config = get_full_config(config, os.path.dirname(config_path))
+        log_prefix = (
+            os.environ["OUTPUT_DIR"] if "OUTPUT_DIR" in os.environ else "../log"
+        )
+        config["log_dir"] = os.path.join(log_prefix, config["log_dir"])
+        run(config)
+    else:
+        print("The config path should be a relative path from EXP_PATH")