code: def run_storage_store_restore_sharded_test(
      pex: parallel.Execution,
      storage_manager: storage.StorageManager,
      clean_up: Optional[Callable] = None,
  ) -> None:
      checkpoint_context = create_checkpoint_context(
          pex, storage_manager
      )
      selector: Optional[Callable[[str], bool]]  # Making lint happy.
  
      metadata = {"steps_completed": 1}
  
      with checkpoint_context.store_path(metadata, shard=True) as (
          path,
          storage_id,
      ):
          if pex.distributed.rank in [0, 1]:
              create_checkpoint(path, EXPECTED_FILES_N0)
          else:
              create_checkpoint(path, EXPECTED_FILES_N1)
  
      # 1. Test all sharded files were uploaded to ckpt.
      with checkpoint_context.restore_path(storage_id) as path:
          validate_checkpoint(
              path,
              expected_files={**EXPECTED_FILES_N0, **EXPECTED_FILES_N1},
          )
          # Make sure all ranks finish validating checkpoint before exiting the context.
          pex.distributed.allgather(None)
  
      if isinstance(
          storage_manager, storage.shared.SharedFSStorageManager
      ):
          # SharedFS does not support downloading partial checkpoints.
          # Skip all remaining tests.
          return
  
      # 2. Test downloading with selector: every rank gets selected files.
      if pex.distributed.rank in [0, 1]:
  
          def selector(x: str) -> bool:
              return x == "subdir/file3_0"
  
      else:
  
          def selector(x: str) -> bool:
              return x == "file1_1"
  
      with checkpoint_context.restore_path(
          storage_id, selector=selector
      ) as path:
          if pex.distributed.rank in [0, 1]:
              validate_checkpoint(
                  path,
                  expected_files={
                      "subdir/file3_0": "nested file node 0",
                      "subdir/": None,
                  },
              )
          else:
              validate_checkpoint(
                  path, expected_files={"file1_1": "file 1 node 1"}
              )
          pex.distributed.allgather(None)
  
      # 3. Test downloading with selector: local ranks get union of selected files.
      if pex.distributed.local_rank == 0:
  
          def selector(x: str) -> bool:
              return x == "subdir/file3_0"
  
      else:
  
          def selector(x: str) -> bool:
              return x == "file1_1"
  
      with checkpoint_context.restore_path(
          storage_id, selector=selector
      ) as path:
          validate_checkpoint(
              path,
              expected_files={
                  "subdir/file3_0": "nested file node 0",
                  "subdir/": None,
                  "file1_1": "file 1 node 1",
              },
          )
          pex.distributed.allgather(None)
  
      # 4. Test downloading with and w/o selector: node 1 gets a selection of files;
      # node 2 gets all the files.
      if pex.distributed.rank in [0, 1]:
  
          def selector(x: str) -> bool:
              return x in list(EXPECTED_FILES_N0.keys())
  
      else:
          selector = None
  
      with checkpoint_context.restore_path(
          storage_id, selector=selector
      ) as path:
          if pex.distributed.rank in [0, 1]:
              validate_checkpoint(
                  path, expected_files=EXPECTED_FILES_N0
              )
          else:
              validate_checkpoint(
                  path,
                  expected_files={
                      **EXPECTED_FILES_N0,
                      **EXPECTED_FILES_N1,
                  },
              )
          pex.distributed.allgather(None)
  
      if pex.distributed.rank == 0 and clean_up is not None:
          clean_up(storage_id, storage_manager)
  
      # 5. Test uploading two files with the same name and different content
      # raises an error.
      with parallel.raises_when(
          True,
          RuntimeError,
          match=r"refusing to upload with files conflicts:.*",
      ):
          with checkpoint_context.store_path(metadata, shard=True) as (
              path,
              storage_id,
          ):
              if pex.distributed.rank in [0, 1]:
                  create_checkpoint(path, {"filename": "content 1"})
              else:
                  create_checkpoint(path, {"filename": "content 2"})
  
      # 6 Test uploading file and directory with the same name.
      with parallel.raises_when(
          True,
          RuntimeError,
          match=r"refusing to upload with files conflicts:.*",
      ):
          with checkpoint_context.store_path(metadata, shard=True) as (
              path,
              storage_id,
          ):
              if pex.distributed.rank in [0, 1]:
                  os.makedirs(os.path.join(path, "abc"))
              else:
                  with open(os.path.join(path, "abc"), "w") as f:
                      f.write("test")
  
imports: import os
  import pathlib
  import shutil
  import time
  import uuid
  from typing import Any, Callable, Dict, List, Optional, Tuple
  from unittest import mock
  import pytest
  import requests
  from determined import core, errors
  from determined.common import storage
  from determined.tensorboard.fetchers import base
  from tests import parallel
label_desc: Use image node_id 7 for calling a global function with example usage: run_storage_store_restore_sharded_test(pex, storage_manager, clean_up) and returns: unknown, unknown, unknown, unknown, unknown