code: def test_load_balancing_under_constrained_memory(
      enable_mac_large_object_store, ray_start_cluster
  ):
      # This test ensures that tasks are being assigned to all raylets in a
      # roughly equal manner even when the tasks have dependencies.
      cluster = ray_start_cluster
      num_nodes = 3
      num_cpus = 4
      object_size = 4e7
      num_tasks = 100
      for _ in range(num_nodes):
          cluster.add_node(
              num_cpus=num_cpus,
              memory=(num_cpus - 2) * object_size,
              object_store_memory=(num_cpus - 2) * object_size,
          )
      cluster.add_node(
          num_cpus=0,
          resources={"custom": 1},
          memory=(num_tasks + 1) * object_size,
          object_store_memory=(num_tasks + 1) * object_size,
      )
      ray.init(address=cluster.address)
  
      @ray.remote(num_cpus=0, resources={"custom": 1})
      def create_object():
          return np.zeros(int(object_size), dtype=np.uint8)
  
      @ray.remote
      def f(i, x):
          print(i, ray._private.worker.global_worker.node.unique_id)
          time.sleep(0.1)
          return ray._private.worker.global_worker.node.unique_id
  
      deps = [create_object.remote() for _ in range(num_tasks)]
      for i, dep in enumerate(deps):
          print(i, dep)
  
      # TODO(swang): Actually test load balancing. Load balancing is currently
      # flaky on Travis, probably due to the scheduling policy ping-ponging
      # waiting tasks.
      deps = [create_object.remote() for _ in range(num_tasks)]
      tasks = [f.remote(i, dep) for i, dep in enumerate(deps)]
      for i, dep in enumerate(deps):
          print(i, dep)
      ray.get(tasks)
  
imports: import os
  import platform
  import sys
  import time
  import numpy
  import pytest
  import ray
  import ray._private.gcs_utils
  import ray.experimental.internal_kv
  from ray._private.test_utils import make_global_state_accessor, wait_for_condition, get_metric_check_condition, MetricSamplePattern
  from ray.util.client.ray_client_helpers import connect_to_client_or_not
  from ray.util.placement_group import placement_group
  from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy, PlacementGroupSchedulingStrategy
  from ray._private.test_utils import SignalActor
label_desc: Use image node_id 1 for calling a global function with example usage: test_load_balancing_under_constrained_memory(enable_mac_large_object_store, ray_start_cluster) and returns: np, ray