initial scaffold

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

Author: luarss

tools/AutoTuner/src/autotuner/distributed.py

@@ -40,6 +40,7 @@
 from uuid import uuid4 as uuid
 
 import numpy as np
+import torch
 
 import ray
 from ray import tune
@@ -805,7 +806,7 @@ def parse_arguments():
         help="Perturbation interval for PopulationBasedTraining.",
     )
     tune_parser.add_argument(
-        "--seed", type=int, metavar="<int>", default=42, help="Random seed."
+        "--seed", type=int, metavar="<int>", default=42, help="Random seed. (0 means no seed.)"
     )
 
     # Workload
@@ -870,10 +871,23 @@ def set_algorithm(experiment_name, config):
     """
     Configure search algorithm.
     """
+    # Pre-set seed if user sets seed to 0
+    if args.seed == 0:
+        print("Warning: you have chosen not to set a seed. Do you wish to continue? (y/n)")
+        if input().lower() != "y":
+            sys.exit(0)
+        args.seed = None
+
     if args.algorithm == "hyperopt":
-        algorithm = HyperOptSearch(points_to_evaluate=best_params)
+        algorithm = HyperOptSearch(
+            points_to_evaluate=best_params,
+            random_state_seed=args.seed,
+        )
     elif args.algorithm == "ax":
-        ax_client = AxClient(enforce_sequential_optimization=False)
+        ax_client = AxClient(
+            enforce_sequential_optimization=False,
+            random_seed=args.seed,
+        )
         AxClientMetric = namedtuple("AxClientMetric", "minimize")
         ax_client.create_experiment(
             name=experiment_name,
@@ -882,18 +896,33 @@ def set_algorithm(experiment_name, config):
         )
         algorithm = AxSearch(ax_client=ax_client, points_to_evaluate=best_params)
     elif args.algorithm == "optuna":
-        algorithm = OptunaSearch(points_to_evaluate=best_params, seed=args.seed)
+        algorithm = OptunaSearch(
+            points_to_evaluate=best_params,
+            seed=args.seed
+        )
     elif args.algorithm == "pbt":
+        print("Warning: PBT does not support seed values. args.seed will be ignored.")
         algorithm = PopulationBasedTraining(
             time_attr="training_iteration",
             perturbation_interval=args.perturbation,
             hyperparam_mutations=config,
             synch=True,
         )
     elif args.algorithm == "random":
-        algorithm = BasicVariantGenerator(max_concurrent=args.jobs)
+        algorithm = BasicVariantGenerator(
+            max_concurrent=args.jobs,
+            random_state=args.seed,)
+
+    # A wrapper algorithm for limiting the number of concurrent trials.
     if args.algorithm not in ["random", "pbt"]:
         algorithm = ConcurrencyLimiter(algorithm, max_concurrent=args.jobs)
+
+    # Self seed
+    if args.seed is not None:
+        torch.manual_seed(args.seed)
+        np.random.seed(args.seed)
+        random.seed(args.seed)
+
     return algorithm