code: def test_get_dataset_shard(shutdown_only):
      dataset = ray.data.from_items([1, 2, 3])
      init_session(
          training_func=lambda: 1,
          world_rank=0,
          local_rank=0,
          node_rank=0,
          local_world_size=1,
          world_size=1,
          dataset_shard=dataset,
          storage=storage,
      )
      assert get_dataset_shard() == dataset
      shutdown_session()
  
imports: import tempfile
  import time
  import warnings
  import pytest
  import ray
  from ray.air._internal.util import StartTraceback
  from ray.air.constants import SESSION_MISUSE_LOG_ONCE_KEY
  from ray.air.session import get_checkpoint, get_dataset_shard, get_local_rank, get_world_rank, get_world_size, report
  from ray.train._internal.accelerator import Accelerator
  from ray.train._internal.session import get_accelerator, get_session, init_session, set_accelerator, shutdown_session
  from ray.train._internal.storage import StorageContext
  from ray.train.error import SessionMisuseError
  from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
label_desc: Use image node_id 9 for calling a global function with example usage: test_get_dataset_shard(shutdown_only) without return types