code: def test_data_parallel_trainer_restore(ray_start_4_cpus, tmpdir):
      """Restoring a DataParallelTrainer with object refs captured in the train fn
      or config works by re-specifying them.
      Success criteria:
      - Restored to the correct iteration. (1 iteration before crash, 1 after restore).
      - Results are being logged to the same directory as before.
      """
      dataset_size = 10
      num_workers = 2
  
      def create_train_fn_and_config():
          obj_ref = ray.put({"test": 1})
  
          def train_fn(config):
              assert ray.get(obj_ref)["test"] == 1
              assert ray.get(config["obj_ref"])["test"] == 1
              ds = train.get_dataset_shard("train")
              assert (
                  sum(
                      [
                          len(batch["feature"])
                          for batch in ds.iter_batches()
                      ]
                  )
                  == dataset_size // num_workers
              )
              _failing_train_fn(config)
  
          train_loop_config = {"obj_ref": obj_ref}
          return train_fn, train_loop_config
  
      datasets = {
          "train": ray.data.from_items(
              [{"feature": i} for i in range(10)]
          )
      }
      train_fn, train_loop_config = create_train_fn_and_config()
      trainer = DataParallelTrainer(
          train_loop_per_worker=train_fn,
          train_loop_config=train_loop_config,
          datasets=datasets,
          scaling_config=ScalingConfig(num_workers=num_workers),
          run_config=RunConfig(
              name="data_parallel_restore_test",
              storage_path=str(tmpdir),
              checkpoint_config=CheckpointConfig(num_to_keep=1),
          ),
      )
      with pytest.raises(TrainingFailedError) as exc_info:
          result = trainer.fit()
      assert isinstance(exc_info.value.__cause__, _TestSpecificError)
  
      # Include an explicit cluster shutdown.
      # Otherwise, the previously registered object references will still exist,
      # and the test may trivially pass.
      ray.shutdown()
      ray.init(num_cpus=4)
  
      train_fn, train_loop_config = create_train_fn_and_config()
      datasets = {
          "train": ray.data.from_items(
              [{"feature": i} for i in range(10)]
          )
      }
      trainer = DataParallelTrainer.restore(
          str(tmpdir / "data_parallel_restore_test"),
          train_loop_per_worker=train_fn,
          train_loop_config=train_loop_config,
          datasets=datasets,
      )
      result = trainer.fit()
      assert not result.error
      assert result.metrics["training_iteration"] == 2
      assert result.metrics["iterations_since_restore"] == 1
      assert (
          tmpdir / "data_parallel_restore_test"
          in Path(result.path).parents
      )
  
imports: import warnings
  from pathlib import Path
  from typing import Dict, List
  import pytest
  import ray
  from ray import train
  from ray.air._internal.uri_utils import URI
  from ray.train import CheckpointConfig, RunConfig, ScalingConfig
  from ray.train.base_trainer import BaseTrainer
  from ray.train.data_parallel_trainer import DataParallelTrainer
  from ray.train.lightgbm import LightGBMTrainer
  from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
  from ray.train.torch import TorchTrainer
  from ray.train.trainer import TrainingFailedError
  from ray.train.xgboost import XGBoostTrainer
  from ray.tune import Callback
  from ray.tune.experiment import Trial
label_desc: Use image node_id 4 for calling a global function with example usage: test_data_parallel_trainer_restore(ray_start_4_cpus, tmpdir) and returns: train_fn,  train_loop_config