add initial scaffold for distributed config

Signed-off-by: Jack Luar <[email protected]>

Author: luarss

tools/AutoTuner/distributed/README.md

@@ -76,31 +76,41 @@ make dashboard
 4. Ray CLI API
 
 ```bash
+# Run this once the dashboard is up
+HEAD_SERVER=$(ray job submit --address http://localhost:8265 -- python3 distributed/scripts/show_main_ip.py | grep "Main IP address" | awk '{print $4}')
+echo "HEAD_SERVER is at $HEAD_SERVER"
+export RAY_ADDRESS=$HEAD_SERVER:6379
+
 # Commands on machine (assume files/commands are present on cluster)
 ray job submit --address http://localhost:8265 ls
 
 # Case 1: 1 job
 ray job submit --address http://localhost:8265 -- python3 -m autotuner.distributed --design gcd --platform asap7 --config ../../flow/designs/asap7/gcd/autotuner.json --cloud_dir gs://autotuner_test tune --samples 1
 
 # Case 2A: 2 job, with resource spec.
-HEAD_SERVER=10.138.0.13
 ray job submit --address http://localhost:8265 --entrypoint-num-cpus 2 -- python3 -m autotuner.distributed --design gcd --platform asap7 --server $HEAD_SERVER --config ../../flow/designs/asap7/gcd/autotuner.json --cloud_dir gs://autotuner_test tune --samples 1
 ray job submit --address http://localhost:8265 --entrypoint-num-cpus 2 -- python3 -m autotuner.distributed --design gcd --platform asap7 --server $HEAD_SERVER --config ../../flow/designs/asap7/gcd/autotuner.json --cloud_dir gs://autotuner_test tune --samples 1
 
 # Case 2B: 2 job, with resource spec (sweep)
-HEAD_SERVER=10.138.0.13
 ray job submit --address http://localhost:8265 --entrypoint-num-cpus 2 -- python3 -m autotuner.distributed --design gcd --platform asap7 --server $HEAD_SERVER --config ./src/autotuner/distributed-sweep-example.json --cloud_dir gs://autotuner_test sweep
 ray job submit --address http://localhost:8265 --entrypoint-num-cpus 2 -- python3 -m autotuner.distributed --design gcd --platform asap7 --server $HEAD_SERVER --config ./src/autotuner/distributed-sweep-example.json --cloud_dir gs://autotuner_test sweep
 
 # Case 3: Overprovisioned resource spec (should fail because the cluster cannot meet this demand.)
-HEAD_SERVER=10.138.0.13
 ray job submit --address http://localhost:8265 --entrypoint-num-cpus 4 -- python3 -m autotuner.distributed --design gcd --platform asap7 --server $HEAD_SERVER --config ../../flow/designs/asap7/gcd/autotuner.json --cloud_dir gs://autotuner_test tune --samples 1
 
 # Commands on machine (sync local working dir, note the dir is stored as some /tmp dir)
 ray job submit --address http://localhost:8265 \
     --working-dir scripts -- python3 hello_world.py
 ```
 
+5. Ray Wrapper
+
+Alternatively, we also provided convenience scripts as an abstraction to the Ray Job Submission API.
+
+```bash
+python ray_wrapper.py --ray-dashboard-address localhost --ray-cluster-head-address $HEAD_ADDRESS
+```
+
 ## Useful commands
 
 ```bash

tools/AutoTuner/distributed/config.yaml

@@ -0,0 +1,20 @@
+#############################################################
+# This file is used to configure the distributed autotuner.
+#############################################################
+# Ray specific variables
+entrypoint_num_cpus: 2
+
+# Autotuner specific - common variables
+cloud_dir: gs://autotuner_test
+
+# Autotuner specific - tune variables
+# note for "config" - please specify a relative path to ./tools/AutoTuner
+design: gcd
+platform: asap7
+config: ../../flow/designs/asap7/gcd/autotuner.json
+mode_arg: tune
+samples: 1
+
+# Autotuner specific - sweep variables
+# config: ./src/autotuner/distributed-sweep-example.json
+# mode_arg: sweep

tools/AutoTuner/distributed/ray_wrapper.py

@@ -0,0 +1,91 @@
+import argparse
+import os
+import datetime
+from ray.job_submission import JobSubmissionClient
+
+from autotuner.utils.config_helper import parse_config
+
+# Constants
+FILE_DIR = os.path.dirname(os.path.abspath(__file__))
+EXPERIMENT_DATETIME = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+
+TUNE_TEMPLATE = """python3 -m autotuner.distributed \
+--design {design} \
+--platform {platform} \
+--config {config} \
+--server {server} \
+--cloud_dir {cloud_dir} \
+tune \
+--samples {samples}"""
+
+SWEEP_TEMPLATE = """python3 -m autotuner.distributed \
+--design {design} \
+--platform {platform} \
+--config {config} \
+--server {server} \
+--cloud_dir {cloud_dir} \
+sweep"""
+
+# CLI Arguments
+arguments = argparse.ArgumentParser()
+arguments.add_argument(
+    "--ray-dashboard-address",
+    type=str,
+    help="IP address of the Ray dashboard. If not provided, it will be set to the IP address of the current machine.",
+    default="localhost",
+)
+arguments.add_argument(
+    "--ray-cluster-head-address",
+    type=str,
+    help="IP address of the Ray cluster head. If not provided, it will be set to the IP address of the current machine.",
+    default="10.138.0.21",
+)
+
+# Parse configs
+# TODO: Validations + Schemas
+config_args = parse_config(os.path.join(FILE_DIR, "config.yaml"))
+
+
+if __name__ == "__main__":
+    args = arguments.parse_args()
+    ray_dashboard_address = args.ray_dashboard_address
+    ray_cluster_head_address = args.ray_cluster_head_address
+
+    # Initialize JobSubmissionClient
+    client = JobSubmissionClient(address=f"http://{ray_dashboard_address}:8265")
+
+    # Get ray parameters
+    entrypoint_num_cpus = config_args.get("entrypoint_num_cpus", 1)
+
+    # Get experiment parameters
+    cloud_dir = config_args.get("cloud_dir", "gs://autotuner_test")
+    design = config_args.get("design", "gcd")
+    platform = config_args.get("platform", "asap7")
+    config = config_args.get("config", "../../flow/designs/asap7/gcd/autotuner.json")
+    mode_arg = config_args.get("mode_arg", "tune")
+    samples = config_args.get("samples", 1)
+
+    # Submit a job
+    entrypoint = (
+        TUNE_TEMPLATE.format(
+            design=design,
+            platform=platform,
+            config=config,
+            server=ray_cluster_head_address,
+            cloud_dir=cloud_dir,
+            samples=samples,
+        )
+        if mode_arg == "tune"
+        else SWEEP_TEMPLATE.format(
+            design=design,
+            platform=platform,
+            config=config,
+            server=ray_cluster_head_address,
+            cloud_dir=cloud_dir,
+        )
+    )
+    job_id = client.submit_job(
+        entrypoint=entrypoint,
+        entrypoint_num_cpus=entrypoint_num_cpus,
+    )
+    print(f"Job submitted with ID: {job_id}")

tools/AutoTuner/distributed/scripts/show_main_ip.py

@@ -0,0 +1,30 @@
+import subprocess
+import re
+
+
+def get_main_ip():
+    """
+    Get the main IP address of the Ray head node.
+    """
+    result = subprocess.run(["ip", "a"], capture_output=True, text=True)
+    ip_output = result.stdout
+
+    # Regex to find all inet addresses (excluding loopback 127.x.x.x and docker/bridge)
+    matches = re.findall(r"inet (\d+\.\d+\.\d+\.\d+)/\d+", ip_output)
+    for ip in matches:
+        # TODO: Add more checks to filter out unwanted IPs
+        if (
+            not ip.startswith("127.")
+            and not ip.startswith("169.")
+            and not ip.startswith("172.17.")
+        ):
+            return ip
+    return "No valid IP found"
+
+
+def main():
+    print(f"Main IP address: {get_main_ip()}")
+
+
+if __name__ == "__main__":
+    main()

tools/AutoTuner/src/autotuner/utils/__init__.py

@@ -0,0 +1,18 @@
+import yaml
+
+
+def parse_config(file_path: str):
+    """
+    Parses a YAML configuration file and returns the content as a dictionary.
+
+    :param file_path: Path to the YAML configuration file.
+    :return: Dictionary containing the parsed configuration.
+    """
+    with open(file_path) as file:
+        config = yaml.safe_load(file)
+    return config
+
+
+if __name__ == "__main__":
+    config = parse_config("config.yaml")
+    print(config)