code: def train_fn(config):
      in_trainer = config.get("in_trainer", False)
      if in_trainer:
          from ray.air._internal.session import _get_session
          from ray.train._internal.session import _TrainSession
  
          train_session = _get_session()
  
          assert isinstance(train_session, _TrainSession)
          assert train_session.storage
          assert train_session.storage.checkpoint_fs_path
  
          # Check that the working dir for each worker is the shared trial dir.
          assert os.getcwd() == train_session.storage.trial_local_path
  
      start = 0
  
      checkpoint = train.get_checkpoint()
      if checkpoint:
          custom_restore_fn = config.get("custom_restore_fn")
          if custom_restore_fn:
              state = custom_restore_fn(checkpoint)
          else:
              with checkpoint.as_directory() as checkpoint_dir:
                  with open(
                      os.path.join(checkpoint_dir, "checkpoint.pkl"),
                      "rb",
                  ) as f:
                      state = pickle.load(f)
          print("Loaded back state from checkpoint:", state)
          start = state["iter"] + 1
  
      for i in range(start, config.get("num_iterations", 5)):
          time.sleep(config.get("time_per_iter", 0.25))
  
          metrics = {"iter": i, TestConstants.SCORE_KEY: i}
  
          # Save an artifact in the local trial dir.
          rank = train.get_context().get_world_rank()
          artifact_file_name = (
              f"artifact-rank={rank}-iter={i}.txt"
              if in_trainer
              else f"artifact-iter={i}.txt"
          )
          with open(artifact_file_name, "w") as f:
              f.write(f"{i}")
  
          if (
              in_trainer
              and train.get_context().get_world_rank()
              in config.get("no_checkpoint_ranks", [])
          ):
              train.report(metrics)
          else:
              with tempfile.TemporaryDirectory() as temp_dir:
                  with open(
                      os.path.join(temp_dir, "checkpoint.pkl"), "wb"
                  ) as f:
                      pickle.dump({"iter": i}, f)
  
                  if in_trainer:
                      checkpoint_file_name = (
                          _create_checkpoint_shard_filename(str(rank))
                      )
                      with open(
                          os.path.join(temp_dir, checkpoint_file_name),
                          "wb",
                      ) as f:
                          pickle.dump({"iter": i}, f)
  
                  with config.get(
                      "custom_save_fn", dummy_context_manager
                  )(temp_dir):
                      train.report(
                          metrics,
                          checkpoint=Checkpoint.from_directory(
                              temp_dir
                          ),
                      )
                  # `train.report` should not have deleted this!
                  assert os.path.exists(temp_dir)
  
          if i in config.get("fail_iters", []):
              raise RuntimeError(f"Failing on iter={i}!!")
  
imports: import logging
  import os
  import pickle
  import re
  import shutil
  import tempfile
  import time
  from contextlib import contextmanager
  from pathlib import Path
  from typing import List, Optional, Tuple
  import pyarrow.fs
  import pytest
  import ray
  from ray import train, tune
  from ray._private.test_utils import simulate_storage
  from ray.air._internal.uri_utils import URI
  from ray.air.constants import EXPR_RESULT_FILE
  from ray.train._checkpoint import Checkpoint
  from ray.train._internal.storage import StorageContext, _delete_fs_path, _download_from_fs_path
  from ray.train.base_trainer import TrainingFailedError
  from ray.train.data_parallel_trainer import DataParallelTrainer
  from ray.tune.trainable.trainable import _DICT_CHECKPOINT_FILE_NAME
label_desc: Use image node_id 11 for calling a global function with example usage: train_fn(config) without return types