Change gpu_batch_predict release test benchmark
diff --git a/doc/source/ray-air/benchmarks.rst b/doc/source/ray-air/benchmarks.rst
index 4c91161d9d81a..b5e30f905c211 100644
--- a/doc/source/ray-air/benchmarks.rst
+++ b/doc/source/ray-air/benchmarks.rst
@@ -113,8 +113,8 @@ using a Pytorch pre-trained ResNet model.
 We test out the performance across different cluster sizes and data sizes.
 
 - `GPU image batch prediction script`_
-- `GPU training small cluster configuration`_
-- `GPU training large cluster configuration`_
+- `GPU prediction small cluster configuration`_
+- `GPU prediction large cluster configuration`_
 
 .. list-table::
 
@@ -122,17 +122,17 @@ We test out the performance across different cluster sizes and data sizes.
       - **Data Size**
       - **Performance**
       - **Command**
-    * - 1 g3.8xlarge node
+    * - 1 g4dn.8xlarge node
       - 1 GB (1623 images)
-      - 72.59 s (22.3 images/sec)
+      - 46.12 s (35.19 images/sec)
       - `python gpu_batch_prediction.py --data-size-gb=1`
-    * - 1 g3.8xlarge node
+    * - 1 g4dn.8xlarge node
       - 20 GB (32460 images)
-      - 1213.48 s (26.76 images/sec)
+      - 285.2 s (113.81 images/sec)
       - `python gpu_batch_prediction.py --data-size-gb=20`
-    * - 4 g3.16xlarge nodes
+    * - 4 g4dn.12xlarge nodes
       - 100 GB (162300 images)
-      - 885.98 s (183.19 images/sec)
+      - 304.01 s (533.86 images/sec)
       - `python gpu_batch_prediction.py --data-size-gb=100`
 
 .. _pytorch_gpu_training_benchmark:
@@ -151,7 +151,7 @@ We test out the performance across different cluster sizes and data sizes.
 
 .. note::
 
-    For multi-host distributed training, on AWS we need to ensure ec2 instances are in the same VPC and 
+    For multi-host distributed training, on AWS we need to ensure ec2 instances are in the same VPC and
     all ports are open in the secure group.
 
 
@@ -259,6 +259,8 @@ overhead of a few seconds for both methods that is negligible for longer trainin
 .. _`XGBoost Cluster Configuration`: https://github.com/ray-project/ray/blob/a241e6a0f5a630d6ed5b84cce30c51963834d15b/release/air_tests/air_benchmarks/xgboost_compute_tpl.yaml#L6-L24
 .. _`GPU image batch prediction script`: https://github.com/ray-project/ray/blob/cec82a1ced631525a4d115e4dc0c283fa4275a7f/release/air_tests/air_benchmarks/workloads/gpu_batch_prediction.py#L18-L49
 .. _`GPU image training script`: https://github.com/ray-project/ray/blob/cec82a1ced631525a4d115e4dc0c283fa4275a7f/release/air_tests/air_benchmarks/workloads/pytorch_training_e2e.py#L95-L106
+.. _`GPU prediction small cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_gpu_1_g4_8xl.yaml#L6-L15
+.. _`GPU prediction large cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_gpu_4_g4_12xl.yaml#L6-L15
 .. _`GPU training small cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_gpu_1.yaml#L6-L24
 .. _`GPU training large cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_gpu_16.yaml#L5-L25
 .. _`Pytorch comparison training script`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/workloads/torch_benchmark.py
diff --git a/release/air_tests/air_benchmarks/compute_gpu_1_g4_8xl.yaml b/release/air_tests/air_benchmarks/compute_gpu_1_g4_8xl.yaml
new file mode 100644
index 0000000000000..ba1bf7e7432d7
--- /dev/null
+++ b/release/air_tests/air_benchmarks/compute_gpu_1_g4_8xl.yaml
@@ -0,0 +1,15 @@
+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:
+    - name: worker_node
+      instance_type: m5.2xlarge
+      max_workers: 0
+      min_workers: 0
+      use_spot: false
diff --git a/release/air_tests/air_benchmarks/workloads/gpu_batch_prediction.py b/release/air_tests/air_benchmarks/workloads/gpu_batch_prediction.py
index fc5092492aacb..7d413006ceb71 100644
--- a/release/air_tests/air_benchmarks/workloads/gpu_batch_prediction.py
+++ b/release/air_tests/air_benchmarks/workloads/gpu_batch_prediction.py
@@ -3,8 +3,8 @@
 import json
 import os
 import numpy as np
-import pandas as pd
 
+import torch
 from torchvision import transforms
 from torchvision.models import resnet18
 
@@ -14,19 +14,20 @@
 from ray.data.preprocessors import BatchMapper
 
 
-def preprocess(batch: np.ndarray) -> pd.DataFrame:
+def preprocess(image_batch: np.ndarray) -> np.ndarray:
     """
-    User Pytorch code to transform user image.
+    User Pytorch code to transform user image with outer dimension of batch size.
     """
     preprocess = transforms.Compose(
         [
-            transforms.ToTensor(),
-            transforms.Resize(256),
+            # Torchvision's ToTensor does not accept outer batch dimension
             transforms.CenterCrop(224),
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
         ]
     )
-    return pd.DataFrame({"image": [preprocess(image) for image in batch]})
+    # Outer dimension is batch size such as (10, 256, 256, 3) -> (10, 3, 256, 256)
+    transposed_torch_tensor = torch.Tensor(image_batch.transpose(0, 3, 1, 2))
+    return preprocess(transposed_torch_tensor).numpy()
 
 
 @click.command(help="Run Batch prediction on Pytorch ResNet models.")
@@ -58,7 +59,6 @@ def main(data_size_gb: int, smoke_test: bool = False):
     predictor.predict(
         dataset,
         num_gpus_per_worker=int(not smoke_test),
-        feature_columns=["image"],
         batch_size=512,
     )
     total_time_s = round(time.time() - start, 2)
diff --git a/release/release_tests.yaml b/release/release_tests.yaml
index 29f4f6450b5d5..ae14d91092eb3 100644
--- a/release/release_tests.yaml
+++ b/release/release_tests.yaml
@@ -304,7 +304,7 @@
 
   cluster:
     cluster_env: app_config.yaml
-    cluster_compute: compute_gpu_1.yaml
+    cluster_compute: compute_gpu_1_g4_8xl.yaml
 
   run:
     timeout: 3600
@@ -316,30 +316,28 @@
   alert: default
 
 
-# TODO(jiaodong): This leads to CUDA OOM on release tool but succeeds on
-# workspaces, re-enabled it once inline conversions are fixed
-# - name: air_benchmark_torch_batch_prediction_gpu_4x4_100gb
-#   group: AIR tests
-#   working_dir: air_tests/air_benchmarks
+- name: air_benchmark_torch_batch_prediction_gpu_4x4_100gb
+  group: AIR tests
+  working_dir: air_tests/air_benchmarks
 
-#   frequency: nightly
-#   team: ml
-#   env: staging
+  frequency: nightly
+  team: ml
+  env: staging
 
-#   stable: false
+  stable: false
 
-#   cluster:
-#     cluster_env: app_config.yaml
-#     cluster_compute: compute_gpu_16.yaml
+  cluster:
+    cluster_env: app_config.yaml
+    cluster_compute: compute_gpu_4_g4_12xl.yaml
 
-#   run:
-#     timeout: 10800
-#     script: python workloads/gpu_batch_prediction.py --data-size-gb 100
+  run:
+    timeout: 10800
+    script: python workloads/gpu_batch_prediction.py --data-size-gb 100
 
-#     type: sdk_command
-#     file_manager: job
+    type: sdk_command
+    file_manager: job
 
-#   alert: default
+  alert: default
 
 
 - name: air_benchmark_torch_mnist_cpu_4x4
@@ -1796,7 +1794,7 @@
 
   run:
     timeout: 86400
-    prepare: ray stop 
+    prepare: ray stop
     script: python workloads/actor_deaths.py
     long_running: true
     type: sdk_command
@@ -1886,7 +1884,7 @@
 
   run:
     timeout: 86400
-    prepare: ray stop 
+    prepare: ray stop
     script: python workloads/many_actor_tasks.py
     long_running: true
     type: sdk_command
@@ -1915,7 +1913,7 @@
 
   run:
     timeout: 86400
-    prepare: ray stop 
+    prepare: ray stop
     script: python workloads/many_drivers.py --iteration-num=4000
     long_running: true
     type: sdk_command
@@ -1978,7 +1976,7 @@
 
   run:
     timeout: 86400
-    prepare: ray stop 
+    prepare: ray stop
     script: python workloads/many_tasks.py
     long_running: true
     type: sdk_command
@@ -2007,7 +2005,7 @@
 
   run:
     timeout: 86400
-    prepare: ray stop 
+    prepare: ray stop
     script: python workloads/many_tasks_serialized_ids.py
     long_running: true
     type: sdk_command
@@ -2036,7 +2034,7 @@
 
   run:
     timeout: 86400
-    prepare: ray stop 
+    prepare: ray stop
     script: python workloads/node_failures.py
     long_running: true
     type: sdk_command
@@ -2127,7 +2125,7 @@
 
   run:
     timeout: 86400
-    prepare: ray stop 
+    prepare: ray stop
     script: python workloads/serve_failure.py
     long_running: true
     type: sdk_command
@@ -4837,7 +4835,7 @@
 
 - name: aws_cluster_launcher
   group: cluster-launcher-test
-  working_dir: cluster_tests 
+  working_dir: cluster_tests
   legacy:
     test_name: aws_cluster_launcher
     test_suite: cluster-launcher-test

## Why are these changes needed?

This should use latest numpy udf with numpy prediction now, also update single column selection code in new numpy predictor as it's not necessary.

Also this PR changes 
 - g3 compute template with g4dn 
 - re-enabled 100gb gpu prediction test -- we can now run it faster with less memory footprint !

```

4 x g4dn.12xlarge, 100GB data

'total_time_s', 'perf_metric_value': 304.01
'throughout_MB_s', 'perf_metric_value': 336.83102529522057

1 x g4dn.8xlarge, 20GB data

'total_time_s', 'perf_metric_value': 285.2
'throughout_MB_s', 'perf_metric_value': 71.80925666199158

1 x g4dn.8xlarge, 1GB data

'total_time_s', 'perf_metric_value': 46.12, 
'perf_metric_name': 'throughout_MB_s', 'perf_metric_value': 22.20294882914137
```

new runs on CI:

20GB image data https://buildkite.com/ray-project/release-tests-pr/builds/21577#01848692-1ff0-4ec5-bb83-7fdea0ca4e89
100GB image data https://buildkite.com/ray-project/release-tests-pr/builds/21577#01848692-1ff3-4f18-852f-9f660dbd1487


## Related issue number

<!-- For example: "Closes #1234" -->

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [x] I've run `scripts/format.sh` to lint the changes in this PR.
- [x] I've included any doc changes needed for https://docs.ray.io/en/master/.
- [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [x] Release tests
   - [ ] This PR is not tested :(
