chore: import upstream snapshot with attribution

release/air_examples/dolly_v2_lightning_fsdp_finetuning/dolly_v2_fsdp_compute_aws.yaml

@@ -0,0 +1,17 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: g4dn.12xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND
+
+advanced_instance_config:
+  TagSpecifications:
+    - ResourceType: "instance"
+      Tags:
+        - Key: ttl-hours
+          Value: '24'

release/air_examples/dolly_v2_lightning_fsdp_finetuning/dolly_v2_fsdp_env.yaml

@@ -0,0 +1,11 @@
+base_image: {{ env["RAY_IMAGE_ML_NIGHTLY_GPU"] | default("anyscale/ray:nightly-py38-cu118") }}
+env_vars: {}
+debian_packages:
+  - curl
+
+post_build_cmds:
+  - pip uninstall -y ray || true && pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
+  - {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}
+  - pip3 uninstall -y pytorch-lightning pytorch_lightning
+  - pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+  - pip3 install "lightning==2.0.2" "transformers==4.29.2" "accelerate==0.19.0"

release/air_examples/dolly_v2_lightning_fsdp_finetuning/lightning-llm-finetuning-7b.ipynb

@@ -0,0 +1 @@
+../../../doc/source/train/examples/lightning/dolly_lightning_fsdp_finetuning.ipynb

release/air_examples/dolly_v2_lightning_fsdp_finetuning/test_myst_doc.py

@@ -0,0 +1 @@
+../../../doc/test_myst_doc.py

release/air_examples/dreambooth/dreambooth

@@ -0,0 +1 @@
+../../../doc/source/templates/05_dreambooth_finetuning/dreambooth

release/air_examples/dreambooth/dreambooth_compute_aws.yaml

@@ -0,0 +1,6 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g5.12xlarge
+
+worker_nodes: []

release/air_examples/dreambooth/dreambooth_env.yaml

@@ -0,0 +1,16 @@
+# NOTE:
+# - This test runs with py38 (see the entry in release_tests.yaml)
+# - This test installs dependencies on top of a base ray image
+#   instead of using the default ray-ml image. See dreambooth/requirements.txt.
+base_image: "anyscale/ray:nightly-py38-cu118"
+env_vars: {}
+debian_packages:
+  - curl
+
+python:
+  pip_packages: []
+  conda_packages: []
+
+post_build_cmds:
+  - pip uninstall -y ray || true && pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
+  - {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}

release/air_examples/dreambooth/dreambooth_run.sh

@@ -0,0 +1 @@
+../../../doc/source/templates/05_dreambooth_finetuning/dreambooth_run.sh

release/air_examples/gptj_deepspeed_finetuning/gptj_deepspeed_compute_aws.yaml

@@ -0,0 +1,17 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: g4dn.4xlarge
+      min_nodes: 8
+      max_nodes: 8
+      market_type: ON_DEMAND
+
+advanced_instance_config:
+  TagSpecifications:
+    - ResourceType: "instance"
+      Tags:
+        - Key: ttl-hours
+          Value: '24'

release/air_examples/gptj_deepspeed_finetuning/gptj_deepspeed_compute_gce.yaml

@@ -0,0 +1,22 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west1
+allowed_azs:
+    - us-west1-b
+
+head_node_type:
+    name: head_node
+    instance_type: n2-standard-8
+
+worker_node_types:
+    - name: worker_node
+      instance_type: n1-standard-16-nvidia-tesla-t4-1
+      min_workers: 8
+      max_workers: 8
+      use_spot: false
+
+#advanced_configurations_json:
+#  TagSpecifications:
+#    - ResourceType: "instance"
+#      Tags:
+#        - Key: ttl-hours
+#          Value: '24'

release/air_examples/gptj_deepspeed_finetuning/gptj_deepspeed_env.yaml

@@ -0,0 +1,8 @@
+base_image: {{ env["RAY_IMAGE_ML_NIGHTLY_GPU"] }}
+env_vars: {}
+debian_packages:
+  - curl
+
+post_build_cmds:
+  - pip uninstall -y ray || true && pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
+  - {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}

release/air_examples/gptj_deepspeed_finetuning/gptj_deepspeed_fine_tuning.ipynb

@@ -0,0 +1 @@
+../../../doc/source/train/examples/deepspeed/gptj_deepspeed_fine_tuning.ipynb

release/air_examples/gptj_deepspeed_finetuning/test_myst_doc.py

@@ -0,0 +1,76 @@
+"""Convert a jupytext-compliant format in to a python script
+and execute it with parsed arguments.
+
+Any cell with 'remove-cell-ci' tag in metadata will not be included
+in the converted python script.
+"""
+
+import argparse
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import jupytext
+
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument(
+    "--path",
+    help="path to the jupytext-compatible file",
+)
+parser.add_argument(
+    "--find-recursively",
+    action="store_true",
+    help="if true, will attempt to find path recursively in cwd",
+)
+parser.add_argument(
+    "--no-postprocess",
+    action="store_true",
+    help="if true, will not postprocess the notebook",
+)
+
+
+def filter_out_cells_with_remove_cell_ci_tag(cells: list):
+    """Filters out cells which contain the 'remove-cell-ci' tag in metadata"""
+
+    def should_keep_cell(cell):
+        tags = cell.metadata.get("tags")
+        if tags:
+            # Both - and _ for consistent behavior with built-in tags
+            return "remove_cell_ci" not in tags and "remove-cell-ci" not in tags
+        return True
+
+    return [cell for cell in cells if should_keep_cell(cell)]
+
+
+def postprocess_notebook(notebook):
+    notebook.cells = filter_out_cells_with_remove_cell_ci_tag(notebook.cells)
+    return notebook
+
+
+if __name__ == "__main__":
+
+    args, remainder = parser.parse_known_args()
+
+    path = Path(args.path)
+    cwd = Path.cwd()
+    if args.find_recursively and not path.exists():
+        path = next((p for p in cwd.rglob("*") if str(p).endswith(args.path)), None)
+    assert path and path.exists()
+
+    with open(path, "r") as f:
+        notebook = jupytext.read(f)
+
+    if not args.no_postprocess:
+        notebook = postprocess_notebook(notebook)
+
+    name = ""
+    with tempfile.NamedTemporaryFile("w", delete=False) as f:
+        jupytext.write(notebook, f, fmt="py:percent")
+        name = f.name
+
+    remainder.insert(0, name)
+    remainder.insert(0, sys.executable)
+
+    # Run the notebook
+    subprocess.run(remainder, check=True)

release/air_examples/opt_deepspeed_batch_inference/30b_deepspeed_compute.yaml

@@ -0,0 +1,15 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west-2
+
+head_node_type:
+    name: head_node
+    instance_type: p3.16xlarge
+
+worker_node_types: []
+
+advanced_configurations_json:
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+            DeleteOnTermination: true
+            VolumeSize: 500

release/air_examples/opt_deepspeed_batch_inference/30b_deepspeed_env.yaml

@@ -0,0 +1,8 @@
+base_image: {{ env["RAY_IMAGE_ML_NIGHTLY_GPU"] }}
+env_vars: {}
+debian_packages:
+  - curl
+
+post_build_cmds:
+  - pip uninstall -y ray || true && pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
+  - {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}

release/air_examples/opt_deepspeed_batch_inference/test_myst_doc.py

@@ -0,0 +1 @@
+../../../doc/test_myst_doc.py

release/air_examples/vicuna_13b_lightning_deepspeed_finetuning/test_myst_doc.py

@@ -0,0 +1 @@
+../../../doc/test_myst_doc.py

release/air_examples/vicuna_13b_lightning_deepspeed_finetuning/vicuna_13b_deepspeed_compute_aws.yaml

@@ -0,0 +1,17 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g5.16xlarge
+
+worker_nodes:
+    - instance_type: g5.4xlarge
+      min_nodes: 15
+      max_nodes: 15
+      market_type: ON_DEMAND
+
+advanced_instance_config:
+  TagSpecifications:
+    - ResourceType: "instance"
+      Tags:
+        - Key: ttl-hours
+          Value: '24'

release/air_examples/vicuna_13b_lightning_deepspeed_finetuning/vicuna_13b_lightning_deepspeed_finetune.ipynb

@@ -0,0 +1 @@
+../../../doc/source/train/examples/lightning/vicuna_13b_lightning_deepspeed_finetune.ipynb

release/air_tests/air_benchmarks/compute_cpu_1_aws.yaml

@@ -0,0 +1,10 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west-2
+
+max_workers: 0
+
+head_node_type:
+    name: head_node
+    instance_type: m5.2xlarge
+
+worker_node_types: []

release/air_tests/air_benchmarks/compute_cpu_1_gce.yaml

@@ -0,0 +1,12 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west1
+allowed_azs:
+    - us-west1-b
+
+max_workers: 0
+
+head_node_type:
+    name: head_node
+    instance_type: n1-standard-8
+
+worker_node_types: []

release/air_tests/air_benchmarks/compute_cpu_4_aws.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: m5.2xlarge
+      min_nodes: 3
+      max_nodes: 3
+      market_type: ON_DEMAND

release/air_tests/air_benchmarks/compute_cpu_4_gce.yaml

@@ -0,0 +1,12 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-b
+
+head_node:
+    instance_type: n1-standard-8
+
+worker_nodes:
+    - instance_type: n1-standard-8
+      min_nodes: 3
+      max_nodes: 3
+      market_type: ON_DEMAND

release/air_tests/air_benchmarks/compute_cpu_8_aws.yaml

@@ -0,0 +1,15 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west-2
+
+max_workers: 7
+
+head_node_type:
+    name: head_node
+    instance_type: m5.2xlarge
+
+worker_node_types:
+    - name: worker_node
+      instance_type: m5.2xlarge
+      max_workers: 7
+      min_workers: 7
+      use_spot: false

release/air_tests/air_benchmarks/compute_cpu_8_gce.yaml

@@ -0,0 +1,17 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west1
+allowed_azs:
+    - us-west1-b
+
+max_workers: 7
+
+head_node_type:
+    name: head_node
+    instance_type: n1-standard-8
+
+worker_node_types:
+    - name: worker_node
+      instance_type: n1-standard-8
+      max_workers: 7
+      min_workers: 7
+      use_spot: false

release/air_tests/air_benchmarks/compute_gpu_1_aws.yaml

@@ -0,0 +1,20 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west-2
+
+max_workers: 0
+
+head_node_type:
+    name: head_node
+    instance_type: g4dn.8xlarge
+
+worker_node_types: []
+
+advanced_configurations_json:
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+            DeleteOnTermination: true
+            Iops: 5000
+            Throughput: 1000
+            VolumeSize: 1000
+            VolumeType: gp3

release/air_tests/air_benchmarks/compute_gpu_1_gce.yaml

@@ -0,0 +1,20 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west1
+allowed_azs:
+    - us-west1-b
+
+max_workers: 0
+
+head_node_type:
+    name: head_node
+    instance_type: n1-standard-32-nvidia-tesla-t4-2
+
+worker_node_types: []
+
+gcp_advanced_configurations_json:
+  instance_properties:
+    disks:
+      - boot: true
+        auto_delete: true
+        initialize_params:
+          disk_size_gb: 1000

release/air_tests/air_benchmarks/compute_gpu_2x2_aws.yaml

@@ -0,0 +1,13 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g4dn.8xlarge
+    resources:
+        CPU: 32
+        GPU: 1
+
+worker_nodes:
+    - instance_type: g4dn.8xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/air_tests/air_benchmarks/compute_gpu_2x2_gce.yaml

@@ -0,0 +1,17 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west1
+allowed_azs:
+    - us-west1-b
+
+max_workers: 1
+
+head_node_type:
+    name: head_node
+    instance_type: n1-standard-32-nvidia-tesla-t4-2
+
+worker_node_types:
+    - name: worker_node
+      instance_type: n1-standard-32-nvidia-tesla-t4-2
+      max_workers: 1
+      min_workers: 1
+      use_spot: false

release/air_tests/air_benchmarks/compute_gpu_4x4_aws.yaml

@@ -0,0 +1,23 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g4dn.12xlarge
+    resources:
+        CPU: 48
+        GPU: 4
+
+worker_nodes:
+    - instance_type: g4dn.12xlarge
+      min_nodes: 3
+      max_nodes: 3
+      market_type: ON_DEMAND
+
+advanced_instance_config:
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+            DeleteOnTermination: true
+            VolumeSize: 800
+            Iops: 5000
+            Throughput: 1000
+            VolumeType: gp3

release/air_tests/air_benchmarks/compute_gpu_4x4_gce.yaml

@@ -0,0 +1,23 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-b
+
+head_node:
+    instance_type: n1-standard-64-nvidia-tesla-t4-4
+    resources:
+        CPU: 64
+        GPU: 4
+
+worker_nodes:
+    - instance_type: n1-standard-64-nvidia-tesla-t4-4
+      min_nodes: 3
+      max_nodes: 3
+      market_type: ON_DEMAND
+
+advanced_instance_config:
+  instance_properties:
+    disks:
+      - boot: true
+        auto_delete: true
+        initialize_params:
+          disk_size_gb: 800

release/air_tests/air_benchmarks/workloads/_tensorflow_prepare.py

@@ -0,0 +1,3 @@
+import tensorflow as tf
+
+tf.keras.datasets.fashion_mnist.load_data()

release/air_tests/air_benchmarks/workloads/_torch_prepare.py

@@ -0,0 +1,3 @@
+import torchvision
+
+torchvision.datasets.FashionMNIST("/tmp/data_fashion_mnist", download=True)

release/air_tests/air_benchmarks/workloads/benchmark_util.py

@@ -0,0 +1,155 @@
+import os
+import socket
+import subprocess
+from collections import defaultdict
+from contextlib import closing
+from pathlib import Path
+from ray.air.util.node import _force_on_node
+
+import ray
+from typing import List, Dict, Union, Callable
+
+
+def schedule_remote_fn_on_all_nodes(
+    remote_fn, exclude_head: bool = False, *args, **kwargs
+):
+    head_ip = ray.util.get_node_ip_address()
+
+    futures = []
+    for node in ray.nodes():
+        if not node["Alive"]:
+            continue
+
+        node_ip = node["NodeManagerAddress"]
+
+        if exclude_head and node_ip == head_ip:
+            continue
+
+        node_id = node["NodeID"]
+
+        future = _force_on_node(node_id, remote_fn).remote(
+            *args,
+            **kwargs,
+        )
+        futures.append(future)
+    return futures
+
+
+@ray.remote
+def _write(stream: bytes, path: str):
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+
+    with open(path, "wb") as f:
+        f.write(stream)
+
+
+def upload_file_to_all_nodes(path: str):
+    path = os.path.abspath(path)
+
+    with open(path, "rb") as f:
+        stream = f.read()
+
+    futures = schedule_remote_fn_on_all_nodes(
+        _write, exclude_head=True, stream=stream, path=path
+    )
+    return ray.get(futures)
+
+
+@ray.remote
+def _run_command(cmd: str):
+    return subprocess.check_call(cmd)
+
+
+def run_command_on_all_nodes(cmd: List[str]):
+    futures = schedule_remote_fn_on_all_nodes(_run_command, cmd=cmd)
+    return ray.get(futures)
+
+
+@ray.remote
+class CommandRunner:
+    def run_command(self, cmd: str):
+        return subprocess.check_call(cmd)
+
+    def run_fn(self, fn: Callable, *args, **kwargs):
+        return fn(*args, **kwargs)
+
+
+def create_actors_with_options(
+    num_actors: int,
+    resources: Dict[str, Union[float, int]],
+) -> List[ray.actor.ActorHandle]:
+    num_cpus = resources.pop("CPU", 1)
+    num_gpus = resources.pop("GPU", 0)
+
+    options = {"num_cpus": num_cpus, "num_gpus": num_gpus, "resources": resources}
+
+    return [CommandRunner.options(**options).remote() for _ in range(num_actors)]
+
+
+def run_commands_on_actors(actors: List[ray.actor.ActorHandle], cmds: List[List[str]]):
+    assert len(actors) == len(cmds)
+    futures = []
+    for actor, cmd in zip(actors, cmds):
+        futures.append(actor.run_command.remote(cmd))
+    return ray.get(futures)
+
+
+def run_fn_on_actors(
+    actors: List[ray.actor.ActorHandle], fn: Callable, *args, **kwargs
+):
+    futures = []
+    for actor in actors:
+        futures.append(actor.run_fn.remote(fn, *args, **kwargs))
+    return ray.get(futures)
+
+
+def get_ip_port_actors(actors: List[ray.actor.ActorHandle]) -> List[str]:
+    # We need this wrapper to avoid deserialization issues with benchmark_util.py
+
+    def get_ip_port():
+        ip = ray.util.get_node_ip_address()
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+            s.bind(("localhost", 0))
+            port = s.getsockname()[1]
+        return ip, port
+
+    return run_fn_on_actors(actors=actors, fn=get_ip_port)
+
+
+def get_gpu_ids_actors(actors: List[ray.actor.ActorHandle]) -> List[List[int]]:
+    # We need this wrapper to avoid deserialization issues with benchmark_util.py
+
+    def get_gpu_ids():
+        return ray.get_gpu_ids()
+
+    return run_fn_on_actors(actors=actors, fn=get_gpu_ids)
+
+
+def map_ips_to_gpus(ips: List[str], gpus: List[List[int]]):
+    assert len(ips) == len(gpus)
+
+    map = defaultdict(set)
+    for ip, gpu in zip(ips, gpus):
+        map[ip].update(set(gpu))
+    return {ip: sorted(gpus) for ip, gpus in map.items()}
+
+
+def set_cuda_visible_devices(
+    actors: List[ray.actor.ActorHandle],
+    actor_ips: List[str],
+    ip_to_gpus: Dict[str, set],
+):
+    assert len(actors) == len(actor_ips)
+
+    def set_env(key: str, val: str):
+        os.environ[key] = val
+
+    futures = []
+    for actor, ip in zip(actors, actor_ips):
+        assert ip in ip_to_gpus
+
+        gpu_str = ",".join([str(device) for device in sorted(ip_to_gpus[ip])])
+        future = actor.run_fn.remote(set_env, "CUDA_VISIBLE_DEVICES", gpu_str)
+        futures.append(future)
+
+    ray.get(futures)

release/air_tests/air_benchmarks/workloads/pytorch_training_e2e.py

@@ -0,0 +1,149 @@
+import click
+import time
+import json
+import os
+import tempfile
+from typing import Dict
+
+import numpy as np
+from torchvision import transforms
+from torchvision.models import resnet18
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+import ray
+from ray import train
+from ray.train import Checkpoint, RunConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def add_fake_labels(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+    batch_size = len(batch["image"])
+    batch["label"] = np.zeros([batch_size], dtype=int)
+    return batch
+
+
+def transform_image(
+    batch: Dict[str, np.ndarray], transform: torch.nn.Module
+) -> Dict[str, np.ndarray]:
+    transformed_tensors = [transform(image).numpy() for image in batch["image"]]
+    batch["image"] = transformed_tensors
+    return batch
+
+
+def train_loop_per_worker(config):
+    raw_model = resnet18(pretrained=True)
+    model = train.torch.prepare_model(raw_model)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
+
+    train_dataset_shard = train.get_dataset_shard("train")
+
+    for epoch in range(config["num_epochs"]):
+        running_loss = 0.0
+        for i, data in enumerate(
+            train_dataset_shard.iter_torch_batches(batch_size=config["batch_size"])
+        ):
+            # get the inputs; data is a list of [inputs, labels]
+            inputs = data["image"].to(device=train.torch.get_device())
+            labels = data["label"].to(device=train.torch.get_device())
+            # zero the parameter gradients
+            optimizer.zero_grad()
+
+            # forward + backward + optimize
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            # print statistics
+            running_loss += loss.item()
+            if i % 2000 == 1999:  # print every 2000 mini-batches
+                print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}")
+                running_loss = 0.0
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(model.state_dict(), os.path.join(tmpdir, "model.pt"))
+            train.report(
+                dict(running_loss=running_loss),
+                checkpoint=Checkpoint.from_directory(tmpdir),
+            )
+
+
+@click.command(help="Run Batch prediction on Pytorch ResNet models.")
+@click.option("--data-size-gb", type=int, default=1)
+@click.option("--num-epochs", type=int, default=2)
+@click.option("--num-workers", type=int, default=1)
+@click.option("--smoke-test", is_flag=True, default=False)
+def main(data_size_gb: int, num_epochs=2, num_workers=1, smoke_test: bool = False):
+    data_url = (
+        f"s3://anonymous@air-example-data-2/{data_size_gb}G-image-data-synthetic-raw"
+    )
+    print(
+        "Running Pytorch image model training with "
+        f"{data_size_gb}GB data from {data_url}"
+    )
+    print(f"Training for {num_epochs} epochs with {num_workers} workers.")
+    start = time.time()
+
+    if smoke_test:
+        # Only read one image
+        data_url = [data_url + "/dog.jpg"]
+        print("Running smoke test on CPU with a single example")
+    else:
+        print(f"Running GPU training with {data_size_gb}GB data from {data_url}")
+
+    dataset = ray.data.read_images(data_url, size=(256, 256))
+
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    dataset = dataset.map_batches(add_fake_labels)
+    dataset = dataset.map_batches(transform_image, fn_kwargs={"transform": transform})
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop_per_worker,
+        train_loop_config={"batch_size": 64, "num_epochs": num_epochs},
+        datasets={"train": dataset},
+        scaling_config=ScalingConfig(
+            num_workers=num_workers, use_gpu=int(not smoke_test)
+        ),
+        run_config=RunConfig(storage_path="/mnt/cluster_storage"),
+    )
+    trainer.fit()
+
+    total_time_s = round(time.time() - start, 2)
+
+    # For structured output integration with internal tooling
+    results = {"data_size_gb": data_size_gb, "num_epochs": num_epochs}
+    results["perf_metrics"] = [
+        {
+            "perf_metric_name": "total_time_s",
+            "perf_metric_value": total_time_s,
+            "perf_metric_type": "LATENCY",
+        },
+        {
+            "perf_metric_name": "throughout_MB_s",
+            "perf_metric_value": round(
+                num_epochs * data_size_gb * 1024 / total_time_s, 2
+            ),
+            "perf_metric_type": "THROUGHPUT",
+        },
+    ]
+
+    test_output_json = os.environ.get("TEST_OUTPUT_JSON", "/tmp/release_test_out.json")
+    with open(test_output_json, "wt") as f:
+        json.dump(results, f)
+
+    print(results)
+
+
+if __name__ == "__main__":
+    main()

release/air_tests/air_benchmarks/workloads/tensorflow_benchmark.py

@@ -0,0 +1,440 @@
+import json
+import os
+import time
+from pathlib import Path
+
+import click
+import numpy as np
+import tensorflow as tf
+from typing import List, Tuple
+
+from ray._common.network_utils import build_address
+
+CONFIG = {"lr": 1e-3, "batch_size": 64}
+VANILLA_RESULT_JSON = "/tmp/vanilla_out.json"
+
+
+def mnist_dataset(batch_size: int) -> tf.data.Dataset:
+    (x_train, y_train), _ = tf.keras.datasets.fashion_mnist.load_data()
+    # The `x` arrays are in uint8 and have values in the [0, 255] range.
+    # You need to convert them to float32 with values in the [0, 1] range.
+    x_train = x_train / np.float32(255)
+    y_train = y_train.astype(np.int64)
+    train_dataset = (
+        tf.data.Dataset.from_tensor_slices((x_train, y_train))
+        .shuffle(60000, seed=1234)
+        .batch(batch_size)
+    )
+    return train_dataset
+
+
+def build_cnn_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=(28, 28)),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(512, activation="relu"),
+            tf.keras.layers.Dense(512, activation="relu"),
+            tf.keras.layers.Dense(10),
+        ]
+    )
+    return model
+
+
+def train_func(use_ray: bool, config: dict):
+    local_start_time = time.monotonic()
+
+    per_worker_batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+    steps_per_epoch = config.get("steps_per_epoch", None)
+    learning_rate = config.get("lr", 0.001)
+
+    tf_config = json.loads(os.environ["TF_CONFIG"])
+    num_workers = len(tf_config["cluster"]["worker"])
+    local_rank = tf_config["task"]["index"]
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    global_batch_size = per_worker_batch_size * num_workers
+    multi_worker_dataset = mnist_dataset(global_batch_size)
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_cnn_model()
+        multi_worker_model.compile(
+            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate),
+            metrics=["accuracy"],
+        )
+
+    if use_ray:
+        from ray.air.integrations.keras import ReportCheckpointCallback
+
+        class CustomReportCallback(ReportCheckpointCallback):
+            def _handle(self, logs: dict, when: str = None):
+                logs["local_time_taken"] = time.monotonic() - local_start_time
+                super()._handle(logs, when)
+
+        # NOTE: We shouldn't checkpoint to be identical to the vanilla TF run.
+        callbacks = [CustomReportCallback()]
+    else:
+        callbacks = []
+
+    history = multi_worker_model.fit(
+        multi_worker_dataset,
+        epochs=epochs,
+        steps_per_epoch=steps_per_epoch,
+        callbacks=callbacks,
+        verbose=2,  # Disables progress bar in remote actors.
+    )
+    results = history.history
+    loss = results["loss"][-1]
+
+    if not use_ray:
+        local_time_taken = time.monotonic() - local_start_time
+        print(f"Reporting loss: {loss:.4f}")
+        if local_rank == 0:
+            with open(VANILLA_RESULT_JSON, "w") as f:
+                json.dump({"loss": loss, "local_time_taken": local_time_taken}, f)
+
+    return results
+
+
+def train_tf_ray_air(
+    *,
+    config: dict,
+    num_workers: int = 4,
+    cpus_per_worker: int = 8,
+    use_gpu: bool = False,
+) -> Tuple[float, float, float]:
+    from ray.train.tensorflow import TensorflowTrainer
+    from ray.train import ScalingConfig
+
+    def train_loop(config):
+        train_func(use_ray=True, config=config)
+
+    start_time = time.monotonic()
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_loop,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(
+            num_workers=num_workers,
+            resources_per_worker={"CPU": cpus_per_worker},
+            use_gpu=use_gpu,
+        ),
+    )
+    result = trainer.fit()
+    time_taken = time.monotonic() - start_time
+
+    print(f"Last result: {result.metrics}")
+    return time_taken, result.metrics["local_time_taken"], result.metrics["loss"]
+
+
+def train_tf_vanilla_worker(
+    *,
+    config: dict,
+    rank: int,
+    world_size: int,
+    worker_ip_port_list: List[str],
+    use_gpu: bool = False,
+):
+    # This function is kicked off by the main() function and runs the vanilla
+    # training script on a single worker.
+    assert world_size == len(worker_ip_port_list)
+
+    tf_config = {
+        "cluster": {"worker": worker_ip_port_list},
+        "task": {"type": "worker", "index": rank},
+    }
+    os.environ["TF_CONFIG"] = json.dumps(tf_config)
+
+    train_func(use_ray=False, config=config)
+
+
+def train_tf_vanilla(
+    *,
+    config: dict,
+    num_workers: int = 4,
+    cpus_per_worker: int = 8,
+    use_gpu: bool = False,
+) -> Tuple[float, float, float]:
+    # This function is kicked off by the main() function and subsequently kicks
+    # off tasks that run train_tf_vanilla_worker() on the worker nodes.
+    from benchmark_util import (
+        upload_file_to_all_nodes,
+        create_actors_with_options,
+        run_commands_on_actors,
+        run_fn_on_actors,
+        get_ip_port_actors,
+    )
+
+    path = os.path.abspath(__file__)
+    upload_file_to_all_nodes(path)
+
+    num_epochs = config["epochs"]
+
+    actors = create_actors_with_options(
+        num_actors=num_workers,
+        resources={
+            "CPU": cpus_per_worker,
+            "GPU": int(use_gpu),
+        },
+    )
+
+    run_fn_on_actors(actors=actors, fn=lambda: os.environ.pop("OMP_NUM_THREADS", None))
+
+    ips_ports = get_ip_port_actors(actors=actors)
+    ip_port_list = [build_address(ip, port) for ip, port in ips_ports]
+    ip_port_str = ",".join(ip_port_list)
+
+    cmds = [
+        [
+            "python",
+            path,
+            "worker",
+            "--num-epochs",
+            str(num_epochs),
+            "--num-workers",
+            str(num_workers),
+            "--rank",
+            str(rank),
+            "--worker-ip-ports",
+            ip_port_str,
+            "--batch-size",
+            str(config["batch_size"]),
+        ]
+        + (["--use-gpu"] if use_gpu else [])
+        for rank in range(num_workers)
+    ]
+
+    run_fn_on_actors(
+        actors=actors, fn=lambda: os.environ.setdefault("OMP_NUM_THREADS", "1")
+    )
+
+    start_time = time.monotonic()
+    run_commands_on_actors(actors=actors, cmds=cmds)
+    time_taken = time.monotonic() - start_time
+
+    loss = local_time_taken = 0.0
+    if os.path.exists(VANILLA_RESULT_JSON):
+        with open(VANILLA_RESULT_JSON, "r") as f:
+            result = json.load(f)
+        loss = result["loss"]
+        local_time_taken = result["local_time_taken"]
+
+    return time_taken, local_time_taken, loss
+
+
+@click.group(help="Run Tensorflow benchmarks")
+def cli():
+    pass
+
+
+@cli.command(help="Kick off Ray and vanilla benchmarks")
+@click.option("--num-runs", type=int, default=1)
+@click.option("--num-epochs", type=int, default=4)
+@click.option("--num-workers", type=int, default=4)
+@click.option("--cpus-per-worker", type=int, default=8)
+@click.option("--use-gpu", is_flag=True, default=False)
+@click.option("--batch-size", type=int, default=64)
+@click.option("--smoke-test", is_flag=True, default=False)
+@click.option("--local", is_flag=True, default=False)
+def run(
+    num_runs: int = 1,
+    num_epochs: int = 4,
+    num_workers: int = 4,
+    cpus_per_worker: int = 8,
+    use_gpu: bool = False,
+    batch_size: int = 64,
+    smoke_test: bool = False,
+    local: bool = False,
+):
+    # Note: smoke_test is ignored as we just adjust the batch size.
+    # The parameter is passed by the release test pipeline.
+    import ray
+    from benchmark_util import upload_file_to_all_nodes, run_command_on_all_nodes
+
+    config = CONFIG.copy()
+    config["epochs"] = num_epochs
+    config["batch_size"] = batch_size
+
+    if local:
+        ray.init(num_cpus=4)
+    else:
+        ray.init("auto")
+
+    print("Preparing Tensorflow benchmark: Downloading MNIST")
+
+    path = str((Path(__file__).parent / "_tensorflow_prepare.py").absolute())
+
+    upload_file_to_all_nodes(path)
+    run_command_on_all_nodes(["python", path])
+
+    times_ray = []
+    times_local_ray = []
+    losses_ray = []
+    times_vanilla = []
+    times_local_vanilla = []
+    losses_vanilla = []
+    for run in range(1, num_runs + 1):
+        time.sleep(2)
+
+        print(f"[Run {run}/{num_runs}] Running Tensorflow Ray benchmark")
+
+        time_ray, time_local_ray, loss_ray = train_tf_ray_air(
+            num_workers=num_workers,
+            cpus_per_worker=cpus_per_worker,
+            use_gpu=use_gpu,
+            config=config,
+        )
+
+        print(
+            f"[Run {run}/{num_runs}] Finished Ray training ({num_epochs} epochs) in "
+            f"{time_ray:.2f} seconds (local training time: {time_local_ray:.2f}s). "
+            f"Observed loss = {loss_ray:.4f}"
+        )
+
+        time.sleep(2)
+
+        print(f"[Run {run}/{num_runs}] Running Tensorflow vanilla benchmark")
+
+        # Todo: Vanilla runs are sometimes failing. We just retry here, but we should
+        # get to the bottom of it.
+        time_vanilla = time_local_vanilla = loss_vanilla = 0.0
+        for i in range(3):
+            try:
+                time_vanilla, time_local_vanilla, loss_vanilla = train_tf_vanilla(
+                    num_workers=num_workers,
+                    cpus_per_worker=cpus_per_worker,
+                    use_gpu=use_gpu,
+                    config=config,
+                )
+            except Exception as e:
+                if i >= 2:
+                    raise RuntimeError("Vanilla TF run failed 3 times") from e
+                print("Vanilla TF run failed:", e)
+                continue
+            break
+
+        print(
+            f"[Run {run}/{num_runs}] Finished vanilla training ({num_epochs} epochs) "
+            f"in {time_vanilla:.2f} seconds "
+            f"(local training time: {time_local_vanilla:.2f}s). "
+            f"Observed loss = {loss_vanilla:.4f}"
+        )
+
+        print(
+            f"[Run {run}/{num_runs}] Observed results: ",
+            {
+                "tensorflow_mnist_ray_time_s": time_ray,
+                "tensorflow_mnist_ray_local_time_s": time_local_ray,
+                "tensorflow_mnist_ray_loss": loss_ray,
+                "tensorflow_mnist_vanilla_time_s": time_vanilla,
+                "tensorflow_mnist_vanilla_local_time_s": time_local_vanilla,
+                "tensorflow_mnist_vanilla_loss": loss_vanilla,
+            },
+        )
+
+        times_ray.append(time_ray)
+        times_local_ray.append(time_local_ray)
+        losses_ray.append(loss_ray)
+
+        times_vanilla.append(time_vanilla)
+        times_local_vanilla.append(time_local_vanilla)
+        losses_vanilla.append(loss_vanilla)
+
+    times_ray_mean = np.mean(times_ray)
+    times_ray_sd = np.std(times_ray)
+
+    times_local_ray_mean = np.mean(times_local_ray)
+    times_local_ray_sd = np.std(times_local_ray)
+
+    times_vanilla_mean = np.mean(times_vanilla)
+    times_vanilla_sd = np.std(times_vanilla)
+
+    times_local_vanilla_mean = np.mean(times_local_vanilla)
+    times_local_vanilla_sd = np.std(times_local_vanilla)
+
+    result = {
+        "tensorflow_mnist_ray_num_runs": num_runs,
+        "tensorflow_mnist_ray_time_s_all": times_ray,
+        "tensorflow_mnist_ray_time_s_mean": times_ray_mean,
+        "tensorflow_mnist_ray_time_s_sd": times_ray_sd,
+        "tensorflow_mnist_ray_time_local_s_all": times_local_ray,
+        "tensorflow_mnist_ray_time_local_s_mean": times_local_ray_mean,
+        "tensorflow_mnist_ray_time_local_s_sd": times_local_ray_sd,
+        "tensorflow_mnist_ray_loss_mean": np.mean(losses_ray),
+        "tensorflow_mnist_ray_loss_sd": np.std(losses_ray),
+        "tensorflow_mnist_vanilla_time_s_all": times_vanilla,
+        "tensorflow_mnist_vanilla_time_s_mean": times_vanilla_mean,
+        "tensorflow_mnist_vanilla_time_s_sd": times_vanilla_sd,
+        "tensorflow_mnist_vanilla_local_time_s_all": times_local_vanilla,
+        "tensorflow_mnist_vanilla_local_time_s_mean": times_local_vanilla_mean,
+        "tensorflow_mnist_vanilla_local_time_s_sd": times_local_vanilla_sd,
+        "tensorflow_mnist_vanilla_loss_mean": np.mean(losses_vanilla),
+        "tensorflow_mnist_vanilla_loss_std": np.std(losses_vanilla),
+    }
+
+    print("Results:", result)
+    test_output_json = os.environ.get("TEST_OUTPUT_JSON", "/tmp/result.json")
+    with open(test_output_json, "wt") as f:
+        json.dump(result, f)
+
+    target_ratio = 1.2
+    ratio = (
+        (times_local_ray_mean / times_local_vanilla_mean)
+        if times_local_vanilla_mean != 0.0
+        else 1.0
+    )
+    if ratio > target_ratio:
+        raise RuntimeError(
+            f"Training on Ray took an average of {times_local_ray_mean:.2f} seconds, "
+            f"which is more than {target_ratio:.2f}x of the average vanilla training "
+            f"time of {times_local_vanilla_mean:.2f} seconds ({ratio:.2f}x). FAILED"
+        )
+
+    print(
+        f"Training on Ray took an average of {times_local_ray_mean:.2f} seconds, "
+        f"which is less than {target_ratio:.2f}x of the average vanilla training "
+        f"time of {times_local_vanilla_mean:.2f} seconds ({ratio:.2f}x). PASSED"
+    )
+
+
+@cli.command(help="Run Tensorflow vanilla worker")
+@click.option("--num-epochs", type=int, default=4)
+@click.option("--num-workers", type=int, default=4)
+@click.option("--rank", type=int, default=0)
+@click.option("--worker-ip-ports", type=str, default="")
+@click.option("--batch-size", type=int, default=64)
+@click.option("--use-gpu", is_flag=True, default=False)
+def worker(
+    num_epochs: int = 4,
+    num_workers: int = 4,
+    rank: int = 0,
+    worker_ip_ports: str = "",
+    batch_size: int = 64,
+    use_gpu: bool = False,
+):
+    config = CONFIG.copy()
+    config["epochs"] = num_epochs
+    config["batch_size"] = batch_size
+
+    # Parse worker ip ports
+    worker_ip_port_list = worker_ip_ports.split(",")
+
+    # Then we kick off the training function on every worker.
+    return train_tf_vanilla_worker(
+        config=config,
+        rank=rank,
+        world_size=num_workers,
+        worker_ip_port_list=worker_ip_port_list,
+        use_gpu=use_gpu,
+    )
+
+
+def main():
+    return cli()
+
+
+if __name__ == "__main__":
+    main()

release/air_tests/air_benchmarks/workloads/torch_benchmark.py

@@ -0,0 +1,592 @@
+import json
+import os
+import time
+from pathlib import Path
+from typing import Dict, Tuple
+import tempfile
+
+import click
+import numpy as np
+import torch
+from torch import nn, distributed
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.utils.data.dataloader import default_collate
+from torchvision import datasets
+from torchvision.transforms import ToTensor
+
+
+CONFIG = {"lr": 1e-3, "batch_size": 64}
+VANILLA_RESULT_JSON = "/tmp/vanilla_out.json"
+
+
+# Define model
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super(NeuralNetwork, self).__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(28 * 28, 512),
+            nn.ReLU(),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Linear(512, 10),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        x = self.flatten(x)
+        logits = self.linear_relu_stack(x)
+        return logits
+
+
+def train_epoch(
+    dataloader, model, loss_fn, optimizer, world_size: int, local_rank: int
+):
+    size = len(dataloader.dataset) // world_size
+    model.train()
+    for batch, (X, y) in enumerate(dataloader):
+        # Compute prediction error
+        pred = model(X)
+        loss = loss_fn(pred, y)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if batch % 100 == 0:
+            loss, current = loss.item(), batch * len(X)
+            print(f"[rank={local_rank}] loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
+
+
+def validate_epoch(dataloader, model, loss_fn, world_size: int, local_rank: int):
+    size = len(dataloader.dataset) // world_size
+    num_batches = len(dataloader)
+    model.eval()
+    test_loss, correct = 0, 0
+    with torch.no_grad():
+        for X, y in dataloader:
+            pred = model(X)
+            test_loss += loss_fn(pred, y).item()
+            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
+    test_loss /= num_batches
+    correct /= size
+    print(
+        f"[rank={local_rank}] Test Error: \n "
+        f"Accuracy: {(100 * correct):>0.1f}%, "
+        f"Avg loss: {test_loss:>8f} \n"
+    )
+    return test_loss
+
+
+def train_func(use_ray: bool, config: Dict):
+    local_start_time = time.monotonic()
+
+    if use_ray:
+        import ray.train as train
+
+    batch_size = config["batch_size"]
+    lr = config["lr"]
+    epochs = config["epochs"]
+    shuffle = config.get("shuffle", False)
+
+    if use_ray:
+        world_size = train.get_context().get_world_size()
+        local_rank = distributed.get_rank()
+    else:
+        world_size = distributed.get_world_size()
+        local_rank = distributed.get_rank()
+
+    worker_batch_size = batch_size // world_size
+
+    # Load datasets. Use download=False to catch errors in preparation, as the
+    # data should have already been downloaded.
+    training_data = datasets.FashionMNIST(
+        root="/tmp/data_fashion_mnist",
+        train=True,
+        download=False,
+        transform=ToTensor(),
+    )
+
+    test_data = datasets.FashionMNIST(
+        root="/tmp/data_fashion_mnist",
+        train=False,
+        download=False,
+        transform=ToTensor(),
+    )
+
+    if use_ray:
+        # Ray adds DistributedSampler in train.torch.prepare_data_loader below
+        training_sampler = None
+        test_sampler = None
+    else:
+        # In vanilla PyTorch we create the distributed sampler here
+        training_sampler = DistributedSampler(training_data, shuffle=shuffle)
+        test_sampler = DistributedSampler(test_data, shuffle=shuffle)
+
+    if not use_ray and config.get("use_gpu", False):
+        assert torch.cuda.is_available(), "No GPUs available"
+        gpu_id = config.get("gpu_id", 0)
+        vanilla_device = torch.device(f"cuda:{gpu_id}")
+        torch.cuda.set_device(vanilla_device)
+
+        print(
+            "Setting GPU ID to",
+            gpu_id,
+            "with visible devices",
+            os.environ.get("CUDA_VISIBLE_DEVICES"),
+        )
+
+        def collate_fn(x):
+            return tuple(x_.to(vanilla_device) for x_ in default_collate(x))
+
+    else:
+        vanilla_device = torch.device("cpu")
+        collate_fn = None
+
+    # Create data loaders and potentially pass distributed sampler
+    train_dataloader = DataLoader(
+        training_data,
+        shuffle=shuffle,
+        batch_size=worker_batch_size,
+        sampler=training_sampler,
+        collate_fn=collate_fn,
+    )
+    test_dataloader = DataLoader(
+        test_data,
+        shuffle=shuffle,
+        batch_size=worker_batch_size,
+        sampler=test_sampler,
+        collate_fn=collate_fn,
+    )
+
+    if use_ray:
+        # In Ray, we now retrofit the DistributedSampler
+        train_dataloader = train.torch.prepare_data_loader(train_dataloader)
+        test_dataloader = train.torch.prepare_data_loader(test_dataloader)
+
+    # Create model.
+    model = NeuralNetwork()
+
+    # Prepare model
+    if use_ray:
+        model = train.torch.prepare_model(model)
+    else:
+        model = model.to(vanilla_device)
+
+        if config.get("use_gpu", False):
+            model = nn.parallel.DistributedDataParallel(
+                model, device_ids=[gpu_id], output_device=gpu_id
+            )
+        else:
+            model = nn.parallel.DistributedDataParallel(model)
+
+    loss_fn = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+
+    for epoch in range(epochs):
+        if world_size > 1:
+            train_dataloader.sampler.set_epoch(epoch)
+
+        train_epoch(
+            train_dataloader,
+            model,
+            loss_fn,
+            optimizer,
+            world_size=world_size,
+            local_rank=local_rank,
+        )
+        loss = validate_epoch(
+            test_dataloader,
+            model,
+            loss_fn,
+            world_size=world_size,
+            local_rank=local_rank,
+        )
+
+        local_time_taken = time.monotonic() - local_start_time
+
+    if use_ray:
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            if train.get_context().get_world_rank() == 0:
+                torch.save(
+                    model.state_dict(),
+                    os.path.join(temp_checkpoint_dir, "model.pt"),
+                )
+
+            train.report(
+                dict(loss=loss, local_time_taken=local_time_taken),
+                checkpoint=train.Checkpoint.from_directory(temp_checkpoint_dir),
+            )
+    else:
+        print(f"Reporting loss: {loss:.4f}")
+        if local_rank == 0:
+            with open(VANILLA_RESULT_JSON, "w") as f:
+                json.dump({"loss": loss, "local_time_taken": local_time_taken}, f)
+
+
+def train_torch_ray_air(
+    *,
+    config: dict,
+    num_workers: int = 4,
+    cpus_per_worker: int = 8,
+    use_gpu: bool = False,
+) -> Tuple[float, float, float]:
+    # This function is kicked off by the main() function and runs a full training
+    # run using Ray Train.
+    from ray.train import ScalingConfig, RunConfig
+    from ray.train.torch import TorchTrainer
+
+    def train_loop(config):
+        train_func(use_ray=True, config=config)
+
+    start_time = time.monotonic()
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(
+            num_workers=num_workers,
+            resources_per_worker={"CPU": cpus_per_worker},
+            use_gpu=use_gpu,
+        ),
+        run_config=RunConfig(storage_path="/mnt/cluster_storage"),
+    )
+    result = trainer.fit()
+    time_taken = time.monotonic() - start_time
+
+    print(f"Last result: {result.metrics}")
+    return time_taken, result.metrics["local_time_taken"], result.metrics["loss"]
+
+
+def train_torch_vanilla_worker(
+    *,
+    config: dict,
+    rank: int,
+    world_size: int,
+    master_addr: str,
+    master_port: int,
+    use_gpu: bool = False,
+    gpu_id: int = 0,
+):
+    # This function is kicked off by the main() function and runs the vanilla
+    # training script on a single worker.
+    backend = "nccl" if use_gpu else "gloo"
+
+    os.environ["MASTER_ADDR"] = master_addr
+    os.environ["MASTER_PORT"] = str(master_port)
+    os.environ["NCCL_BLOCKING_WAIT"] = "1"
+    distributed.init_process_group(
+        backend=backend, rank=rank, world_size=world_size, init_method="env://"
+    )
+
+    config["use_gpu"] = use_gpu
+    config["gpu_id"] = gpu_id
+    train_func(use_ray=False, config=config)
+
+    distributed.destroy_process_group()
+
+
+def train_torch_vanilla(
+    *,
+    config: dict,
+    num_workers: int = 4,
+    cpus_per_worker: int = 8,
+    use_gpu: bool = False,
+) -> Tuple[float, float, float]:
+    # This function is kicked off by the main() function and subsequently kicks
+    # off tasks that run train_torch_vanilla_worker() on the worker nodes.
+    from benchmark_util import (
+        upload_file_to_all_nodes,
+        create_actors_with_options,
+        run_commands_on_actors,
+        run_fn_on_actors,
+        get_ip_port_actors,
+        get_gpu_ids_actors,
+        map_ips_to_gpus,
+        set_cuda_visible_devices,
+    )
+
+    path = os.path.abspath(__file__)
+    upload_file_to_all_nodes(path)
+
+    num_epochs = config["epochs"]
+
+    actors = create_actors_with_options(
+        num_actors=num_workers,
+        resources={
+            "CPU": cpus_per_worker,
+            "GPU": int(use_gpu),
+        },
+    )
+
+    run_fn_on_actors(actors=actors, fn=lambda: os.environ.pop("OMP_NUM_THREADS", None))
+
+    # Get IPs and ports for all actors
+    ip_ports = get_ip_port_actors(actors=actors)
+
+    # Rank 0 is the master addr/port
+    master_addr, master_port = ip_ports[0]
+
+    if use_gpu:
+        # Extract IPs
+        actor_ips = [ipp[0] for ipp in ip_ports]
+
+        # Get allocated GPU IDs for all actors
+        gpu_ids = get_gpu_ids_actors(actors=actors)
+
+        # Build a map of IP to all allocated GPUs on that machine
+        ip_to_gpu_map = map_ips_to_gpus(ips=actor_ips, gpus=gpu_ids)
+
+        # Set the environment variables on the workers
+        set_cuda_visible_devices(
+            actors=actors, actor_ips=actor_ips, ip_to_gpus=ip_to_gpu_map
+        )
+
+        use_gpu_ids = [gi[0] for gi in gpu_ids]
+    else:
+        use_gpu_ids = [0] * num_workers
+
+    cmds = [
+        [
+            "python",
+            path,
+            "worker",
+            "--num-epochs",
+            str(num_epochs),
+            "--num-workers",
+            str(num_workers),
+            "--rank",
+            str(rank),
+            "--master-addr",
+            master_addr,
+            "--master-port",
+            str(master_port),
+            "--batch-size",
+            str(config["batch_size"]),
+        ]
+        + (["--use-gpu"] if use_gpu else [])
+        + (["--gpu-id", str(use_gpu_ids[rank])] if use_gpu else [])
+        for rank in range(num_workers)
+    ]
+
+    run_fn_on_actors(
+        actors=actors, fn=lambda: os.environ.setdefault("OMP_NUM_THREADS", "1")
+    )
+
+    start_time = time.monotonic()
+    run_commands_on_actors(actors=actors, cmds=cmds)
+    time_taken = time.monotonic() - start_time
+
+    loss = 0.0
+    if os.path.exists(VANILLA_RESULT_JSON):
+        with open(VANILLA_RESULT_JSON, "r") as f:
+            result = json.load(f)
+        loss = result["loss"]
+        local_time_taken = result["local_time_taken"]
+
+    return time_taken, local_time_taken, loss
+
+
+@click.group(help="Run Torch benchmarks")
+def cli():
+    pass
+
+
+@cli.command(help="Kick off Ray and vanilla benchmarks")
+@click.option("--num-runs", type=int, default=1)
+@click.option("--num-epochs", type=int, default=4)
+@click.option("--num-workers", type=int, default=4)
+@click.option("--cpus-per-worker", type=int, default=8)
+@click.option("--use-gpu", is_flag=True, default=False)
+@click.option("--batch-size", type=int, default=64)
+@click.option("--smoke-test", is_flag=True, default=False)
+@click.option("--local", is_flag=True, default=False)
+def run(
+    num_runs: int = 1,
+    num_epochs: int = 4,
+    num_workers: int = 4,
+    cpus_per_worker: int = 8,
+    use_gpu: bool = False,
+    batch_size: int = 64,
+    smoke_test: bool = False,
+    local: bool = False,
+):
+    # Note: smoke_test is ignored as we just adjust the batch size.
+    # The parameter is passed by the release test pipeline.
+    import ray
+    from benchmark_util import upload_file_to_all_nodes, run_command_on_all_nodes
+
+    config = CONFIG.copy()
+    config["epochs"] = num_epochs
+    config["batch_size"] = batch_size
+
+    if local:
+        ray.init(num_cpus=4)
+    else:
+        ray.init("auto")
+
+    print("Preparing Torch benchmark: Downloading MNIST")
+
+    path = str((Path(__file__).parent / "_torch_prepare.py").absolute())
+    upload_file_to_all_nodes(path)
+    run_command_on_all_nodes(["python", path])
+
+    times_ray = []
+    times_local_ray = []
+    losses_ray = []
+    times_vanilla = []
+    times_local_vanilla = []
+    losses_vanilla = []
+    for run in range(1, num_runs + 1):
+        time.sleep(2)
+
+        print(f"[Run {run}/{num_runs}] Running Torch Ray benchmark")
+
+        time_ray, time_local_ray, loss_ray = train_torch_ray_air(
+            num_workers=num_workers,
+            cpus_per_worker=cpus_per_worker,
+            use_gpu=use_gpu,
+            config=config,
+        )
+
+        print(
+            f"[Run {run}/{num_runs}] Finished Ray training ({num_epochs} epochs) in "
+            f"{time_ray:.2f} seconds (local training time: {time_local_ray:.2f}s). "
+            f"Observed loss = {loss_ray:.4f}"
+        )
+
+        time.sleep(2)
+
+        print(f"[Run {run}/{num_runs}] Running Torch vanilla benchmark")
+
+        time_vanilla, time_local_vanilla, loss_vanilla = train_torch_vanilla(
+            num_workers=num_workers,
+            cpus_per_worker=cpus_per_worker,
+            use_gpu=use_gpu,
+            config=config,
+        )
+
+        print(
+            f"[Run {run}/{num_runs}] Finished vanilla training ({num_epochs} epochs) "
+            f"in {time_vanilla:.2f} seconds "
+            f"(local training time: {time_local_vanilla:.2f}s). "
+            f"Observed loss = {loss_vanilla:.4f}"
+        )
+
+        print(
+            f"[Run {run}/{num_runs}] Observed results: ",
+            {
+                "tensorflow_mnist_ray_time_s": time_ray,
+                "tensorflow_mnist_ray_local_time_s": time_local_ray,
+                "tensorflow_mnist_ray_loss": loss_ray,
+                "tensorflow_mnist_vanilla_time_s": time_vanilla,
+                "tensorflow_mnist_vanilla_local_time_s": time_local_vanilla,
+                "tensorflow_mnist_vanilla_loss": loss_vanilla,
+            },
+        )
+
+        times_ray.append(time_ray)
+        times_local_ray.append(time_local_ray)
+        losses_ray.append(loss_ray)
+        times_vanilla.append(time_vanilla)
+        times_local_vanilla.append(time_local_vanilla)
+        losses_vanilla.append(loss_vanilla)
+
+    times_ray_mean = np.mean(times_ray)
+    times_ray_sd = np.std(times_ray)
+
+    times_local_ray_mean = np.mean(times_local_ray)
+    times_local_ray_sd = np.std(times_local_ray)
+
+    times_vanilla_mean = np.mean(times_vanilla)
+    times_vanilla_sd = np.std(times_vanilla)
+
+    times_local_vanilla_mean = np.mean(times_local_vanilla)
+    times_local_vanilla_sd = np.std(times_local_vanilla)
+
+    result = {
+        "torch_mnist_ray_num_runs": num_runs,
+        "torch_mnist_ray_time_s_all": times_ray,
+        "torch_mnist_ray_time_s_mean": times_ray_mean,
+        "torch_mnist_ray_time_s_sd": times_ray_sd,
+        "torch_mnist_ray_time_local_s_all": times_local_ray,
+        "torch_mnist_ray_time_local_s_mean": times_local_ray_mean,
+        "torch_mnist_ray_time_local_s_sd": times_local_ray_sd,
+        "torch_mnist_ray_loss_mean": np.mean(losses_ray),
+        "torch_mnist_ray_loss_sd": np.std(losses_ray),
+        "torch_mnist_vanilla_time_s_all": times_vanilla,
+        "torch_mnist_vanilla_time_s_mean": times_vanilla_mean,
+        "torch_mnist_vanilla_time_s_sd": times_vanilla_sd,
+        "torch_mnist_vanilla_local_time_s_all": times_local_vanilla,
+        "torch_mnist_vanilla_local_time_s_mean": times_local_vanilla_mean,
+        "torch_mnist_vanilla_local_time_s_sd": times_local_vanilla_sd,
+        "torch_mnist_vanilla_loss_mean": np.mean(losses_vanilla),
+        "torch_mnist_vanilla_loss_std": np.std(losses_vanilla),
+    }
+
+    print("Results:", result)
+    test_output_json = os.environ.get("TEST_OUTPUT_JSON", "/tmp/result.json")
+    with open(test_output_json, "wt") as f:
+        json.dump(result, f)
+
+    target_ratio = 1.15
+    ratio = (
+        (times_local_ray_mean / times_local_vanilla_mean)
+        if times_local_vanilla_mean != 0.0
+        else 1.0
+    )
+    if ratio > target_ratio:
+        raise RuntimeError(
+            f"Training on Ray took an average of {times_local_ray_mean:.2f} seconds, "
+            f"which is more than {target_ratio:.2f}x of the average vanilla training "
+            f"time of {times_local_vanilla_mean:.2f} seconds ({ratio:.2f}x). FAILED"
+        )
+
+    print(
+        f"Training on Ray took an average of {times_local_ray_mean:.2f} seconds, "
+        f"which is less than {target_ratio:.2f}x of the average vanilla training "
+        f"time of {times_local_vanilla_mean:.2f} seconds ({ratio:.2f}x). PASSED"
+    )
+
+
+@cli.command(help="Run PyTorch vanilla worker")
+@click.option("--num-epochs", type=int, default=4)
+@click.option("--num-workers", type=int, default=4)
+@click.option("--rank", type=int, default=0)
+@click.option("--master-addr", type=str, default="")
+@click.option("--master-port", type=int, default=0)
+@click.option("--batch-size", type=int, default=64)
+@click.option("--use-gpu", is_flag=True, default=False)
+@click.option("--gpu-id", type=int, default=0)
+def worker(
+    num_epochs: int = 4,
+    num_workers: int = 4,
+    rank: int = 0,
+    master_addr: str = "",
+    master_port: int = 0,
+    batch_size: int = 64,
+    use_gpu: bool = False,
+    gpu_id: int = 0,
+):
+    config = CONFIG.copy()
+    config["epochs"] = num_epochs
+    config["batch_size"] = batch_size
+
+    # Then we kick off the training function on every worker.
+    return train_torch_vanilla_worker(
+        config=config,
+        rank=rank,
+        world_size=num_workers,
+        master_addr=master_addr,
+        master_port=master_port,
+        use_gpu=use_gpu,
+        gpu_id=gpu_id,
+    )
+
+
+def main():
+    return cli()
+
+
+if __name__ == "__main__":
+    main()

release/air_tests/air_benchmarks/workloads/tune_torch_benchmark.py

@@ -0,0 +1,232 @@
+import json
+import os
+import time
+from typing import Optional, Dict, List
+
+import click
+import numpy as np
+
+import ray
+from ray.train import ScalingConfig, RunConfig
+from ray.train.torch import TorchTrainer
+from ray.tune.integration.ray_train import TuneReportCallback
+
+
+CONFIG = {"lr": 1e-3, "batch_size": 64, "epochs": 20}
+
+
+def prepare_mnist():
+    # Pre-download the data onto each node.
+    from benchmark_util import schedule_remote_fn_on_all_nodes
+
+    print("Preparing Torch benchmark: Downloading MNIST")
+
+    @ray.remote(num_cpus=0)
+    def _download_data():
+        import torchvision
+
+        torchvision.datasets.FashionMNIST("/tmp/data_fashion_mnist", download=True)
+        return True
+
+    ray.get(schedule_remote_fn_on_all_nodes(_download_data))
+
+
+def train_loop(config: Dict):
+    from torch_benchmark import train_func
+
+    train_func(use_ray=True, config=config)
+
+
+def get_trainer(
+    num_workers: int = 4,
+    use_gpu: bool = False,
+    config: Optional[Dict] = None,
+):
+    """Get the trainer to be used across train and tune to ensure consistency."""
+    # We are using STRICT_PACK here to do an apples to apples comparison.
+    # PyTorch defaults to using multithreading, so if the workers are spread,
+    # they are able to utilize more resources. We would effectively be comparing
+    # X tune runs with 2 CPUs per worker vs. 1 tune run with up to 8 CPUs per
+    # worker. Using STRICT_PACK avoids this by forcing all workers to be
+    # co-located.
+    config = config or CONFIG
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(
+            num_workers=num_workers,
+            resources_per_worker={"CPU": 2},
+            use_gpu=use_gpu,
+            placement_strategy="STRICT_PACK",
+        ),
+        run_config=RunConfig(
+            name="train_torch_benchmark",
+            storage_path="/mnt/cluster_storage/ray-train-results",
+        ),
+    )
+    return trainer
+
+
+def train_torch(
+    num_workers: int, use_gpu: bool = False, config: Optional[Dict] = None
+) -> float:
+    trainer = get_trainer(num_workers=num_workers, use_gpu=use_gpu, config=config)
+    result = trainer.fit()
+    return result.metrics["local_time_taken"]
+
+
+def train_driver_fn(config: Dict):
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop,
+        train_loop_config=config["train_loop_config"],
+        run_config=RunConfig(
+            name="tune_torch_benchmark",
+            storage_path="/mnt/cluster_storage/ray-tune-results",
+            callbacks=[TuneReportCallback()],
+        ),
+        scaling_config=ScalingConfig(
+            num_workers=config["num_workers"],
+            resources_per_worker={"CPU": 2},
+            use_gpu=config["use_gpu"],
+            placement_strategy="STRICT_PACK",
+        ),
+    )
+
+    trainer.fit()
+
+
+def tune_torch(
+    num_workers: int = 4,
+    num_trials: int = 8,
+    use_gpu: bool = False,
+    config: Optional[Dict] = None,
+) -> List[float]:
+    """Making sure that tuning multiple trials in parallel is not
+    taking significantly longer than training each one individually.
+
+    Some overhead is expected.
+    """
+
+    from ray import tune
+    from ray.tune.tuner import Tuner
+    from ray.tune.tune_config import TuneConfig
+
+    param_space = {
+        "train_loop_config": {
+            "lr": tune.loguniform(1e-4, 1e-1),
+        },
+        "num_workers": num_workers,
+        "use_gpu": use_gpu,
+    }
+
+    param_space["train_loop_config"].update(config or {})
+
+    tuner = Tuner(
+        trainable=train_driver_fn,
+        param_space=param_space,
+        tune_config=TuneConfig(mode="min", metric="loss", num_samples=num_trials),
+    )
+    results = tuner.fit()
+    return [result.metrics["local_time_taken"] for result in results]
+
+
+@click.command(help="Run Benchmark comparing Train to Tune.")
+@click.option("--num-runs", type=int, default=1)
+@click.option("--num-trials", type=int, default=8)
+@click.option("--num-workers", type=int, default=4)
+@click.option("--use-gpu", is_flag=True)
+@click.option("--smoke-test", is_flag=True, default=False)
+def main(
+    num_runs: int = 1,
+    num_trials: int = 8,
+    num_workers: int = 4,
+    use_gpu: bool = False,
+    smoke_test: bool = False,
+):
+    ray.init(
+        runtime_env={
+            "working_dir": os.path.dirname(__file__),
+        }
+    )
+    prepare_mnist()
+
+    config = CONFIG.copy()
+
+    if smoke_test:
+        config["epochs"] = 1
+
+    train_times = []
+    tune_times = []
+
+    train_computes = []
+    tune_trial_computes = []
+
+    for i in range(num_runs):
+        print(f"Run {i+1} / {num_runs}")
+
+        time.sleep(2)
+
+        train_start = time.monotonic()
+        train_compute = train_torch(
+            num_workers=num_workers, use_gpu=use_gpu, config=config
+        )
+        train_time = time.monotonic() - train_start
+        train_times.append(train_time)
+        train_computes.append(train_compute)
+
+        time.sleep(2)
+
+        tune_start = time.monotonic()
+        trial_computes = tune_torch(
+            num_workers=num_workers,
+            num_trials=num_trials,
+            use_gpu=use_gpu,
+            config=config,
+        )
+        tune_time = time.monotonic() - tune_start
+        tune_times.append(tune_time)
+        tune_trial_computes.append(trial_computes)
+
+        result = {
+            "train_time": train_time,
+            "train_compute": train_compute,
+            "train_overhead": train_time - train_compute,
+            "tune_time": tune_time,
+            "tune_overhead": tune_time - max(trial_computes),
+        }
+        print(f"Results run {i+1}: {result}")
+
+    mean_train_time = np.mean(train_times)
+    mean_tune_time = np.mean(tune_times)
+
+    full_results = {
+        "train_times": train_times,
+        "train_mean": mean_train_time,
+        "train_sd": np.std(train_times),
+        "tune_times": tune_times,
+        "tune_mean": mean_tune_time,
+        "tune_sd": np.std(tune_times),
+        "train_computes": train_computes,
+        "tune_trial_computes": tune_trial_computes,
+    }
+
+    print("Full results:", full_results)
+
+    # NOTE: The value of `factor` is mostly arbitrary. It was previously `1.2`, but
+    # that value turned out to be too low. For more context, see #29682.
+    factor = 1.35
+    threshold = mean_train_time * factor
+
+    test_output_json = os.environ.get("TEST_OUTPUT_JSON", "/tmp/result.json")
+    with open(test_output_json, "wt") as f:
+        json.dump(full_results, f)
+
+    assert (
+        mean_tune_time <= threshold
+    ), f"{mean_tune_time:.2f} > {threshold:.2f} = {factor:.1f} * {mean_train_time:.2f}"
+
+
+if __name__ == "__main__":
+    main()

release/autoscaling_tests/aws.yaml

@@ -0,0 +1,17 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west-2
+
+max_workers: 20
+
+head_node_type:
+  name: head_node
+  instance_type: m5.16xlarge
+  resources:
+    cpu: 0
+
+
+worker_node_types:
+  - name: worker_node
+    instance_type: m5.xlarge
+    min_workers: 0
+    max_workers: 20

release/autoscaling_tests/logger.py

@@ -0,0 +1,18 @@
+import logging
+from rich.logging import RichHandler
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+
+def add_handlers(logger: logging.Logger):
+    handler = RichHandler()
+    formatter = logging.Formatter(
+        fmt="[%(levelname)s %(asctime)s] %(filename)s:%(lineno)d - %(message)s"
+    )
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+
+
+if not logger.hasHandlers():
+    add_handlers(logger)

release/autoscaling_tests/run.py

@@ -0,0 +1,82 @@
+import subprocess
+import click
+import json
+import os
+import time
+
+from logger import logger
+
+WORKLOAD_SCRIPTS = [
+    "test_core.py",
+]
+
+
+def setup_cluster():
+    from ray.cluster_utils import AutoscalingCluster
+
+    cluster = AutoscalingCluster(
+        head_resources={"CPU": 0},
+        worker_node_types={
+            "type-1": {
+                "resources": {"CPU": 4},
+                "node_config": {},
+                "min_workers": 0,
+                "max_workers": 10,
+            },
+        },
+        idle_timeout_minutes=1 * 0.1,
+    )
+
+    cluster.start(_system_config={"enable_autoscaler_v2": True})
+    return cluster
+
+
+def run_test():
+    failed_workloads = []
+    for workload in WORKLOAD_SCRIPTS:
+        # Run the python script.
+        logger.info(f"Running workload {workload}:")
+        try:
+            subprocess.check_call(["python", workload])
+        except subprocess.CalledProcessError as e:
+            failed_workloads.append((workload, e))
+
+    if failed_workloads:
+        for workload, e in failed_workloads:
+            logger.error(f"Workload {workload} failed with {e}")
+        raise RuntimeError(f"{len(failed_workloads)} workloads failed.")
+    else:
+        logger.info("All workloads passed!")
+
+
+@click.command()
+@click.option("--local", is_flag=True, help="Run locally.", default=False)
+def run(local):
+    start_time = time.time()
+    cluster = None
+    try:
+        if local:
+            cluster = setup_cluster()
+            run_test()
+            cluster.shutdown()
+        else:
+            run_test()
+    except Exception as e:
+        logger.error(f"Test failed with {e}")
+        raise e
+    finally:
+        if cluster:
+            cluster.shutdown()
+
+    results = {
+        "time": time.time() - start_time,
+    }
+    if "TEST_OUTPUT_JSON" in os.environ:
+        with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+            json.dump(results, out_file)
+
+    print(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    run()

release/autoscaling_tests/test_core.py

@@ -0,0 +1,131 @@
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray.autoscaler.v2.sdk import get_cluster_status
+import time
+from logger import logger
+from typing import Dict
+
+ray.init("auto")
+
+# Sync with the compute config.
+HEAD_NODE_CPU = 0
+WORKER_NODE_CPU = 4
+IDLE_TERMINATION_S = 60 * 5  # 5 min
+DEFAULT_RETRY_INTERVAL_MS = 15 * 1000  # 15 sec
+
+
+def check_cluster(target_num_nodes: int, target_resources: Dict[str, float]):
+    gcs_address = ray.get_runtime_context().gcs_address
+    cluster_status = get_cluster_status(gcs_address)
+
+    assert (
+        len(cluster_status.active_nodes) + len(cluster_status.idle_nodes)
+    ) == target_num_nodes
+
+    for k, v in target_resources.items():
+        assert cluster_status.total_resources().get(k, 0) == v
+
+    return True
+
+
+ctx = {
+    "num_cpus": 0,
+    "num_nodes": 1,
+}
+logger.info(f"Starting cluster with {ctx['num_nodes']} nodes, {ctx['num_cpus']} cpus")
+check_cluster(
+    target_num_nodes=ctx["num_nodes"], target_resources={"CPU": ctx["num_cpus"]}
+)
+
+
+# Request for cluster resources
+def test_request_cluster_resources(ctx: dict):
+    from ray.autoscaler._private.commands import request_resources
+
+    request_resources(num_cpus=8)
+
+    ctx["num_cpus"] += 8
+    ctx["num_nodes"] += 8 // WORKER_NODE_CPU
+
+    # Assert on number of worker nodes.
+    logger.info(
+        f"Requesting cluster constraints: {ctx['num_nodes']} nodes, "
+        f"{ctx['num_cpus']} cpus"
+    )
+    wait_for_condition(
+        check_cluster,
+        timeout=60 * 5,  # 5min
+        retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
+        target_num_nodes=ctx["num_nodes"],
+        target_resources={"CPU": ctx["num_cpus"]},
+    )
+
+    # Reset the cluster constraints.
+    request_resources(num_cpus=0)
+
+    ctx["num_cpus"] -= 8
+    ctx["num_nodes"] -= 8 // WORKER_NODE_CPU
+    logger.info(
+        f"Waiting for cluster go idle after constraint cleared: {ctx['num_nodes']} "
+        f"nodes, {ctx['num_cpus']} cpus"
+    )
+    wait_for_condition(
+        check_cluster,
+        timeout=60 + IDLE_TERMINATION_S,  # 1min + idle timeout
+        retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
+        target_num_nodes=ctx["num_nodes"],
+        target_resources={"CPU": ctx["num_cpus"]},
+    )
+
+
+# Run actors/tasks that exceed the cluster resources should upscale the cluster
+def test_run_tasks_concurrent(ctx: dict):
+    num_tasks = 2
+    num_actors = 2
+
+    @ray.remote(num_cpus=WORKER_NODE_CPU)
+    def f():
+        while True:
+            time.sleep(1)
+
+    @ray.remote(num_cpus=WORKER_NODE_CPU)
+    class Actor:
+        def __init__(self):
+            pass
+
+    tasks = [f.remote() for _ in range(num_tasks)]
+    actors = [Actor.remote() for _ in range(num_actors)]
+
+    ctx["num_cpus"] += (num_tasks + num_actors) * WORKER_NODE_CPU
+    ctx["num_nodes"] += num_tasks + num_actors
+
+    logger.info(f"Waiting for {ctx['num_nodes']} nodes, {ctx['num_cpus']} cpus")
+    wait_for_condition(
+        check_cluster,
+        timeout=60 * 5,  # 5min
+        retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
+        target_num_nodes=ctx["num_nodes"],
+        target_resources={"CPU": ctx["num_cpus"]},
+    )
+
+    [ray.cancel(task) for task in tasks]
+    [ray.kill(actor) for actor in actors]
+
+    ctx["num_cpus"] -= (num_actors + num_tasks) * WORKER_NODE_CPU
+    ctx["num_nodes"] -= num_actors + num_tasks
+
+    logger.info(
+        f"Waiting for cluster to scale down to {ctx['num_nodes']} nodes, "
+        f"{ctx['num_cpus']} cpus"
+    )
+    wait_for_condition(
+        check_cluster,
+        timeout=60 + IDLE_TERMINATION_S,
+        retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
+        target_num_nodes=ctx["num_nodes"],
+        target_resources={"CPU": ctx["num_cpus"]},
+    )
+
+
+test_request_cluster_resources(ctx)
+test_run_tasks_concurrent(ctx)

release/aws2gce_iam.json

@@ -0,0 +1,13 @@
+{
+  "type": "external_account",
+  "audience": "//iam.googleapis.com/projects/498773744730/locations/global/workloadIdentityPools/anyscale-ci-aws/providers/aws-anyscale",
+  "subject_token_type": "urn:ietf:params:aws:token-type:aws4_request",
+  "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/498773744730-compute@developer.gserviceaccount.com:generateAccessToken",
+  "token_url": "https://sts.googleapis.com/v1/token",
+  "credential_source": {
+    "environment_id": "aws1",
+    "region_url": "http://169.254.169.254/latest/meta-data/placement/availability-zone",
+    "url": "http://169.254.169.254/latest/meta-data/iam/security-credentials",
+    "regional_cred_verification_url": "https://sts.{region}.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15"
+  }
+}

release/azure_docker_login.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+# This script is used to login to azure docker registry using azure cli
+
+set -euo pipefail
+
+# Retrieve credentials from Secrets Manager
+SECRET="$(aws secretsmanager get-secret-value \
+  --secret-id azure-service-principal-oss-release \
+  --query SecretString \
+  --region us-west-2 \
+  --output text)"
+
+CLIENT_ID="$(echo "$SECRET" | jq -r '.client_id')"
+TENANT_ID="$(echo "$SECRET" | jq -r '.tenant_id')"
+
+temp_dir=$(mktemp -d)
+
+aws secretsmanager get-secret-value \
+--secret-id azure-service-principal-certificate \
+--query SecretString \
+--region us-west-2 \
+--output text > "${temp_dir}/azure_cert.pem"
+
+# Login to azure
+az login --service-principal --username "$CLIENT_ID" --certificate "${temp_dir}/azure_cert.pem" --tenant "$TENANT_ID"

release/benchmark-worker-startup/app_config_gpu.yaml

@@ -0,0 +1,10 @@
+base_image: {{ env["RAY_IMAGE_ML_NIGHTLY_GPU"] | default("anyscale/ray-ml:nightly-py39-gpu") }}
+env_vars: {}
+
+python:
+  pip_packages: []
+  conda_packages: []
+
+post_build_cmds:
+  - pip uninstall -y ray || true && pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
+  - {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}

release/benchmark-worker-startup/benchmark_worker_startup.py

@@ -0,0 +1,365 @@
+#!/usr/bin/env python3
+
+"""
+$ ./benchmark_worker_startup.py --help
+usage: benchmark_worker_startup.py [-h] --num_gpus_in_cluster
+                                   NUM_GPUS_IN_CLUSTER
+                                   --num_cpus_in_cluster
+                                   NUM_CPUS_IN_CLUSTER
+                                   --num_tasks_or_actors_per_run
+                                   NUM_TASKS_OR_ACTORS_PER_RUN
+                                   --num_measurements_per_configuration
+                                   NUM_MEASUREMENTS_PER_CONFIGURATION
+
+This release test measures Ray worker startup time. Specifically, it
+measures the time to start N different tasks or actors, where each task or
+actor imports a large library (currently PyTorch). N is configurable. The
+test runs under a few different configurations: {task, actor} x {runtime
+env, no runtime env} x {GPU, no GPU} x {cold start, warm start} x {import
+torch, no imports}.
+
+options:
+  -h, --help            show this help message and exit
+  --num_gpus_in_cluster NUM_GPUS_IN_CLUSTER
+                        The number of GPUs in the cluster. This determines
+                        how many GPU resources each actor/task requests.
+  --num_cpus_in_cluster NUM_CPUS_IN_CLUSTER
+                        The number of CPUs in the cluster. This determines
+                        how many CPU resources each actor/task requests.
+  --num_tasks_or_actors_per_run NUM_TASKS_OR_ACTORS_PER_RUN
+                        The number of tasks or actors per 'run'. A run
+                        starts this many tasks/actors and consitutes a
+                        single measurement. Several runs can be composed
+                        within a single job for measure warm start, or
+                        spread across different jobs to measure cold start.
+  --num_measurements_per_configuration NUM_MEASUREMENTS_PER_CONFIGURATION
+                        The number of measurements to record per
+                        configuration.
+
+This script uses test_single_configuration.py to run the actual
+measurements.
+"""
+
+import argparse
+import asyncio
+import random
+import statistics
+import subprocess
+import sys
+from collections import defaultdict
+from dataclasses import dataclass
+
+import ray
+from ray._private.test_utils import safe_write_to_results_json
+from ray.job_submission import JobStatus, JobSubmissionClient
+
+
+def main(
+    num_cpus_in_cluster: int,
+    num_gpus_in_cluster: int,
+    num_tasks_or_actors_per_run: int,
+    num_measurements_per_configuration: int,
+):
+    """
+    Generate test cases, then run them in random order via run_and_stream_logs.
+    """
+    metrics_actor_name = "metrics_actor"
+    metrics_actor_namespace = "metrics_actor_namespace"
+    metrics_actor = MetricsActor.options(  # noqa: F841
+        name=metrics_actor_name,
+        namespace=metrics_actor_namespace,
+    ).remote(
+        expected_measurements_per_test=num_measurements_per_configuration,
+    )
+
+    print_disk_config()
+
+    run_matrix = generate_test_matrix(
+        num_cpus_in_cluster,
+        num_gpus_in_cluster,
+        num_tasks_or_actors_per_run,
+        num_measurements_per_configuration,
+    )
+    print(f"List of tests: {run_matrix}")
+
+    for test in random.sample(list(run_matrix), k=len(run_matrix)):
+        print(f"Running test {test}")
+        asyncio.run(
+            run_and_stream_logs(
+                metrics_actor_name,
+                metrics_actor_namespace,
+                test,
+            )
+        )
+
+
+@ray.remote(num_cpus=0)
+class MetricsActor:
+    """
+    Actor which tests will report metrics to.
+    """
+
+    def __init__(self, expected_measurements_per_test: int):
+        self.measurements = defaultdict(list)
+        self.expected_measurements_per_test = expected_measurements_per_test
+
+    def submit(self, test_name: str, latency: float):
+        print(f"got latency {latency} s for test {test_name}")
+        self.measurements[test_name].append(latency)
+        results = self.create_results_dict_from_measurements(
+            self.measurements, self.expected_measurements_per_test
+        )
+        safe_write_to_results_json(results)
+
+        assert (
+            len(self.measurements[test_name]) <= self.expected_measurements_per_test
+        ), (
+            f"Expected {self.measurements[test_name]} to not have more elements than "
+            f"{self.expected_measurements_per_test}"
+        )
+
+    @staticmethod
+    def create_results_dict_from_measurements(
+        all_measurements, expected_measurements_per_test
+    ):
+        results = {}
+        perf_metrics = []
+
+        for test_name, measurements in all_measurements.items():
+            test_summary = {
+                "measurements": measurements,
+            }
+
+            if len(measurements) == expected_measurements_per_test:
+                median = statistics.median(measurements)
+                test_summary["p50"] = median
+                perf_metrics.append(
+                    {
+                        "perf_metric_name": f"p50.{test_name}",
+                        "perf_metric_value": median,
+                        "perf_metric_type": "LATENCY",
+                    }
+                )
+
+            results[test_name] = test_summary
+
+        results["perf_metrics"] = perf_metrics
+        return results
+
+
+def print_disk_config():
+    print("Getting disk sizes via df -h")
+    subprocess.check_call("df -h", shell=True)
+
+
+def generate_test_matrix(
+    num_cpus_in_cluster: int,
+    num_gpus_in_cluster: int,
+    num_tasks_or_actors_per_run: int,
+    num_measurements_per_test: int,
+):
+
+    num_repeated_jobs_or_runs = num_measurements_per_test
+    total_num_tasks_or_actors = num_tasks_or_actors_per_run * num_repeated_jobs_or_runs
+
+    num_jobs_per_type = {
+        "cold_start": num_repeated_jobs_or_runs,
+        "warm_start": 1,
+    }
+
+    imports_to_try = ["torch", "none"]
+
+    tests = set()
+
+    for with_tasks in [True, False]:
+        for with_gpu in [True, False]:
+            # Do not run without runtime env. TODO(cade) Infra team added cgroups to
+            # default runtime env, need to find some way around that if we want
+            # "pure" (non-runtime-env) measurements.
+            for with_runtime_env in [True]:
+                for import_to_try in imports_to_try:
+                    for num_jobs in num_jobs_per_type.values():
+
+                        num_tasks_or_actors_per_job = (
+                            total_num_tasks_or_actors // num_jobs
+                        )
+                        num_runs_per_job = (
+                            num_tasks_or_actors_per_job // num_tasks_or_actors_per_run
+                        )
+
+                        test = TestConfiguration(
+                            num_jobs=num_jobs,
+                            num_runs_per_job=num_runs_per_job,
+                            num_tasks_or_actors_per_run=num_tasks_or_actors_per_run,
+                            with_tasks=with_tasks,
+                            with_gpu=with_gpu,
+                            with_runtime_env=with_runtime_env,
+                            import_to_try=import_to_try,
+                            num_cpus_in_cluster=num_cpus_in_cluster,
+                            num_gpus_in_cluster=num_gpus_in_cluster,
+                            num_nodes_in_cluster=1,
+                        )
+                        tests.add(test)
+
+    return tests
+
+
+@dataclass(eq=True, frozen=True)
+class TestConfiguration:
+    num_jobs: int
+    num_runs_per_job: int
+    num_tasks_or_actors_per_run: int
+    with_gpu: bool
+    with_tasks: bool
+    with_runtime_env: bool
+    import_to_try: str
+    num_cpus_in_cluster: int
+    num_gpus_in_cluster: int
+    num_nodes_in_cluster: int
+
+    def __repr__(self):
+        with_gpu_str = "with_gpu" if self.with_gpu else "without_gpu"
+        executable_unit = "tasks" if self.with_tasks else "actors"
+        cold_or_warm_start = "cold" if self.num_jobs > 1 else "warm"
+        with_runtime_env_str = (
+            "with_runtime_env" if self.with_runtime_env else "without_runtime_env"
+        )
+        single_node_or_multi_node = (
+            "single_node" if self.num_nodes_in_cluster == 1 else "multi_node"
+        )
+        import_torch_or_none = (
+            "import_torch" if self.import_to_try == "torch" else "no_import"
+        )
+
+        return "-".join(
+            [
+                f"seconds_to_{cold_or_warm_start}_start_"
+                f"{self.num_tasks_or_actors_per_run}_{executable_unit}",
+                import_torch_or_none,
+                with_gpu_str,
+                single_node_or_multi_node,
+                with_runtime_env_str,
+                f"{self.num_cpus_in_cluster}_CPU_{self.num_gpus_in_cluster}"
+                "_GPU_cluster",
+            ]
+        )
+
+
+async def run_and_stream_logs(
+    metrics_actor_name, metrics_actor_namespace, test: TestConfiguration
+):
+    """
+    Run a particular test configuration by invoking ./test_single_configuration.py.
+    """
+    client = JobSubmissionClient("http://127.0.0.1:8265")
+    entrypoint = generate_entrypoint(metrics_actor_name, metrics_actor_namespace, test)
+
+    for _ in range(test.num_jobs):
+        print(f"Running {entrypoint}")
+
+        if not test.with_runtime_env:
+            # On non-workspaces, this will run as a job but without a runtime env.
+            subprocess.check_call(entrypoint, shell=True)
+        else:
+            job_id = client.submit_job(
+                entrypoint=entrypoint,
+                runtime_env={"working_dir": "./"},
+            )
+
+            try:
+                async for lines in client.tail_job_logs(job_id):
+                    print(lines, end="")
+            except KeyboardInterrupt:
+                print(f"Stopping job {job_id}")
+                client.stop_job(job_id)
+                raise
+
+            job_status = client.get_job_status(job_id)
+            if job_status != JobStatus.SUCCEEDED:
+                raise ValueError(
+                    f"Job {job_id} was not successful; status is {job_status}"
+                )
+
+
+def generate_entrypoint(
+    metrics_actor_name: str, metrics_actor_namespace: str, test: TestConfiguration
+):
+
+    task_or_actor_arg = "--with_tasks" if test.with_tasks else "--with_actors"
+    with_gpu_arg = "--with_gpu" if test.with_gpu else "--without_gpu"
+    with_runtime_env_arg = (
+        "--with_runtime_env" if test.with_runtime_env else "--without_runtime_env"
+    )
+    return " ".join(
+        [
+            "python ./test_single_configuration.py",
+            f"--metrics_actor_name {metrics_actor_name}",
+            f"--metrics_actor_namespace {metrics_actor_namespace}",
+            f"--test_name {test}",
+            f"--num_runs {test.num_runs_per_job} ",
+            f"--num_tasks_or_actors_per_run {test.num_tasks_or_actors_per_run}",
+            f"--num_cpus_in_cluster {test.num_cpus_in_cluster}",
+            f"--num_gpus_in_cluster {test.num_gpus_in_cluster}",
+            task_or_actor_arg,
+            with_gpu_arg,
+            with_runtime_env_arg,
+            f"--library_to_import {test.import_to_try}",
+        ]
+    )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="This release test measures Ray worker startup time. "
+        "Specifically, it measures the time to start N different tasks or"
+        " actors, where each task or actor imports a large library ("
+        "currently PyTorch). N is configurable.\nThe test runs under a "
+        "few different configurations: {task, actor} x {runtime env, "
+        "no runtime env} x {GPU, no GPU} x {cold start, warm start} x "
+        "{import torch, no imports}.",
+        epilog="This script uses test_single_configuration.py to run the "
+        "actual measurements.",
+    )
+    parser.add_argument(
+        "--num_gpus_in_cluster",
+        type=int,
+        required=True,
+        help="The number of GPUs in the cluster. This determines how many "
+        "GPU resources each actor/task requests.",
+    )
+    parser.add_argument(
+        "--num_cpus_in_cluster",
+        type=int,
+        required=True,
+        help="The number of CPUs in the cluster. This determines how many "
+        "CPU resources each actor/task requests.",
+    )
+    parser.add_argument(
+        "--num_tasks_or_actors_per_run",
+        type=int,
+        required=True,
+        help="The number of tasks or actors per 'run'. A run starts this "
+        "many tasks/actors and consitutes a single measurement. Several "
+        "runs can be composed within a single job for measure warm start, "
+        "or spread across different jobs to measure cold start.",
+    )
+    parser.add_argument(
+        "--num_measurements_per_configuration",
+        type=int,
+        required=True,
+        help="The number of measurements to record per configuration.",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    sys.exit(
+        main(
+            args.num_cpus_in_cluster,
+            args.num_gpus_in_cluster,
+            args.num_tasks_or_actors_per_run,
+            args.num_measurements_per_configuration,
+        )
+    )

release/benchmark-worker-startup/only_head_node_1gpu_64cpu.yaml

@@ -0,0 +1,21 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+advanced_instance_config:
+    # Fix the volume size so that IOPS is constant even if the default changes.
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+            DeleteOnTermination: true
+            # 150GB is the default in Anyscale.
+            VolumeSize: 150
+
+head_node:
+    instance_type: g5.16xlarge
+    # GPU: 64 is an intentional override (instance has 1 physical GPU);
+    # the benchmark script uses --num_gpus_in_cluster 64 to test
+    # resource scheduling with 64 logical GPU slots.
+    resources:
+        CPU: 64
+        GPU: 64
+
+worker_nodes: []

release/benchmark-worker-startup/only_head_node_1gpu_64cpu_gce.yaml

@@ -0,0 +1,21 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones: [us-west1-b]
+
+advanced_instance_config:
+    instance_properties:
+        disks:
+            - boot: true
+              auto_delete: true
+              initialize_params:
+                  disk_size_gb: 150
+
+head_node:
+    instance_type: n2-standard-64
+    # GPU: 64 is an intentional override (instance has no physical GPU);
+    # the benchmark script uses --num_gpus_in_cluster 64 to test
+    # resource scheduling with 64 logical GPU slots.
+    resources:
+        CPU: 64
+        GPU: 64
+
+worker_nodes: []

release/benchmark-worker-startup/test_single_configuration.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+"""
+Helper file for benchmark_worker_startup.py. This file runs a particular test
+configuration.
+"""
+
+import argparse
+import sys
+import time
+
+import ray
+
+
+@ray.remote
+class Actor:
+    def run_code(self, should_import_torch: bool):
+        if should_import_torch:
+            import torch  # noqa: F401
+
+
+@ray.remote
+def task(should_import_torch: bool):
+    if should_import_torch:
+        import torch  # noqa: F401
+
+
+def main(
+    metrics_actor,
+    test_name: str,
+    num_runs: int,
+    num_tasks_or_actors_per_run: int,
+    num_cpus_in_cluster: int,
+    num_gpus_in_cluster: int,
+    library_to_import: str,
+    use_actors: bool,
+    with_gpu: bool,
+    with_runtime_env: bool,
+):
+
+    num_gpus = (num_gpus_in_cluster / num_tasks_or_actors_per_run) if with_gpu else 0
+    num_cpus = num_cpus_in_cluster / num_tasks_or_actors_per_run
+
+    print(f"Assigning each task/actor {num_cpus} num_cpus and {num_gpus} num_gpus")
+    actor_with_resources = Actor.options(num_gpus=num_gpus, num_cpus=num_cpus)
+    task_with_resources = task.options(num_gpus=num_gpus, num_cpus=num_cpus)
+
+    should_import_torch = library_to_import == "torch"
+    print(f"should_import_torch: {should_import_torch}")
+
+    fail_if_incorrect_runtime_env(expect_runtime_env=with_runtime_env)
+
+    def with_actors():
+        actors = [
+            actor_with_resources.remote() for _ in range(num_tasks_or_actors_per_run)
+        ]
+        ray.get([actor.run_code.remote(should_import_torch) for actor in actors])
+
+    def with_tasks():
+        ray.get(
+            [
+                task_with_resources.remote(should_import_torch)
+                for _ in range(num_tasks_or_actors_per_run)
+            ]
+        )
+
+    func_to_measure = with_actors if use_actors else with_tasks
+
+    for run in range(num_runs):
+        print(f"Starting measurement for run {run}")
+        start = time.time()
+        func_to_measure()
+        dur_s = time.time() - start
+        ray.get(metrics_actor.submit.remote(test_name, dur_s))
+
+
+def fail_if_incorrect_runtime_env(expect_runtime_env: bool):
+    ctx = ray.runtime_context.get_runtime_context()
+    print(f"Found runtime_env={ctx.runtime_env}")
+    if expect_runtime_env and ctx.runtime_env == {}:
+        raise AssertionError(
+            f"Expected a runtime environment but found runtime_env={ctx.runtime_env}"
+        )
+
+    if not expect_runtime_env and ctx.runtime_env != {}:
+        raise AssertionError(
+            f"Expected no runtime environment but found runtime_env={ctx.runtime_env}"
+        )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--metrics_actor_name", type=str, required=True)
+    parser.add_argument("--metrics_actor_namespace", type=str, required=True)
+    parser.add_argument("--test_name", type=str, required=True)
+    parser.add_argument("--num_runs", type=int, required=True)
+    parser.add_argument("--num_tasks_or_actors_per_run", type=int, required=True)
+    parser.add_argument("--num_cpus_in_cluster", type=int, required=True)
+    parser.add_argument("--num_gpus_in_cluster", type=int, required=True)
+    parser.add_argument(
+        "--library_to_import", type=str, required=True, choices=["torch", "none"]
+    )
+
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--with_actors", action="store_true")
+    group.add_argument("--with_tasks", action="store_true")
+
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--with_gpu", action="store_true")
+    group.add_argument("--without_gpu", action="store_true")
+
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--with_runtime_env", action="store_true")
+    group.add_argument("--without_runtime_env", action="store_true")
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    metrics_actor = ray.get_actor(
+        args.metrics_actor_name,
+        args.metrics_actor_namespace,
+    )
+
+    sys.exit(
+        main(
+            metrics_actor,
+            args.test_name,
+            args.num_runs,
+            args.num_tasks_or_actors_per_run,
+            args.num_cpus_in_cluster,
+            args.num_gpus_in_cluster,
+            args.library_to_import,
+            args.with_actors,
+            args.with_gpu,
+            args.with_runtime_env,
+        )
+    )

release/benchmarks/README.md

@@ -0,0 +1,33 @@
+# Ray Scalability Envelope
+
+**NOTE**: the Ray scalability benchmarks are in the process of being refreshed. If you have questions about a specific workload or limit, please get in touch by filing a [GitHub issue](https://github.com/ray-project/ray/issues).
+
+## Distributed Benchmarks
+
+All distributed tests are run on 64 nodes with 64 cores/node. Maximum number of nodes is achieved by adding 4 core nodes.
+
+| Dimension                                       | Quantity |
+| ---------                                       | -------- |
+| # nodes in cluster (with trivial task workload) | 2k+      |
+| # actors in cluster (with trivial workload)     | 40k+     |
+| # simultaneously running tasks                  | 10k+     |
+| # simultaneously running placement groups       | 1k+      |
+
+## Object Store Benchmarks
+
+| Dimension                           | Quantity |
+| ---------                           | -------- |
+| 1 GiB object broadcast (# of nodes) | 50+      |
+
+
+## Single Node Benchmarks.
+
+All single node benchmarks are run on a single m4.16xlarge.
+
+| Dimension                                      | Quantity   |
+| ---------                                      | --------   |
+| # of object arguments to a single task         | 10000+     |
+| # of objects returned from a single task       | 3000+      |
+| # of plasma objects in a single `ray.get` call | 10000+     |
+| # of tasks queued on a single node             | 1,000,000+ |
+| Maximum `ray.get` numpy object size            | 100GiB+    |

release/benchmarks/distributed.yaml

@@ -0,0 +1,22 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: r5dn.16xlarge # Network optimized.
+    resources:
+      CPU: 0
+      node: 1
+      small: 1
+
+worker_nodes:
+    - instance_type: m5.16xlarge
+      min_nodes: 32
+      max_nodes: 32
+      market_type: ON_DEMAND
+      resources:
+        node: 1
+    - instance_type: r5.16xlarge
+      min_nodes: 32
+      max_nodes: 32
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/distributed/many_nodes_tests/actor_test.py

@@ -0,0 +1,145 @@
+import argparse
+import os
+import math
+from time import sleep, perf_counter
+import json
+import ray
+
+from dashboard_test import DashboardTestAtScale
+
+
+def test_max_actors_launch(cpus_per_actor, total_actors):
+    @ray.remote(num_cpus=cpus_per_actor)
+    class Actor:
+        def foo(self):
+            pass
+
+    print("Start launch actors")
+    actors = [Actor.options(max_restarts=-1).remote() for _ in range(total_actors)]
+    return actors
+
+
+def parse_script_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cpus-per-actor", type=float, default=0.2)
+    parser.add_argument("--total-actors", nargs="+", type=int, required=True)
+    parser.add_argument("--no-report", default=False, action="store_true")
+    parser.add_argument("--no-wait", default=False, action="store_true")
+    return parser.parse_known_args()
+
+
+def scale_cluster_up(num_cpus):
+    print(f"Start to scale up to {num_cpus} cpus")
+
+    def get_curr_cpus():
+        return int(sum([r.get("Resources", {}).get("CPU", 0) for r in ray.nodes()]))
+
+    step = 1000
+    curr_cpus = get_curr_cpus()
+    target_cpus = curr_cpus
+
+    while curr_cpus < num_cpus:
+        curr_cpus = get_curr_cpus()
+        new_target_cpus = min(curr_cpus + step, num_cpus)
+        if new_target_cpus != target_cpus:
+            target_cpus = new_target_cpus
+            ray.autoscaler.sdk.request_resources(num_cpus=target_cpus)
+        print(f"Waiting for cluster to be up: {curr_cpus}->{target_cpus}->{num_cpus}")
+        sleep(10)
+
+
+def get_head_node_cpus():
+    head_ip = ray.util.get_node_ip_address()
+    for node in ray.nodes():
+        if node["Alive"] and node["NodeManagerAddress"] == head_ip:
+            return int(node.get("Resources", {}).get("CPU", 0))
+    return 0
+
+
+def run_one(total_actors, cpus_per_actor, no_wait):
+    total_cpus = cpus_per_actor * total_actors + get_head_node_cpus()
+    total_cpus = int(math.ceil(total_cpus))
+    scale_cluster_up(total_cpus)
+
+    actor_launch_start = perf_counter()
+    actors = test_max_actors_launch(cpus_per_actor, total_actors)
+    actor_launch_end = perf_counter()
+    actor_launch_time = actor_launch_end - actor_launch_start
+    actor_ready_start = perf_counter()
+    total_actors = len(actors)
+    objs = [actor.foo.remote() for actor in actors]
+
+    while len(objs) != 0:
+        timeout = None if no_wait else 30
+        objs_ready, objs = ray.wait(objs, num_returns=len(objs), timeout=timeout)
+        print(
+            f"Status: {total_actors - len(objs)}/{total_actors}, "
+            f"{perf_counter() - actor_ready_start}"
+        )
+    actor_ready_end = perf_counter()
+    actor_ready_time = actor_ready_end - actor_ready_start
+
+    throughput = total_actors / (actor_ready_time + actor_launch_time)
+    print(f"Actor launch time: {actor_launch_time} ({total_actors} actors)")
+    print(f"Actor ready time: {actor_ready_time} ({total_actors} actors)")
+    print(
+        f"Total time: {actor_launch_time + actor_ready_time}"
+        f" ({total_actors} actors)"
+    )
+    print(f"Through put: {throughput}")
+
+    return {
+        "actor_launch_time": actor_launch_time,
+        "actor_ready_time": actor_ready_time,
+        "total_time": actor_launch_time + actor_ready_time,
+        "num_actors": total_actors,
+        "throughput": throughput,
+    }
+
+
+def main():
+    args, unknown = parse_script_args()
+    args.total_actors.sort()
+
+    addr = ray.init(address="auto")
+    dashboard_test = DashboardTestAtScale(addr)
+
+    result = {}
+    for i in args.total_actors:
+        result[f"many_nodes_actor_tests_{i}"] = run_one(
+            i, args.cpus_per_actor, args.no_wait
+        )
+
+    # Print the results early so if failed in the future, we still
+    # can see it in the log.
+    print(f"Result: {json.dumps(result, indent=2)}")
+
+    if "TEST_OUTPUT_JSON" in os.environ and not args.no_report:
+        with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+            perf = [
+                {
+                    "perf_metric_name": name,
+                    "perf_metric_value": r["throughput"],
+                    "perf_metric_type": "THROUGHPUT",
+                }
+                for (name, r) in result.items()
+            ]
+            result["perf_metrics"] = perf
+            dashboard_test.update_release_test_result(result)
+
+            print(f"Writing data into file: {os.environ['TEST_OUTPUT_JSON']}")
+            json.dump(result, out_file)
+
+    print("Test finished successfully!")
+    ray.shutdown()
+
+    # We need to make sure GCS cool down otherwise, testing infra
+    # might get timeout when fetching the result because when the driver
+    # got shutdown, many actors needs to be terminated which will
+    # overload GCS.
+    print("Sleep for 60s, waiting for the cluster to cool down.")
+    sleep(60)
+
+
+if __name__ == "__main__":
+    main()

release/benchmarks/distributed/many_nodes_tests/compute_config.yaml

@@ -0,0 +1,25 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+# NFS needs to be disabled for this test, since the test spawns too many nodes
+# and may hit the limit on the # of clients.
+advanced_instance_config:
+  TagSpecifications:
+    - ResourceType: "instance"
+      Tags:
+        - Key: as-feature-disable-nfs-mount
+          Value: "true"
+
+  BlockDeviceMappings:
+    - DeviceName: /dev/sda1
+      Ebs:
+        DeleteOnTermination: true
+        VolumeSize: 30
+
+head_node:
+  instance_type: m5.16xlarge
+
+worker_nodes:
+  - instance_type: m6i.large
+    min_nodes: 500
+    max_nodes: 2000
+    market_type: ON_DEMAND

release/benchmarks/distributed/many_nodes_tests/compute_config_gce.yaml

@@ -0,0 +1,25 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+  - us-west1-c
+
+advanced_instance_config:
+  instance_properties:
+    disks:
+      - boot: true
+        auto_delete: true
+        initialize_params:
+          disk_size_gb: 30
+
+    # NFS needs to be disabled for this test, since the test spawns too many nodes
+    # and may hit the limit on the # of clients.
+    labels:
+      as-feature-disable-nfs-mount: "true"
+
+head_node:
+  instance_type: n2-standard-64
+
+worker_nodes:
+  - instance_type: n2-standard-2
+    min_nodes: 500
+    max_nodes: 2000
+    market_type: ON_DEMAND

release/benchmarks/distributed/many_nodes_tests/dashboard_test.py

@@ -0,0 +1,185 @@
+import asyncio
+import time
+import urllib
+from typing import Dict, Optional, List
+from pprint import pprint
+
+import requests
+import ray
+import logging
+import os
+
+from collections import defaultdict
+from ray.util.state import list_nodes
+from ray._private.test_utils import get_system_metric_for_component
+from pydantic import BaseModel
+from ray.dashboard.utils import get_address_for_submission_client
+from ray.dashboard.modules.metrics.metrics_head import (
+    DEFAULT_PROMETHEUS_HOST,
+    PROMETHEUS_HOST_ENV_VAR,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def calc_p(latencies, percent):
+    if len(latencies) == 0:
+        return 0
+
+    return round(sorted(latencies)[int(len(latencies) / 100 * percent)] * 1000, 3)
+
+
+class Result(BaseModel):
+    success: bool
+    # endpoints -> list of latencies
+    result: Dict[str, List[float]]
+    # Dashboard memory usage in MB.
+    memory_mb: Optional[float]
+
+
+# Currently every endpoint is GET endpoints.
+endpoints = [
+    "/logical/actors",
+    "/nodes?view=summary",
+    "/",
+    "/api/cluster_status",
+    "/events",
+    "/api/jobs/",
+    "/api/v0/logs",
+    "/api/prometheus_health",
+]
+
+
+@ray.remote(num_cpus=0)
+class DashboardTester:
+    def __init__(self, interval_s: int = 1):
+        self.dashboard_url = get_address_for_submission_client(None)
+        # Ping interval for all endpoints.
+        self.interval_s = interval_s
+        # endpoint -> a list of latencies
+        self.result = defaultdict(list)
+
+    async def run(self):
+        await asyncio.gather(*[self.ping(endpoint) for endpoint in endpoints])
+
+    async def ping(self, endpoint):
+        """Synchronously call an endpoint."""
+        node_id = ray.get_runtime_context().get_node_id()
+        while True:
+            start = time.monotonic()
+            # for logs API, we should append node ID and glob.
+            if "/api/v0/logs" in endpoint:
+                glob_filter = "*"
+
+                options_dict = {"node_id": node_id, "glob": glob_filter}
+                url = (
+                    f"{self.dashboard_url}{endpoint}?"
+                    f"{urllib.parse.urlencode(options_dict)}"
+                )
+            else:
+                url = f"{self.dashboard_url}{endpoint}"
+
+            resp = requests.get(url, timeout=30)
+            elapsed = time.monotonic() - start
+
+            if resp.status_code == 200:
+                self.result[endpoint].append(time.monotonic() - start)
+            else:
+                try:
+                    resp.raise_for_status()
+                except Exception as e:
+                    logger.exception(e)
+            await asyncio.sleep(max(0, self.interval_s, elapsed))
+
+    def get_result(self):
+        return self.result
+
+
+class DashboardTestAtScale:
+    """This is piggybacked into existing scalability tests."""
+
+    def __init__(self, addr: ray._private.worker.RayContext):
+        self.addr = addr
+
+        # Schedule the actor on the current node (which is a head node).
+        current_node_ip = ray._private.worker.global_worker.node_ip_address
+        nodes = list_nodes(filters=[("node_ip", "=", current_node_ip)])
+        assert len(nodes) > 0, f"{current_node_ip} not found in the cluster"
+        node = nodes[0]
+        # Schedule on a head node.
+        self.tester = DashboardTester.options(
+            label_selector={ray._raylet.RAY_NODE_ID_KEY: node["node_id"]}
+        ).remote()
+
+        self.tester.run.remote()
+
+    def get_result(self):
+        """Get the result from the test.
+
+        Returns:
+            A tuple of success, and the result (Result object).
+        """
+        try:
+            result = ray.get(self.tester.get_result.remote(), timeout=60)
+        except ray.exceptions.GetTimeoutError:
+            return Result(success=False)
+
+        # Get the memory usage.
+        memories = get_system_metric_for_component(
+            "ray_component_uss_bytes",
+            "dashboard",
+            os.environ.get(PROMETHEUS_HOST_ENV_VAR, DEFAULT_PROMETHEUS_HOST),
+        )
+
+        return Result(
+            success=True,
+            result=result,
+            memory_mb=max(memories) / 1.0e6 if memories else None,
+        )
+
+    def update_release_test_result(self, release_result: dict):
+        test_result = self.get_result()
+
+        def calc_endpoints_p(result, percent):
+            return {
+                # sort -> get PX -> convert second to ms -> round up.
+                endpoint: calc_p(latencies, percent)
+                for endpoint, latencies in result.items()
+            }
+
+        print("======Print per dashboard endpoint latencies======")
+        print("=====================P50==========================")
+        pprint(calc_endpoints_p(test_result.result, 50))
+        print("=====================P95==========================")
+        pprint(calc_endpoints_p(test_result.result, 95))
+        print("=====================P99==========================")
+        pprint(calc_endpoints_p(test_result.result, 99))
+
+        latencies = []
+        for per_endpoint_latencies in test_result.result.values():
+            latencies.extend(per_endpoint_latencies)
+        aggregated_metrics = {
+            "p50": calc_p(latencies, 50),
+            "p95": calc_p(latencies, 95),
+            "p99": calc_p(latencies, 99),
+        }
+
+        print("=====================Aggregated====================")
+        pprint(aggregated_metrics)
+
+        release_result["_dashboard_test_success"] = test_result.success
+        if test_result.success:
+            if "perf_metrics" not in release_result:
+                release_result["perf_metrics"] = []
+
+            release_result["perf_metrics"].extend(
+                [
+                    {
+                        "perf_metric_name": f"dashboard_{p}_latency_ms",
+                        "perf_metric_value": value,
+                        "perf_metric_type": "LATENCY",
+                    }
+                    for p, value in aggregated_metrics.items()
+                ]
+            )
+            release_result["_dashboard_memory_usage_mb"] = test_result.memory_mb

release/benchmarks/distributed/many_nodes_tests/multi_master_test.py

@@ -0,0 +1,90 @@
+import argparse
+import os
+from time import sleep, perf_counter
+import json
+import ray
+
+
+def test_max_actors_launch(cpus_per_actor, total_actors, num_masters):
+    # By default, there are 50 groups, each group has 1 master and 99 slaves.
+    num_slaves_per_master = total_actors / num_masters - 1
+
+    @ray.remote(num_cpus=cpus_per_actor)
+    class Actor:
+        def foo(self):
+            pass
+
+        def create(self):
+            return [
+                Actor.options(max_restarts=-1).remote()
+                for _ in range(num_slaves_per_master)
+            ]
+
+    print("Start launch actors")
+    # The 50 masters are spreaded.
+    actors = [
+        Actor.options(max_restarts=-1, scheduling_strategy="SPREAD").remote()
+        for _ in range(num_masters)
+    ]
+    slaves_per_master = []
+    for master in actors:
+        slaves_per_master.append(master.create.remote())
+    for slaves in slaves_per_master:
+        actors.extend(ray.get(slaves))
+    return actors
+
+
+def test_actor_ready(actors):
+    remaining = [actor.foo.remote() for actor in actors]
+    ray.get(remaining)
+
+
+def parse_script_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cpus-per-actor", type=float, default=0.2)
+    parser.add_argument("--total-actors", type=int, default=5000)
+    parser.add_argument("--num-masters", type=int, default=50)
+    parser.add_argument("--no-report", default=False, action="store_true")
+    parser.add_argument("--fail", default=False, action="store_true")
+    return parser.parse_known_args()
+
+
+def main():
+    args, unknown = parse_script_args()
+
+    ray.init(address="auto")
+
+    actor_launch_start = perf_counter()
+    actors = test_max_actors_launch(
+        args.cpus_per_actor, args.total_actors, args.num_masters
+    )
+    actor_launch_end = perf_counter()
+    actor_launch_time = actor_launch_end - actor_launch_start
+    if args.fail:
+        sleep(10)
+        return
+    actor_ready_start = perf_counter()
+    test_actor_ready(actors)
+    actor_ready_end = perf_counter()
+    actor_ready_time = actor_ready_end - actor_ready_start
+
+    print(f"Actor launch time: {actor_launch_time} ({args.total_actors} actors)")
+    print(f"Actor ready time: {actor_ready_time} ({args.total_actors} actors)")
+    print(
+        f"Total time: {actor_launch_time + actor_ready_time}"
+        f" ({args.total_actors} actors)"
+    )
+
+    if "TEST_OUTPUT_JSON" in os.environ and not args.no_report:
+        with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+            results = {
+                "actor_launch_time": actor_launch_time,
+                "actor_ready_time": actor_ready_time,
+                "total_time": actor_launch_time + actor_ready_time,
+                "num_actors": args.total_actors,
+            }
+            json.dump(results, out_file)
+
+
+if __name__ == "__main__":
+    main()

release/benchmarks/distributed/test_many_actors.py

@@ -0,0 +1,90 @@
+import os
+import time
+
+import tqdm
+from many_nodes_tests.dashboard_test import DashboardTestAtScale
+
+import ray
+import ray._common.test_utils
+import ray._private.test_utils as test_utils
+from ray._private.state_api_test_utils import summarize_worker_startup_time
+
+is_smoke_test = True
+if "SMOKE_TEST" in os.environ:
+    MAX_ACTORS_IN_CLUSTER = 100
+else:
+    MAX_ACTORS_IN_CLUSTER = 10000
+    is_smoke_test = False
+
+
+def test_max_actors():
+    # TODO (Alex): Dynamically set this based on number of cores
+    cpus_per_actor = 0.25
+
+    @ray.remote(num_cpus=cpus_per_actor)
+    class Actor:
+        def foo(self):
+            pass
+
+    actors = [
+        Actor.remote()
+        for _ in tqdm.trange(MAX_ACTORS_IN_CLUSTER, desc="Launching actors")
+    ]
+
+    done = ray.get([actor.foo.remote() for actor in actors])
+    for result in done:
+        assert result is None
+
+
+def no_resource_leaks():
+    return test_utils.no_resource_leaks_excluding_node_resources()
+
+
+addr = ray.init(address="auto")
+
+ray._common.test_utils.wait_for_condition(no_resource_leaks)
+monitor_actor = test_utils.monitor_memory_usage()
+dashboard_test = DashboardTestAtScale(addr)
+
+start_time = time.time()
+test_max_actors()
+end_time = time.time()
+
+ray.get(monitor_actor.stop_run.remote())
+used_gb, usage = ray.get(monitor_actor.get_peak_memory_info.remote())
+print(f"Peak memory usage: {round(used_gb, 2)}GB")
+print(f"Peak memory usage per processes:\n {usage}")
+del monitor_actor
+
+# Get the dashboard result
+ray._common.test_utils.wait_for_condition(no_resource_leaks)
+
+rate = MAX_ACTORS_IN_CLUSTER / (end_time - start_time)
+try:
+    summarize_worker_startup_time()
+except Exception as e:
+    print("Failed to summarize worker startup time.")
+    print(e)
+
+print(
+    f"Success! Started {MAX_ACTORS_IN_CLUSTER} actors in "
+    f"{end_time - start_time}s. ({rate} actors/s)"
+)
+
+results = {
+    "actors_per_second": rate,
+    "num_actors": MAX_ACTORS_IN_CLUSTER,
+    "time": end_time - start_time,
+    "_peak_memory": round(used_gb, 2),
+    "_peak_process_memory": usage,
+}
+if not is_smoke_test:
+    results["perf_metrics"] = [
+        {
+            "perf_metric_name": "actors_per_second",
+            "perf_metric_value": rate,
+            "perf_metric_type": "THROUGHPUT",
+        }
+    ]
+dashboard_test.update_release_test_result(results)
+test_utils.safe_write_to_results_json(results)

release/benchmarks/distributed/test_many_pgs.py

@@ -0,0 +1,121 @@
+import os
+import time
+
+import tqdm
+from many_nodes_tests.dashboard_test import DashboardTestAtScale
+
+import ray
+import ray._common.test_utils
+import ray._private.test_utils as test_utils
+from ray.util.placement_group import placement_group, remove_placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+is_smoke_test = True
+if "SMOKE_TEST" in os.environ:
+    MAX_PLACEMENT_GROUPS = 20
+else:
+    MAX_PLACEMENT_GROUPS = 1000
+    is_smoke_test = False
+
+
+def test_many_placement_groups():
+    # @ray.remote(num_cpus=1, resources={"node": 0.02})
+    @ray.remote
+    class C1:
+        def ping(self):
+            return "pong"
+
+    # @ray.remote(num_cpus=1)
+    @ray.remote
+    class C2:
+        def ping(self):
+            return "pong"
+
+    # @ray.remote(resources={"node": 0.02})
+    @ray.remote
+    class C3:
+        def ping(self):
+            return "pong"
+
+    bundle1 = {"node": 0.02, "CPU": 1}
+    bundle2 = {"CPU": 1}
+    bundle3 = {"node": 0.02}
+
+    pgs = []
+    for _ in tqdm.trange(MAX_PLACEMENT_GROUPS, desc="Creating pgs"):
+        pg = placement_group(bundles=[bundle1, bundle2, bundle3])
+        pgs.append(pg)
+
+    for pg in tqdm.tqdm(pgs, desc="Waiting for pgs to be ready"):
+        ray.get(pg.ready())
+
+    actors = []
+    for pg in tqdm.tqdm(pgs, desc="Scheduling tasks"):
+        actors.append(
+            C1.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
+            ).remote()
+        )
+        actors.append(
+            C2.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
+            ).remote()
+        )
+        actors.append(
+            C3.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
+            ).remote()
+        )
+
+    not_ready = [actor.ping.remote() for actor in actors]
+    for _ in tqdm.trange(len(actors)):
+        ready, not_ready = ray.wait(not_ready)
+        assert ray.get(*ready) == "pong"
+
+    for pg in tqdm.tqdm(pgs, desc="Cleaning up pgs"):
+        remove_placement_group(pg)
+
+
+def no_resource_leaks():
+    return test_utils.no_resource_leaks_excluding_node_resources()
+
+
+addr = ray.init(address="auto")
+
+ray._common.test_utils.wait_for_condition(no_resource_leaks)
+monitor_actor = test_utils.monitor_memory_usage()
+dashboard_test = DashboardTestAtScale(addr)
+
+start_time = time.time()
+test_many_placement_groups()
+end_time = time.time()
+ray.get(monitor_actor.stop_run.remote())
+used_gb, usage = ray.get(monitor_actor.get_peak_memory_info.remote())
+print(f"Peak memory usage: {round(used_gb, 2)}GB")
+print(f"Peak memory usage per processes:\n {usage}")
+del monitor_actor
+ray._common.test_utils.wait_for_condition(no_resource_leaks)
+
+rate = MAX_PLACEMENT_GROUPS / (end_time - start_time)
+print(
+    f"Success! Started {MAX_PLACEMENT_GROUPS} pgs in "
+    f"{end_time - start_time}s. ({rate} pgs/s)"
+)
+
+results = {
+    "pgs_per_second": rate,
+    "num_pgs": MAX_PLACEMENT_GROUPS,
+    "time": end_time - start_time,
+    "_peak_memory": round(used_gb, 2),
+    "_peak_process_memory": usage,
+}
+if not is_smoke_test:
+    results["perf_metrics"] = [
+        {
+            "perf_metric_name": "pgs_per_second",
+            "perf_metric_value": rate,
+            "perf_metric_type": "THROUGHPUT",
+        }
+    ]
+dashboard_test.update_release_test_result(results)
+test_utils.safe_write_to_results_json(results)

release/benchmarks/distributed/test_many_tasks.py

@@ -0,0 +1,138 @@
+import time
+
+import click
+import tqdm
+from many_nodes_tests.dashboard_test import DashboardTestAtScale
+
+import ray
+import ray._common.test_utils
+import ray._private.test_utils as test_utils
+from ray._private.state_api_test_utils import (
+    StateAPICallSpec,
+    periodic_invoke_state_apis_with_actor,
+    summarize_worker_startup_time,
+)
+from ray.util.state import summarize_tasks
+
+sleep_time = 300
+
+
+def test_max_running_tasks(num_tasks):
+    cpus_per_task = 0.25
+
+    @ray.remote(num_cpus=cpus_per_task)
+    def task():
+        time.sleep(sleep_time)
+
+    def time_up(start_time):
+        return time.time() - start_time >= sleep_time
+
+    refs = [task.remote() for _ in tqdm.trange(num_tasks, desc="Launching tasks")]
+
+    max_cpus = ray.cluster_resources()["CPU"]
+    min_cpus_available = max_cpus
+    start_time = time.time()
+    for _ in tqdm.trange(int(sleep_time / 0.1), desc="Waiting"):
+        try:
+            cur_cpus = ray.available_resources().get("CPU", 0)
+            min_cpus_available = min(min_cpus_available, cur_cpus)
+        except Exception:
+            # There are race conditions `.get` can fail if a new heartbeat
+            # comes at the same time.
+            pass
+        if time_up(start_time):
+            print(f"Time up for sleeping {sleep_time} seconds")
+            break
+        time.sleep(0.1)
+
+    # There are some relevant magic numbers in this check. 10k tasks each
+    # require 1/4 cpus. Therefore, ideally 2.5k cpus will be used.
+    used_cpus = max_cpus - min_cpus_available
+    err_str = f"Only {used_cpus}/{max_cpus} cpus used."
+    # 1500 tasks. Note that it is a pretty low threshold, and the
+    # performance should be tracked via perf dashboard.
+    threshold = num_tasks * cpus_per_task * 0.60
+    print(f"{used_cpus}/{max_cpus} used.")
+    assert used_cpus > threshold, err_str
+
+    for _ in tqdm.trange(num_tasks, desc="Ensuring all tasks have finished"):
+        done, refs = ray.wait(refs)
+        assert ray.get(done[0]) is None
+
+    return used_cpus
+
+
+def no_resource_leaks():
+    return test_utils.no_resource_leaks_excluding_node_resources()
+
+
+@click.command()
+@click.option("--num-tasks", required=True, type=int, help="Number of tasks to launch.")
+def test(num_tasks):
+    addr = ray.init(address="auto")
+
+    ray._common.test_utils.wait_for_condition(no_resource_leaks)
+    monitor_actor = test_utils.monitor_memory_usage()
+    dashboard_test = DashboardTestAtScale(addr)
+
+    def not_none(res):
+        return res is not None
+
+    api_caller = periodic_invoke_state_apis_with_actor(
+        apis=[StateAPICallSpec(summarize_tasks, not_none)],
+        call_interval_s=4,
+        print_result=True,
+    )
+
+    start_time = time.time()
+    used_cpus = test_max_running_tasks(num_tasks)
+    end_time = time.time()
+    ray.get(monitor_actor.stop_run.remote())
+    used_gb, usage = ray.get(monitor_actor.get_peak_memory_info.remote())
+    print(f"Peak memory usage: {round(used_gb, 2)}GB")
+    print(f"Peak memory usage per processes:\n {usage}")
+    ray.get(api_caller.stop.remote())
+
+    del api_caller
+    del monitor_actor
+    ray._common.test_utils.wait_for_condition(no_resource_leaks)
+
+    try:
+        summarize_worker_startup_time()
+    except Exception as e:
+        print("Failed to summarize worker startup time.")
+        print(e)
+
+    rate = num_tasks / (end_time - start_time - sleep_time)
+    print(
+        f"Success! Started {num_tasks} tasks in {end_time - start_time}s. "
+        f"({rate} tasks/s)"
+    )
+
+    results = {
+        "tasks_per_second": rate,
+        "num_tasks": num_tasks,
+        "time": end_time - start_time,
+        "used_cpus": used_cpus,
+        "_peak_memory": round(used_gb, 2),
+        "_peak_process_memory": usage,
+        "perf_metrics": [
+            {
+                "perf_metric_name": "tasks_per_second",
+                "perf_metric_value": rate,
+                "perf_metric_type": "THROUGHPUT",
+            },
+            {
+                "perf_metric_name": "used_cpus_by_deadline",
+                "perf_metric_value": used_cpus,
+                "perf_metric_type": "THROUGHPUT",
+            },
+        ],
+    }
+
+    dashboard_test.update_release_test_result(results)
+    test_utils.safe_write_to_results_json(results)
+
+
+if __name__ == "__main__":
+    test()

release/benchmarks/distributed/test_scheduling.py

@@ -0,0 +1,134 @@
+import argparse
+from math import floor
+from time import sleep, time
+
+import ray
+import ray._private.test_utils as test_utils
+from ray._private.test_utils import safe_write_to_results_json
+
+
+@ray.remote
+def simple_task(t):
+    sleep(t)
+
+
+@ray.remote
+class SimpleActor:
+    def __init__(self, job=None):
+        self._job = job
+
+    def ready(self):
+        return
+
+    def do_job(self):
+        if self._job is not None:
+            self._job()
+
+
+def start_tasks(num_task, num_cpu_per_task, task_duration):
+    ray.get(
+        [
+            simple_task.options(num_cpus=num_cpu_per_task).remote(task_duration)
+            for _ in range(num_task)
+        ]
+    )
+
+
+def measure(f):
+    start = time()
+    ret = f()
+    end = time()
+    return (end - start, ret)
+
+
+def start_actor(num_actors, num_actors_per_nodes, job):
+    resources = {"node": floor(1.0 / num_actors_per_nodes)}
+    submission_cost, actors = measure(
+        lambda: [
+            SimpleActor.options(resources=resources, num_cpus=0).remote(job)
+            for _ in range(num_actors)
+        ]
+    )
+    ready_cost, _ = measure(lambda: ray.get([actor.ready.remote() for actor in actors]))
+    actor_job_cost, _ = measure(
+        lambda: ray.get([actor.do_job.remote() for actor in actors])
+    )
+    return (submission_cost, ready_cost, actor_job_cost)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(prog="Test Scheduling")
+    # Task workloads
+    parser.add_argument(
+        "--total-num-task", type=int, help="Total number of tasks.", required=False
+    )
+    parser.add_argument(
+        "--num-cpu-per-task",
+        type=int,
+        help="Resources needed for tasks.",
+        required=False,
+    )
+    parser.add_argument(
+        "--task-duration-s",
+        type=int,
+        help="How long does each task execute.",
+        required=False,
+        default=1,
+    )
+
+    # Actor workloads
+    parser.add_argument(
+        "--total-num-actors", type=int, help="Total number of actors.", required=True
+    )
+    parser.add_argument(
+        "--num-actors-per-nodes",
+        type=int,
+        help="How many actors to allocate for each nodes.",
+        required=True,
+    )
+
+    ray.init(address="auto")
+    monitor_actor = test_utils.monitor_memory_usage()
+
+    total_cpus_per_node = [node["Resources"].get("CPU", 0) for node in ray.nodes()]
+    num_nodes = len(total_cpus_per_node)
+    total_cpus = sum(total_cpus_per_node)
+
+    args = parser.parse_args()
+    job = None
+    if args.total_num_task is not None:
+        if args.num_cpu_per_task is None:
+            args.num_cpu_per_task = floor(1.0 * total_cpus / args.total_num_task)
+        job = lambda: start_tasks(  # noqa: E731
+            args.total_num_task, args.num_cpu_per_task, args.task_duration_s
+        )
+
+    submission_cost, ready_cost, actor_job_cost = start_actor(
+        args.total_num_actors, args.num_actors_per_nodes, job
+    )
+
+    ray.get(monitor_actor.stop_run.remote())
+    used_gb, usage = ray.get(monitor_actor.get_peak_memory_info.remote())
+    print(f"Peak memory usage: {round(used_gb, 2)}GB")
+    print(f"Peak memory usage per processes:\n {usage}")
+    del monitor_actor
+
+    result = {
+        "total_num_task": args.total_num_task,
+        "num_cpu_per_task": args.num_cpu_per_task,
+        "task_duration_s": args.task_duration_s,
+        "total_num_actors": args.total_num_actors,
+        "num_actors_per_nodes": args.num_actors_per_nodes,
+        "num_nodes": num_nodes,
+        "total_cpus": total_cpus,
+        "submission_cost": submission_cost,
+        "ready_cost": ready_cost,
+        "actor_job_cost": actor_job_cost,
+        "_peak_memory": round(used_gb, 2),
+        "_peak_process_memory": usage,
+        "_runtime": submission_cost + ready_cost + actor_job_cost,
+    }
+
+    safe_write_to_results_json(result)
+
+    print(result)

release/benchmarks/distributed_gce.yaml

@@ -0,0 +1,26 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-c
+
+head_node:
+    instance_type: n2-standard-64 # Network optimized.
+    resources:
+      CPU: 0
+      node: 1
+      small: 1
+
+worker_nodes:
+    - name: worker_node_m
+      instance_type: n2-standard-64
+      min_nodes: 32
+      max_nodes: 32
+      market_type: ON_DEMAND
+      resources:
+        node: 1
+    - name: worker_node_r
+      instance_type: n2-standard-64
+      min_nodes: 32
+      max_nodes: 32
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/distributed_smoke_test.yaml

@@ -0,0 +1,16 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: r5dn.16xlarge # Network optimized.
+    resources:
+      CPU: 0
+      node: 1
+      small: 1
+
+worker_nodes:
+    - instance_type: m5.16xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/many_nodes.yaml

@@ -0,0 +1,15 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: r6idn.16xlarge # Network optimized.
+    resources:
+      node: 1
+      small: 1
+
+worker_nodes:
+    - instance_type: m6i.2xlarge
+      min_nodes: 249
+      max_nodes: 249
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/many_nodes_gce.yaml

@@ -0,0 +1,17 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-c
+
+head_node:
+    instance_type: n2-standard-64 # Network optimized.
+    resources:
+      node: 1
+      small: 1
+
+worker_nodes:
+    - instance_type: n2-standard-8
+      min_nodes: 249
+      max_nodes: 249
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/object_store.yaml

@@ -0,0 +1,14 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m6i.16xlarge
+    resources:
+      node: 1
+
+worker_nodes:
+    - instance_type: m6i.2xlarge
+      min_nodes: 49
+      max_nodes: 49
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/object_store/callback_throughput.yaml

@@ -0,0 +1,14 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m6i.2xlarge
+    resources:
+      node: 1
+
+worker_nodes:
+    - instance_type: m6i.2xlarge
+      min_nodes: 10
+      max_nodes: 10
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/object_store/large_objects.yaml

@@ -0,0 +1,14 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m6i.16xlarge
+    resources:
+      node: 1
+
+worker_nodes:
+    - instance_type: m6i.4xlarge
+      min_nodes: 9
+      max_nodes: 9
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/object_store/small_objects.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m6i.16xlarge
+
+worker_nodes:
+    - instance_type: m6i.4xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND

release/benchmarks/object_store/test_callback_throughput.py

@@ -0,0 +1,108 @@
+import json
+import os
+import threading
+import time
+
+import numpy as np
+
+import ray
+import ray._private.worker
+
+NUM_WORKERS = 10
+OBJECT_SIZE = 1024 * 1024  # 1 MiB, above the 100 KB inlining threshold
+
+
+@ray.remote(num_cpus=1)
+def produce_block():
+    return np.zeros(OBJECT_SIZE, dtype=np.uint8)
+
+
+@ray.remote(num_cpus=1)
+def consume_block(block_ref):
+    return len(block_ref)
+
+
+def test_callback_pipeline(num_blocks, timeout_s=60):
+    core_worker = ray._private.worker.global_worker.core_worker
+
+    latencies = []
+    drop_times = {}
+    lock = threading.Lock()
+    done = threading.Event()
+
+    def on_freed(id_bytes):
+        with lock:
+            latencies.append(time.perf_counter() - drop_times[id_bytes])
+            if len(latencies) == num_blocks:
+                done.set()
+
+    refs = [
+        produce_block.options(scheduling_strategy="SPREAD").remote()
+        for _ in range(num_blocks)
+    ]
+    ray.wait(refs, num_returns=len(refs))
+
+    # live_refs keeps each block ref alive until its consumer completes.
+    live_refs = {}
+    for ref in refs:
+        assert core_worker.add_object_out_of_scope_callback(ref, on_freed)
+        consumer = consume_block.remote(ref)
+        live_refs[consumer] = ref
+    del refs
+
+    # Release each ref as its consumer completes.
+    pending = list(live_refs.keys())
+    while pending:
+        done_list, pending = ray.wait(pending, num_returns=1)
+        for consumer in done_list:
+            ref = live_refs.pop(consumer)
+            drop_times[ref.binary()] = time.perf_counter()
+            del ref
+
+    if not done.wait(timeout=timeout_s):
+        raise TimeoutError(
+            f"Only {len(latencies)}/{num_blocks} callbacks fired within {timeout_s}s"
+        )
+
+    latencies.sort()
+    p95 = latencies[int(len(latencies) * 0.95)]
+    print(f"  {num_blocks} blocks: p95={p95:.4f}s")
+    return p95
+
+
+ray.init(address="auto")
+
+# Warm up gRPC connections and worker pools.
+ray.get(
+    [
+        produce_block.options(scheduling_strategy="SPREAD").remote()
+        for _ in range(NUM_WORKERS)
+    ]
+)
+
+p95_100 = test_callback_pipeline(100)
+p95_1k = test_callback_pipeline(1000)
+
+print("\nSummary:")
+print(f"  100 blocks: p95={p95_100:.4f}s")
+print(f"  1k blocks: p95={p95_1k:.4f}s")
+
+if "TEST_OUTPUT_JSON" in os.environ:
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+        results = {
+            "p95_100": p95_100,
+            "p95_1k": p95_1k,
+            "perf_metrics": [
+                {
+                    "perf_metric_name": "callback_p95_latency_100_blocks_s",
+                    "perf_metric_value": p95_100,
+                    "perf_metric_type": "LATENCY",
+                },
+                {
+                    "perf_metric_name": "callback_p95_latency_1k_blocks_s",
+                    "perf_metric_value": p95_1k,
+                    "perf_metric_type": "LATENCY",
+                },
+            ],
+        }
+        json.dump(results, out_file)

release/benchmarks/object_store/test_large_objects.py

@@ -0,0 +1,101 @@
+import json
+import os
+from time import perf_counter
+
+import numpy as np
+from tqdm import tqdm
+
+import ray
+
+NUM_NODES = 9
+OBJECT_SIZE = 2**32
+
+
+def test_object_many_to_one():
+    @ray.remote(num_cpus=1, resources={"node": 1})
+    class Actor:
+        def foo(self):
+            pass
+
+        def send_objects(self):
+            return np.ones(OBJECT_SIZE, dtype=np.uint8)
+
+    actors = [Actor.remote() for _ in range(NUM_NODES)]
+
+    for actor in tqdm(actors, desc="Ensure all actors have started."):
+        ray.get(actor.foo.remote())
+
+    start = perf_counter()
+    result_refs = []
+    for actor in tqdm(actors, desc="Tasks kickoff"):
+        result_refs.append(actor.send_objects.remote())
+
+    results = ray.get(result_refs)
+    end = perf_counter()
+
+    for result in results:
+        assert len(result) == OBJECT_SIZE
+
+    return end - start
+
+
+def test_object_one_to_many():
+    @ray.remote(num_cpus=1, resources={"node": 1})
+    class Actor:
+        def foo(self):
+            pass
+
+        def data_len(self, arr):
+            return len(arr)
+
+    actors = [Actor.remote() for _ in range(NUM_NODES)]
+
+    arr = np.ones(OBJECT_SIZE, dtype=np.uint8)
+    ref = ray.put(arr)
+
+    for actor in tqdm(actors, desc="Ensure all actors have started."):
+        ray.get(actor.foo.remote())
+
+    start = perf_counter()
+    result_refs = []
+    for actor in tqdm(actors, desc="Tasks kickoff"):
+        result_refs.append(actor.data_len.remote(ref))
+
+    results = ray.get(result_refs)
+    end = perf_counter()
+
+    for result in results:
+        assert result == OBJECT_SIZE
+
+    return end - start
+
+
+ray.init(address="auto")
+many_to_one_duration = test_object_many_to_one()
+print(f"many_to_one time: {many_to_one_duration} ({OBJECT_SIZE} B x {NUM_NODES} nodes)")
+one_to_many_duration = test_object_one_to_many()
+print(f"one_to_many time: {one_to_many_duration} ({OBJECT_SIZE} B x {NUM_NODES} nodes)")
+
+
+if "TEST_OUTPUT_JSON" in os.environ:
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+        results = {
+            "many_to_one_time": many_to_one_duration,
+            "one_to_many_time": one_to_many_duration,
+            "object_size": OBJECT_SIZE,
+            "num_nodes": NUM_NODES,
+        }
+        results["perf_metrics"] = [
+            {
+                "perf_metric_name": f"time_many_to_one_{OBJECT_SIZE}_bytes_from_{NUM_NODES}_nodes",
+                "perf_metric_value": many_to_one_duration,
+                "perf_metric_type": "LATENCY",
+            },
+            {
+                "perf_metric_name": f"time_one_to_many_{OBJECT_SIZE}_bytes_to_{NUM_NODES}_nodes",
+                "perf_metric_value": one_to_many_duration,
+                "perf_metric_type": "LATENCY",
+            },
+        ]
+
+        json.dump(results, out_file)

release/benchmarks/object_store/test_object_store.py

@@ -0,0 +1,75 @@
+import json
+import os
+from time import perf_counter
+
+import numpy as np
+from tqdm import tqdm
+
+import ray
+import ray.autoscaler.sdk
+
+NUM_NODES = 50
+OBJECT_SIZE = 2**30
+
+
+def num_alive_nodes():
+    n = 0
+    for node in ray.nodes():
+        if node["Alive"]:
+            n += 1
+    return n
+
+
+def test_object_broadcast():
+    assert num_alive_nodes() == NUM_NODES
+
+    @ray.remote(num_cpus=1, resources={"node": 1})
+    class Actor:
+        def foo(self):
+            pass
+
+        def data_len(self, arr):
+            return len(arr)
+
+    actors = [Actor.remote() for _ in range(NUM_NODES)]
+
+    arr = np.ones(OBJECT_SIZE, dtype=np.uint8)
+    ref = ray.put(arr)
+
+    for actor in tqdm(actors, desc="Ensure all actors have started."):
+        ray.get(actor.foo.remote())
+
+    start = perf_counter()
+    result_refs = []
+    for actor in tqdm(actors, desc="Broadcasting objects"):
+        result_refs.append(actor.data_len.remote(ref))
+
+    results = ray.get(result_refs)
+    end = perf_counter()
+
+    for result in results:
+        assert result == OBJECT_SIZE
+
+    return end - start
+
+
+ray.init(address="auto")
+duration = test_object_broadcast()
+print(f"Broadcast time: {duration} ({OBJECT_SIZE} B x {NUM_NODES} nodes)")
+
+if "TEST_OUTPUT_JSON" in os.environ:
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+        results = {
+            "broadcast_time": duration,
+            "object_size": OBJECT_SIZE,
+            "num_nodes": NUM_NODES,
+        }
+        perf_metric_name = f"time_to_broadcast_{OBJECT_SIZE}_bytes_to_{NUM_NODES}_nodes"
+        results["perf_metrics"] = [
+            {
+                "perf_metric_name": perf_metric_name,
+                "perf_metric_value": duration,
+                "perf_metric_type": "LATENCY",
+            }
+        ]
+        json.dump(results, out_file)

release/benchmarks/object_store/test_small_objects.py

@@ -0,0 +1,81 @@
+import json
+import os
+import time
+
+import numpy as np
+
+import ray
+
+
+def test_small_objects_many_to_one():
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def send(self, _, actor_idx):
+            # this size is chosen because it's >100kb so big enough to be stored in plasma
+            numpy_arr = np.ones((20, 1024))
+            return (numpy_arr, actor_idx)
+
+    actors = [Actor.remote() for _ in range(64)]
+    not_ready = []
+    for index, actor in enumerate(actors):
+        not_ready.append(actor.send.remote(0, index))
+    num_messages = 0
+    start_time = time.time()
+    while time.time() - start_time < 60:
+        ready, not_ready = ray.wait(not_ready, num_returns=10)
+        for ready_ref in ready:
+            _, actor_idx = ray.get(ready_ref)
+            not_ready.append(actors[actor_idx].send.remote(0, actor_idx))
+        num_messages += 10
+    return num_messages / 60
+
+
+def test_small_objects_one_to_many():
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def receive(self, numpy_arr, actor_idx):
+            return actor_idx
+
+    actors = [Actor.remote() for _ in range(64)]
+    numpy_arr_ref = ray.put(np.ones((20, 1024)))
+    not_ready = []
+
+    num_messages = 0
+    start_time = time.time()
+    for idx, actor in enumerate(actors):
+        not_ready.append(actor.receive.remote(numpy_arr_ref, idx))
+    while time.time() - start_time < 60:
+        ready, not_ready = ray.wait(not_ready, num_returns=10)
+        actor_idxs = ray.get(ready)
+        for actor_idx in actor_idxs:
+            not_ready.append(actors[actor_idx].receive.remote(numpy_arr_ref, actor_idx))
+        num_messages += 10
+    return num_messages / 60
+
+
+ray.init(address="auto")
+many_to_one_throughput = test_small_objects_many_to_one()
+print(f"Number of messages per second many_to_one: {many_to_one_throughput}")
+one_to_many_throughput = test_small_objects_one_to_many()
+print(f"Number of messages per second one_to_many: {one_to_many_throughput}")
+
+
+if "TEST_OUTPUT_JSON" in os.environ:
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+        results = {
+            "num_messages_many_to_one": many_to_one_throughput,
+            "num_messages_one_to_many": one_to_many_throughput,
+        }
+        results["perf_metrics"] = [
+            {
+                "perf_metric_name": "num_small_objects_many_to_one",
+                "perf_metric_value": many_to_one_throughput,
+                "perf_metric_type": "THROUGHPUT",
+            },
+            {
+                "perf_metric_name": "num_small_objects_one_to_many_per_second",
+                "perf_metric_value": one_to_many_throughput,
+                "perf_metric_type": "THROUGHPUT",
+            },
+        ]
+        json.dump(results, out_file)

release/benchmarks/object_store_gce.yaml

@@ -0,0 +1,16 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-c
+
+head_node:
+    instance_type: n2-standard-64
+    resources:
+      node: 1
+
+worker_nodes:
+    - instance_type: n2-standard-8
+      min_nodes: 49
+      max_nodes: 49
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/scheduling.yaml

@@ -0,0 +1,19 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.4xlarge
+    resources:
+      # Assume the node has 64 CPU instead of 16.
+      # This should be fine since each task has little
+      # computation in scheduling tests.
+      CPU: 64
+      node: 1
+
+worker_nodes:
+    - instance_type: m5.4xlarge
+      min_nodes: 31
+      max_nodes: 31
+      market_type: ON_DEMAND
+      resources:
+        CPU: 64
+        node: 1

release/benchmarks/scheduling_gce.yaml

@@ -0,0 +1,21 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-c
+
+head_node:
+    instance_type: n2-standard-16
+    resources:
+      # Assume the node has 64 CPU instead of 16.
+      # This should be fine since each task has little
+      # computation in scheduling tests.
+      CPU: 64
+      node: 1
+
+worker_nodes:
+    - instance_type: n2-standard-16
+      min_nodes: 31
+      max_nodes: 31
+      market_type: ON_DEMAND
+      resources:
+        CPU: 64
+        node: 1

release/benchmarks/single_node.yaml

@@ -0,0 +1,16 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+advanced_instance_config:
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+            DeleteOnTermination: true
+            VolumeSize: 500
+
+head_node:
+    instance_type: m4.16xlarge
+    resources:
+      # 128 GB
+      object_store_memory: 128000000000
+
+worker_nodes: []

release/benchmarks/single_node/test_single_node.py

@@ -0,0 +1,237 @@
+import json
+import os
+import time
+from time import perf_counter
+
+import numpy as np
+from tqdm import tqdm, trange
+
+import ray
+import ray.autoscaler.sdk
+from ray._common.test_utils import Semaphore
+
+MAX_ARGS = 10000
+MAX_RETURNS = 3000
+MAX_RAY_GET_ARGS = 10000
+MAX_QUEUED_TASKS = 1_000_000
+MAX_RAY_GET_SIZE = 100 * 2**30
+
+
+def assert_no_leaks():
+    total = ray.cluster_resources()
+    current = ray.available_resources()
+    total.pop("memory")
+    total.pop("object_store_memory")
+    current.pop("memory")
+    current.pop("object_store_memory")
+    assert total == current, (total, current)
+
+
+def test_many_args():
+    @ray.remote
+    def sum_args(*args):
+        return sum(sum(arg) for arg in args)
+
+    args = [[1 for _ in range(10000)] for _ in range(MAX_ARGS)]
+    result = ray.get(sum_args.remote(*args))
+    assert result == MAX_ARGS * 10000
+
+
+def test_many_returns():
+    @ray.remote(num_returns=MAX_RETURNS)
+    def f():
+        to_return = []
+        for _ in range(MAX_RETURNS):
+            obj = list(range(10000))
+            to_return.append(obj)
+
+        return tuple(to_return)
+
+    returned_refs = f.remote()
+    assert len(returned_refs) == MAX_RETURNS
+
+    for ref in returned_refs:
+        expected = list(range(10000))
+        obj = ray.get(ref)
+        assert obj == expected
+
+
+def test_ray_get_args():
+    def with_dese():
+        print("Putting test objects:")
+        refs = []
+        for _ in trange(MAX_RAY_GET_ARGS):
+            obj = list(range(10000))
+            refs.append(ray.put(obj))
+
+        print("Getting objects")
+        results = ray.get(refs)
+        assert len(results) == MAX_RAY_GET_ARGS
+
+        print("Asserting correctness")
+        for obj in tqdm(results):
+            expected = list(range(10000))
+            assert obj == expected
+
+    def with_zero_copy():
+        print("Putting test objects:")
+        refs = []
+        for _ in trange(MAX_RAY_GET_ARGS):
+            obj = np.arange(10000)
+            refs.append(ray.put(obj))
+
+        print("Getting objects")
+        results = ray.get(refs)
+        assert len(results) == MAX_RAY_GET_ARGS
+
+        print("Asserting correctness")
+        for obj in tqdm(results):
+            expected = np.arange(10000)
+            assert (obj == expected).all()
+
+    with_dese()
+    print("Done with dese")
+    with_zero_copy()
+    print("Done with zero copy")
+
+
+def test_many_queued_tasks():
+    sema = Semaphore.remote(0)
+
+    @ray.remote(num_cpus=1)
+    def block():
+        ray.get(sema.acquire.remote())
+
+    @ray.remote(num_cpus=1)
+    def f():
+        pass
+
+    num_cpus = int(ray.cluster_resources()["CPU"])
+    blocked_tasks = []
+    for _ in range(num_cpus):
+        blocked_tasks.append(block.remote())
+
+    print("Submitting many tasks")
+    pending_tasks = []
+    for _ in trange(MAX_QUEUED_TASKS):
+        pending_tasks.append(f.remote())
+
+    # Make sure all the tasks can actually run.
+    for _ in range(num_cpus):
+        sema.release.remote()
+
+    print("Unblocking tasks")
+    for ref in tqdm(pending_tasks):
+        assert ray.get(ref) is None
+
+
+def test_large_object():
+    print("Generating object")
+    obj = np.zeros(MAX_RAY_GET_SIZE, dtype=np.int8)
+    print("Putting object")
+    ref = ray.put(obj)
+    del obj
+    print("Getting object")
+    big_obj = ray.get(ref)
+
+    assert big_obj[0] == 0
+    assert big_obj[-1] == 0
+
+
+ray.init(address="auto")
+
+args_start = perf_counter()
+test_many_args()
+args_end = perf_counter()
+
+time.sleep(5)
+assert_no_leaks()
+print("Finished many args")
+
+returns_start = perf_counter()
+test_many_returns()
+returns_end = perf_counter()
+
+time.sleep(5)
+assert_no_leaks()
+print("Finished many returns")
+
+get_start = perf_counter()
+test_ray_get_args()
+get_end = perf_counter()
+
+time.sleep(5)
+assert_no_leaks()
+print("Finished ray.get on many objects")
+
+queued_start = perf_counter()
+test_many_queued_tasks()
+queued_end = perf_counter()
+
+time.sleep(5)
+assert_no_leaks()
+print("Finished queueing many tasks")
+
+large_object_start = perf_counter()
+test_large_object()
+large_object_end = perf_counter()
+
+time.sleep(5)
+assert_no_leaks()
+print("Done")
+
+args_time = args_end - args_start
+returns_time = returns_end - returns_start
+get_time = get_end - get_start
+queued_time = queued_end - queued_start
+large_object_time = large_object_end - large_object_start
+
+print(f"Many args time: {args_time} ({MAX_ARGS} args)")
+print(f"Many returns time: {returns_time} ({MAX_RETURNS} returns)")
+print(f"Ray.get time: {get_time} ({MAX_RAY_GET_ARGS} args)")
+print(f"Queued task time: {queued_time} ({MAX_QUEUED_TASKS} tasks)")
+print(f"Ray.get large object time: {large_object_time} " f"({MAX_RAY_GET_SIZE} bytes)")
+
+if "TEST_OUTPUT_JSON" in os.environ:
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+        results = {
+            "args_time": args_time,
+            "num_args": MAX_ARGS,
+            "returns_time": returns_time,
+            "num_returns": MAX_RETURNS,
+            "get_time": get_time,
+            "num_get_args": MAX_RAY_GET_ARGS,
+            "queued_time": queued_time,
+            "num_queued": MAX_QUEUED_TASKS,
+            "large_object_time": large_object_time,
+            "large_object_size": MAX_RAY_GET_SIZE,
+            "success": "1",
+        }
+        results["perf_metrics"] = [
+            {
+                "perf_metric_name": f"{MAX_ARGS}_args_time",
+                "perf_metric_value": args_time,
+                "perf_metric_type": "LATENCY",
+            },
+            {
+                "perf_metric_name": f"{MAX_RETURNS}_returns_time",
+                "perf_metric_value": returns_time,
+                "perf_metric_type": "LATENCY",
+            },
+            {
+                "perf_metric_name": f"{MAX_RAY_GET_ARGS}_get_time",
+                "perf_metric_value": get_time,
+                "perf_metric_type": "LATENCY",
+            },
+            {
+                "perf_metric_name": f"{MAX_QUEUED_TASKS}_queued_time",
+                "perf_metric_value": queued_time,
+                "perf_metric_type": "LATENCY",
+            },
+            {
+                "perf_metric_name": f"{MAX_RAY_GET_SIZE}_large_object_time",
+                "perf_metric_value": large_object_time,
+                "perf_metric_type": "LATENCY",
+            },
+        ]
+        json.dump(results, out_file)

release/benchmarks/single_node_gce.yaml

@@ -0,0 +1,19 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-c
+
+advanced_instance_config:
+  instance_properties:
+    disks:
+      - boot: true
+        auto_delete: true
+        initialize_params:
+          disk_size_gb: 500
+
+head_node:
+    instance_type: n2-standard-64
+    resources:
+      # 128 GB
+      object_store_memory: 128000000000
+
+worker_nodes: []

release/cluster_tests/cpt_autoscaling_1-3_aws.yaml

@@ -0,0 +1,12 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.xlarge  # 4 CPUs
+    resources:
+      CPU: 4
+
+worker_nodes:
+    - instance_type: m5.xlarge
+      min_nodes: 0
+      max_nodes: 2
+      market_type: ON_DEMAND

release/cluster_tests/cpt_autoscaling_1-3_gce.yaml

@@ -0,0 +1,14 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-b
+
+head_node:
+    instance_type: n2-standard-4  # 4 CPUs
+    resources:
+      CPU: 4
+
+worker_nodes:
+    - instance_type: n2-standard-4
+      min_nodes: 0
+      max_nodes: 2
+      market_type: ON_DEMAND

release/cluster_tests/cpt_autoscaling_1-3_kuberay.yaml

@@ -0,0 +1,26 @@
+
+head_node_type:
+    name: head_node
+    instance_type: n2-standard-4  # 4 CPUs
+    resources:
+      limits:
+        cpu: "4"
+        memory: "16Gi"
+      requests:
+        cpu: "4"
+        memory: "16Gi"
+
+worker_node_types:
+    - name: worker_node
+      instance_type: n2-standard-4
+      resources:
+        limits:
+          cpu: "4"
+          memory: "16Gi"
+        requests:
+          cpu: "4"
+          memory: "16Gi"
+      min_workers: 0
+      max_workers: 2
+      use_spot: false
+autoscaler_version: v2

release/cluster_tests/workloads/tune_scale_up_down.py

@@ -0,0 +1,81 @@
+"""Test cluster up/down scaling behavior.
+
+This test should run on a cluster with autoscaling enabled. It assumes 1-3 nodes
+with 4 CPUs each.
+
+We start a Ray Tune run with 3 trials. Each trial uses 4 CPUs, so fills up a node
+completely. This means we will trigger autoscaling after starting up.
+
+The trial on the head node will run for 30 minutes. This is to make sure that
+we have enough time that the nodes for the other two trials come up, complete
+training, and come down before the first trial finishes.
+
+The other two trials will run once their nodes are up, and take 3 minutes each
+to finish. The three minutes have been chosen to make sure that both trials
+run in parallel for some time, i.e. to avoid that both additional trials run on
+only one node.
+
+We keep track of the number of nodes we observe at any point during the run.
+
+Test owner: krfricke
+
+Acceptance criteria: Should have scaled to 3 nodes at some point during the run.
+Should have scaled down to 1 node at the end.
+"""
+from collections import Counter
+import time
+
+import ray
+
+from ray import tune
+
+
+def train_fn(config):
+    this_node_ip = ray.util.get_node_ip_address()
+    if config["head_node_ip"] == this_node_ip:
+        # On the head node, run for 30 minutes
+        for i in range(30):
+            tune.report({"metric": i})
+            time.sleep(60)
+    else:
+        # On worker nodes, run for 3 minutes
+        for i in range(3):
+            tune.report({"metric": i})
+            time.sleep(60)
+
+
+class NodeCountCallback(tune.Callback):
+    def __init__(self):
+        self.node_counts = []
+
+    def on_step_begin(self, iteration, trials, **info):
+        node_count = len([n for n in ray.nodes() if n["Alive"]])
+        self.node_counts.append(node_count)
+
+
+def main():
+    ray.init()
+
+    head_node_ip = ray.util.get_node_ip_address()
+
+    assert (
+        len([n for n in ray.nodes() if n["Alive"]]) == 1
+    ), "Too many nodes available at start of script"
+
+    node_counter = NodeCountCallback()
+
+    tune.run(
+        train_fn,
+        num_samples=3,
+        config={"head_node_ip": head_node_ip},
+        callbacks=[node_counter],
+        resources_per_trial={"cpu": 4},
+    )
+
+    node_counts = Counter(node_counter.node_counts)
+    assert node_counts[3] > 0, "Cluster never scaled to 3 nodes"
+    assert node_counter.node_counts[-1] == 1, "Cluster didn't scale down to 1 node."
+
+
+if __name__ == "__main__":
+    main()

release/dashboard/agent_stress_compute.yaml

@@ -0,0 +1,15 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+advanced_instance_config:
+  TagSpecifications:
+    - ResourceType: "instance"
+      Tags:
+        - Key: ttl-hours
+          Value: '24'
+
+head_node:
+  instance_type: m5.16xlarge
+  resources:
+    CPU: 85
+
+worker_nodes: []

release/dashboard/agent_stress_compute_gce.yaml

@@ -0,0 +1,9 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones: [us-west1-c]
+
+head_node:
+  instance_type: n2-standard-64
+  resources:
+    CPU: 85
+
+worker_nodes: []

release/dashboard/mem_check.py

@@ -0,0 +1,93 @@
+import argparse
+import json
+import os
+import time
+
+import ray
+from ray._private.memory_monitor import MemoryMonitor, get_top_n_memory_usage
+from ray._private.test_utils import get_system_metric_for_component
+from ray.dashboard.modules.metrics.metrics_head import (
+    DEFAULT_PROMETHEUS_HOST,
+    PROMETHEUS_HOST_ENV_VAR,
+)
+from ray.job_submission import JobStatus, JobSubmissionClient
+
+# Initialize ray to avoid autosuspend.
+addr = ray.init()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--working-dir",
+        required=True,
+        help="working_dir to use for the job within this test.",
+    )
+    args = parser.parse_args()
+    client = JobSubmissionClient("http://127.0.0.1:8265")
+    job_id = client.submit_job(
+        # Entrypoint shell command to execute
+        entrypoint="python workload.py",
+        runtime_env={"working_dir": args.working_dir},
+    )
+    print(job_id)
+
+    # If using a remote cluster, replace 127.0.0.1 with the head node's IP address.
+    client = JobSubmissionClient("http://127.0.0.1:8265")
+    m = MemoryMonitor()
+    start = time.time()
+    # Run for 3 hours
+    initial_used_gb = m.get_memory_usage()[0]
+
+    terminal_states = {JobStatus.SUCCEEDED, JobStatus.STOPPED, JobStatus.FAILED}
+
+    while time.time() - start < 3600 * 3:
+        print(f"{round((time.time() - start) / 60, 2)}m passed...")
+        m.raise_if_low_memory()
+        used_gb = m.get_memory_usage()[0]
+        print("Used GB: ", used_gb)
+        print(get_top_n_memory_usage())
+        print("\n\n")
+
+        # Terminate the test if the job is failed.
+        status = client.get_job_status(job_id)
+        print(f"Job status: {status}")
+        if status in terminal_states:
+            break
+        time.sleep(15)
+
+    ending_used_gb = m.get_memory_usage()[0]
+
+    mem_growth = ending_used_gb - initial_used_gb
+    top_n_mem_usage = get_top_n_memory_usage()
+    print(top_n_mem_usage)
+    print(f"Memory growth: {mem_growth} GB")
+
+    if status == JobStatus.FAILED or status == JobStatus.STOPPED:
+        print(client.get_job_logs(job_id))
+        assert False, "Job has failed."
+
+    uss_bytes_for_agent_component = get_system_metric_for_component(
+        "ray_component_uss_bytes",
+        "agent",
+        os.environ.get(PROMETHEUS_HOST_ENV_VAR, DEFAULT_PROMETHEUS_HOST),
+    )
+    assert (
+        len(uss_bytes_for_agent_component) > 0
+    ), "Agent component memory metrics are not found."
+    for bytes in uss_bytes_for_agent_component:
+        print(f"Agent component memory usage: {bytes} bytes")
+        assert bytes < 500 * 1024 * 1024, "Agent component memory usage is too high."
+
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as f:
+        results = {
+            "memory_growth_gb": mem_growth,
+        }
+        results["perf_metrics"] = [
+            {
+                "perf_metric_name": "memory_growth_gb",
+                "perf_metric_value": mem_growth,
+                "perf_metric_type": "LATENCY",
+            }
+        ]
+
+        f.write(json.dumps(results))

release/dashboard/workload.py

@@ -0,0 +1,19 @@
+import time
+
+import ray
+
+ray.init("auto")
+
+
+@ray.remote(num_cpus=1)
+class A:
+    def f(self):
+        return 1
+
+
+actors = [A.remote() for _ in range(85)]
+
+# Keep calling actor methods which will generate lots of metrics.
+while True:
+    time.sleep(0.1)
+    ray.get([actor.f.remote() for actor in actors])

release/gcloud_docker_login.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+# This script is used to login to gcloud docker registry using GCP workload identity
+# federation service account
+
+set -euo pipefail
+
+# Only install gcloud if not already available
+if ! command -v gcloud &> /dev/null; then
+  ARCH=$(uname -m)
+  case "$ARCH" in
+      x86_64)
+          GCLOUD_ARCH="x86_64"
+          ;;
+      aarch64|arm64)
+          GCLOUD_ARCH="arm"
+          ;;
+      *)
+          echo "Unsupported architecture: $ARCH"
+          exit 1
+          ;;
+  esac
+
+  GCLOUD_VERSION="550.0.0"
+  GCLOUD_TARBALL="google-cloud-cli-${GCLOUD_VERSION}-linux-${GCLOUD_ARCH}.tar.gz"
+
+  curl -fO "https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/${GCLOUD_TARBALL}"
+  tar -xf "$GCLOUD_TARBALL"
+  ./google-cloud-sdk/install.sh -q
+  PATH="$(pwd)/google-cloud-sdk/bin:$PATH"
+  export PATH
+fi
+
+gcloud auth login --cred-file="$1" --quiet
+gcloud auth configure-docker us-west1-docker.pkg.dev --quiet

release/golden_notebook_tests/compute_tpl.yaml

@@ -0,0 +1,15 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west-2
+
+max_workers: 3
+
+head_node_type:
+    name: head_node
+    instance_type: m5.4xlarge
+
+worker_node_types:
+    - name: worker_node
+      instance_type: m5.4xlarge
+      min_workers: 3
+      max_workers: 3
+      use_spot: false

release/golden_notebook_tests/gpu_tpl_aws.yaml

@@ -0,0 +1,18 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west-2
+
+max_workers: 2
+
+head_node_type:
+    name: head_node
+    instance_type: m5.xlarge
+
+worker_node_types:
+    - name: worker_node
+      instance_type: g4dn.12xlarge
+      min_workers: 2
+      max_workers: 2
+      use_spot: true
+      resources:
+        custom_resources:
+          worker: 1

release/golden_notebook_tests/gpu_tpl_gce.yaml

@@ -0,0 +1,17 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west1
+allowed_azs:
+    - us-west1-b
+
+max_workers: 2
+
+head_node_type:
+    name: head_node
+    instance_type: n1-standard-4
+
+worker_node_types:
+    - name: worker_node
+      instance_type: n1-standard-32-nvidia-tesla-t4-2
+      min_workers: 2
+      max_workers: 2
+      use_spot: true

release/golden_notebook_tests/workloads/torch_tune_serve_test.py

@@ -0,0 +1,302 @@
+import argparse
+import atexit
+import json
+import os
+import tempfile
+import time
+import subprocess
+
+import ray
+from ray.train import Checkpoint, ScalingConfig, RunConfig
+from ray.tune.tune_config import TuneConfig
+import requests
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from filelock import FileLock
+from ray import serve, tune, train
+from ray.train.torch import TorchTrainer
+from ray.tune import Tuner
+from torch.utils.data import DataLoader, Subset
+from torchvision.datasets import MNIST
+from torchvision.models import resnet18
+
+
+def load_mnist_data(train: bool, download: bool):
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+    )
+
+    with FileLock(os.path.expanduser("~/.ray.lock")):
+        return MNIST(
+            root=os.path.expanduser("~/data"),
+            train=train,
+            download=download,
+            transform=transform,
+        )
+
+
+def train_epoch(epoch, dataloader, model, loss_fn, optimizer):
+    if ray.train.get_context().get_world_size() > 1:
+        dataloader.sampler.set_epoch(epoch)
+
+    for X, y in dataloader:
+        # Compute prediction error
+        pred = model(X)
+        loss = loss_fn(pred, y)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+
+def validate_epoch(dataloader, model, loss_fn):
+    num_batches = len(dataloader)
+    model.eval()
+    loss = 0
+    with torch.no_grad():
+        for X, y in dataloader:
+            pred = model(X)
+            loss += loss_fn(pred, y).item()
+    loss /= num_batches
+    result = {"val_loss": loss}
+    return result
+
+
+def training_loop(config):
+    # Create model.
+    model = resnet18()
+    model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=1, padding=3, bias=False)
+    model = train.torch.prepare_model(model)
+
+    # Create optimizer.
+    optimizer = torch.optim.SGD(
+        model.parameters(),
+        lr=config.get("lr", 0.1),
+        momentum=config.get("momentum", 0.9),
+    )
+
+    # Load in training and validation data.
+    train_dataset = load_mnist_data(True, True)
+    validation_dataset = load_mnist_data(False, False)
+
+    if config["test_mode"]:
+        train_dataset = Subset(train_dataset, list(range(64)))
+        validation_dataset = Subset(validation_dataset, list(range(64)))
+
+    train_loader = DataLoader(
+        train_dataset, batch_size=config["batch_size"], num_workers=2, shuffle=True
+    )
+    validation_loader = DataLoader(
+        validation_dataset, batch_size=config["batch_size"], num_workers=2
+    )
+
+    train_loader = train.torch.prepare_data_loader(train_loader)
+    validation_loader = train.torch.prepare_data_loader(validation_loader)
+
+    # Create loss.
+    criterion = nn.CrossEntropyLoss()
+
+    for epoch_idx in range(2):
+        train_epoch(epoch_idx, train_loader, model, criterion, optimizer)
+        validation_loss = validate_epoch(validation_loader, model, criterion)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(model.module.state_dict(), os.path.join(tmpdir, "model.pt"))
+            train.report(validation_loss, checkpoint=Checkpoint.from_directory(tmpdir))
+
+
+def train_mnist(test_mode=False, num_workers=1, use_gpu=False):
+    trainer = TorchTrainer(
+        training_loop,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+    )
+    tuner = Tuner(
+        trainer,
+        param_space={
+            "train_loop_config": {
+                "lr": tune.grid_search([1e-4, 1e-3]),
+                "test_mode": test_mode,
+                "batch_size": 128,
+            }
+        },
+        tune_config=TuneConfig(
+            metric="val_loss",
+            mode="min",
+            num_samples=1,
+        ),
+        run_config=RunConfig(
+            verbose=1,
+            storage_path=(
+                "/mnt/cluster_storage"
+                if os.path.exists("/mnt/cluster_storage")
+                else None
+            ),
+        ),
+    )
+
+    return tuner.fit()
+
+
+def get_model(checkpoint_dir: str):
+    model = resnet18()
+    model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=1, padding=3, bias=False)
+
+    model_state_dict = torch.load(
+        os.path.join(checkpoint_dir, "model.pt"), map_location="cpu"
+    )
+    model.load_state_dict(model_state_dict)
+
+    return model
+
+
+@serve.deployment(name="mnist")
+class MnistDeployment:
+    def __init__(self, model):
+        use_cuda = torch.cuda.is_available()
+        self.device = torch.device("cuda" if use_cuda else "cpu")
+        model = model.to(self.device)
+        self.model = model
+
+    async def __call__(self, request):
+        json_input = await request.json()
+        prediction = await self.my_batch_handler(json_input["image"])
+        return {"result": prediction.cpu().numpy().tolist()}
+
+    @serve.batch(max_batch_size=4, batch_wait_timeout_s=1)
+    async def my_batch_handler(self, images):
+        print(f"Processing request with batch size {len(images)}.")
+        image_tensors = torch.tensor(images)
+        image_tensors = image_tensors.to(self.device)
+        outputs = self.model(image_tensors)
+        predictions = torch.max(outputs.data, 1)[1]
+        return predictions
+
+
+def setup_serve(model, use_gpu: bool = False):
+    serve.start(
+        http_options={"location": "EveryNode"}
+    )  # Start on every node so `predict` can hit localhost.
+    serve.run(
+        MnistDeployment.options(
+            num_replicas=2,
+            ray_actor_options={"num_gpus": 1, "resources": {"worker": 1}}
+            if use_gpu
+            else {},
+        ).bind(model),
+        route_prefix="/mnist",
+    )
+
+
+@ray.remote
+def predict_and_validate(index, image, label):
+    def predict(image):
+        response = requests.post(
+            "http://localhost:8000/mnist", json={"image": image.numpy().tolist()}
+        )
+        try:
+            return response.json()["result"]
+        except:  # noqa: E722
+            return -1
+
+    prediction = predict(image)
+    print(
+        "Querying model with example #{}. "
+        "Label = {}, Prediction = {}, Correct = {}".format(
+            index, label, prediction, label == prediction
+        )
+    )
+    return prediction
+
+
+def test_predictions(test_mode=False):
+    # Load in data
+    dataset = load_mnist_data(False, True)
+    num_to_test = 10 if test_mode else 1000
+    filtered_dataset = [dataset[i] for i in range(num_to_test)]
+    images, labels = zip(*filtered_dataset)
+
+    # Remote function calls are done here for parallelism.
+    # As a byproduct `predict` can hit localhost.
+    predictions = ray.get(
+        [
+            predict_and_validate.remote(i, images[i], labels[i])
+            for i in range(num_to_test)
+        ]
+    )
+
+    correct = 0
+    for label, prediction in zip(labels, predictions):
+        if label == prediction:
+            correct += 1
+
+    print(
+        "Labels = {}. Predictions = {}. {} out of {} are correct.".format(
+            list(labels), predictions, correct, num_to_test
+        )
+    )
+
+    return correct / float(num_to_test)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+    args = parser.parse_args()
+
+    if os.environ.get("IS_SMOKE_TEST"):
+        args.smoke_test = True
+        proc = subprocess.Popen(["ray", "start", "--head"])
+        proc.wait()
+        os.environ["RAY_ADDRESS"] = "ray://localhost:10001"
+
+        def stop_ray():
+            subprocess.Popen(["ray", "stop", "--force"]).wait()
+
+        atexit.register(stop_ray)
+
+    start = time.time()
+
+    addr = os.environ.get("RAY_ADDRESS")
+    job_name = os.environ.get("RAY_JOB_NAME", "torch_tune_serve_test")
+    if addr is not None and addr.startswith("anyscale://"):
+        client = ray.init(address=addr, job_name=job_name)
+    else:
+        client = ray.init()
+
+    num_workers = 2
+    use_gpu = not args.smoke_test
+
+    print("Training model.")
+    result_grid = train_mnist(args.smoke_test, num_workers, use_gpu)
+
+    print("Retrieving best model.")
+    best_result = result_grid.get_best_result()
+    best_checkpoint = best_result.get_best_checkpoint(metric="val_loss", mode="min")
+    with best_checkpoint.as_directory() as checkpoint_dir:
+        model = get_model(checkpoint_dir)
+
+    print("Setting up Serve.")
+    setup_serve(model, use_gpu)
+
+    print("Testing Prediction Service.")
+    accuracy = test_predictions(args.smoke_test)
+
+    taken = time.time() - start
+    result = {
+        "time_taken": taken,
+        "accuracy": accuracy,
+    }
+    test_output_json = os.environ.get(
+        "TEST_OUTPUT_JSON", "/tmp/torch_tune_serve_test.json"
+    )
+    with open(test_output_json, "wt") as f:
+        json.dump(result, f)
+
+    print("Test Successful!")

release/golden_notebook_tests/workloads/util.py

@@ -0,0 +1,17 @@
+import os.path
+from pathlib import Path
+import importlib.util
+
+
+def import_and_execute_test_script(relative_path_to_test_script: str):
+    """Imports and executes a module from a path relative to Ray repo root."""
+    # get the ray folder
+    ray_path = Path(
+        os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+    )
+    notebook_path = ray_path.joinpath(relative_path_to_test_script)
+    assert notebook_path.exists()
+
+    spec = importlib.util.spec_from_file_location("notebook_test", notebook_path)
+    notebook_test_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(notebook_test_module)

release/hello_world_tests/hello_world.py

@@ -0,0 +1,14 @@
+import ray
+
+
+@ray.remote
+def hello_world():
+    return "Hello, world!"
+
+
+def main():
+    print(ray.get(hello_world.remote()))
+
+
+if __name__ == "__main__":
+    main()

release/hello_world_tests/hello_world_compute_config.yaml

@@ -0,0 +1,6 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.xlarge
+
+worker_nodes: []

release/hello_world_tests/hello_world_compute_config_azure.yaml

@@ -0,0 +1,6 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: 4CPU-16GB
+
+worker_nodes: []

release/hello_world_tests/hello_world_compute_config_gce.yaml

@@ -0,0 +1,8 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-c
+
+head_node:
+    instance_type: n2-standard-4
+
+worker_nodes: []

release/hello_world_tests/hello_world_emoji.py

@@ -0,0 +1,15 @@
+import ray
+import emoji
+
+
+@ray.remote
+def hello_world_emoji():
+    return emoji.emojize(":globe_showing_Americas:")
+
+
+def main():
+    print(ray.get(hello_world_emoji.remote()))
+
+
+if __name__ == "__main__":
+    main()