code: def test_distributed_context(
      cross_size: int, local_size: int, force_tcp: bool
  ) -> None:
      size = cross_size * local_size
  
      # Make sure `make test` doesn't hang on macbook's default values.  Avoid skipping on linux
      # because it's not a common default, and to avoid false positives in CI.
      if sys.platform == "darwin" and size == 16:
          import resource
  
          if resource.getrlimit(resource.RLIMIT_NOFILE)[0] < 1024:
              pytest.skip(
                  "increase the open fd limit with `ulimit -n 1024` or greater to run this test"
              )
  
      with parallel.Execution(
          size, local_size=local_size, make_distributed_context=False
      ) as pex:
  
          @pex.run
          def contexts() -> core.DistributedContext:
              return core.DistributedContext(
                  rank=pex.rank,
                  size=pex.size,
                  local_rank=pex.local_rank,
                  local_size=pex.local_size,
                  cross_rank=pex.cross_rank,
                  cross_size=pex.cross_size,
                  chief_ip="localhost",
                  force_tcp=force_tcp,
              )
  
          # Perform a broadcast.
          results = pex.run(lambda: contexts[pex.rank].broadcast(pex.rank))  # type: ignore
          assert (
              results == [0] * size
          ), "not all threads ran broadcast correctly"
  
          # Perform a local broadcast.
          results = pex.run(
              lambda: contexts[pex.rank].broadcast_local(pex.rank)
          )
          expect = [
              rank - (rank % local_size) for rank in range(size)
          ]  # type: Any
  
          assert (
              results == expect
          ), "not all threads ran broadcast_local correctly"
  
          # Perform a gather.
          results = pex.run(
              lambda: set(contexts[pex.rank].gather(pex.rank) or [])
          )
          chief = set(range(size))
          expect = [
              set(range(size)) if rank == 0 else set()
              for rank in range(size)
          ]
          assert results == [chief] + [set()] * (
              size - 1
          ), "not all threads ran gather correctly"
  
          # Perform a local gather.
          results = pex.run(
              lambda: set(
                  contexts[pex.rank].gather_local(pex.rank) or []
              )
          )
          expect = [
              set(range(rank, rank + local_size))
              if rank % local_size == 0
              else set()
              for rank in range(size)
          ]
          assert (
              results == expect
          ), "not all threads ran gather correctly"
  
          # Perform an allgather.
          results = pex.run(
              lambda: set(contexts[pex.rank].allgather(pex.rank))
          )
          expect = set(range(size))
          assert (
              results == [expect] * size
          ), "not all threads ran allgather correctly"
  
          # Perform a local allgather.
          results = pex.run(
              lambda: set(contexts[pex.rank].allgather_local(pex.rank))
          )
          expect = [
              set(
                  range(
                      cross_rank * local_size,
                      (cross_rank + 1) * local_size,
                  )
              )
              for cross_rank, _ in itertools.product(
                  range(cross_size), range(local_size)
              )
          ]
          assert (
              results == expect
          ), "not all threads ran allgather_local correctly"
  
          # Close all contexts.
          for context in contexts:
              context.close()
  
imports: import abc
  import itertools
  import multiprocessing
  import sys
  import textwrap
  import time
  import traceback
  from typing import Any, List, Optional, cast
  import pytest
  import determined
  from determined import core, ipc
  from tests import parallel
label_desc: Use image node_id 2 for calling a global function with example usage: test_distributed_context(cross_size, local_size, force_tcp) and returns: core