code: def run_rpc_feature_test(
      world_size: int,
      rank: int,
      feature: LocalFeatureStore,
      partition_book: torch.Tensor,
      master_port: int,
  ):
      # 1) Initialize the context info:
      current_ctx = DistContext(
          rank=rank,
          global_rank=rank,
          world_size=world_size,
          global_world_size=world_size,
          group_name="dist-feature-test",
      )
  
      rpc.init_rpc(
          current_ctx=current_ctx,
          master_addr="localhost",
          master_port=master_port,
      )
  
      # 2) Collect all workers:
      partition_to_workers = rpc.rpc_partition_to_workers(
          current_ctx, world_size, rank
      )
  
      assert partition_to_workers == [
          ["dist-feature-test-0"],
          ["dist-feature-test-1"],
      ]
  
      # 3) Find the mapping between worker and partition ID:
      rpc_router = RPCRouter(partition_to_workers)
  
      assert (
          rpc_router.get_to_worker(partition_idx=0)
          == "dist-feature-test-0"
      )
      assert (
          rpc_router.get_to_worker(partition_idx=1)
          == "dist-feature-test-1"
      )
  
      meta = {
          "edge_types": None,
          "is_hetero": False,
          "node_types": None,
          "num_parts": 2,
      }
  
      feature.num_partitions = world_size
      feature.partition_idx = rank
      feature.node_feat_pb = partition_book
      feature.meta = meta
      feature.local_only = False
      feature.set_rpc_router(rpc_router)
  
      # Global node IDs:
      global_id0 = torch.arange(128 * 2)
      global_id1 = torch.arange(128 * 2) + 128 * 2
  
      # Lookup the features from stores including locally and remotely:
      tensor0 = feature.lookup_features(global_id0)
      tensor1 = feature.lookup_features(global_id1)
  
      # Expected searched results:
      cpu_tensor0 = torch.cat(
          [torch.ones(128, 1024), torch.ones(128, 1024) * 2]
      )
      cpu_tensor1 = torch.cat(
          [torch.zeros(128, 1024), torch.zeros(128, 1024)]
      )
  
      # Verify..
      assert torch.allclose(cpu_tensor0, tensor0.wait())
      assert torch.allclose(cpu_tensor1, tensor1.wait())
  
      rpc.shutdown_rpc()
      assert rpc.rpc_is_initialized() is False
  
file_path: pytorch_geometric/test/distributed/test_rpc.py
imports: import socket
  import torch
  import torch_geometric.distributed.rpc
  from torch_geometric.distributed import LocalFeatureStore
  from torch_geometric.distributed.dist_context import DistContext
  from torch_geometric.distributed.rpc import RPCRouter
  from torch_geometric.testing import onlyLinux
label_desc: Use image node_id 1 for calling a global function with example usage: run_rpc_feature_test(world_size, rank, feature, partition_book, master_port) without return types