code: def run_storage_upload_download_sharded_test(
      pex: parallel.Execution,
      storage_manager: storage.StorageManager,
      tmp_path: pathlib.Path,
      clean_up: Optional[Callable] = None,
  ) -> None:
      checkpoint_context = create_checkpoint_context(
          pex, storage_manager
      )
      selector: Optional[Callable[[str], bool]]  # Making lint happy.
  
      # Create "local" file structure.
      ckpt_dir = tmp_path.joinpath(
          f"ckpt_dir_{pex.distributed.rank // 2}"
      )
      if pex.distributed.rank == 0:
          create_checkpoint(ckpt_dir, EXPECTED_FILES_N0)
      elif pex.distributed.rank == 2:
          create_checkpoint(ckpt_dir, EXPECTED_FILES_N1)
  
      # Wait for all ranks to save files.
      pex.distributed.allgather(None)
  
      metadata = {"steps_completed": 1}
  
      # Upload sharded data.
      storage_id = checkpoint_context.upload(
          ckpt_dir, metadata, shard=True
      )
  
      # 1. Test all sharded files were uploaded to ckpt.
      # Since DownloadMode = LocalWorkerShareDownload, let's make sure that ranks=[0,1]
      # download data to one directory, and ranks=[2,3] use another directory.
      download_dir = tmp_path.joinpath(
          f"test1_download{pex.distributed.rank // 2}"
      )
      try:
          checkpoint_context.download(storage_id, download_dir)
          validate_checkpoint(
              download_dir,
              expected_files={**EXPECTED_FILES_N0, **EXPECTED_FILES_N1},
          )
          pex.distributed.allgather(None)
      finally:
          shutil.rmtree(download_dir, ignore_errors=True)
  
      # 2. Test downloading with selector: every rank gets selected files.
      download_dir = tmp_path.joinpath(
          f"test2_download_{pex.distributed.rank // 2}"
      )
      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"
  
      try:
          checkpoint_context.download(
              storage_id, download_dir, selector=selector
          )
  
          if pex.distributed.rank in [0, 1]:
              validate_checkpoint(
                  download_dir,
                  expected_files={
                      "subdir/file3_0": "nested file node 0",
                      "subdir/": None,
                  },
              )
          else:
              validate_checkpoint(
                  download_dir,
                  expected_files={"file1_1": "file 1 node 1"},
              )
  
          pex.distributed.allgather(None)
      finally:
          shutil.rmtree(download_dir, ignore_errors=True)
  
      # 3.Test downloading with and w/o selector.
      download_dir = tmp_path.joinpath(
          f"test3_download_{pex.distributed.rank // 2}"
      )
      if pex.distributed.rank in [0, 1]:
  
          def selector(x: str) -> bool:
              return x in EXPECTED_FILES_N0
  
      else:
          selector = None
  
      try:
          checkpoint_context.download(
              storage_id, download_dir, selector=selector
          )
          if pex.distributed.rank in [0, 1]:
              validate_checkpoint(
                  download_dir, expected_files=EXPECTED_FILES_N0
              )
          else:
              validate_checkpoint(
                  download_dir,
                  expected_files={
                      **EXPECTED_FILES_N0,
                      **EXPECTED_FILES_N1,
                  },
              )
          pex.distributed.allgather(None)
      finally:
          shutil.rmtree(download_dir, ignore_errors=True)
  
      # 4.Test downloading with and w/o selector in NoSharedDownload mode.
      download_dir = tmp_path.joinpath(
          f"test4_download_{pex.distributed.rank}"
      )
      if pex.distributed.rank in [0, 1]:
  
          def selector(x: str) -> bool:
              return x in EXPECTED_FILES_N0
  
      else:
          selector = None
  
      try:
          checkpoint_context.download(
              storage_id,
              download_dir,
              selector=selector,
              download_mode=core.DownloadMode.NoSharedDownload,
          )
          if pex.distributed.rank in [0, 1]:
              validate_checkpoint(
                  download_dir, expected_files=EXPECTED_FILES_N0
              )
          else:
              validate_checkpoint(
                  download_dir,
                  expected_files={
                      **EXPECTED_FILES_N0,
                      **EXPECTED_FILES_N1,
                  },
              )
          pex.distributed.allgather(None)
      finally:
          shutil.rmtree(download_dir, ignore_errors=True)
  
      # 5.Test downloading with selector excluding all files.
      download_dir = tmp_path.joinpath(
          f"test5_download_{pex.distributed.rank//2}"
      )
  
      def selector1(x: str) -> bool:
          return False
  
      try:
          checkpoint_context.download(
              storage_id, download_dir, selector=selector1
          )
          assert (
              not download_dir.exists()
          ), f"{download_dir} should not exist"
      finally:
          shutil.rmtree(download_dir, ignore_errors=True)
  
      sync_and_clean(pex, clean_up, storage_id, storage_manager)
  
      # 6. Upload sharded data from rank 2 only.
      # Metadata should be saved and uploaded as well.
      storage_id = checkpoint_context.upload(
          ckpt_dir if pex.distributed.rank == 2 else None,
          metadata,
          shard=True,
      )
      download_dir = tmp_path.joinpath(
          f"test6_download_{pex.distributed.rank // 2}"
      )
      checkpoint_context.download(storage_id, download_dir)
      validate_checkpoint(
          download_dir,
          expected_files={
              **EXPECTED_FILES_N1,
              **{"metadata.json": '{\n  "steps_completed": 1\n}'},
          },
      )
      sync_and_clean(pex, clean_up, storage_id, storage_manager)
  
      # 7. Test uploading selected paths: rank 0 and rank 1 uploads a subset of files,
      # rank 2 and 3 uploads all files.
      if pex.distributed.rank == 0:
  
          def selector(x: str) -> bool:
              return x in x in ["file_repeated", "subdir/file_repeated"]
  
      elif pex.distributed.rank == 1:
  
          def selector(x: str) -> bool:
              return x in x == "file0_0"
  
      else:
          selector = None
  
      storage_id = checkpoint_context.upload(
          ckpt_dir, metadata, shard=True, selector=selector
      )
  
      download_dir = tmp_path.joinpath(
          f"test7_download_{pex.distributed.rank // 2}"
      )
      checkpoint_context.download(storage_id, download_dir)
      validate_checkpoint(
          download_dir,
          expected_files={
              **EXPECTED_FILES_N1,
              **{
                  "file_repeated": "file repeated",
                  "file0_0": "file 0 node 0",
                  "subdir/": None,
                  "subdir/file_repeated": "nested file repeated",
              },
          },
      )
      sync_and_clean(pex, clean_up, storage_id, storage_manager)
  
      # 8. Test uploading selected paths: rank=[0, 1] uploads nothing,
      # rank=2 uploads a directory from ckpt_dir,
      # rank=3 uploads selected files in a ckpt_dir directory
      if pex.distributed.rank in [0, 1]:
          storage_id = checkpoint_context.upload(
              None, metadata, shard=True, selector=None
          )
      elif pex.distributed.rank == 2:
          dir_in_ckpt_dir = ckpt_dir.joinpath("another_subdir")
          storage_id = checkpoint_context.upload(
              dir_in_ckpt_dir, metadata, shard=True, selector=None
          )
      else:
  
          def selector(x: str) -> bool:
              return x == "subdir/file_repeated"
  
          storage_id = checkpoint_context.upload(
              ckpt_dir, metadata, shard=True, selector=selector
          )
  
      download_dir = tmp_path.joinpath(
          f"test8_download_{pex.distributed.rank // 2}"
      )
      checkpoint_context.download(storage_id, download_dir)
      validate_checkpoint(
          download_dir,
          expected_files={
              "file": "file in another subdir",
              "subdir/": None,
              "subdir/file_repeated": "nested file repeated",
              "metadata.json": '{\n  "steps_completed": 1\n}',
          },
      )
      sync_and_clean(pex, clean_up, storage_id, storage_manager)
  
      # 9. Test uploading selected paths: ranks try to upload not existing paths.
      def selector_upload(x: str) -> bool:
          return False
  
      storage_id = checkpoint_context.upload(
          ckpt_dir, metadata, shard=True, selector=selector_upload
      )
      download_dir = tmp_path.joinpath(
          f"test9_download_{pex.distributed.rank // 2}"
      )
      checkpoint_context.download(storage_id, download_dir)
      validate_checkpoint(
          download_dir,
          expected_files={
              "metadata.json": '{\n  "steps_completed": 1\n}'
          },
      )
      sync_and_clean(pex, clean_up, storage_id, storage_manager)
  
      # 10. 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:.*",
      ):
          if pex.distributed.rank == 0:
              create_checkpoint(ckpt_dir, {"filename": "content 1"})
          elif pex.distributed.rank == 2:
              create_checkpoint(ckpt_dir, {"filename": "content 2"})
  
          pex.distributed.allgather(None)
          checkpoint_context.upload(ckpt_dir, metadata, shard=True)
  
      # 11. Test uploading directory and file with the same name.
      with parallel.raises_when(
          True,
          RuntimeError,
          match=r"refusing to upload with files conflicts:.*",
      ):
          if pex.distributed.rank == 0:
              os.makedirs(os.path.join(ckpt_dir, "abc"))
          elif pex.distributed.rank == 2:
              with open(os.path.join(ckpt_dir, "abc"), "w") as f:
                  f.write("test")
          pex.distributed.allgather(None)
          checkpoint_context.upload(ckpt_dir, metadata, shard=True)
  
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 6 for calling a global function with example usage: run_storage_upload_download_sharded_test(pex, storage_manager, tmp_path, clean_up) and returns: False, False, unknown, unknown, unknown, unknown, unknown, unknown, unknown