Add training service (#13)

Author: HYLcool

environments/data.yaml

@@ -5,9 +5,11 @@ dependencies:
   - python=3.10
   - pip:
     - py-data-juicer
+    - agentscope
     - flask
     - omegaconf
     - sqlalchemy
     - psycopg2
     - networkx
     - transformers
+    - "-e ..[dev]"

environments/env_mapping.json

@@ -3,5 +3,10 @@
     "env_name": "trinity_data",
     "env_yaml": "environments/data.yaml",
     "env_entry": "trinity/data/server.py"
+  },
+  "trinity.training": {
+    "env_name": "trinity",
+    "env_yaml": "environments/training.yaml",
+    "env_entry": "trinity/cli/server.py"
   }
 }

environments/training.yaml

@@ -0,0 +1,7 @@
+name: trinity
+channels:
+  - defaults
+dependencies:
+  - python=3.10
+  - pip:
+    - "-e ..[dev]"

pyproject.toml

@@ -32,6 +32,8 @@ dependencies = [
     "math_verify",
     "ninja",
     "fire",
+    "flask",
+    "requests",
 ]
 
 [project.scripts]

scripts/install.py

@@ -24,13 +24,13 @@ def main():
         env_mapping = json.load(f)
     for env_path, env_config in env_mapping.items():
         env_name = env_config["env_name"]
-        print(f"Installing dependencies for module {env_name}...")
+        print(f"Installing dependencies for module [{env_name}]...")
         # check if it's existing
         res = subprocess.run(
             f"{env_mng} env list | grep {env_name}", shell=True, text=True, stdout=subprocess.PIPE
         )
         if res.returncode == 0 and env_name in res.stdout:
-            print(f"Environment {env_name} already exists. Skipping...")
+            print(f"Environment [{env_name}] already exists. Skipping...")
         else:
             res = subprocess.run(
                 f'{env_mng} env create -f {env_config["env_yaml"]}'
@@ -39,9 +39,9 @@ def main():
                 shell=True,
             )
             if res.returncode == 0:
-                print(f"Environment {env_name} created successfully.")
+                print(f"Environment [{env_name}] created successfully.")
             else:
-                print(f"Failed to create environment {env_name} with exit code {res.returncode}.")
+                print(f"Failed to create environment [{env_name}] with exit code {res.returncode}.")
 
 
 if __name__ == "__main__":

scripts/start_servers.py

@@ -28,7 +28,7 @@ def main():
         env_mapping = json.load(f)
     for env_path, env_config in env_mapping.items():
         env_name = env_config["env_name"]
-        print(f"Starting server for module {env_name}...")
+        print(f"Starting server for module [{env_name}]...")
         timestamp = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
         with open(os.path.join(args.log_dir, f"{env_name}_{timestamp}_log.txt"), "w") as log_file:
             server = subprocess.Popen(
@@ -38,7 +38,7 @@ def main():
                 shell=True,
             )
             servers.append(server)
-            print(f"Server of module {env_name} is started with PID {server.pid}")
+            print(f"Server of module [{env_name}] is started with PID {server.pid}")
     for server in servers:
         server.wait()
 

trinity/data/client.py trinity/cli/client.pytrinity/data/client.py renamed to trinity/cli/client.py

@@ -1,9 +1,7 @@
 import requests
 
-LOCAL_SERVER_URL = "http://127.0.0.1:5000/data_workflow"
 
-
-def send_get_request(url: str, params: dict) -> None:
+def send_get_request(url: str, params: dict):
     """
     Send GET request with parameters.
 
@@ -32,8 +30,15 @@ def request(url, **kwargs):
 
 
 if __name__ == "__main__":
+    # --- only for local testing
+    LOCAL_DATA_WORKFLOW_SERVER_URL = "http://127.0.0.1:5005/data_workflow"
+    LOCAL_TRINITY_TRAINING_SERVER_URL = "http://127.0.0.1:5006/trinity_rft"
+    # --- only for local testing
+
     res = request(
-        url=LOCAL_SERVER_URL,
+        url=LOCAL_DATA_WORKFLOW_SERVER_URL,
         configPath="examples/grpo_gsm8k/gsm8k.yaml",
     )
-    print(res)
+    if res:
+        print(res)
+        print(res["message"])

trinity/cli/launcher.py

@@ -107,20 +107,38 @@ def both(config: Config) -> None:
         ray.get(trainer.log_finalize.remote(step=train_iter_num))
 
 
-def activate_data_module(config_path: str):
+def activate_data_module(data_workflow_url: str, config_path: str):
     """Check whether to activate data module and preprocess datasets."""
-    from trinity.data.client import LOCAL_SERVER_URL, request
+    from trinity.cli.client import request
 
     logger.info("Activating data module...")
     res = request(
-        url=LOCAL_SERVER_URL,
+        url=data_workflow_url,
         configPath=config_path,
     )
     if res["return_code"] != 0:
         logger.error(f"Failed to activate data module: {res['return_msg']}.")
         return
 
 
+def run(config_path: str):
+    config = load_config(config_path)
+    config.check_and_update()
+    # try to activate data module
+    data_config = config.data
+    if data_config.data_workflow_url and (
+        data_config.dj_config_path or data_config.dj_process_desc
+    ):
+        activate_data_module(data_config.data_workflow_url, config_path)
+    ray.init()
+    if config.mode == "explore":
+        explore(config)
+    elif config.mode == "train":
+        train(config)
+    elif config.mode == "both":
+        both(config)
+
+
 def main() -> None:
     """The main entrypoint."""
     parser = argparse.ArgumentParser()
@@ -135,19 +153,7 @@ def main() -> None:
     args = parser.parse_args()
     if args.command == "run":
         # TODO: support parse all args from command line
-        config = load_config(args.config)
-        config.check_and_update()
-        # try to activate data module
-        data_config = config.data
-        if data_config.dj_config_path or data_config.dj_process_desc:
-            activate_data_module(args.config)
-        ray.init()
-        if config.mode == "explore":
-            explore(config)
-        elif config.mode == "train":
-            train(config)
-        elif config.mode == "both":
-            both(config)
+        run(args.config)
 
 
 if __name__ == "__main__":

trinity/cli/server.py

@@ -0,0 +1,32 @@
+import traceback
+
+import fire
+from flask import Flask, jsonify, request
+
+app = Flask(__name__)
+
+APP_NAME = "trinity_rft"
+
+
+@app.route(f"/{APP_NAME}", methods=["GET"])
+def trinity_training():
+    config_path = request.args.get("configPath")
+    try:
+        from trinity.cli.launcher import run
+
+        run(config_path)
+        ret = 0
+        msg = "Training Success."
+    except:  # noqa: E722
+        traceback.print_exc()
+        msg = traceback.format_exc()
+        ret = 1
+    return jsonify({"return_code": ret, "message": msg})
+
+
+def main(port=5006):
+    app.run(port=port, debug=True)
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

trinity/common/config.py

@@ -41,7 +41,8 @@ class FormatConfig:
 class DataConfig:
     """Data config"""
 
-    # TODO: add more
+    data_workflow_url: Optional[str] = None
+
     dataset_path: str = ""
     train_split: str = "train"
     eval_split: Optional[str] = None  # TODO: check data format