code: def get_mock_distributed_context(
      rank: int = 0,
      all_gather_return_value: Optional[Any] = None,
      gather_return_value: Optional[Any] = None,
  ) -> unittest.mock.MagicMock:
      mock_distributed_context = unittest.mock.MagicMock()
      mock_distributed_context.get_rank.return_value = rank
      mock_distributed_context.broadcast.return_value = (
          "mock_checkpoint_uuid"
      )
      mock_distributed_context.allgather.return_value = (
          all_gather_return_value
      )
      mock_distributed_context.gather.return_value = gather_return_value
      return mock_distributed_context
  
imports: import importlib
  import os
  import pathlib
  import re
  import unittest.mock
  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type
  import numpy
  import pytest
  from mypy_extensions import DefaultNamedArg
  import determined
  from determined import core, gpu, workload
  from determined.common import util
label_desc: Use image node_id 18 for calling a global function with example usage: get_mock_distributed_context(rank, all_gather_return_value, gather_return_value) and returns: mock_distributed_context