code: def test_trainer(
      root_path_storage_filesystem_label, tmp_path, monkeypatch
  ):
      """Tests that a data parallel trainer can save and restore checkpoints to
      various storage types properly. Also records checkpoint save/restore timing.
  
      Here's the rundown of what this test does:
      1. Passes in a `custom_save_fn` and `custom_restore_fn` to the trainer to
         record how long the operations take, as well as save a large checkpoint.
         See `create_checkpoint` for details on the checkpoint contents.
      2. Configures the training loop to fail 3 times.
      3. Runs the trainer, which will fail 2 times and recover via FailureConfig.
         This first run will exit on the 3rd failure.
      4. Manually restores the trainer, which will restore from the 3rd failure and
         run to completion.
      5. Downloads the results from the storage path and asserts that the contents
         are all correct. See `ray.train.test_new_persistence` for the expected filetree.
      6. Tests a new run with `resume_from_checkpoint`.
      """
      monkeypatch.setenv(
          "RAY_AIR_LOCAL_CACHE_DIR", str(tmp_path / "ray_results")
      )
  
      ray.init(
          runtime_env={"working_dir": "."}, ignore_reinit_error=True
      )
  
      (
          root_path,
          storage_filesystem,
          label,
      ) = root_path_storage_filesystem_label
      storage_path = root_path + label
      checkpoint_config = train.CheckpointConfig(
          num_to_keep=TestConstants.NUM_ITERATIONS // 2
      )
      exp_name = "test_trainer"
  
      print(
          "\nSaving results under (storage_path, exp_name) = "
          f"({storage_path}, {exp_name})\n"
      )
  
      trainer = TorchTrainer(
          train_fn,
          train_loop_config={
              "in_trainer": True,
              "fail_iters": [3, 6, 8],
              "time_per_iter": 1.0,
              "num_iterations": TestConstants.NUM_ITERATIONS,
              "custom_save_fn": custom_save_fn,
              "custom_restore_fn": custom_restore_fn,
          },
          scaling_config=train.ScalingConfig(
              num_workers=TestConstants.NUM_WORKERS,
              trainer_resources={"CPU": 0},
              resources_per_worker={
                  "CPU": TestConstants.NUM_CPUS_PER_WORKER
              },
          ),
          run_config=train.RunConfig(
              failure_config=train.FailureConfig(max_failures=2),
              name=exp_name,
              storage_path=storage_path,
              storage_filesystem=storage_filesystem,
              checkpoint_config=checkpoint_config,
              sync_config=train.SyncConfig(sync_artifacts=True),
          ),
      )
      print("\nStarting initial run.\n")
      with pytest.raises(TrainingFailedError):
          result = trainer.fit()
  
      print("\nStarting manually restored run.\n")
      restored_trainer = TorchTrainer.restore(
          path=str(URI(storage_path) / exp_name),
          storage_filesystem=storage_filesystem,
      )
      result = restored_trainer.fit()
      print(result)
  
      print("\nAsserting contents of uploaded results.\n")
      local_inspect_dir = tmp_path / "inspect_dir"
      local_inspect_dir.mkdir()
      # Download the results from storage
      if "cloud" in label:
          # NOTE: Use the CLI to download, since the python libraries
          # (pyarrow, fsspec) aren't consistent across cloud platforms (s3, gs).
          cloud_uri = CLOUD_TEST_DIR + label
          print("\nDownloading from cloud URI:", cloud_uri, "\n")
          download_from_uri(cloud_uri, str(local_inspect_dir))
      elif label == "nfs":
          local_inspect_dir = Path(storage_path)
      else:
          raise NotImplementedError(f"Invalid storage type: {label}")
  
      _assert_storage_contents(
          local_inspect_dir,
          exp_name,
          checkpoint_config,
          "TorchTrainer",
          test_trainer=True,
          constants=TestConstants,
      )
  
      # Test `resume_from_checkpoint`
      _resume_from_checkpoint(
          result.checkpoint,
          expected_state={"iter": TestConstants.NUM_ITERATIONS - 1},
          storage_path=storage_path,
          storage_filesystem=storage_filesystem,
      )
  
      # Upload checkpoint save and restore timing release test metrics
      all_checkpoint_timing_metrics = collections.defaultdict(list)
      for checkpoint, _ in result.best_checkpoints:
          metadata = checkpoint.get_metadata()
          for metric, value in metadata.items():
              all_checkpoint_timing_metrics[metric].append(value)
  
      aggregated_metrics = {
          key: np.mean(values)
          for key, values in all_checkpoint_timing_metrics.items()
      }
      checkpoint_size_mb = (
          TestConstants.NUM_GB * 1000
          + TestConstants.NUM_MB
          + TestConstants.NUM_KB / 1000
      )
      speeds = {
          key + "_speed_mbps": checkpoint_size_mb / time_s
          for key, time_s in aggregated_metrics.items()
      }
      # Add units as the suffix
      aggregated_metrics = {
          key + "_avg_s": time_s
          for key, time_s in aggregated_metrics.items()
      }
      aggregated_metrics.update(speeds)
      aggregated_metrics["checkpoint_size_mb"] = checkpoint_size_mb
  
      print(aggregated_metrics)
      update_output_json(flatten_dict({label: aggregated_metrics}))
  
      print("Deleting files from the run...")
      if "cloud" in label:
          # NOTE: Use the CLI to delete files on cloud, since the python libraries
          # (pyarrow, fsspec) aren't consistent across cloud platforms (s3, gs).
          delete_at_uri(CLOUD_TEST_DIR)
      elif label == "nfs":
          shutil.rmtree(NFS_TEST_DIR, ignore_errors=True)
      else:
          raise NotImplementedError(f"Invalid storage type: {label}")
  
imports: import collections
  from contextlib import contextmanager
  from datetime import datetime
  import json
  import os
  from pathlib import Path
  import pickle
  import shutil
  import subprocess
  import time
  from typing import Any, Dict
  import uuid
  import fsspec
  import numpy
  import pyarrow.fs
  import pytest
  import torch
  import torch.distributed
  import ray
  from ray import train
  from ray._private.dict import flatten_dict
  from ray.air.constants import TRAINING_ITERATION
  from ray.air._internal.uri_utils import URI
  from ray.train import Checkpoint
  from ray.train.base_trainer import TrainingFailedError
  from ray.train.torch import TorchTrainer
  from test_new_persistence import train_fn, _assert_storage_contents, _resume_from_checkpoint
label_desc: Use image node_id 9 for calling a global function with example usage: test_trainer(root_path_storage_filesystem_label, tmp_path, monkeypatch) without return types