Add files using upload-large-folder tool

.venv/lib/python3.11/site-packages/ray/tune/cli/commands.py

@@ -0,0 +1,306 @@
+import logging
+import operator
+import os
+import shutil
+import subprocess
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional
+
+import click
+import pandas as pd
+from pandas.api.types import is_numeric_dtype, is_string_dtype
+
+from ray._private.thirdparty.tabulate.tabulate import tabulate
+from ray.air.constants import EXPR_RESULT_FILE
+from ray.tune import TuneError
+from ray.tune.analysis import ExperimentAnalysis
+from ray.tune.result import (
+    CONFIG_PREFIX,
+    DEFAULT_EXPERIMENT_INFO_KEYS,
+    DEFAULT_RESULT_KEYS,
+)
+
+logger = logging.getLogger(__name__)
+
+EDITOR = os.getenv("EDITOR", "vim")
+
+TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S (%A)"
+
+DEFAULT_CLI_KEYS = DEFAULT_EXPERIMENT_INFO_KEYS + DEFAULT_RESULT_KEYS
+
+DEFAULT_PROJECT_INFO_KEYS = (
+    "name",
+    "total_trials",
+    "last_updated",
+)
+
+TERM_WIDTH, TERM_HEIGHT = shutil.get_terminal_size(fallback=(100, 100))
+
+OPERATORS = {
+    "<": operator.lt,
+    "<=": operator.le,
+    "==": operator.eq,
+    "!=": operator.ne,
+    ">=": operator.ge,
+    ">": operator.gt,
+}
+
+
+def _check_tabulate():
+    """Checks whether tabulate is installed."""
+    if tabulate is None:
+        raise ImportError("Tabulate not installed. Please run `pip install tabulate`.")
+
+
+def print_format_output(dataframe):
+    """Prints output of given dataframe to fit into terminal.
+
+    Returns:
+        table: Final outputted dataframe.
+        dropped_cols: Columns dropped due to terminal size.
+        empty_cols: Empty columns (dropped on default).
+    """
+    print_df = pd.DataFrame()
+    dropped_cols = []
+    empty_cols = []
+    # column display priority is based on the info_keys passed in
+    for i, col in enumerate(dataframe):
+        if dataframe[col].isnull().all():
+            # Don't add col to print_df if is fully empty
+            empty_cols += [col]
+            continue
+
+        print_df[col] = dataframe[col]
+        test_table = tabulate(print_df, headers="keys", tablefmt="psql")
+        if str(test_table).index("\n") > TERM_WIDTH:
+            # Drop all columns beyond terminal width
+            print_df.drop(col, axis=1, inplace=True)
+            dropped_cols += list(dataframe.columns)[i:]
+            break
+
+    table = tabulate(print_df, headers="keys", tablefmt="psql", showindex="never")
+
+    print(table)
+    if dropped_cols:
+        click.secho("Dropped columns: {}".format(dropped_cols), fg="yellow")
+        click.secho("Please increase your terminal size to view remaining columns.")
+    if empty_cols:
+        click.secho("Empty columns: {}".format(empty_cols), fg="yellow")
+
+    return table, dropped_cols, empty_cols
+
+
+def list_trials(
+    experiment_path: str,
+    sort: Optional[List[str]] = None,
+    output: Optional[str] = None,
+    filter_op: Optional[str] = None,
+    info_keys: Optional[List[str]] = None,
+    limit: int = None,
+    desc: bool = False,
+):
+    """Lists trials in the directory subtree starting at the given path.
+
+    Args:
+        experiment_path: Directory where trials are located.
+            Like Experiment.local_dir/Experiment.name/experiment*.json.
+        sort: Keys to sort by.
+        output: Name of file where output is saved.
+        filter_op: Filter operation in the format
+            "<column> <operator> <value>".
+        info_keys: Keys that are displayed.
+        limit: Number of rows to display.
+        desc: Sort ascending vs. descending.
+    """
+    _check_tabulate()
+
+    try:
+        checkpoints_df = ExperimentAnalysis(experiment_path).dataframe()  # last result
+    except TuneError as e:
+        raise click.ClickException("No trial data found!") from e
+
+    config_prefix = CONFIG_PREFIX + "/"
+
+    def key_filter(k):
+        return k in DEFAULT_CLI_KEYS or k.startswith(config_prefix)
+
+    col_keys = [k for k in checkpoints_df.columns if key_filter(k)]
+
+    if info_keys:
+        for k in info_keys:
+            if k not in checkpoints_df.columns:
+                raise click.ClickException(
+                    "Provided key invalid: {}. "
+                    "Available keys: {}.".format(k, checkpoints_df.columns)
+                )
+        col_keys = [k for k in checkpoints_df.columns if k in info_keys]
+
+    if not col_keys:
+        raise click.ClickException("No columns to output.")
+
+    checkpoints_df = checkpoints_df[col_keys]
+    if "last_update_time" in checkpoints_df:
+        with pd.option_context("mode.use_inf_as_null", True):
+            datetime_series = checkpoints_df["last_update_time"].dropna()
+
+        datetime_series = datetime_series.apply(
+            lambda t: datetime.fromtimestamp(t).strftime(TIMESTAMP_FORMAT)
+        )
+        checkpoints_df["last_update_time"] = datetime_series
+
+    if "logdir" in checkpoints_df:
+        # logdir often too long to view in table, so drop experiment_path
+        checkpoints_df["logdir"] = checkpoints_df["logdir"].str.replace(
+            experiment_path, ""
+        )
+
+    if filter_op:
+        col, op, val = filter_op.split(" ")
+        col_type = checkpoints_df[col].dtype
+        if is_numeric_dtype(col_type):
+            val = float(val)
+        elif is_string_dtype(col_type):
+            val = str(val)
+        # TODO(Andrew): add support for datetime and boolean
+        else:
+            raise click.ClickException(
+                "Unsupported dtype for {}: {}".format(val, col_type)
+            )
+        op = OPERATORS[op]
+        filtered_index = op(checkpoints_df[col], val)
+        checkpoints_df = checkpoints_df[filtered_index]
+
+    if sort:
+        for key in sort:
+            if key not in checkpoints_df:
+                raise click.ClickException(
+                    "{} not in: {}".format(key, list(checkpoints_df))
+                )
+        ascending = not desc
+        checkpoints_df = checkpoints_df.sort_values(by=sort, ascending=ascending)
+
+    if limit:
+        checkpoints_df = checkpoints_df[:limit]
+
+    print_format_output(checkpoints_df)
+
+    if output:
+        file_extension = os.path.splitext(output)[1].lower()
+        if file_extension in (".p", ".pkl", ".pickle"):
+            checkpoints_df.to_pickle(output)
+        elif file_extension == ".csv":
+            checkpoints_df.to_csv(output, index=False)
+        else:
+            raise click.ClickException("Unsupported filetype: {}".format(output))
+        click.secho("Output saved at {}".format(output), fg="green")
+
+
+def list_experiments(
+    project_path: str,
+    sort: Optional[List[str]] = None,
+    output: str = None,
+    filter_op: str = None,
+    info_keys: Optional[List[str]] = None,
+    limit: int = None,
+    desc: bool = False,
+):
+    """Lists experiments in the directory subtree.
+
+    Args:
+        project_path: Directory where experiments are located.
+            Corresponds to Experiment.local_dir.
+        sort: Keys to sort by.
+        output: Name of file where output is saved.
+        filter_op: Filter operation in the format
+            "<column> <operator> <value>".
+        info_keys: Keys that are displayed.
+        limit: Number of rows to display.
+        desc: Sort ascending vs. descending.
+    """
+    _check_tabulate()
+    base, experiment_folders, _ = next(os.walk(project_path))
+
+    experiment_data_collection = []
+
+    for experiment_dir in experiment_folders:
+        num_trials = sum(
+            EXPR_RESULT_FILE in files
+            for _, _, files in os.walk(os.path.join(base, experiment_dir))
+        )
+
+        experiment_data = {"name": experiment_dir, "total_trials": num_trials}
+        experiment_data_collection.append(experiment_data)
+
+    if not experiment_data_collection:
+        raise click.ClickException("No experiments found!")
+
+    info_df = pd.DataFrame(experiment_data_collection)
+    if not info_keys:
+        info_keys = DEFAULT_PROJECT_INFO_KEYS
+    col_keys = [k for k in list(info_keys) if k in info_df]
+    if not col_keys:
+        raise click.ClickException(
+            "None of keys {} in experiment data!".format(info_keys)
+        )
+    info_df = info_df[col_keys]
+
+    if filter_op:
+        col, op, val = filter_op.split(" ")
+        col_type = info_df[col].dtype
+        if is_numeric_dtype(col_type):
+            val = float(val)
+        elif is_string_dtype(col_type):
+            val = str(val)
+        # TODO(Andrew): add support for datetime and boolean
+        else:
+            raise click.ClickException(
+                "Unsupported dtype for {}: {}".format(val, col_type)
+            )
+        op = OPERATORS[op]
+        filtered_index = op(info_df[col], val)
+        info_df = info_df[filtered_index]
+
+    if sort:
+        for key in sort:
+            if key not in info_df:
+                raise click.ClickException("{} not in: {}".format(key, list(info_df)))
+        ascending = not desc
+        info_df = info_df.sort_values(by=sort, ascending=ascending)
+
+    if limit:
+        info_df = info_df[:limit]
+
+    print_format_output(info_df)
+
+    if output:
+        file_extension = os.path.splitext(output)[1].lower()
+        if file_extension in (".p", ".pkl", ".pickle"):
+            info_df.to_pickle(output)
+        elif file_extension == ".csv":
+            info_df.to_csv(output, index=False)
+        else:
+            raise click.ClickException("Unsupported filetype: {}".format(output))
+        click.secho("Output saved at {}".format(output), fg="green")
+
+
+def add_note(path: str, filename: str = "note.txt"):
+    """Opens a txt file at the given path where user can add and save notes.
+
+    Args:
+        path: Directory where note will be saved.
+        filename: Name of note. Defaults to "note.txt"
+    """
+    path = Path(path).expanduser()
+    assert path.is_dir(), "{} is not a valid directory.".format(path)
+
+    filepath = path / filename
+
+    try:
+        subprocess.call([EDITOR, filepath.as_posix()])
+    except Exception as exc:
+        click.secho("Editing note failed: {}".format(str(exc)), fg="red")
+    if filepath.exists():
+        print("Note updated at:", filepath.as_posix())
+    else:
+        print("Note created at:", filepath.as_posix())

.venv/lib/python3.11/site-packages/ray/tune/cli/scripts.py

@@ -0,0 +1,101 @@
+import click
+
+import ray.tune.cli.commands as commands
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command()
+@click.argument("experiment_path", required=True, type=str)
+@click.option("--sort", default=None, type=str, help="Select which column to sort on.")
+@click.option(
+    "--output",
+    "-o",
+    default=None,
+    type=str,
+    help="Select file to output information to.",
+)
+@click.option(
+    "--filter",
+    "filter_op",
+    default=None,
+    type=str,
+    help="Select filter in the format '<column> <operator> <value>'.",
+)
+@click.option(
+    "--columns", default=None, type=str, help="Select columns to be displayed."
+)
+@click.option(
+    "--limit", default=None, type=int, help="Select number of rows to display."
+)
+@click.option("--desc", default=False, type=bool, help="Sort ascending vs. descending.")
+def list_trials(experiment_path, sort, output, filter_op, columns, limit, desc):
+    """Lists trials in the directory subtree starting at the given path."""
+    if sort:
+        sort = sort.split(",")
+    if columns:
+        columns = columns.split(",")
+    commands.list_trials(experiment_path, sort, output, filter_op, columns, limit, desc)
+
+
+@cli.command()
+@click.argument("project_path", required=True, type=str)
+@click.option("--sort", default=None, type=str, help="Select which column to sort on.")
+@click.option(
+    "--output",
+    "-o",
+    default=None,
+    type=str,
+    help="Select file to output information to.",
+)
+@click.option(
+    "--filter",
+    "filter_op",
+    default=None,
+    type=str,
+    help="Select filter in the format '<column> <operator> <value>'.",
+)
+@click.option(
+    "--columns", default=None, type=str, help="Select columns to be displayed."
+)
+@click.option(
+    "--limit", default=None, type=int, help="Select number of rows to display."
+)
+@click.option("--desc", default=False, type=bool, help="Sort ascending vs. descending.")
+def list_experiments(project_path, sort, output, filter_op, columns, limit, desc):
+    """Lists experiments in the directory subtree."""
+    if sort:
+        sort = sort.split(",")
+    if columns:
+        columns = columns.split(",")
+    commands.list_experiments(
+        project_path, sort, output, filter_op, columns, limit, desc
+    )
+
+
+@cli.command()
+@click.argument("path", required=True, type=str)
+@click.option(
+    "--filename", default="note.txt", type=str, help="Specify filename for note."
+)
+def add_note(path, filename):
+    """Adds user notes as a text file at the given path."""
+    commands.add_note(path, filename)
+
+
+cli.add_command(list_trials, name="ls")
+cli.add_command(list_trials, name="list-trials")
+cli.add_command(list_experiments, name="lsx")
+cli.add_command(list_experiments, name="list-experiments")
+cli.add_command(add_note, name="add-note")
+
+
+def main():
+    return cli()
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,285 @@
+import os
+
+import matplotlib.animation as animation
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.utils.data
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+import torchvision.utils as vutils
+from scipy.stats import entropy
+from torch.autograd import Variable
+from torch.nn import functional as F
+
+import ray
+
+# Training parameters
+workers = 2
+batch_size = 64
+image_size = 32
+
+# Number of channels in the training images. For color images this is 3
+nc = 1
+
+# Size of z latent vector (i.e. size of generator input)
+nz = 100
+
+# Size of feature maps in generator
+ngf = 32
+
+# Size of feature maps in discriminator
+ndf = 32
+
+# Beta1 hyperparam for Adam optimizers
+beta1 = 0.5
+
+# iterations of actual training in each Trainable _train
+train_iterations_per_step = 5
+
+MODEL_PATH = os.path.expanduser("~/.ray/models/mnist_cnn.pt")
+
+
+def get_data_loader(data_dir="~/data"):
+    dataset = dset.MNIST(
+        root=data_dir,
+        download=True,
+        transform=transforms.Compose(
+            [
+                transforms.Resize(image_size),
+                transforms.ToTensor(),
+                transforms.Normalize((0.5,), (0.5,)),
+            ]
+        ),
+    )
+
+    # Create the dataloader
+    dataloader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=True, num_workers=workers
+    )
+
+    return dataloader
+
+
+# __GANmodel_begin__
+# custom weights initialization called on netG and netD
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find("BatchNorm") != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+
+# Generator Code
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.main = nn.Sequential(
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(nz, ngf * 4, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(ngf * 4),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf * 2),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
+            nn.Tanh(),
+        )
+
+    def forward(self, input):
+        return self.main(input)
+
+
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+        self.main = nn.Sequential(
+            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 2),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 4),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, input):
+        return self.main(input)
+
+
+# __GANmodel_end__
+
+
+# __INCEPTION_SCORE_begin__
+class Net(nn.Module):
+    """
+    LeNet for MNist classification, used for inception_score
+    """
+
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def inception_score(imgs, mnist_model_ref, batch_size=32, splits=1):
+    N = len(imgs)
+    dtype = torch.FloatTensor
+    dataloader = torch.utils.data.DataLoader(imgs, batch_size=batch_size)
+    cm = ray.get(mnist_model_ref)  # Get the mnist model from Ray object store.
+    up = nn.Upsample(size=(28, 28), mode="bilinear").type(dtype)
+
+    def get_pred(x):
+        x = up(x)
+        x = cm(x)
+        return F.softmax(x).data.cpu().numpy()
+
+    preds = np.zeros((N, 10))
+    for i, batch in enumerate(dataloader, 0):
+        batch = batch.type(dtype)
+        batchv = Variable(batch)
+        batch_size_i = batch.size()[0]
+        preds[i * batch_size : i * batch_size + batch_size_i] = get_pred(batchv)
+
+    # Now compute the mean kl-div
+    split_scores = []
+    for k in range(splits):
+        part = preds[k * (N // splits) : (k + 1) * (N // splits), :]
+        py = np.mean(part, axis=0)
+        scores = []
+        for i in range(part.shape[0]):
+            pyx = part[i, :]
+            scores.append(entropy(pyx, py))
+        split_scores.append(np.exp(np.mean(scores)))
+
+    return np.mean(split_scores), np.std(split_scores)
+
+
+# __INCEPTION_SCORE_end__
+
+
+def train_func(
+    netD,
+    netG,
+    optimG,
+    optimD,
+    criterion,
+    dataloader,
+    iteration,
+    device,
+    mnist_model_ref,
+):
+    real_label = 1
+    fake_label = 0
+
+    for i, data in enumerate(dataloader, 0):
+        if i >= train_iterations_per_step:
+            break
+
+        netD.zero_grad()
+        real_cpu = data[0].to(device)
+        b_size = real_cpu.size(0)
+        label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
+        output = netD(real_cpu).view(-1)
+        errD_real = criterion(output, label)
+        errD_real.backward()
+        D_x = output.mean().item()
+
+        noise = torch.randn(b_size, nz, 1, 1, device=device)
+        fake = netG(noise)
+        label.fill_(fake_label)
+        output = netD(fake.detach()).view(-1)
+        errD_fake = criterion(output, label)
+        errD_fake.backward()
+        D_G_z1 = output.mean().item()
+        errD = errD_real + errD_fake
+        optimD.step()
+
+        netG.zero_grad()
+        label.fill_(real_label)
+        output = netD(fake).view(-1)
+        errG = criterion(output, label)
+        errG.backward()
+        D_G_z2 = output.mean().item()
+        optimG.step()
+
+        is_score, is_std = inception_score(fake, mnist_model_ref)
+
+        # Output training stats
+        if iteration % 10 == 0:
+            print(
+                "[%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z))"
+                ": %.4f / %.4f \tInception score: %.4f"
+                % (
+                    iteration,
+                    len(dataloader),
+                    errD.item(),
+                    errG.item(),
+                    D_x,
+                    D_G_z1,
+                    D_G_z2,
+                    is_score,
+                )
+            )
+
+    return errG.item(), errD.item(), is_score
+
+
+def plot_images(dataloader):
+    # Plot some training images
+    real_batch = next(iter(dataloader))
+    plt.figure(figsize=(8, 8))
+    plt.axis("off")
+    plt.title("Original Images")
+    plt.imshow(
+        np.transpose(
+            vutils.make_grid(real_batch[0][:64], padding=2, normalize=True).cpu(),
+            (1, 2, 0),
+        )
+    )
+
+    plt.show()
+
+
+def demo_gan(checkpoint_paths):
+    img_list = []
+    fixed_noise = torch.randn(64, nz, 1, 1)
+    for path in checkpoint_paths:
+        checkpoint_dict = torch.load(os.path.join(path, "checkpoint.pt"))
+
+        loadedG = Generator()
+        loadedG.load_state_dict(checkpoint_dict["netGmodel"])
+        with torch.no_grad():
+            fake = loadedG(fixed_noise).detach().cpu()
+        img_list.append(vutils.make_grid(fake, padding=2, normalize=True))
+
+    fig = plt.figure(figsize=(8, 8))
+    plt.axis("off")
+    ims = [[plt.imshow(np.transpose(i, (1, 2, 0)), animated=True)] for i in img_list]
+    ani = animation.ArtistAnimation(
+        fig, ims, interval=1000, repeat_delay=1000, blit=True
+    )
+    ani.save("./generated.gif", writer="imagemagick", dpi=72)
+    plt.show()

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

@@ -0,0 +1,191 @@
+#!/usr/bin/env python
+"""
+Example of training DCGAN on MNIST using PBT with Tune's function API.
+"""
+import argparse
+import os
+import tempfile
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.optim as optim
+import torch.utils.data
+from filelock import FileLock
+
+import ray
+from ray import train, tune
+from ray.train import Checkpoint
+from ray.tune.examples.pbt_dcgan_mnist.common import (
+    MODEL_PATH,
+    Discriminator,
+    Generator,
+    Net,
+    beta1,
+    demo_gan,
+    get_data_loader,
+    plot_images,
+    train_func,
+    weights_init,
+)
+from ray.tune.schedulers import PopulationBasedTraining
+
+
+# __Train_begin__
+def dcgan_train(config):
+    use_cuda = config.get("use_gpu") and torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+    netD = Discriminator().to(device)
+    netD.apply(weights_init)
+    netG = Generator().to(device)
+    netG.apply(weights_init)
+    criterion = nn.BCELoss()
+    optimizerD = optim.Adam(
+        netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+    )
+    optimizerG = optim.Adam(
+        netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+    )
+    with FileLock(os.path.expanduser("~/ray_results/.data.lock")):
+        dataloader = get_data_loader()
+
+    step = 1
+    checkpoint = train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            checkpoint_dict = torch.load(os.path.join(checkpoint_dir, "checkpoint.pt"))
+        netD.load_state_dict(checkpoint_dict["netDmodel"])
+        netG.load_state_dict(checkpoint_dict["netGmodel"])
+        optimizerD.load_state_dict(checkpoint_dict["optimD"])
+        optimizerG.load_state_dict(checkpoint_dict["optimG"])
+        # Note: Make sure to increment the loaded step by 1 to get the
+        # current step.
+        last_step = checkpoint_dict["step"]
+        step = last_step + 1
+
+        # NOTE: It's important to set the optimizer learning rates
+        # again, since we want to explore the parameters passed in by PBT.
+        # Without this, we would continue using the exact same
+        # configuration as the trial whose checkpoint we are exploiting.
+        if "netD_lr" in config:
+            for param_group in optimizerD.param_groups:
+                param_group["lr"] = config["netD_lr"]
+        if "netG_lr" in config:
+            for param_group in optimizerG.param_groups:
+                param_group["lr"] = config["netG_lr"]
+
+    while True:
+        lossG, lossD, is_score = train_func(
+            netD,
+            netG,
+            optimizerG,
+            optimizerD,
+            criterion,
+            dataloader,
+            step,
+            device,
+            config["mnist_model_ref"],
+        )
+        metrics = {"lossg": lossG, "lossd": lossD, "is_score": is_score}
+
+        if step % config["checkpoint_interval"] == 0:
+            with tempfile.TemporaryDirectory() as tmpdir:
+                torch.save(
+                    {
+                        "netDmodel": netD.state_dict(),
+                        "netGmodel": netG.state_dict(),
+                        "optimD": optimizerD.state_dict(),
+                        "optimG": optimizerG.state_dict(),
+                        "step": step,
+                    },
+                    os.path.join(tmpdir, "checkpoint.pt"),
+                )
+                train.report(metrics, checkpoint=Checkpoint.from_directory(tmpdir))
+        else:
+            train.report(metrics)
+
+        step += 1
+
+
+# __Train_end__
+
+
+def download_mnist_cnn():
+    import urllib.request
+
+    # Download a pre-trained MNIST model for inception score calculation.
+    # This is a tiny model (<100kb).
+    if not os.path.exists(MODEL_PATH):
+        print("downloading model")
+        os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
+        urllib.request.urlretrieve(
+            "https://github.com/ray-project/ray/raw/master/python/ray/tune/"
+            "examples/pbt_dcgan_mnist/mnist_cnn.pt",
+            MODEL_PATH,
+        )
+    return MODEL_PATH
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    parser.add_argument(
+        "--data-dir", type=str, default="~/data/", help="Set the path of the dataset."
+    )
+    args, _ = parser.parse_known_args()
+    ray.init()
+
+    download_mnist_cnn()
+
+    dataloader = get_data_loader(args.data_dir)
+    if not args.smoke_test:
+        plot_images(dataloader)
+
+    # __tune_begin__
+
+    # load the pretrained mnist classification model for inception_score
+    mnist_cnn = Net()
+    mnist_cnn.load_state_dict(torch.load(MODEL_PATH))
+    mnist_cnn.eval()
+    # Put the model in Ray object store.
+    mnist_model_ref = ray.put(mnist_cnn)
+
+    scheduler = PopulationBasedTraining(
+        perturbation_interval=5,
+        hyperparam_mutations={
+            # distribution for resampling
+            "netG_lr": lambda: np.random.uniform(1e-2, 1e-5),
+            "netD_lr": lambda: np.random.uniform(1e-2, 1e-5),
+        },
+    )
+
+    tune_iter = 5 if args.smoke_test else 300
+    tuner = tune.Tuner(
+        dcgan_train,
+        run_config=train.RunConfig(
+            name="pbt_dcgan_mnist",
+            stop={"training_iteration": tune_iter},
+            verbose=1,
+        ),
+        tune_config=tune.TuneConfig(
+            metric="is_score",
+            mode="max",
+            num_samples=8,
+            scheduler=scheduler,
+        ),
+        param_space={
+            "netG_lr": tune.choice([0.0001, 0.0002, 0.0005]),
+            "netD_lr": tune.choice([0.0001, 0.0002, 0.0005]),
+            "mnist_model_ref": mnist_model_ref,
+        },
+    )
+    results = tuner.fit()
+    # __tune_end__
+
+    # demo of the trained Generators
+    if not args.smoke_test:
+        checkpoint_paths = [result.checkpoint.to_directory() for result in results]
+        demo_gan(checkpoint_paths)

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

@@ -0,0 +1,185 @@
+#!/usr/bin/env python
+"""
+Example of training DCGAN on MNIST using PBT with Tune's Trainable Class
+API.
+"""
+import argparse
+import os
+import random
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.optim as optim
+import torch.utils.data
+from filelock import FileLock
+
+import ray
+from ray import train, tune
+from ray.tune.examples.pbt_dcgan_mnist.common import (
+    MODEL_PATH,
+    Discriminator,
+    Generator,
+    Net,
+    beta1,
+    demo_gan,
+    get_data_loader,
+    plot_images,
+    train_func,
+    weights_init,
+)
+from ray.tune.schedulers import PopulationBasedTraining
+
+
+# __Trainable_begin__
+class PytorchTrainable(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.netD = Discriminator().to(self.device)
+        self.netD.apply(weights_init)
+        self.netG = Generator().to(self.device)
+        self.netG.apply(weights_init)
+        self.criterion = nn.BCELoss()
+        self.optimizerD = optim.Adam(
+            self.netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+        )
+        self.optimizerG = optim.Adam(
+            self.netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)
+        )
+        with FileLock(os.path.expanduser("~/.data.lock")):
+            self.dataloader = get_data_loader(config.get("data_dir", "~/data"))
+        self.mnist_model_ref = config["mnist_model_ref"]
+
+    def step(self):
+        lossG, lossD, is_score = train_func(
+            self.netD,
+            self.netG,
+            self.optimizerG,
+            self.optimizerD,
+            self.criterion,
+            self.dataloader,
+            self._iteration,
+            self.device,
+            self.mnist_model_ref,
+        )
+        return {"lossg": lossG, "lossd": lossD, "is_score": is_score}
+
+    def save_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint.pt")
+        torch.save(
+            {
+                "netDmodel": self.netD.state_dict(),
+                "netGmodel": self.netG.state_dict(),
+                "optimD": self.optimizerD.state_dict(),
+                "optimG": self.optimizerG.state_dict(),
+            },
+            path,
+        )
+
+        return checkpoint_dir
+
+    def load_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "checkpoint.pt")
+        checkpoint = torch.load(path)
+        self.netD.load_state_dict(checkpoint["netDmodel"])
+        self.netG.load_state_dict(checkpoint["netGmodel"])
+        self.optimizerD.load_state_dict(checkpoint["optimD"])
+        self.optimizerG.load_state_dict(checkpoint["optimG"])
+
+    def reset_config(self, new_config):
+        if "netD_lr" in new_config:
+            for param_group in self.optimizerD.param_groups:
+                param_group["lr"] = new_config["netD_lr"]
+        if "netG_lr" in new_config:
+            for param_group in self.optimizerG.param_groups:
+                param_group["lr"] = new_config["netG_lr"]
+
+        self.config = new_config
+        return True
+
+
+# __Trainable_end__
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    parser.add_argument(
+        "--data-dir", type=str, default="~/data/", help="Set the path of the dataset."
+    )
+    args, _ = parser.parse_known_args()
+    ray.init()
+
+    import urllib.request
+
+    # Download a pre-trained MNIST model for inception score calculation.
+    # This is a tiny model (<100kb).
+    if not os.path.exists(MODEL_PATH):
+        print("downloading model")
+        os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
+        urllib.request.urlretrieve(
+            "https://github.com/ray-project/ray/raw/master/python/ray/tune/"
+            "examples/pbt_dcgan_mnist/mnist_cnn.pt",
+            MODEL_PATH,
+        )
+
+    dataloader = get_data_loader()
+    if not args.smoke_test:
+        plot_images(dataloader)
+
+    # load the pretrained mnist classification model for inception_score
+    mnist_cnn = Net()
+    mnist_cnn.load_state_dict(torch.load(MODEL_PATH))
+    mnist_cnn.eval()
+    mnist_model_ref = ray.put(mnist_cnn)
+
+    # __tune_begin__
+    scheduler = PopulationBasedTraining(
+        time_attr="training_iteration",
+        perturbation_interval=5,
+        hyperparam_mutations={
+            # distribution for resampling
+            "netG_lr": lambda: np.random.uniform(1e-2, 1e-5),
+            "netD_lr": lambda: np.random.uniform(1e-2, 1e-5),
+        },
+    )
+
+    tune_iter = 10 if args.smoke_test else 300
+    tuner = tune.Tuner(
+        PytorchTrainable,
+        run_config=train.RunConfig(
+            name="pbt_dcgan_mnist",
+            stop={"training_iteration": tune_iter},
+            verbose=1,
+            checkpoint_config=train.CheckpointConfig(checkpoint_at_end=True),
+        ),
+        tune_config=tune.TuneConfig(
+            metric="is_score",
+            mode="max",
+            num_samples=8,
+            scheduler=scheduler,
+            reuse_actors=True,
+        ),
+        param_space={
+            "netG_lr": tune.sample_from(
+                lambda spec: random.choice([0.0001, 0.0002, 0.0005])
+            ),
+            "netD_lr": tune.sample_from(
+                lambda spec: random.choice([0.0001, 0.0002, 0.0005])
+            ),
+            "mnist_model_ref": mnist_model_ref,
+            "data_dir": args.data_dir,
+        },
+    )
+    results = tuner.fit()
+
+    # export_formats=[ExportFormat.MODEL]
+    # __tune_end__
+
+    # demo of the trained Generators
+    if not args.smoke_test:
+        checkpoint_paths = [result.checkpoint.to_directory() for result in results]
+        demo_gan(checkpoint_paths)

.venv/lib/python3.11/site-packages/ray/tune/experimental/output.py

@@ -0,0 +1,1043 @@
+import argparse
+import collections
+import datetime
+import logging
+import math
+import numbers
+import os
+import sys
+import textwrap
+import time
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Any, Collection, Dict, Iterable, List, Optional, Tuple, Union
+
+import numpy as np
+import pandas as pd
+
+import ray
+from ray._private.dict import flatten_dict, unflattened_lookup
+from ray._private.thirdparty.tabulate.tabulate import (
+    DataRow,
+    Line,
+    TableFormat,
+    tabulate,
+)
+from ray.air._internal.usage import AirEntrypoint
+from ray.air.constants import TRAINING_ITERATION
+from ray.train import Checkpoint
+from ray.tune.callback import Callback
+from ray.tune.experiment.trial import Trial
+from ray.tune.result import (
+    AUTO_RESULT_KEYS,
+    EPISODE_REWARD_MEAN,
+    MEAN_ACCURACY,
+    MEAN_LOSS,
+    TIME_TOTAL_S,
+    TIMESTEPS_TOTAL,
+)
+from ray.tune.search.sample import Domain
+from ray.tune.utils.log import Verbosity
+
+try:
+    import rich
+    import rich.layout
+    import rich.live
+except ImportError:
+    rich = None
+
+
+logger = logging.getLogger(__name__)
+
+# defines the mapping of the key in result and the key to be printed in table.
+# Note this is ordered!
+DEFAULT_COLUMNS = collections.OrderedDict(
+    {
+        MEAN_ACCURACY: "acc",
+        MEAN_LOSS: "loss",
+        TRAINING_ITERATION: "iter",
+        TIME_TOTAL_S: "total time (s)",
+        TIMESTEPS_TOTAL: "ts",
+        EPISODE_REWARD_MEAN: "reward",
+    }
+)
+
+# These keys are blacklisted for printing out training/tuning intermediate/final result!
+BLACKLISTED_KEYS = {
+    "config",
+    "date",
+    "done",
+    "hostname",
+    "iterations_since_restore",
+    "node_ip",
+    "pid",
+    "time_since_restore",
+    "timestamp",
+    "trial_id",
+    "experiment_tag",
+    "should_checkpoint",
+    "_report_on",  # LIGHTNING_REPORT_STAGE_KEY
+}
+
+VALID_SUMMARY_TYPES = {
+    int,
+    float,
+    np.float32,
+    np.float64,
+    np.int32,
+    np.int64,
+    type(None),
+}
+
+# The order of summarizing trials.
+ORDER = [
+    Trial.RUNNING,
+    Trial.TERMINATED,
+    Trial.PAUSED,
+    Trial.PENDING,
+    Trial.ERROR,
+]
+
+
+class AirVerbosity(IntEnum):
+    SILENT = 0
+    DEFAULT = 1
+    VERBOSE = 2
+
+    def __repr__(self):
+        return str(self.value)
+
+
+IS_NOTEBOOK = ray.widgets.util.in_notebook()
+
+
+def get_air_verbosity(
+    verbose: Union[int, AirVerbosity, Verbosity]
+) -> Optional[AirVerbosity]:
+    if os.environ.get("RAY_AIR_NEW_OUTPUT", "1") == "0":
+        return None
+
+    if isinstance(verbose, AirVerbosity):
+        return verbose
+
+    verbose_int = verbose if isinstance(verbose, int) else verbose.value
+
+    # Verbosity 2 and 3 both map to AirVerbosity 2
+    verbose_int = min(2, verbose_int)
+
+    return AirVerbosity(verbose_int)
+
+
+def _infer_params(config: Dict[str, Any]) -> List[str]:
+    params = []
+    flat_config = flatten_dict(config)
+    for key, val in flat_config.items():
+        if isinstance(val, Domain):
+            params.append(key)
+        # Grid search is a special named field. Because we flattened
+        # the whole config, we look it up per string
+        if key.endswith("/grid_search"):
+            # Truncate `/grid_search`
+            params.append(key[:-12])
+    return params
+
+
+def _get_time_str(start_time: float, current_time: float) -> Tuple[str, str]:
+    """Get strings representing the current and elapsed time.
+
+    Args:
+        start_time: POSIX timestamp of the start of the tune run
+        current_time: POSIX timestamp giving the current time
+
+    Returns:
+        Current time and elapsed time for the current run
+    """
+    current_time_dt = datetime.datetime.fromtimestamp(current_time)
+    start_time_dt = datetime.datetime.fromtimestamp(start_time)
+    delta: datetime.timedelta = current_time_dt - start_time_dt
+
+    rest = delta.total_seconds()
+    days = int(rest // (60 * 60 * 24))
+
+    rest -= days * (60 * 60 * 24)
+    hours = int(rest // (60 * 60))
+
+    rest -= hours * (60 * 60)
+    minutes = int(rest // 60)
+
+    seconds = int(rest - minutes * 60)
+
+    running_for_str = ""
+    if days > 0:
+        running_for_str += f"{days:d}d "
+
+    if hours > 0 or running_for_str:
+        running_for_str += f"{hours:d}hr "
+
+    if minutes > 0 or running_for_str:
+        running_for_str += f"{minutes:d}min "
+
+    running_for_str += f"{seconds:d}s"
+
+    return f"{current_time_dt:%Y-%m-%d %H:%M:%S}", running_for_str
+
+
+def _get_trials_by_state(trials: List[Trial]) -> Dict[str, List[Trial]]:
+    trials_by_state = collections.defaultdict(list)
+    for t in trials:
+        trials_by_state[t.status].append(t)
+    return trials_by_state
+
+
+def _get_trials_with_error(trials: List[Trial]) -> List[Trial]:
+    return [t for t in trials if t.error_file]
+
+
+def _infer_user_metrics(trials: List[Trial], limit: int = 4) -> List[str]:
+    """Try to infer the metrics to print out.
+
+    By default, only the first 4 meaningful metrics in `last_result` will be
+    inferred as user implied metrics.
+    """
+    # Using OrderedDict for OrderedSet.
+    result = collections.OrderedDict()
+    for t in trials:
+        if not t.last_result:
+            continue
+        for metric, value in t.last_result.items():
+            if metric not in DEFAULT_COLUMNS:
+                if metric not in AUTO_RESULT_KEYS:
+                    if type(value) in VALID_SUMMARY_TYPES:
+                        result[metric] = ""  # not important
+
+            if len(result) >= limit:
+                return list(result.keys())
+    return list(result.keys())
+
+
+def _current_best_trial(
+    trials: List[Trial], metric: Optional[str], mode: Optional[str]
+) -> Tuple[Optional[Trial], Optional[str]]:
+    """
+    Returns the best trial and the metric key. If anything is empty or None,
+    returns a trivial result of None, None.
+
+    Args:
+        trials: List of trials.
+        metric: Metric that trials are being ranked.
+        mode: One of "min" or "max".
+
+    Returns:
+         Best trial and the metric key.
+    """
+    if not trials or not metric or not mode:
+        return None, None
+
+    metric_op = 1.0 if mode == "max" else -1.0
+    best_metric = float("-inf")
+    best_trial = None
+    for t in trials:
+        if not t.last_result:
+            continue
+        metric_value = unflattened_lookup(metric, t.last_result, default=None)
+        if pd.isnull(metric_value):
+            continue
+        if not best_trial or metric_value * metric_op > best_metric:
+            best_metric = metric_value * metric_op
+            best_trial = t
+    return best_trial, metric
+
+
+@dataclass
+class _PerStatusTrialTableData:
+    trial_infos: List[List[str]]
+    more_info: str
+
+
+@dataclass
+class _TrialTableData:
+    header: List[str]
+    data: List[_PerStatusTrialTableData]
+
+
+def _max_len(value: Any, max_len: int = 20, wrap: bool = False) -> Any:
+    """Abbreviate a string representation of an object to `max_len` characters.
+
+    For numbers, booleans and None, the original value will be returned for
+    correct rendering in the table formatting tool.
+
+    Args:
+        value: Object to be represented as a string.
+        max_len: Maximum return string length.
+    """
+    if value is None or isinstance(value, (int, float, numbers.Number, bool)):
+        return value
+
+    string = str(value)
+    if len(string) <= max_len:
+        return string
+
+    if wrap:
+        # Maximum two rows.
+        # Todo: Make this configurable in the refactor
+        if len(value) > max_len * 2:
+            value = "..." + string[(3 - (max_len * 2)) :]
+
+        wrapped = textwrap.wrap(value, width=max_len)
+        return "\n".join(wrapped)
+
+    result = "..." + string[(3 - max_len) :]
+    return result
+
+
+def _get_trial_info(
+    trial: Trial, param_keys: List[str], metric_keys: List[str]
+) -> List[str]:
+    """Returns the following information about a trial:
+
+    name | status | metrics...
+
+    Args:
+        trial: Trial to get information for.
+        param_keys: Names of parameters to include.
+        metric_keys: Names of metrics to include.
+    """
+    result = trial.last_result
+    trial_info = [str(trial), trial.status]
+
+    # params
+    trial_info.extend(
+        [
+            _max_len(
+                unflattened_lookup(param, trial.config, default=None),
+            )
+            for param in param_keys
+        ]
+    )
+    # metrics
+    trial_info.extend(
+        [
+            _max_len(
+                unflattened_lookup(metric, result, default=None),
+            )
+            for metric in metric_keys
+        ]
+    )
+    return trial_info
+
+
+def _get_trial_table_data_per_status(
+    status: str,
+    trials: List[Trial],
+    param_keys: List[str],
+    metric_keys: List[str],
+    force_max_rows: bool = False,
+) -> Optional[_PerStatusTrialTableData]:
+    """Gather all information of trials pertained to one `status`.
+
+    Args:
+        status: The trial status of interest.
+        trials: all the trials of that status.
+        param_keys: *Ordered* list of parameters to be displayed in the table.
+        metric_keys: *Ordered* list of metrics to be displayed in the table.
+            Including both default and user defined.
+        force_max_rows: Whether or not to enforce a max row number for this status.
+            If True, only a max of `5` rows will be shown.
+
+    Returns:
+        All information of trials pertained to the `status`.
+    """
+    # TODO: configure it.
+    max_row = 5 if force_max_rows else math.inf
+    if not trials:
+        return None
+
+    trial_infos = list()
+    more_info = None
+    for t in trials:
+        if len(trial_infos) >= max_row:
+            remaining = len(trials) - max_row
+            more_info = f"{remaining} more {status}"
+            break
+        trial_infos.append(_get_trial_info(t, param_keys, metric_keys))
+    return _PerStatusTrialTableData(trial_infos, more_info)
+
+
+def _get_trial_table_data(
+    trials: List[Trial],
+    param_keys: List[str],
+    metric_keys: List[str],
+    all_rows: bool = False,
+    wrap_headers: bool = False,
+) -> _TrialTableData:
+    """Generate a table showing the current progress of tuning trials.
+
+    Args:
+        trials: List of trials for which progress is to be shown.
+        param_keys: Ordered list of parameters to be displayed in the table.
+        metric_keys: Ordered list of metrics to be displayed in the table.
+            Including both default and user defined.
+            Will only be shown if at least one trial is having the key.
+        all_rows: Force to show all rows.
+        wrap_headers: If True, header columns can be wrapped with ``\n``.
+
+    Returns:
+        Trial table data, including header and trial table per each status.
+    """
+    # TODO: configure
+    max_trial_num_to_show = 20
+    max_column_length = 20
+    trials_by_state = _get_trials_by_state(trials)
+
+    # get the right metric to show.
+    metric_keys = [
+        k
+        for k in metric_keys
+        if any(
+            unflattened_lookup(k, t.last_result, default=None) is not None
+            for t in trials
+        )
+    ]
+
+    # get header from metric keys
+    formatted_metric_columns = [
+        _max_len(k, max_len=max_column_length, wrap=wrap_headers) for k in metric_keys
+    ]
+
+    formatted_param_columns = [
+        _max_len(k, max_len=max_column_length, wrap=wrap_headers) for k in param_keys
+    ]
+
+    metric_header = [
+        DEFAULT_COLUMNS[metric] if metric in DEFAULT_COLUMNS else formatted
+        for metric, formatted in zip(metric_keys, formatted_metric_columns)
+    ]
+
+    param_header = formatted_param_columns
+
+    # Map to the abbreviated version if necessary.
+    header = ["Trial name", "status"] + param_header + metric_header
+
+    trial_data = list()
+    for t_status in ORDER:
+        trial_data_per_status = _get_trial_table_data_per_status(
+            t_status,
+            trials_by_state[t_status],
+            param_keys=param_keys,
+            metric_keys=metric_keys,
+            force_max_rows=not all_rows and len(trials) > max_trial_num_to_show,
+        )
+        if trial_data_per_status:
+            trial_data.append(trial_data_per_status)
+    return _TrialTableData(header, trial_data)
+
+
+def _best_trial_str(
+    trial: Trial,
+    metric: str,
+):
+    """Returns a readable message stating the current best trial."""
+    # returns something like
+    # Current best trial: 18ae7_00005 with loss=0.5918508041056858 and params={'train_loop_config': {'lr': 0.059253447253394785}}. # noqa
+    val = unflattened_lookup(metric, trial.last_result, default=None)
+    config = trial.last_result.get("config", {})
+    parameter_columns = list(config.keys())
+    params = {p: unflattened_lookup(p, config) for p in parameter_columns}
+    return (
+        f"Current best trial: {trial.trial_id} with {metric}={val} and "
+        f"params={params}"
+    )
+
+
+def _render_table_item(
+    key: str, item: Any, prefix: str = ""
+) -> Iterable[Tuple[str, str]]:
+    key = prefix + key
+
+    if isinstance(item, argparse.Namespace):
+        item = item.__dict__
+
+    if isinstance(item, float):
+        # tabulate does not work well with mixed-type columns, so we format
+        # numbers ourselves.
+        yield key, f"{item:.5f}".rstrip("0")
+    elif isinstance(item, dict):
+        flattened = flatten_dict(item)
+        for k, v in sorted(flattened.items()):
+            yield key + "/" + str(k), _max_len(v)
+    else:
+        yield key, _max_len(item, 20)
+
+
+def _get_dict_as_table_data(
+    data: Dict,
+    include: Optional[Collection] = None,
+    exclude: Optional[Collection] = None,
+    upper_keys: Optional[Collection] = None,
+):
+    """Get ``data`` dict as table rows.
+
+    If specified, excluded keys are removed. Excluded keys can either be
+    fully specified (e.g. ``foo/bar/baz``) or specify a top-level dictionary
+    (e.g. ``foo``), but no intermediate levels (e.g. ``foo/bar``). If this is
+    needed, we can revisit the logic at a later point.
+
+    The same is true for included keys. If a top-level key is included (e.g. ``foo``)
+    then all sub keys will be included, too, except if they are excluded.
+
+    If keys are both excluded and included, exclusion takes precedence. Thus, if
+    ``foo`` is excluded but ``foo/bar`` is included, it won't show up in the output.
+    """
+    include = include or set()
+    exclude = exclude or set()
+    upper_keys = upper_keys or set()
+
+    upper = []
+    lower = []
+
+    for key, value in sorted(data.items()):
+        # Exclude top-level keys
+        if key in exclude:
+            continue
+
+        for k, v in _render_table_item(str(key), value):
+            # k is now the full subkey, e.g. config/nested/key
+
+            # We can exclude the full key
+            if k in exclude:
+                continue
+
+            # If we specify includes, top-level includes should take precedence
+            # (e.g. if `config` is in include, include config always).
+            if include and key not in include and k not in include:
+                continue
+
+            if key in upper_keys:
+                upper.append([k, v])
+            else:
+                lower.append([k, v])
+
+    if not upper:
+        return lower
+    elif not lower:
+        return upper
+    else:
+        return upper + lower
+
+
+if sys.stdout and sys.stdout.encoding and sys.stdout.encoding.startswith("utf"):
+    # Copied/adjusted from tabulate
+    AIR_TABULATE_TABLEFMT = TableFormat(
+        lineabove=Line("╭", "─", "─", "╮"),
+        linebelowheader=Line("├", "─", "─", "┤"),
+        linebetweenrows=None,
+        linebelow=Line("╰", "─", "─", "╯"),
+        headerrow=DataRow("│", " ", "│"),
+        datarow=DataRow("│", " ", "│"),
+        padding=1,
+        with_header_hide=None,
+    )
+else:
+    # For non-utf output, use ascii-compatible characters.
+    # This prevents errors e.g. when legacy windows encoding is used.
+    AIR_TABULATE_TABLEFMT = TableFormat(
+        lineabove=Line("+", "-", "-", "+"),
+        linebelowheader=Line("+", "-", "-", "+"),
+        linebetweenrows=None,
+        linebelow=Line("+", "-", "-", "+"),
+        headerrow=DataRow("|", " ", "|"),
+        datarow=DataRow("|", " ", "|"),
+        padding=1,
+        with_header_hide=None,
+    )
+
+
+def _print_dict_as_table(
+    data: Dict,
+    header: Optional[str] = None,
+    include: Optional[Collection[str]] = None,
+    exclude: Optional[Collection[str]] = None,
+    division: Optional[Collection[str]] = None,
+):
+    table_data = _get_dict_as_table_data(
+        data=data, include=include, exclude=exclude, upper_keys=division
+    )
+
+    headers = [header, ""] if header else []
+
+    if not table_data:
+        return
+
+    print(
+        tabulate(
+            table_data,
+            headers=headers,
+            colalign=("left", "right"),
+            tablefmt=AIR_TABULATE_TABLEFMT,
+        )
+    )
+
+
+class ProgressReporter(Callback):
+    """Periodically prints out status update."""
+
+    # TODO: Make this configurable
+    _heartbeat_freq = 30  # every 30 sec
+    # to be updated by subclasses.
+    _heartbeat_threshold = None
+    _start_end_verbosity = None
+    _intermediate_result_verbosity = None
+    _addressing_tmpl = None
+
+    def __init__(
+        self,
+        verbosity: AirVerbosity,
+        progress_metrics: Optional[Union[List[str], List[Dict[str, str]]]] = None,
+    ):
+        """
+
+        Args:
+            verbosity: AirVerbosity level.
+        """
+        self._verbosity = verbosity
+        self._start_time = time.time()
+        self._last_heartbeat_time = float("-inf")
+        self._start_time = time.time()
+        self._progress_metrics = progress_metrics
+        self._trial_last_printed_results = {}
+
+        self._in_block = None
+
+    @property
+    def verbosity(self) -> AirVerbosity:
+        return self._verbosity
+
+    def setup(
+        self,
+        start_time: Optional[float] = None,
+        **kwargs,
+    ):
+        self._start_time = start_time
+
+    def _start_block(self, indicator: Any):
+        if self._in_block != indicator:
+            self._end_block()
+        self._in_block = indicator
+
+    def _end_block(self):
+        if self._in_block:
+            print("")
+        self._in_block = None
+
+    def on_experiment_end(self, trials: List["Trial"], **info):
+        self._end_block()
+
+    def experiment_started(
+        self,
+        experiment_name: str,
+        experiment_path: str,
+        searcher_str: str,
+        scheduler_str: str,
+        total_num_samples: int,
+        tensorboard_path: Optional[str] = None,
+        **kwargs,
+    ):
+        self._start_block("exp_start")
+        print(f"\nView detailed results here: {experiment_path}")
+
+        if tensorboard_path:
+            print(
+                f"To visualize your results with TensorBoard, run: "
+                f"`tensorboard --logdir {tensorboard_path}`"
+            )
+
+    @property
+    def _time_heartbeat_str(self):
+        current_time_str, running_time_str = _get_time_str(
+            self._start_time, time.time()
+        )
+        return (
+            f"Current time: {current_time_str}. Total running time: " + running_time_str
+        )
+
+    def print_heartbeat(self, trials, *args, force: bool = False):
+        if self._verbosity < self._heartbeat_threshold:
+            return
+        if force or time.time() - self._last_heartbeat_time >= self._heartbeat_freq:
+            self._print_heartbeat(trials, *args, force=force)
+            self._last_heartbeat_time = time.time()
+
+    def _print_heartbeat(self, trials, *args, force: bool = False):
+        raise NotImplementedError
+
+    def _print_result(self, trial, result: Optional[Dict] = None, force: bool = False):
+        """Only print result if a different result has been reported, or force=True"""
+        result = result or trial.last_result
+
+        last_result_iter = self._trial_last_printed_results.get(trial.trial_id, -1)
+        this_iter = result.get(TRAINING_ITERATION, 0)
+
+        if this_iter != last_result_iter or force:
+            _print_dict_as_table(
+                result,
+                header=f"{self._addressing_tmpl.format(trial)} result",
+                include=self._progress_metrics,
+                exclude=BLACKLISTED_KEYS,
+                division=AUTO_RESULT_KEYS,
+            )
+            self._trial_last_printed_results[trial.trial_id] = this_iter
+
+    def _print_config(self, trial):
+        _print_dict_as_table(
+            trial.config, header=f"{self._addressing_tmpl.format(trial)} config"
+        )
+
+    def on_trial_result(
+        self,
+        iteration: int,
+        trials: List[Trial],
+        trial: Trial,
+        result: Dict,
+        **info,
+    ):
+        if self.verbosity < self._intermediate_result_verbosity:
+            return
+        self._start_block(f"trial_{trial}_result_{result[TRAINING_ITERATION]}")
+        curr_time_str, running_time_str = _get_time_str(self._start_time, time.time())
+        print(
+            f"{self._addressing_tmpl.format(trial)} "
+            f"finished iteration {result[TRAINING_ITERATION]} "
+            f"at {curr_time_str}. Total running time: " + running_time_str
+        )
+        self._print_result(trial, result)
+
+    def on_trial_complete(
+        self, iteration: int, trials: List[Trial], trial: Trial, **info
+    ):
+        if self.verbosity < self._start_end_verbosity:
+            return
+        curr_time_str, running_time_str = _get_time_str(self._start_time, time.time())
+        finished_iter = 0
+        if trial.last_result and TRAINING_ITERATION in trial.last_result:
+            finished_iter = trial.last_result[TRAINING_ITERATION]
+
+        self._start_block(f"trial_{trial}_complete")
+        print(
+            f"{self._addressing_tmpl.format(trial)} "
+            f"completed after {finished_iter} iterations "
+            f"at {curr_time_str}. Total running time: " + running_time_str
+        )
+        self._print_result(trial)
+
+    def on_trial_error(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        curr_time_str, running_time_str = _get_time_str(self._start_time, time.time())
+        finished_iter = 0
+        if trial.last_result and TRAINING_ITERATION in trial.last_result:
+            finished_iter = trial.last_result[TRAINING_ITERATION]
+
+        self._start_block(f"trial_{trial}_error")
+        print(
+            f"{self._addressing_tmpl.format(trial)} "
+            f"errored after {finished_iter} iterations "
+            f"at {curr_time_str}. Total running time: {running_time_str}\n"
+            f"Error file: {trial.error_file}"
+        )
+        self._print_result(trial)
+
+    def on_trial_recover(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.on_trial_error(iteration=iteration, trials=trials, trial=trial, **info)
+
+    def on_checkpoint(
+        self,
+        iteration: int,
+        trials: List[Trial],
+        trial: Trial,
+        checkpoint: Checkpoint,
+        **info,
+    ):
+        if self._verbosity < self._intermediate_result_verbosity:
+            return
+        # don't think this is supposed to happen but just to be safe.
+        saved_iter = "?"
+        if trial.last_result and TRAINING_ITERATION in trial.last_result:
+            saved_iter = trial.last_result[TRAINING_ITERATION]
+
+        self._start_block(f"trial_{trial}_result_{saved_iter}")
+
+        loc = f"({checkpoint.filesystem.type_name}){checkpoint.path}"
+
+        print(
+            f"{self._addressing_tmpl.format(trial)} "
+            f"saved a checkpoint for iteration {saved_iter} "
+            f"at: {loc}"
+        )
+
+    def on_trial_start(self, iteration: int, trials: List[Trial], trial: Trial, **info):
+        if self.verbosity < self._start_end_verbosity:
+            return
+        has_config = bool(trial.config)
+
+        self._start_block(f"trial_{trial}_start")
+        if has_config:
+            print(
+                f"{self._addressing_tmpl.format(trial)} " f"started with configuration:"
+            )
+            self._print_config(trial)
+        else:
+            print(
+                f"{self._addressing_tmpl.format(trial)} "
+                f"started without custom configuration."
+            )
+
+
+def _detect_reporter(
+    verbosity: AirVerbosity,
+    num_samples: int,
+    entrypoint: Optional[AirEntrypoint] = None,
+    metric: Optional[str] = None,
+    mode: Optional[str] = None,
+    config: Optional[Dict] = None,
+    progress_metrics: Optional[Union[List[str], List[Dict[str, str]]]] = None,
+):
+    if entrypoint in {
+        AirEntrypoint.TUNE_RUN,
+        AirEntrypoint.TUNE_RUN_EXPERIMENTS,
+        AirEntrypoint.TUNER,
+    }:
+        reporter = TuneTerminalReporter(
+            verbosity,
+            num_samples=num_samples,
+            metric=metric,
+            mode=mode,
+            config=config,
+            progress_metrics=progress_metrics,
+        )
+    else:
+        reporter = TrainReporter(verbosity, progress_metrics=progress_metrics)
+    return reporter
+
+
+class TuneReporterBase(ProgressReporter):
+    _heartbeat_threshold = AirVerbosity.DEFAULT
+    _wrap_headers = False
+    _intermediate_result_verbosity = AirVerbosity.VERBOSE
+    _start_end_verbosity = AirVerbosity.DEFAULT
+    _addressing_tmpl = "Trial {}"
+
+    def __init__(
+        self,
+        verbosity: AirVerbosity,
+        num_samples: int = 0,
+        metric: Optional[str] = None,
+        mode: Optional[str] = None,
+        config: Optional[Dict] = None,
+        progress_metrics: Optional[Union[List[str], List[Dict[str, str]]]] = None,
+    ):
+        self._num_samples = num_samples
+        self._metric = metric
+        self._mode = mode
+        # will be populated when first result comes in.
+        self._inferred_metric = None
+        self._inferred_params = _infer_params(config or {})
+        super(TuneReporterBase, self).__init__(
+            verbosity=verbosity, progress_metrics=progress_metrics
+        )
+
+    def setup(
+        self,
+        start_time: Optional[float] = None,
+        total_samples: Optional[int] = None,
+        **kwargs,
+    ):
+        super().setup(start_time=start_time)
+        self._num_samples = total_samples
+
+    def _get_overall_trial_progress_str(self, trials):
+        result = " | ".join(
+            [
+                f"{len(trials)} {status}"
+                for status, trials in _get_trials_by_state(trials).items()
+            ]
+        )
+        return f"Trial status: {result}"
+
+    # TODO: Return a more structured type to share code with Jupyter flow.
+    def _get_heartbeat(
+        self, trials, *sys_args, force_full_output: bool = False
+    ) -> Tuple[List[str], _TrialTableData]:
+        result = list()
+        # Trial status: 1 RUNNING | 7 PENDING
+        result.append(self._get_overall_trial_progress_str(trials))
+        # Current time: 2023-02-24 12:35:39 (running for 00:00:37.40)
+        result.append(self._time_heartbeat_str)
+        # Logical resource usage: 8.0/64 CPUs, 0/0 GPUs
+        result.extend(sys_args)
+        # Current best trial: TRIAL NAME, metrics: {...}, parameters: {...}
+        current_best_trial, metric = _current_best_trial(
+            trials, self._metric, self._mode
+        )
+        if current_best_trial:
+            result.append(_best_trial_str(current_best_trial, metric))
+        # Now populating the trial table data.
+        if not self._inferred_metric:
+            # try inferring again.
+            self._inferred_metric = _infer_user_metrics(trials)
+
+        all_metrics = list(DEFAULT_COLUMNS.keys()) + self._inferred_metric
+
+        trial_table_data = _get_trial_table_data(
+            trials,
+            param_keys=self._inferred_params,
+            metric_keys=all_metrics,
+            all_rows=force_full_output,
+            wrap_headers=self._wrap_headers,
+        )
+        return result, trial_table_data
+
+    def _print_heartbeat(self, trials, *sys_args, force: bool = False):
+        raise NotImplementedError
+
+
+class TuneTerminalReporter(TuneReporterBase):
+    def experiment_started(
+        self,
+        experiment_name: str,
+        experiment_path: str,
+        searcher_str: str,
+        scheduler_str: str,
+        total_num_samples: int,
+        tensorboard_path: Optional[str] = None,
+        **kwargs,
+    ):
+        if total_num_samples > sys.maxsize:
+            total_num_samples_str = "infinite"
+        else:
+            total_num_samples_str = str(total_num_samples)
+
+        print(
+            tabulate(
+                [
+                    ["Search algorithm", searcher_str],
+                    ["Scheduler", scheduler_str],
+                    ["Number of trials", total_num_samples_str],
+                ],
+                headers=["Configuration for experiment", experiment_name],
+                tablefmt=AIR_TABULATE_TABLEFMT,
+            )
+        )
+        super().experiment_started(
+            experiment_name=experiment_name,
+            experiment_path=experiment_path,
+            searcher_str=searcher_str,
+            scheduler_str=scheduler_str,
+            total_num_samples=total_num_samples,
+            tensorboard_path=tensorboard_path,
+            **kwargs,
+        )
+
+    def _print_heartbeat(self, trials, *sys_args, force: bool = False):
+        if self._verbosity < self._heartbeat_threshold and not force:
+            return
+        heartbeat_strs, table_data = self._get_heartbeat(
+            trials, *sys_args, force_full_output=force
+        )
+
+        self._start_block("heartbeat")
+        for s in heartbeat_strs:
+            print(s)
+        # now print the table using Tabulate
+        more_infos = []
+        all_data = []
+        fail_header = table_data.header
+        for sub_table in table_data.data:
+            all_data.extend(sub_table.trial_infos)
+            if sub_table.more_info:
+                more_infos.append(sub_table.more_info)
+
+        print(
+            tabulate(
+                all_data,
+                headers=fail_header,
+                tablefmt=AIR_TABULATE_TABLEFMT,
+                showindex=False,
+            )
+        )
+        if more_infos:
+            print(", ".join(more_infos))
+
+        if not force:
+            # Only print error table at end of training
+            return
+
+        trials_with_error = _get_trials_with_error(trials)
+        if not trials_with_error:
+            return
+
+        self._start_block("status_errored")
+        print(f"Number of errored trials: {len(trials_with_error)}")
+        fail_header = ["Trial name", "# failures", "error file"]
+        fail_table_data = [
+            [
+                str(trial),
+                str(trial.run_metadata.num_failures)
+                + ("" if trial.status == Trial.ERROR else "*"),
+                trial.error_file,
+            ]
+            for trial in trials_with_error
+        ]
+        print(
+            tabulate(
+                fail_table_data,
+                headers=fail_header,
+                tablefmt=AIR_TABULATE_TABLEFMT,
+                showindex=False,
+                colalign=("left", "right", "left"),
+            )
+        )
+        if any(trial.status == Trial.TERMINATED for trial in trials_with_error):
+            print("* The trial terminated successfully after retrying.")
+
+
+class TrainReporter(ProgressReporter):
+    # the minimal verbosity threshold at which heartbeat starts getting printed.
+    _heartbeat_threshold = AirVerbosity.VERBOSE
+    _intermediate_result_verbosity = AirVerbosity.DEFAULT
+    _start_end_verbosity = AirVerbosity.DEFAULT
+    _addressing_tmpl = "Training"
+
+    def _get_heartbeat(self, trials: List[Trial], force_full_output: bool = False):
+        # Training on iteration 1. Current time: 2023-03-22 15:29:25 (running for 00:00:03.24)  # noqa
+        if len(trials) == 0:
+            return
+        trial = trials[0]
+        if trial.status != Trial.RUNNING:
+            return " ".join(
+                [f"Training is in {trial.status} status.", self._time_heartbeat_str]
+            )
+        if not trial.last_result or TRAINING_ITERATION not in trial.last_result:
+            iter_num = 1
+        else:
+            iter_num = trial.last_result[TRAINING_ITERATION] + 1
+        return " ".join(
+            [f"Training on iteration {iter_num}.", self._time_heartbeat_str]
+        )
+
+    def _print_heartbeat(self, trials, *args, force: bool = False):
+        print(self._get_heartbeat(trials, force_full_output=force))
+
+    def on_trial_result(
+        self,
+        iteration: int,
+        trials: List[Trial],
+        trial: Trial,
+        result: Dict,
+        **info,
+    ):
+        self._last_heartbeat_time = time.time()
+        super().on_trial_result(
+            iteration=iteration, trials=trials, trial=trial, result=result, **info
+        )

.venv/lib/python3.11/site-packages/ray/tune/logger/__init__.py

@@ -0,0 +1,32 @@
+from ray.tune.logger.csv import CSVLogger, CSVLoggerCallback
+from ray.tune.logger.json import JsonLogger, JsonLoggerCallback
+from ray.tune.logger.logger import (
+    LegacyLoggerCallback,
+    Logger,
+    LoggerCallback,
+    pretty_print,
+)
+from ray.tune.logger.noop import NoopLogger
+from ray.tune.logger.tensorboardx import TBXLogger, TBXLoggerCallback
+
+DEFAULT_LOGGERS = (JsonLogger, CSVLogger, TBXLogger)
+
+# isort: off
+from ray.tune.logger.unified import UnifiedLogger  # noqa: E402
+
+# isort: on
+
+__all__ = [
+    "Logger",
+    "LoggerCallback",
+    "LegacyLoggerCallback",
+    "pretty_print",
+    "CSVLogger",
+    "CSVLoggerCallback",
+    "JsonLogger",
+    "JsonLoggerCallback",
+    "NoopLogger",
+    "TBXLogger",
+    "TBXLoggerCallback",
+    "UnifiedLogger",
+]

.venv/lib/python3.11/site-packages/ray/tune/logger/aim.py

@@ -0,0 +1,187 @@
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+import numpy as np
+
+from ray.air.constants import TRAINING_ITERATION
+from ray.tune.logger.logger import LoggerCallback
+from ray.tune.result import TIME_TOTAL_S, TIMESTEPS_TOTAL
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial
+
+try:
+    from aim.sdk import Repo, Run
+except ImportError:
+    Repo, Run = None, None
+
+logger = logging.getLogger(__name__)
+
+VALID_SUMMARY_TYPES = [int, float, np.float32, np.float64, np.int32, np.int64]
+
+
+@PublicAPI
+class AimLoggerCallback(LoggerCallback):
+    """Aim Logger: logs metrics in Aim format.
+
+    Aim is an open-source, self-hosted ML experiment tracking tool.
+    It's good at tracking lots (thousands) of training runs, and it allows you to
+    compare them with a performant and well-designed UI.
+
+    Source: https://github.com/aimhubio/aim
+
+    Args:
+        repo: Aim repository directory or a `Repo` object that the Run object will
+            log results to. If not provided, a default repo will be set up in the
+            experiment directory (one level above trial directories).
+        experiment: Sets the `experiment` property of each Run object, which is the
+            experiment name associated with it. Can be used later to query
+            runs/sequences.
+            If not provided, the default will be the Tune experiment name set
+            by `RunConfig(name=...)`.
+        metrics: List of metric names (out of the metrics reported by Tune) to
+            track in Aim. If no metric are specified, log everything that
+            is reported.
+        aim_run_kwargs: Additional arguments that will be passed when creating the
+            individual `Run` objects for each trial. For the full list of arguments,
+            please see the Aim documentation:
+            https://aimstack.readthedocs.io/en/latest/refs/sdk.html
+    """
+
+    VALID_HPARAMS = (str, bool, int, float, list, type(None))
+    VALID_NP_HPARAMS = (np.bool_, np.float32, np.float64, np.int32, np.int64)
+
+    def __init__(
+        self,
+        repo: Optional[Union[str, "Repo"]] = None,
+        experiment_name: Optional[str] = None,
+        metrics: Optional[List[str]] = None,
+        **aim_run_kwargs,
+    ):
+        """
+        See help(AimLoggerCallback) for more information about parameters.
+        """
+        assert Run is not None, (
+            "aim must be installed!. You can install aim with"
+            " the command: `pip install aim`."
+        )
+        self._repo_path = repo
+        self._experiment_name = experiment_name
+        if not (bool(metrics) or metrics is None):
+            raise ValueError(
+                "`metrics` must either contain at least one metric name, or be None, "
+                "in which case all reported metrics will be logged to the aim repo."
+            )
+        self._metrics = metrics
+        self._aim_run_kwargs = aim_run_kwargs
+        self._trial_to_run: Dict["Trial", Run] = {}
+
+    def _create_run(self, trial: "Trial") -> Run:
+        """Initializes an Aim Run object for a given trial.
+
+        Args:
+            trial: The Tune trial that aim will track as a Run.
+
+        Returns:
+            Run: The created aim run for a specific trial.
+        """
+        experiment_dir = trial.local_experiment_path
+        run = Run(
+            repo=self._repo_path or experiment_dir,
+            experiment=self._experiment_name or trial.experiment_dir_name,
+            **self._aim_run_kwargs,
+        )
+        # Attach a few useful trial properties
+        run["trial_id"] = trial.trial_id
+        run["trial_log_dir"] = trial.path
+        trial_ip = trial.get_ray_actor_ip()
+        if trial_ip:
+            run["trial_ip"] = trial_ip
+        return run
+
+    def log_trial_start(self, trial: "Trial"):
+        if trial in self._trial_to_run:
+            # Cleanup an existing run if the trial has been restarted
+            self._trial_to_run[trial].close()
+
+        trial.init_local_path()
+        self._trial_to_run[trial] = self._create_run(trial)
+
+        if trial.evaluated_params:
+            self._log_trial_hparams(trial)
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        tmp_result = result.copy()
+
+        step = result.get(TIMESTEPS_TOTAL, None) or result[TRAINING_ITERATION]
+
+        for k in ["config", "pid", "timestamp", TIME_TOTAL_S, TRAINING_ITERATION]:
+            tmp_result.pop(k, None)  # not useful to log these
+
+        # `context` and `epoch` are special keys that users can report,
+        # which are treated as special aim metrics/configurations.
+        context = tmp_result.pop("context", None)
+        epoch = tmp_result.pop("epoch", None)
+
+        trial_run = self._trial_to_run[trial]
+        path = ["ray", "tune"]
+
+        flat_result = flatten_dict(tmp_result, delimiter="/")
+        valid_result = {}
+
+        for attr, value in flat_result.items():
+            if self._metrics and attr not in self._metrics:
+                continue
+
+            full_attr = "/".join(path + [attr])
+            if isinstance(value, tuple(VALID_SUMMARY_TYPES)) and not (
+                np.isnan(value) or np.isinf(value)
+            ):
+                valid_result[attr] = value
+                trial_run.track(
+                    value=value,
+                    name=full_attr,
+                    epoch=epoch,
+                    step=step,
+                    context=context,
+                )
+            elif (isinstance(value, (list, tuple, set)) and len(value) > 0) or (
+                isinstance(value, np.ndarray) and value.size > 0
+            ):
+                valid_result[attr] = value
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        trial_run = self._trial_to_run.pop(trial)
+        trial_run.close()
+
+    def _log_trial_hparams(self, trial: "Trial"):
+        params = flatten_dict(trial.evaluated_params, delimiter="/")
+        flat_params = flatten_dict(params)
+
+        scrubbed_params = {
+            k: v for k, v in flat_params.items() if isinstance(v, self.VALID_HPARAMS)
+        }
+
+        np_params = {
+            k: v.tolist()
+            for k, v in flat_params.items()
+            if isinstance(v, self.VALID_NP_HPARAMS)
+        }
+
+        scrubbed_params.update(np_params)
+        removed = {
+            k: v
+            for k, v in flat_params.items()
+            if not isinstance(v, self.VALID_HPARAMS + self.VALID_NP_HPARAMS)
+        }
+        if removed:
+            logger.info(
+                "Removed the following hyperparameter values when "
+                "logging to aim: %s",
+                str(removed),
+            )
+
+        run = self._trial_to_run[trial]
+        run["hparams"] = scrubbed_params

.venv/lib/python3.11/site-packages/ray/tune/logger/comet.py

@@ -0,0 +1,3 @@
+from ray.air.integrations.comet import CometLoggerCallback
+
+CometLoggerCallback.__module__ = "ray.tune.logger.comet"

.venv/lib/python3.11/site-packages/ray/tune/logger/csv.py

@@ -0,0 +1,135 @@
+import csv
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, TextIO
+
+from ray.air.constants import EXPR_PROGRESS_FILE
+from ray.tune.logger.logger import _LOGGER_DEPRECATION_WARNING, Logger, LoggerCallback
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import Deprecated, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+@Deprecated(
+    message=_LOGGER_DEPRECATION_WARNING.format(
+        old="CSVLogger", new="ray.tune.csv.CSVLoggerCallback"
+    ),
+    warning=True,
+)
+@PublicAPI
+class CSVLogger(Logger):
+    """Logs results to progress.csv under the trial directory.
+
+    Automatically flattens nested dicts in the result dict before writing
+    to csv:
+
+        {"a": {"b": 1, "c": 2}} -> {"a/b": 1, "a/c": 2}
+
+    """
+
+    def _init(self):
+        self._initialized = False
+
+    def _maybe_init(self):
+        """CSV outputted with Headers as first set of results."""
+        if not self._initialized:
+            progress_file = Path(self.logdir, EXPR_PROGRESS_FILE)
+            self._continuing = (
+                progress_file.exists() and progress_file.stat().st_size > 0
+            )
+            self._file = progress_file.open("a")
+            self._csv_out = None
+            self._initialized = True
+
+    def on_result(self, result: Dict):
+        self._maybe_init()
+
+        tmp = result.copy()
+        if "config" in tmp:
+            del tmp["config"]
+        result = flatten_dict(tmp, delimiter="/")
+        if self._csv_out is None:
+            self._csv_out = csv.DictWriter(self._file, result.keys())
+            if not self._continuing:
+                self._csv_out.writeheader()
+        self._csv_out.writerow(
+            {k: v for k, v in result.items() if k in self._csv_out.fieldnames}
+        )
+        self._file.flush()
+
+    def flush(self):
+        if self._initialized and not self._file.closed:
+            self._file.flush()
+
+    def close(self):
+        if self._initialized:
+            self._file.close()
+
+
+@PublicAPI
+class CSVLoggerCallback(LoggerCallback):
+    """Logs results to progress.csv under the trial directory.
+
+    Automatically flattens nested dicts in the result dict before writing
+    to csv:
+
+        {"a": {"b": 1, "c": 2}} -> {"a/b": 1, "a/c": 2}
+
+    """
+
+    _SAVED_FILE_TEMPLATES = [EXPR_PROGRESS_FILE]
+
+    def __init__(self):
+        self._trial_continue: Dict["Trial", bool] = {}
+        self._trial_files: Dict["Trial", TextIO] = {}
+        self._trial_csv: Dict["Trial", csv.DictWriter] = {}
+
+    def _setup_trial(self, trial: "Trial"):
+        if trial in self._trial_files:
+            self._trial_files[trial].close()
+
+        # Make sure logdir exists
+        trial.init_local_path()
+        local_file_path = Path(trial.local_path, EXPR_PROGRESS_FILE)
+
+        # Resume the file from remote storage.
+        self._restore_from_remote(EXPR_PROGRESS_FILE, trial)
+
+        self._trial_continue[trial] = (
+            local_file_path.exists() and local_file_path.stat().st_size > 0
+        )
+
+        self._trial_files[trial] = local_file_path.open("at")
+        self._trial_csv[trial] = None
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        if trial not in self._trial_files:
+            self._setup_trial(trial)
+
+        tmp = result.copy()
+        tmp.pop("config", None)
+        result = flatten_dict(tmp, delimiter="/")
+
+        if not self._trial_csv[trial]:
+            self._trial_csv[trial] = csv.DictWriter(
+                self._trial_files[trial], result.keys()
+            )
+            if not self._trial_continue[trial]:
+                self._trial_csv[trial].writeheader()
+
+        self._trial_csv[trial].writerow(
+            {k: v for k, v in result.items() if k in self._trial_csv[trial].fieldnames}
+        )
+        self._trial_files[trial].flush()
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        if trial not in self._trial_files:
+            return
+
+        del self._trial_csv[trial]
+        self._trial_files[trial].close()
+        del self._trial_files[trial]

.venv/lib/python3.11/site-packages/ray/tune/logger/json.py

@@ -0,0 +1,128 @@
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, TextIO
+
+import numpy as np
+
+import ray.cloudpickle as cloudpickle
+from ray.air.constants import EXPR_PARAM_FILE, EXPR_PARAM_PICKLE_FILE, EXPR_RESULT_FILE
+from ray.tune.logger.logger import _LOGGER_DEPRECATION_WARNING, Logger, LoggerCallback
+from ray.tune.utils.util import SafeFallbackEncoder
+from ray.util.annotations import Deprecated, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+tf = None
+VALID_SUMMARY_TYPES = [int, float, np.float32, np.float64, np.int32, np.int64]
+
+
+@Deprecated(
+    message=_LOGGER_DEPRECATION_WARNING.format(
+        old="JsonLogger", new="ray.tune.json.JsonLoggerCallback"
+    ),
+    warning=True,
+)
+@PublicAPI
+class JsonLogger(Logger):
+    """Logs trial results in json format.
+
+    Also writes to a results file and param.json file when results or
+    configurations are updated. Experiments must be executed with the
+    JsonLogger to be compatible with the ExperimentAnalysis tool.
+    """
+
+    def _init(self):
+        self.update_config(self.config)
+        local_file = Path(self.logdir, EXPR_RESULT_FILE)
+        self.local_out = local_file.open("a")
+
+    def on_result(self, result: Dict):
+        json.dump(result, self, cls=SafeFallbackEncoder)
+        self.write("\n")
+        self.local_out.flush()
+
+    def write(self, b):
+        self.local_out.write(b)
+
+    def flush(self):
+        if not self.local_out.closed:
+            self.local_out.flush()
+
+    def close(self):
+        self.local_out.close()
+
+    def update_config(self, config: Dict):
+        self.config = config
+        config_out = Path(self.logdir, EXPR_PARAM_FILE)
+        with open(config_out, "w") as f:
+            json.dump(self.config, f, indent=2, sort_keys=True, cls=SafeFallbackEncoder)
+        config_pkl = Path(self.logdir, EXPR_PARAM_PICKLE_FILE)
+        with config_pkl.open("wb") as f:
+            cloudpickle.dump(self.config, f)
+
+
+@PublicAPI
+class JsonLoggerCallback(LoggerCallback):
+    """Logs trial results in json format.
+
+    Also writes to a results file and param.json file when results or
+    configurations are updated. Experiments must be executed with the
+    JsonLoggerCallback to be compatible with the ExperimentAnalysis tool.
+    """
+
+    _SAVED_FILE_TEMPLATES = [EXPR_RESULT_FILE, EXPR_PARAM_FILE, EXPR_PARAM_PICKLE_FILE]
+
+    def __init__(self):
+        self._trial_configs: Dict["Trial", Dict] = {}
+        self._trial_files: Dict["Trial", TextIO] = {}
+
+    def log_trial_start(self, trial: "Trial"):
+        if trial in self._trial_files:
+            self._trial_files[trial].close()
+
+        # Update config
+        self.update_config(trial, trial.config)
+
+        # Make sure logdir exists
+        trial.init_local_path()
+        local_file = Path(trial.local_path, EXPR_RESULT_FILE)
+
+        # Resume the file from remote storage.
+        self._restore_from_remote(EXPR_RESULT_FILE, trial)
+
+        self._trial_files[trial] = local_file.open("at")
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        if trial not in self._trial_files:
+            self.log_trial_start(trial)
+        json.dump(result, self._trial_files[trial], cls=SafeFallbackEncoder)
+        self._trial_files[trial].write("\n")
+        self._trial_files[trial].flush()
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        if trial not in self._trial_files:
+            return
+
+        self._trial_files[trial].close()
+        del self._trial_files[trial]
+
+    def update_config(self, trial: "Trial", config: Dict):
+        self._trial_configs[trial] = config
+
+        config_out = Path(trial.local_path, EXPR_PARAM_FILE)
+        with config_out.open("w") as f:
+            json.dump(
+                self._trial_configs[trial],
+                f,
+                indent=2,
+                sort_keys=True,
+                cls=SafeFallbackEncoder,
+            )
+
+        config_pkl = Path(trial.local_path, EXPR_PARAM_PICKLE_FILE)
+        with config_pkl.open("wb") as f:
+            cloudpickle.dump(self._trial_configs[trial], f)