code: def test_restore_retry(
      ray_start_2_cpus, tmpdir, monkeypatch, retry_num
  ):
      """Test retrying restore on a trial level by setting `TUNE_RESTORE_RETRY_NUM`."""
  
      class MockTrainable(Trainable):
          """A trainable that can generate one failure during training and
          another `config["retry_num_to_fail"]` times during restoring.
          """
  
          def setup(self, config):
              self.idx = 0
              self.tag_file_path = config["tag_file_path"]
              self.retry_num_to_fail = config.get(
                  "retry_num_to_fail", 2
              )
              self._is_restored = False
  
          def step(self):
              time.sleep(1)
              if self.idx == 0 and self._is_restored:
                  raise RuntimeError(
                      "===== Restored trial cannot start from scratch ====="
                  )
              elif self.idx == 2 and not self._is_restored:
                  raise RuntimeError(
                      "===== First run fails at idx=2 ====="
                  )
              self.idx += 1
              return {"score": self.idx}
  
          def save_checkpoint(self, checkpoint_dir):
              path = os.path.join(checkpoint_dir, "checkpoint")
              with open(path, "w") as f:
                  f.write(json.dumps({"idx": self.idx}))
  
          def load_checkpoint(self, checkpoint_dir):
              self._is_restored = True
              with open(self.tag_file_path, "r") as f:
                  retried_num = json.loads(f.read())["retried_num"]
  
              with open(self.tag_file_path, "w") as f:
                  f.write(json.dumps({"retried_num": retried_num + 1}))
  
              if retried_num < self.retry_num_to_fail:
                  raise RuntimeError(
                      f"===== Failing restore #{retried_num + 1} ====="
                  )
              with open(
                  os.path.join(checkpoint_dir, "checkpoint"), "r"
              ) as f:
                  self.idx = json.loads(f.read())["idx"]
  
      # Set environment variable just for this test
      with unittest.mock.patch.dict(
          os.environ, {"TUNE_RESTORE_RETRY_NUM": str(retry_num)}
      ):
          tag_file = os.path.join(tmpdir, "tag")
          # set up tag file
          with open(tag_file, "w") as f:
              f.write(json.dumps({"retried_num": 0}))
          tuner = Tuner(
              MockTrainable,
              run_config=RunConfig(
                  name="tryout_restore",
                  stop={"training_iteration": 5},
                  storage_path=str(tmpdir),
                  failure_config=FailureConfig(max_failures=1),
                  checkpoint_config=CheckpointConfig(
                      checkpoint_frequency=1
                  ),
              ),
              param_space={"tag_file_path": tag_file},
          )
          results = tuner.fit()
          [result] = list(results)
          if retry_num > 0:
              assert result.metrics["score"] == 5
          else:
              assert result.metrics["score"] == 2
  
imports: import json
  import logging
  import os
  import shutil
  import time
  import unittest
  import pytest
  import ray
  from ray import train, tune
  import ray.cloudpickle
  from ray.train import Checkpoint, CheckpointConfig, FailureConfig, RunConfig, ScalingConfig
  from ray.air._internal.uri_utils import URI
  from ray.train.data_parallel_trainer import DataParallelTrainer
  from ray.train._internal.storage import get_fs_and_path, _download_from_fs_path, _upload_to_fs_path
  from ray.tune import Callback, Trainable
  from ray.tune.analysis import ExperimentAnalysis
  from ray.tune.execution.experiment_state import _find_newest_experiment_checkpoint
  from ray.tune.experiment import Trial
  from ray.tune.result_grid import ResultGrid
  from ray.tune.schedulers.async_hyperband import ASHAScheduler
  from ray.tune.search.optuna import OptunaSearch
  from ray.tune.tune_config import TuneConfig
  from ray.tune.tuner import Tuner
  from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
label_desc: Use image node_id 18 for calling a global function with example usage: test_restore_retry(ray_start_2_cpus, tmpdir, monkeypatch, retry_num) and returns: dict