code: def test_checkpoint_save_restore(
      ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir
  ):
      """Test that a checkpoint is saved and can be used to restore a trainable.
  
      The trainable saves a checkpoint and terminates. We then start another trial
      that should restore from the saved checkpoint and assert that it picks up
      the state and continues to run to termination.
  
      Legacy test: test_trial_runner_2.py::TrialRunnerTest::testCheckpointing
      Legacy test: test_trial_runner_2.py::TrialRunnerTest::testRestoreMetricsAfterCheckpointing  # noqa
      """
      runner = TuneController(
          resource_manager_factory=lambda: resource_manager_cls(),
          storage=STORAGE,
      )
      kwargs = {
          "stopping_criterion": {"training_iteration": 1},
          "placement_group_factory": PlacementGroupFactory(
              [{"CPU": 1, "GPU": 1}]
          ),
          "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
          "storage": STORAGE,
      }
      runner.add_trial(Trial("__fake", **kwargs))
      trials = runner.get_trials()
  
      runner.step()  # Start trial
  
      while trials[0].status != Trial.RUNNING:
          runner.step()
  
      # Set some state that will be saved in the checkpoint
      assert (
          ray.get(
              trials[0].temporary_state.ray_actor.set_info.remote(1)
          )
          == 1
      )
  
      while trials[0].status != Trial.TERMINATED:
          runner.step()
  
      assert (
          trials[0].latest_checkpoint_result.metrics[TRAINING_ITERATION]
          == 1
      )
      assert trials[0].last_result[TRAINING_ITERATION] == 1
      assert trials[0].last_result["iterations_since_restore"] == 1
  
      # Prepare new trial
      kwargs["restore_path"] = trials[0].checkpoint.path
      new_trial = Trial("__fake", **kwargs)
      runner.add_trial(new_trial)
      trials = runner.get_trials()
  
      assert trials[1].status == Trial.PENDING
  
      # Start trial, restore, run to termination
      while trials[1].status != Trial.RUNNING:
          runner.step()
  
      # Restore
      runner.step()
  
      assert (
          ray.get(trials[1].temporary_state.ray_actor.get_info.remote())
          == 1
      )
  
      # Run to termination
      while trials[1].status != Trial.TERMINATED:
          runner.step()
  
      assert (
          trials[0].latest_checkpoint_result.metrics[TRAINING_ITERATION]
          == 1
      )
      assert trials[1].last_result[TRAINING_ITERATION] == 1
      assert trials[1].last_result["iterations_since_restore"] == 1
  
imports: import json
  import logging
  import os
  import tempfile
  from unittest import mock
  import pytest
  import sys
  import time
  from functools import partial
  import ray
  from freezegun import freeze_time
  from ray.train import CheckpointConfig
  from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
  from ray.air.constants import TRAINING_ITERATION
  from ray.train import Checkpoint
  from ray.train._internal.session import _TrainingResult
  from ray.train._internal.storage import StorageContext
  from ray.tune import PlacementGroupFactory
  from ray.tune.execution.tune_controller import TuneController
  from ray.tune.experiment import Trial
  from ray.tune.result import DONE
  from ray.tune.schedulers import FIFOScheduler
  from ray.tune.search import BasicVariantGenerator
  from ray.train.tests.util import mock_storage_context
  from ray.tune.tests.tune_test_util import TrialResultObserver
label_desc: Use image node_id 3 for calling a global function with example usage: test_checkpoint_save_restore(ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir) without return types