docstr: Integration test to confirm distributed initialization and consistency with process group backend utilities.
docstr_tok: ["Integration","test","to","confirm","distributed","initialization","and","consistency","with","process","group","backend","utilities","."]
code: def _test_worker_fn(init_pg_explicit: bool) -> torch.device:
      """
      Integration test to confirm distributed initialization and consistency with process group backend utilities.
      """
      if init_pg_explicit:
          torch.distributed.init_process_group(backend="gloo")
      device = init_from_env()
      if not torch.distributed.is_initialized():
          raise AssertionError(
              "Expected torch.distributed to be initialized"
          )
      device_from_env = get_device_from_env()
      if device != device_from_env:
          raise AssertionError(
              f"Expected different device: received {device}, expected {device_from_env}"
          )
      pg_backend = torch.distributed.get_backend()
      expected_pg_backend = (
          get_process_group_backend_from_device(device)
          if not init_pg_explicit
          else "gloo"
      )
      if pg_backend != expected_pg_backend:
          raise AssertionError(
              f"Expected different process group backend: received {pg_backend}, expected {expected_pg_backend}"
          )
      return device
  
imports: import os
  import unittest
  import numpy
  import torch
  import torch.distributed.launcher
  from torchtnt.utils.device import get_device_from_env
  from torchtnt.utils.distributed import get_process_group_backend_from_device
  from torchtnt.utils.env import init_from_env, seed
  from torchtnt.utils.test_utils import get_pet_launch_config
label_desc: Use image node_id 2 for calling the EnvTest obj's underlying member method code with example usage: obj._test_worker_fn(init_pg_explicit) and returns: device