[Data] Add autoscaler_scales_up_when_compute_bound release test

[Hyunoh-Yeo]

Description

Adds autoscaler_scales_up_when_compute_bound under a new "Autoscaling tests" section in release_data_tests.yaml, labeled group: data-autoscaling. It checks that the autoscaler scales up both worker groups when a pipeline is compute bound on its CPU stage.

produce sleeps 5s per emitted block and emits four 128 MiB blocks per input; consume sleeps 1s per batch and holds a GPU. The test asserts the peak worker node counts: exactly 10 GPU nodes, at least 3 CPU nodes, and at most --assert-max-cpu-nodes CPU nodes when passed. The run takes ~30 minutes.

Related issues

Closes #65771

Additional information

ClusterResourceMonitor gains get_peak_node_counts(). Aggregate CPU/GPU totals can't tell the two groups apart.

The node-count thresholds are derived from the sleep values and the compute config.

[gemini-code-assist]

Code Review

This pull request introduces a new nightly test to verify that the autoscaler scales up sufficiently to balance a compute-bound pipeline, along with its cluster configuration and registration in the release tests. It also updates the ClusterResourceMonitor to track peak CPU and GPU node counts by querying ray.nodes(). The review feedback recommends clamping the calculated minimum CPU nodes to zero to prevent negative assertions, and wrapping the node-counting logic in a try-except block within the background thread to prevent silent thread crashes during transient network or GCS RPC hiccups.

[gemini-code-assist]

If the CPU requirements are fully satisfied by the GPU nodes (i.e., _CPUS_NEEDED <= _CPUS_FROM_GPU_NODES), MIN_CPU_NODES can evaluate to a negative number. It is safer and more semantically correct to clamp this value to a minimum of 0 to prevent negative node count assertions.

Suggested change

	_CPUS_FROM_GPU_NODES = EXPECTED_GPU_NODES * CPUS_PER_NODE
	MIN_CPU_NODES = max(0, math.ceil((_CPUS_NEEDED - _CPUS_FROM_GPU_NODES) / CPUS_PER_NODE))

[Hyunoh-Yeo]

agree. However, if that is the case, the test needs changes in sleeps or configs. Therefore, added assert to positive number and print(change numbers) when fails

[gemini-code-assist]

The background monitoring thread calls _count_worker_nodes(), which queries GCS via ray.nodes(). If there are transient network or GCS RPC hiccups during this long-running test (~30 minutes), an unhandled exception will crash the daemon thread silently. Wrapping this call in a try-except block ensures the monitoring thread remains resilient and continues tracking resource metrics.

Suggested change

	node_counts = _count_worker_nodes()
	self._peak_node_counts = NodeCounts(
	cpu=max(self._peak_node_counts.cpu, node_counts.cpu),
	gpu=max(self._peak_node_counts.gpu, node_counts.gpu),
	)
	try:
	node_counts = _count_worker_nodes()
	self._peak_node_counts = NodeCounts(
	cpu=max(self._peak_node_counts.cpu, node_counts.cpu),
	gpu=max(self._peak_node_counts.gpu, node_counts.gpu),
	)
	except Exception as e:
	import logging
	logging.warning(f'Failed to count worker nodes: {e}')

[Hyunoh-Yeo]

agreed, but instead wrapped the whole loop for a consistency with ray.cluster_resources()

[yuhuan130]

Do you think we could move these test-specific values into the function and add a docstring explaining the setup and expected scaling behavior? That would make the test easier to follow without jumping between so many module-level constants.

Suggested change

	# With 1000 inputs this takes ~30 minutes, long enough that node provisioning
	def func(...):
	"""
	docstring
	"""
	num_inputs = 1000
	blocks_per_input = 4
	produce_sleep_s = 5
	consume_sleep_s = 1
	block_shape = (128, 1024, 1024)
	consume_batch_size = 2 * block_shape[0]

[Hyunoh-Yeo]

Since the values are repetitively referenced in produce and consume and used in derivation, moving them to main() will raise the NameError. Therefore, I would rather keep them where they are and add the docstring at the top of the file.

[yuhuan130]

Could we simplify this calculation? There are many intermediate constants, which makes it hard to follow. For example, a short comment could explain that 10 GPU workers require 100 producer CPUs, the GPU nodes provide 80 CPUs, and the remaining 20 CPUs require 3 CPU-only nodes.

[Hyunoh-Yeo]

Personally I'd like to keep the derivation. With hardcoded values, changing any single constant means editing several lines and redoing the arithmetic by hand. I've collapsed the intermediates so the whole calculation reads in one place. Any thoughts about this?

[yuhuan130]

Would it make sense to keep the CPU and GPU counts from the same sample? Since their peaks are tracked separately, we could end up with a combination that never existed.

# First sample
node_counts = NodeCounts(cpu=0, gpu=10)
peak = NodeCounts(cpu=0, gpu=10)

# Later sample
node_counts = NodeCounts(cpu=3, gpu=7)

# The separate max calculations produce:
peak = NodeCounts(cpu=3, gpu=10)

[yuhuan130]

It looks like this may be redundant. Benchmark.run_fn() API automatically initializes Ray.

[yuhuan130]

Should we set num_cpus=0 here? My understanding is that the consumer otherwise reserves 1 CPU as well as 1 GPU. This leaves 7 CPUs per GPU node for the producer, so we would need 4 CPU nodes instead of 3. Setting num_cpus=0 would make the code match the calculation above.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 3c22277. Configure here.}

[cursor]

CPU threshold ignores head node

Medium Severity

MIN_CPU_NODES still assumes only GPU workers supply producer CPUs. After consume was set to num_cpus=0, those 80 CPUs plus the head's 8 already cover 88 of the 100 producer slots, so a correct autoscaler can add only 2 CPU workers and still fail the >= 3 check. Both worker groups start at min_nodes: 0, so produce tasks will use the head.

Additional Locations (1)

• release/nightly_tests/dataset/autoscaling/autoscaler_scales_up_when_compute_bound.py#L89-L90

 

^{Reviewed by Cursor Bugbot for commit 3c22277. Configure here.}

[yuhuan130]

CMIIW I don’t understand the NameError concern. I was imagining the pattern used in backpressure_benchmark.py: keepig the test configuration and calculation in main(), then pass the configuration to the UDFs with functools.partial. Or did you have another structure in mind?

For example:

def produce(_, *, blocks_per_input, sleep_s, block_shape):
    for _ in range(blocks_per_input):
        time.sleep(sleep_s)
        yield {"data": np.zeros(block_shape, dtype=np.uint8)}


def consume(batch, *, sleep_s):
    time.sleep(sleep_s)
    return batch


def main(args):
    num_inputs = 1000
    blocks_per_input = 4
    produce_sleep_s = 5
    consume_sleep_s = 1
    block_shape = (128, 1024, 1024)
    rows_per_block = block_shape[0]
    consume_batch_size = 2 * rows_per_block

    expected_gpu_nodes = 10
    cpus_per_node = 8

    # Each consumer processes 2 blocks/s. Each producer emits 0.2 blocks/s,
    # so one consumer needs 10 producers.
    blocks_per_consumer_batch = consume_batch_size / rows_per_block
    producers_per_consumer = (
        blocks_per_consumer_batch * produce_sleep_s / consume_sleep_s
    )
    
    # Ten consumers need 100 producer CPUs. The GPU nodes provide 80,
    # leaving a 20-CPU shortfall, or 3 CPU-only nodes.
    producer_cpus_needed = expected_gpu_nodes * producers_per_consumer
    producer_cpus_on_gpu_nodes = expected_gpu_nodes * cpus_per_node
    cpu_shortfall = producer_cpus_needed - producer_cpus_on_gpu_nodes
    min_cpu_nodes = math.ceil(cpu_shortfall / cpus_per_node)

    producer = functools.partial(
        produce,
        blocks_per_input=blocks_per_input,
        sleep_s=produce_sleep_s,
        block_shape=block_shape,
    )
    consumer = functools.partial(consume, sleep_s=consume_sleep_s)