Add files using upload-large-folder tool

.venv/lib/python3.11/site-packages/ray/tune/examples/async_hyperband_example.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+
+import argparse
+import time
+
+from ray import train, tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+
+
+def evaluation_fn(step, width, height):
+    time.sleep(0.1)
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be an arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        train.report({"iterations": step, "mean_loss": intermediate_score})
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    # AsyncHyperBand enables aggressive early stopping of bad trials.
+    scheduler = AsyncHyperBandScheduler(grace_period=5, max_t=100)
+
+    # 'training_iteration' is incremented every time `trainable.step` is called
+    stopping_criteria = {"training_iteration": 1 if args.smoke_test else 9999}
+
+    tuner = tune.Tuner(
+        tune.with_resources(easy_objective, {"cpu": 1, "gpu": 0}),
+        run_config=train.RunConfig(
+            name="asynchyperband_test",
+            stop=stopping_criteria,
+            verbose=1,
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_loss", mode="min", scheduler=scheduler, num_samples=20
+        ),
+        param_space={  # Hyperparameter space
+            "steps": 100,
+            "width": tune.uniform(10, 100),
+            "height": tune.uniform(0, 100),
+        },
+    )
+    results = tuner.fit()
+    print("Best hyperparameters found were: ", results.get_best_result().config)

.venv/lib/python3.11/site-packages/ray/tune/examples/ax_example.py

@@ -0,0 +1,97 @@
+"""This example demonstrates the usage of AxSearch with Ray Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the Ax library to be installed (`pip install ax-platform`).
+"""
+
+import time
+
+import numpy as np
+
+from ray import train, tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search.ax import AxSearch
+
+
+def hartmann6(x):
+    alpha = np.array([1.0, 1.2, 3.0, 3.2])
+    A = np.array(
+        [
+            [10, 3, 17, 3.5, 1.7, 8],
+            [0.05, 10, 17, 0.1, 8, 14],
+            [3, 3.5, 1.7, 10, 17, 8],
+            [17, 8, 0.05, 10, 0.1, 14],
+        ]
+    )
+    P = 10 ** (-4) * np.array(
+        [
+            [1312, 1696, 5569, 124, 8283, 5886],
+            [2329, 4135, 8307, 3736, 1004, 9991],
+            [2348, 1451, 3522, 2883, 3047, 6650],
+            [4047, 8828, 8732, 5743, 1091, 381],
+        ]
+    )
+    y = 0.0
+    for j, alpha_j in enumerate(alpha):
+        t = 0
+        for k in range(6):
+            t += A[j, k] * ((x[k] - P[j, k]) ** 2)
+        y -= alpha_j * np.exp(-t)
+    return y
+
+
+def easy_objective(config):
+    for i in range(config["iterations"]):
+        x = np.array([config.get("x{}".format(i + 1)) for i in range(6)])
+        train.report(
+            {
+                "timesteps_total": i,
+                "hartmann6": hartmann6(x),
+                "l2norm": np.sqrt((x**2).sum()),
+            }
+        )
+        time.sleep(0.02)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    algo = AxSearch(
+        parameter_constraints=["x1 + x2 <= 2.0"],  # Optional.
+        outcome_constraints=["l2norm <= 1.25"],  # Optional.
+    )
+    # Limit to 4 concurrent trials
+    algo = tune.search.ConcurrencyLimiter(algo, max_concurrent=4)
+    scheduler = AsyncHyperBandScheduler()
+    tuner = tune.Tuner(
+        easy_objective,
+        run_config=train.RunConfig(
+            name="ax",
+            stop={"timesteps_total": 100},
+        ),
+        tune_config=tune.TuneConfig(
+            metric="hartmann6",  # provided in the 'easy_objective' function
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if args.smoke_test else 50,
+        ),
+        param_space={
+            "iterations": 100,
+            "x1": tune.uniform(0.0, 1.0),
+            "x2": tune.uniform(0.0, 1.0),
+            "x3": tune.uniform(0.0, 1.0),
+            "x4": tune.uniform(0.0, 1.0),
+            "x5": tune.uniform(0.0, 1.0),
+            "x6": tune.uniform(0.0, 1.0),
+        },
+    )
+    results = tuner.fit()
+    print("Best hyperparameters found were: ", results.get_best_result().config)

.venv/lib/python3.11/site-packages/ray/tune/examples/bayesopt_example.py

@@ -0,0 +1,63 @@
+"""This example demonstrates the usage of BayesOpt with Ray Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the BayesOpt library to be installed (`pip install bayesian-optimization`).
+"""
+import time
+
+from ray import train, tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.bayesopt import BayesOptSearch
+
+
+def evaluation_fn(step, width, height):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        train.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    algo = BayesOptSearch(utility_kwargs={"kind": "ucb", "kappa": 2.5, "xi": 0.0})
+    algo = ConcurrencyLimiter(algo, max_concurrent=4)
+    scheduler = AsyncHyperBandScheduler()
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if args.smoke_test else 1000,
+        ),
+        run_config=train.RunConfig(
+            name="my_exp",
+        ),
+        param_space={
+            "steps": 100,
+            "width": tune.uniform(0, 20),
+            "height": tune.uniform(-100, 100),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

.venv/lib/python3.11/site-packages/ray/tune/examples/bohb_example.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python
+
+"""This example demonstrates the usage of BOHB with Ray Tune.
+
+Requires the HpBandSter and ConfigSpace libraries to be installed
+(`pip install hpbandster ConfigSpace`).
+"""
+
+import json
+import os
+import time
+
+import numpy as np
+
+import ray
+from ray import train, tune
+from ray.tune import Trainable
+from ray.tune.schedulers.hb_bohb import HyperBandForBOHB
+from ray.tune.search.bohb import TuneBOHB
+
+
+class MyTrainableClass(Trainable):
+    """Example agent whose learning curve is a random sigmoid.
+
+    The dummy hyperparameters "width" and "height" determine the slope and
+    maximum reward value reached.
+    """
+
+    def setup(self, config):
+        self.timestep = 0
+
+    def step(self):
+        self.timestep += 1
+        v = np.tanh(float(self.timestep) / self.config.get("width", 1))
+        v *= self.config.get("height", 1)
+        time.sleep(0.1)
+        # Here we use `episode_reward_mean`, but you can also report other
+        # objectives such as loss or accuracy.
+        return {"episode_reward_mean": v}
+
+    def save_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint")
+        with open(path, "w") as f:
+            f.write(json.dumps({"timestep": self.timestep}))
+
+    def load_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint")
+        with open(path, "r") as f:
+            self.timestep = json.loads(f.read())["timestep"]
+
+
+if __name__ == "__main__":
+    import sys
+
+    if sys.version_info >= (3, 12):
+        # TuneBOHB is not compatible with Python 3.12
+        sys.exit(0)
+
+    ray.init(num_cpus=8)
+
+    config = {
+        "iterations": 100,
+        "width": tune.uniform(0, 20),
+        "height": tune.uniform(-100, 100),
+        "activation": tune.choice(["relu", "tanh"]),
+    }
+
+    # Optional: Pass the parameter space yourself
+    # import ConfigSpace as CS
+    # config_space = CS.ConfigurationSpace()
+    # config_space.add_hyperparameter(
+    #     CS.UniformFloatHyperparameter("width", lower=0, upper=20))
+    # config_space.add_hyperparameter(
+    #     CS.UniformFloatHyperparameter("height", lower=-100, upper=100))
+    # config_space.add_hyperparameter(
+    #     CS.CategoricalHyperparameter(
+    #         "activation", choices=["relu", "tanh"]))
+
+    max_iterations = 10
+    bohb_hyperband = HyperBandForBOHB(
+        time_attr="training_iteration",
+        max_t=max_iterations,
+        reduction_factor=2,
+        stop_last_trials=False,
+    )
+
+    bohb_search = TuneBOHB(
+        # space=config_space,  # If you want to set the space manually
+    )
+    bohb_search = tune.search.ConcurrencyLimiter(bohb_search, max_concurrent=4)
+
+    tuner = tune.Tuner(
+        MyTrainableClass,
+        run_config=train.RunConfig(
+            name="bohb_test", stop={"training_iteration": max_iterations}
+        ),
+        tune_config=tune.TuneConfig(
+            metric="episode_reward_mean",
+            mode="max",
+            scheduler=bohb_hyperband,
+            search_alg=bohb_search,
+            num_samples=32,
+        ),
+        param_space=config,
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

.venv/lib/python3.11/site-packages/ray/tune/examples/custom_func_checkpointing.py

@@ -0,0 +1,70 @@
+# If want to use checkpointing with a custom training function (not a Ray
+# integration like PyTorch or Tensorflow), your function can read/write
+# checkpoint through the ``ray.train.report(metrics, checkpoint=...)`` API.
+import argparse
+import json
+import os
+import tempfile
+import time
+
+from ray import train, tune
+from ray.train import Checkpoint
+
+
+def evaluation_fn(step, width, height):
+    time.sleep(0.1)
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def train_func(config):
+    step = 0
+    width, height = config["width"], config["height"]
+
+    checkpoint = train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            with open(os.path.join(checkpoint_dir, "checkpoint.json")) as f:
+                state = json.load(f)
+            step = state["step"] + 1
+
+    for current_step in range(step, 100):
+        intermediate_score = evaluation_fn(current_step, width, height)
+
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            with open(os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w") as f:
+                json.dump({"step": current_step}, f)
+            train.report(
+                {"iterations": current_step, "mean_loss": intermediate_score},
+                checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
+            )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    tuner = tune.Tuner(
+        train_func,
+        run_config=train.RunConfig(
+            name="hyperband_test",
+            stop={"training_iteration": 1 if args.smoke_test else 10},
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            num_samples=5,
+        ),
+        param_space={
+            "steps": 10,
+            "width": tune.randint(10, 100),
+            "height": tune.loguniform(10, 100),
+        },
+    )
+    results = tuner.fit()
+    best_result = results.get_best_result()
+    print("Best hyperparameters: ", best_result.config)
+    best_checkpoint = best_result.checkpoint
+    print("Best checkpoint: ", best_checkpoint)

.venv/lib/python3.11/site-packages/ray/tune/examples/hyperband_example.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python
+
+import argparse
+
+import ray
+from ray import train, tune
+from ray.tune.schedulers import HyperBandScheduler
+from ray.tune.utils.mock_trainable import MyTrainableClass
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(num_cpus=4 if args.smoke_test else None)
+
+    # Hyperband early stopping, configured with `episode_reward_mean` as the
+    # objective and `training_iteration` as the time unit,
+    # which is automatically filled by Tune.
+    hyperband = HyperBandScheduler(time_attr="training_iteration", max_t=200)
+
+    tuner = tune.Tuner(
+        MyTrainableClass,
+        run_config=train.RunConfig(
+            name="hyperband_test",
+            stop={"training_iteration": 1 if args.smoke_test else 200},
+            verbose=1,
+            failure_config=train.FailureConfig(
+                fail_fast=True,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            num_samples=20 if args.smoke_test else 200,
+            metric="episode_reward_mean",
+            mode="max",
+            scheduler=hyperband,
+        ),
+        param_space={"width": tune.randint(10, 90), "height": tune.randint(0, 100)},
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)

.venv/lib/python3.11/site-packages/ray/tune/examples/hyperband_function_example.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+
+import argparse
+import json
+import os
+import tempfile
+
+import numpy as np
+
+import ray
+from ray import train, tune
+from ray.train import Checkpoint
+from ray.tune.schedulers import HyperBandScheduler
+
+
+def train_func(config):
+    step = 0
+    checkpoint = train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            with open(os.path.join(checkpoint_dir, "checkpoint.json")) as f:
+                step = json.load(f)["timestep"] + 1
+
+    for timestep in range(step, 100):
+        v = np.tanh(float(timestep) / config.get("width", 1))
+        v *= config.get("height", 1)
+
+        # Checkpoint the state of the training every 3 steps
+        # Note that this is only required for certain schedulers
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            checkpoint = None
+            if timestep % 3 == 0:
+                with open(
+                    os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w"
+                ) as f:
+                    json.dump({"timestep": timestep}, f)
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+            # Here we use `episode_reward_mean`, but you can also report other
+            # objectives such as loss or accuracy.
+            train.report({"episode_reward_mean": v}, checkpoint=checkpoint)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(num_cpus=4 if args.smoke_test else None)
+
+    # Hyperband early stopping, configured with `episode_reward_mean` as the
+    # objective and `training_iteration` as the time unit,
+    # which is automatically filled by Tune.
+    hyperband = HyperBandScheduler(max_t=200)
+
+    tuner = tune.Tuner(
+        train_func,
+        run_config=train.RunConfig(
+            name="hyperband_test",
+            stop={"training_iteration": 10 if args.smoke_test else 99999},
+            failure_config=train.FailureConfig(
+                fail_fast=True,
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            num_samples=20,
+            metric="episode_reward_mean",
+            mode="max",
+            scheduler=hyperband,
+        ),
+        param_space={"height": tune.uniform(0, 100)},
+    )
+    results = tuner.fit()
+    print("Best hyperparameters found were: ", results.get_best_result().config)

.venv/lib/python3.11/site-packages/ray/tune/examples/hyperopt_conditional_search_space_example.py

@@ -0,0 +1,110 @@
+"""This example demonstrates the usage of conditional search spaces with Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the HyperOpt library to be installed (`pip install hyperopt`).
+
+For an example of using a Tune search space, see
+:doc:`/tune/examples/hyperopt_example`.
+"""
+
+import time
+
+from hyperopt import hp
+
+import ray
+from ray import train, tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.hyperopt import HyperOptSearch
+
+
+def f_unpack_dict(dct):
+    """
+    Unpacks all sub-dictionaries in given dictionary recursively.
+    There should be no duplicated keys across all nested
+    subdictionaries, or some instances will be lost without warning
+
+    Source: https://www.kaggle.com/fanvacoolt/tutorial-on-hyperopt
+
+    Parameters:
+    ----------------
+    dct : dictionary to unpack
+
+    Returns:
+    ----------------
+    : unpacked dictionary
+    """
+
+    res = {}
+    for k, v in dct.items():
+        if isinstance(v, dict):
+            res = {**res, **f_unpack_dict(v)}
+        else:
+            res[k] = v
+
+    return res
+
+
+def evaluation_fn(step, width, height, mult=1):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1 * mult
+
+
+def easy_objective(config_in):
+    # Hyperparameters
+    config = f_unpack_dict(config_in)
+    width, height, mult = config["width"], config["height"], config.get("mult", 1)
+    print(config)
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height, mult)
+        # Feed the score back back to Tune.
+        train.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+config_space = {
+    "activation": hp.choice(
+        "activation",
+        [
+            {"activation": "relu", "mult": hp.uniform("mult", 1, 2)},
+            {"activation": "tanh"},
+        ],
+    ),
+    "width": hp.uniform("width", 0, 20),
+    "height": hp.uniform("heright", -100, 100),
+    "steps": 100,
+}
+
+
+def run_hyperopt_tune(config_dict=config_space, smoke_test=False):
+    algo = HyperOptSearch(space=config_dict, metric="mean_loss", mode="min")
+    algo = ConcurrencyLimiter(algo, max_concurrent=4)
+    scheduler = AsyncHyperBandScheduler()
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if smoke_test else 100,
+        ),
+    )
+    results = tuner.fit()
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(configure_logging=False)
+
+    run_hyperopt_tune(smoke_test=args.smoke_test)

.venv/lib/python3.11/site-packages/ray/tune/examples/logging_example.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python
+
+import argparse
+import time
+
+from ray import train, tune
+from ray.tune.logger import LoggerCallback
+
+
+class TestLoggerCallback(LoggerCallback):
+    def on_trial_result(self, iteration, trials, trial, result, **info):
+        print(f"TestLogger for trial {trial}: {result}")
+
+
+def trial_str_creator(trial):
+    return "{}_{}_123".format(trial.trainable_name, trial.trial_id)
+
+
+def evaluation_fn(step, width, height):
+    time.sleep(0.1)
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        train.report({"iterations": step, "mean_loss": intermediate_score})
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    tuner = tune.Tuner(
+        easy_objective,
+        run_config=train.RunConfig(
+            name="hyperband_test",
+            callbacks=[TestLoggerCallback()],
+            stop={"training_iteration": 1 if args.smoke_test else 100},
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            num_samples=5,
+            trial_name_creator=trial_str_creator,
+            trial_dirname_creator=trial_str_creator,
+        ),
+        param_space={
+            "steps": 100,
+            "width": tune.randint(10, 100),
+            "height": tune.loguniform(10, 100),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters: ", results.get_best_result().config)

.venv/lib/python3.11/site-packages/ray/tune/examples/mlflow_example.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+"""Examples using MLfowLoggerCallback and setup_mlflow.
+"""
+import os
+import tempfile
+import time
+
+import mlflow
+
+from ray import train, tune
+from ray.air.integrations.mlflow import MLflowLoggerCallback, setup_mlflow
+
+
+def evaluation_fn(step, width, height):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def train_function(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config.get("steps", 100)):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back to Tune.
+        train.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+def tune_with_callback(mlflow_tracking_uri, finish_fast=False):
+
+    tuner = tune.Tuner(
+        train_function,
+        run_config=train.RunConfig(
+            name="mlflow",
+            callbacks=[
+                MLflowLoggerCallback(
+                    tracking_uri=mlflow_tracking_uri,
+                    experiment_name="example",
+                    save_artifact=True,
+                )
+            ],
+        ),
+        tune_config=tune.TuneConfig(
+            num_samples=5,
+        ),
+        param_space={
+            "width": tune.randint(10, 100),
+            "height": tune.randint(0, 100),
+            "steps": 5 if finish_fast else 100,
+        },
+    )
+    tuner.fit()
+
+
+def train_function_mlflow(config):
+    setup_mlflow(config)
+
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config.get("steps", 100)):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Log the metrics to mlflow
+        mlflow.log_metrics(dict(mean_loss=intermediate_score), step=step)
+        # Feed the score back to Tune.
+        train.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+def tune_with_setup(mlflow_tracking_uri, finish_fast=False):
+    # Set the experiment, or create a new one if does not exist yet.
+    mlflow.set_tracking_uri(mlflow_tracking_uri)
+    mlflow.set_experiment(experiment_name="mixin_example")
+    tuner = tune.Tuner(
+        train_function_mlflow,
+        run_config=train.RunConfig(
+            name="mlflow",
+        ),
+        tune_config=tune.TuneConfig(
+            num_samples=5,
+        ),
+        param_space={
+            "width": tune.randint(10, 100),
+            "height": tune.randint(0, 100),
+            "steps": 5 if finish_fast else 100,
+            "mlflow": {
+                "experiment_name": "mixin_example",
+                "tracking_uri": mlflow.get_tracking_uri(),
+            },
+        },
+    )
+    tuner.fit()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    parser.add_argument(
+        "--tracking-uri",
+        type=str,
+        help="The tracking URI for the MLflow tracking server.",
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        mlflow_tracking_uri = os.path.join(tempfile.gettempdir(), "mlruns")
+    else:
+        mlflow_tracking_uri = args.tracking_uri
+
+    tune_with_callback(mlflow_tracking_uri, finish_fast=args.smoke_test)
+    if not args.smoke_test:
+        df = mlflow.search_runs(
+            [mlflow.get_experiment_by_name("example").experiment_id]
+        )
+        print(df)
+
+    tune_with_setup(mlflow_tracking_uri, finish_fast=args.smoke_test)
+    if not args.smoke_test:
+        df = mlflow.search_runs(
+            [mlflow.get_experiment_by_name("mixin_example").experiment_id]
+        )
+        print(df)

.venv/lib/python3.11/site-packages/ray/tune/examples/mlflow_ptl.py

@@ -0,0 +1,105 @@
+"""An example showing how to use Pytorch Lightning training, Ray Tune
+HPO, and MLflow autologging all together."""
+
+import os
+import tempfile
+
+import mlflow
+import pytorch_lightning as pl
+
+from ray import train, tune
+from ray.air.integrations.mlflow import setup_mlflow
+from ray.tune.examples.mnist_ptl_mini import LightningMNISTClassifier, MNISTDataModule
+from ray.tune.integration.pytorch_lightning import TuneReportCallback
+
+
+def train_mnist_tune(config, data_dir=None, num_epochs=10, num_gpus=0):
+    setup_mlflow(
+        config,
+        experiment_name=config.get("experiment_name", None),
+        tracking_uri=config.get("tracking_uri", None),
+    )
+
+    model = LightningMNISTClassifier(config, data_dir)
+    dm = MNISTDataModule(
+        data_dir=data_dir, num_workers=1, batch_size=config["batch_size"]
+    )
+    metrics = {"loss": "ptl/val_loss", "acc": "ptl/val_accuracy"}
+    mlflow.pytorch.autolog()
+    trainer = pl.Trainer(
+        max_epochs=num_epochs,
+        gpus=num_gpus,
+        progress_bar_refresh_rate=0,
+        callbacks=[TuneReportCallback(metrics, on="validation_end")],
+    )
+    trainer.fit(model, dm)
+
+
+def tune_mnist(
+    num_samples=10,
+    num_epochs=10,
+    gpus_per_trial=0,
+    tracking_uri=None,
+    experiment_name="ptl_autologging_example",
+):
+    data_dir = os.path.join(tempfile.gettempdir(), "mnist_data_")
+    # Download data
+    MNISTDataModule(data_dir=data_dir, batch_size=32).prepare_data()
+
+    # Set the MLflow experiment, or create it if it does not exist.
+    mlflow.set_tracking_uri(tracking_uri)
+    mlflow.set_experiment(experiment_name)
+
+    config = {
+        "layer_1": tune.choice([32, 64, 128]),
+        "layer_2": tune.choice([64, 128, 256]),
+        "lr": tune.loguniform(1e-4, 1e-1),
+        "batch_size": tune.choice([32, 64, 128]),
+        "experiment_name": experiment_name,
+        "tracking_uri": mlflow.get_tracking_uri(),
+        "data_dir": os.path.join(tempfile.gettempdir(), "mnist_data_"),
+        "num_epochs": num_epochs,
+    }
+
+    trainable = tune.with_parameters(
+        train_mnist_tune,
+        data_dir=data_dir,
+        num_epochs=num_epochs,
+        num_gpus=gpus_per_trial,
+    )
+
+    tuner = tune.Tuner(
+        tune.with_resources(trainable, resources={"cpu": 1, "gpu": gpus_per_trial}),
+        tune_config=tune.TuneConfig(
+            metric="loss",
+            mode="min",
+            num_samples=num_samples,
+        ),
+        run_config=train.RunConfig(
+            name="tune_mnist",
+        ),
+        param_space=config,
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        tune_mnist(
+            num_samples=1,
+            num_epochs=1,
+            gpus_per_trial=0,
+            tracking_uri=os.path.join(tempfile.gettempdir(), "mlruns"),
+        )
+    else:
+        tune_mnist(num_samples=10, num_epochs=10, gpus_per_trial=0)

.venv/lib/python3.11/site-packages/ray/tune/examples/mnist_ptl_mini.py

@@ -0,0 +1,166 @@
+import math
+import os
+
+import pytorch_lightning as pl
+import torch
+from datasets import load_dataset
+from filelock import FileLock
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchmetrics import Accuracy
+from torchvision import transforms
+
+from ray import train, tune
+from ray.tune.integration.pytorch_lightning import TuneReportCheckpointCallback
+
+PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
+
+
+class MNISTDataModule(pl.LightningDataModule):
+    def __init__(self, batch_size: int, data_dir: str = PATH_DATASETS):
+        super().__init__()
+        self.data_dir = data_dir
+        self.transform = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize((0.1307,), (0.3081,)),
+            ]
+        )
+        self.batch_size = batch_size
+        self.dims = (1, 28, 28)
+        self.num_classes = 10
+
+    def prepare_data(self):
+        # download
+        with FileLock(os.path.expanduser("~/.data.lock")):
+            load_dataset("ylecun/mnist", cache_dir=self.data_dir)
+
+    def setup(self, stage=None):
+        dataset = load_dataset("ylecun/mnist", cache_dir=self.data_dir)
+
+        def transform_fn(sample):
+            return (self.transform(sample["image"]), sample["label"])
+
+        self.mnist_train = [transform_fn(sample) for sample in dataset["train"]]
+        self.mnist_val = [transform_fn(sample) for sample in dataset["test"]]
+
+    def train_dataloader(self):
+        return DataLoader(self.mnist_train, batch_size=self.batch_size)
+
+    def val_dataloader(self):
+        return DataLoader(self.mnist_val, batch_size=self.batch_size)
+
+
+class LightningMNISTClassifier(pl.LightningModule):
+    def __init__(self, config, data_dir=None):
+        super(LightningMNISTClassifier, self).__init__()
+
+        self.data_dir = data_dir or os.getcwd()
+        self.lr = config["lr"]
+        layer_1, layer_2 = config["layer_1"], config["layer_2"]
+        self.batch_size = config["batch_size"]
+
+        # mnist images are (1, 28, 28) (channels, width, height)
+        self.layer_1 = torch.nn.Linear(28 * 28, layer_1)
+        self.layer_2 = torch.nn.Linear(layer_1, layer_2)
+        self.layer_3 = torch.nn.Linear(layer_2, 10)
+        self.accuracy = Accuracy(task="multiclass", num_classes=10, top_k=1)
+
+    def forward(self, x):
+        batch_size, channels, width, height = x.size()
+        x = x.view(batch_size, -1)
+        x = self.layer_1(x)
+        x = torch.relu(x)
+        x = self.layer_2(x)
+        x = torch.relu(x)
+        x = self.layer_3(x)
+        x = torch.log_softmax(x, dim=1)
+        return x
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.lr)
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        logits = self.forward(x)
+        loss = F.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        self.log("ptl/train_loss", loss)
+        self.log("ptl/train_accuracy", acc)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        logits = self.forward(x)
+        loss = F.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        return {"val_loss": loss, "val_accuracy": acc}
+
+    def validation_epoch_end(self, outputs):
+        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
+        avg_acc = torch.stack([x["val_accuracy"] for x in outputs]).mean()
+        self.log("ptl/val_loss", avg_loss)
+        self.log("ptl/val_accuracy", avg_acc)
+
+
+def train_mnist_tune(config, num_epochs=10, num_gpus=0):
+    data_dir = os.path.abspath("./data")
+    model = LightningMNISTClassifier(config, data_dir)
+    with FileLock(os.path.expanduser("~/.data.lock")):
+        dm = MNISTDataModule(data_dir=data_dir, batch_size=config["batch_size"])
+    metrics = {"loss": "ptl/val_loss", "acc": "ptl/val_accuracy"}
+    trainer = pl.Trainer(
+        max_epochs=num_epochs,
+        # If fractional GPUs passed in, convert to int.
+        gpus=math.ceil(num_gpus),
+        enable_progress_bar=False,
+        callbacks=[
+            TuneReportCheckpointCallback(
+                metrics, on="validation_end", save_checkpoints=False
+            )
+        ],
+    )
+    trainer.fit(model, dm)
+
+
+def tune_mnist(num_samples=10, num_epochs=10, gpus_per_trial=0):
+    config = {
+        "layer_1": tune.choice([32, 64, 128]),
+        "layer_2": tune.choice([64, 128, 256]),
+        "lr": tune.loguniform(1e-4, 1e-1),
+        "batch_size": tune.choice([32, 64, 128]),
+    }
+
+    trainable = tune.with_parameters(
+        train_mnist_tune, num_epochs=num_epochs, num_gpus=gpus_per_trial
+    )
+    tuner = tune.Tuner(
+        tune.with_resources(trainable, resources={"cpu": 1, "gpu": gpus_per_trial}),
+        tune_config=tune.TuneConfig(
+            metric="loss",
+            mode="min",
+            num_samples=num_samples,
+        ),
+        run_config=train.RunConfig(
+            name="tune_mnist",
+        ),
+        param_space=config,
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        tune_mnist(num_samples=1, num_epochs=1, gpus_per_trial=0)
+    else:
+        tune_mnist(num_samples=10, num_epochs=10, gpus_per_trial=0)

.venv/lib/python3.11/site-packages/ray/tune/examples/mnist_pytorch.py

@@ -0,0 +1,161 @@
+# Original Code here:
+# https://github.com/pytorch/examples/blob/master/mnist/main.py
+
+import argparse
+import os
+import tempfile
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from filelock import FileLock
+from torchvision import datasets, transforms
+
+import ray
+from ray import train, tune
+from ray.train import Checkpoint
+from ray.tune.schedulers import AsyncHyperBandScheduler
+
+# Change these values if you want the training to run quicker or slower.
+EPOCH_SIZE = 512
+TEST_SIZE = 256
+
+
+class ConvNet(nn.Module):
+    def __init__(self):
+        super(ConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(1, 3, kernel_size=3)
+        self.fc = nn.Linear(192, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 3))
+        x = x.view(-1, 192)
+        x = self.fc(x)
+        return F.log_softmax(x, dim=1)
+
+
+def train_func(model, optimizer, train_loader, device=None):
+    device = device or torch.device("cpu")
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if batch_idx * len(data) > EPOCH_SIZE:
+            return
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+
+
+def test_func(model, data_loader, device=None):
+    device = device or torch.device("cpu")
+    model.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for batch_idx, (data, target) in enumerate(data_loader):
+            if batch_idx * len(data) > TEST_SIZE:
+                break
+            data, target = data.to(device), target.to(device)
+            outputs = model(data)
+            _, predicted = torch.max(outputs.data, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+    return correct / total
+
+
+def get_data_loaders(batch_size=64):
+    mnist_transforms = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+    )
+
+    # We add FileLock here because multiple workers will want to
+    # download data, and this may cause overwrites since
+    # DataLoader is not threadsafe.
+    with FileLock(os.path.expanduser("~/data.lock")):
+        train_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "~/data", train=True, download=True, transform=mnist_transforms
+            ),
+            batch_size=batch_size,
+            shuffle=True,
+        )
+        test_loader = torch.utils.data.DataLoader(
+            datasets.MNIST(
+                "~/data", train=False, download=True, transform=mnist_transforms
+            ),
+            batch_size=batch_size,
+            shuffle=True,
+        )
+    return train_loader, test_loader
+
+
+def train_mnist(config):
+    should_checkpoint = config.get("should_checkpoint", False)
+    use_cuda = torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+    train_loader, test_loader = get_data_loaders()
+    model = ConvNet().to(device)
+
+    optimizer = optim.SGD(
+        model.parameters(), lr=config["lr"], momentum=config["momentum"]
+    )
+
+    while True:
+        train_func(model, optimizer, train_loader, device)
+        acc = test_func(model, test_loader, device)
+        metrics = {"mean_accuracy": acc}
+
+        # Report metrics (and possibly a checkpoint)
+        if should_checkpoint:
+            with tempfile.TemporaryDirectory() as tempdir:
+                torch.save(model.state_dict(), os.path.join(tempdir, "model.pt"))
+                train.report(metrics, checkpoint=Checkpoint.from_directory(tempdir))
+        else:
+            train.report(metrics)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
+    parser.add_argument(
+        "--cuda", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init(num_cpus=2 if args.smoke_test else None)
+
+    # for early stopping
+    sched = AsyncHyperBandScheduler()
+
+    resources_per_trial = {"cpu": 2, "gpu": int(args.cuda)}  # set this for GPUs
+    tuner = tune.Tuner(
+        tune.with_resources(train_mnist, resources=resources_per_trial),
+        tune_config=tune.TuneConfig(
+            metric="mean_accuracy",
+            mode="max",
+            scheduler=sched,
+            num_samples=1 if args.smoke_test else 50,
+        ),
+        run_config=train.RunConfig(
+            name="exp",
+            stop={
+                "mean_accuracy": 0.98,
+                "training_iteration": 5 if args.smoke_test else 100,
+            },
+        ),
+        param_space={
+            "lr": tune.loguniform(1e-4, 1e-2),
+            "momentum": tune.uniform(0.1, 0.9),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best config is:", results.get_best_result().config)
+
+    assert not results.errors

.venv/lib/python3.11/site-packages/ray/tune/examples/mnist_pytorch_trainable.py

@@ -0,0 +1,98 @@
+# Original Code here:
+# https://github.com/pytorch/examples/blob/master/mnist/main.py
+from __future__ import print_function
+
+import argparse
+import os
+
+import torch
+import torch.optim as optim
+
+import ray
+from ray import train, tune
+from ray.tune.examples.mnist_pytorch import (
+    ConvNet,
+    get_data_loaders,
+    test_func,
+    train_func,
+)
+from ray.tune.schedulers import ASHAScheduler
+
+# Change these values if you want the training to run quicker or slower.
+EPOCH_SIZE = 512
+TEST_SIZE = 256
+
+# Training settings
+parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
+parser.add_argument(
+    "--use-gpu", action="store_true", default=False, help="enables CUDA training"
+)
+parser.add_argument("--ray-address", type=str, help="The Redis address of the cluster.")
+parser.add_argument(
+    "--smoke-test", action="store_true", help="Finish quickly for testing"
+)
+
+
+# Below comments are for documentation purposes only.
+# fmt: off
+# __trainable_example_begin__
+class TrainMNIST(tune.Trainable):
+    def setup(self, config):
+        use_cuda = config.get("use_gpu") and torch.cuda.is_available()
+        self.device = torch.device("cuda" if use_cuda else "cpu")
+        self.train_loader, self.test_loader = get_data_loaders()
+        self.model = ConvNet().to(self.device)
+        self.optimizer = optim.SGD(
+            self.model.parameters(),
+            lr=config.get("lr", 0.01),
+            momentum=config.get("momentum", 0.9))
+
+    def step(self):
+        train_func(
+            self.model, self.optimizer, self.train_loader, device=self.device)
+        acc = test_func(self.model, self.test_loader, self.device)
+        return {"mean_accuracy": acc}
+
+    def save_checkpoint(self, checkpoint_dir):
+        checkpoint_path = os.path.join(checkpoint_dir, "model.pth")
+        torch.save(self.model.state_dict(), checkpoint_path)
+
+    def load_checkpoint(self, checkpoint_dir):
+        checkpoint_path = os.path.join(checkpoint_dir, "model.pth")
+        self.model.load_state_dict(torch.load(checkpoint_path))
+
+
+# __trainable_example_end__
+# fmt: on
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    ray.init(address=args.ray_address, num_cpus=6 if args.smoke_test else None)
+    sched = ASHAScheduler()
+
+    tuner = tune.Tuner(
+        tune.with_resources(TrainMNIST, resources={"cpu": 3, "gpu": int(args.use_gpu)}),
+        run_config=train.RunConfig(
+            stop={
+                "mean_accuracy": 0.95,
+                "training_iteration": 3 if args.smoke_test else 20,
+            },
+            checkpoint_config=train.CheckpointConfig(
+                checkpoint_at_end=True, checkpoint_frequency=3
+            ),
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_accuracy",
+            mode="max",
+            scheduler=sched,
+            num_samples=1 if args.smoke_test else 20,
+        ),
+        param_space={
+            "args": args,
+            "lr": tune.uniform(0.001, 0.1),
+            "momentum": tune.uniform(0.1, 0.9),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best config is:", results.get_best_result().config)

.venv/lib/python3.11/site-packages/ray/tune/examples/nevergrad_example.py

@@ -0,0 +1,77 @@
+"""This example demonstrates the usage of Nevergrad with Ray Tune.
+
+It also checks that it is usable with a separate scheduler.
+
+Requires the Nevergrad library to be installed (`pip install nevergrad`).
+"""
+
+import time
+
+from ray import train, tune
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from ray.tune.search import ConcurrencyLimiter
+from ray.tune.search.nevergrad import NevergradSearch
+
+
+def evaluation_fn(step, width, height):
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        train.report({"iterations": step, "mean_loss": intermediate_score})
+        time.sleep(0.1)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    import nevergrad as ng
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    # Optional: Pass the parameter space yourself
+    # space = ng.p.Dict(
+    #     width=ng.p.Scalar(lower=0, upper=20),
+    #     height=ng.p.Scalar(lower=-100, upper=100),
+    #     activation=ng.p.Choice(choices=["relu", "tanh"])
+    # )
+
+    algo = NevergradSearch(
+        optimizer=ng.optimizers.OnePlusOne,
+        # space=space,  # If you want to set the space manually
+    )
+    algo = ConcurrencyLimiter(algo, max_concurrent=4)
+
+    scheduler = AsyncHyperBandScheduler()
+
+    tuner = tune.Tuner(
+        easy_objective,
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            search_alg=algo,
+            scheduler=scheduler,
+            num_samples=10 if args.smoke_test else 50,
+        ),
+        run_config=train.RunConfig(name="nevergrad"),
+        param_space={
+            "steps": 100,
+            "width": tune.uniform(0, 20),
+            "height": tune.uniform(-100, 100),
+            "activation": tune.choice(["relu", "tanh"]),
+        },
+    )
+    results = tuner.fit()
+
+    print("Best hyperparameters found were: ", results.get_best_result().config)