code: def main():
      args, unknown = parse_script_args()
      logging.info("Received arguments: {}".format(args))
  
      # Create test spec
      test_spec = TestSpec(
          num_workers=args.num_workers,
          worker_obj_store_size_in_gb=args.worker_obj_store_size_in_gb,
          error_rate=args.error_rate,
          trigger_object_spill=args.trigger_object_spill,
      )
      logging.info("Created test spec: {}".format(test_spec))
  
      # Create the data save path if it doesn't exist.
      data_save_path = args.data_save_path
      if not os.path.exists(data_save_path):
          os.makedirs(data_save_path, mode=0o777, exist_ok=True)
  
      # Lazily construct Xarrays
      xarray_filename_pairs = lazy_create_xarray_filename_pairs(
          test_spec
      )
  
      # Connect to the Ray cluster
      ray.init(address="auto")
      monitor_actor = monitor_memory_usage()
  
      # Save all the Xarrays to disk; this will trigger
      # Dask computations on Ray.
      logging.info(
          "Saving {} xarrays..".format(len(xarray_filename_pairs))
      )
      SaveRoutines.save_all_xarrays(
          xarray_filename_pairs=xarray_filename_pairs,
          dirpath=data_save_path,
          batch_size=test_spec.batch_size,
          ray_scheduler=ray_dask_get,
      )
      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}")
      try:
          print(
              ray._private.internal_api.memory_summary(stats_only=True)
          )
      except Exception as e:
          print(f"Warning: query memory summary failed: {e}")
      with open(os.environ["TEST_OUTPUT_JSON"], "w") as f:
          f.write(
              json.dumps(
                  {
                      "success": 1,
                      "_peak_memory": round(used_gb, 2),
                      "_peak_process_memory": usage,
                  }
              )
          )
  
imports: import argparse
  import json
  import logging
  import math
  import os
  import random
  import time
  from copy import copy, deepcopy
  from typing import List, Tuple
  import dask.array
  import numpy
  import xarray
  import ray
  from ray._private.test_utils import monitor_memory_usage
  from ray.util.dask import ray_dask_get
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types