code: 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)
  
file_path: ray/release/air_tests/air_benchmarks/workloads/pytorch_training_e2e.py
imports: import click
  import time
  import json
  import os
  import tempfile
  from typing import Dict
  import numpy
  from torchvision import transforms
  from torchvision.models import resnet18
  import torch
  import torch.nn
  import torch.optim
  import ray
  from ray import train
  from ray.train import Checkpoint, RunConfig, ScalingConfig
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 4 for calling a global function with example usage: main(data_size_gb, num_epochs, num_workers, smoke_test) without return types