Repository: ray-project/ray
File: python/ray/data/tests/test_default_cluster_autoscaler_v2.py (Python)

Code before:
            {"CPU": 4, "GPU": 0, "memory": 1000},
            {"CPU": 8, "GPU": 2, "memory": 2000},
        ]
        fake_coordinator.request_resources(
            requester_id="data-test_execution_id",
            resources=initial_allocation,
            expire_after_s=180,
            request_remaining=True,
        )

        autoscaler = DefaultClusterAutoscalerV2(
            resource_manager=MagicMock(),
            execution_id="test_execution_id",
            cluster_scaling_up_delta=1,
            resource_utilization_calculator=StubUtilizationGauge(low_utilization),
            cluster_scaling_up_util_threshold=scale_up_threshold,
            min_gap_between_autoscaling_requests_s=0,
            autoscaling_coordinator=fake_coordinator,
            get_node_counts=lambda: {
                resource_spec1: 1,
                resource_spec2: 1,
            },
        )

        # Initialization should send []
        _send_resource_request.assert_called_with([])

        # Trigger scaling with low utilization
        autoscaler.try_trigger_scaling()

        # Should send current allocation to preserve resource footprint
        _send_resource_request.assert_called_with(initial_allocation)

    def test_get_node_resource_spec_and_count_skips_max_count_zero(self):
        """Test that node types with max_count=0 are skipped."""
        # Simulate a cluster with only head node (no worker nodes)
        node_table = [
            {
                "Resources": self._head_node,
                "Alive": True,
            },
        ]

        # Create a mock cluster config with one valid node type and one with max_count=0
        cluster_config = autoscaler_pb2.ClusterConfig()

        # Node type 1: 4 CPU, 0 GPU, 1000 memory, max_count=10
        node_group_config1 = autoscaler_pb2.NodeGroupConfig()
        node_group_config1.resources["CPU"] = 4
        node_group_config1.resources["memory"] = 1000
        node_group_config1.max_count = 10

Review comment (suggestion):
As a general practice, I think we should avoid testing against implementation details, but I think it's okay in this case if there's not a clear alternative

Code after:
        # one node rather than no resources.
        utilization = ExecutionResources(cpu=0)
        autoscaler.try_trigger_scaling()
        assert autoscaler.get_total_resources() == ExecutionResources(cpu=1)

    def test_get_node_resource_spec_and_count_skips_max_count_zero(self):
        """Test that node types with max_count=0 are skipped."""
        # Simulate a cluster with only head node (no worker nodes)
        node_table = [
            {
                "Resources": self._head_node,
                "Alive": True,
            },
        ]

        # Create a mock cluster config with one valid node type and one with max_count=0
        cluster_config = autoscaler_pb2.ClusterConfig()

        # Node type 1: 4 CPU, 0 GPU, 1000 memory, max_count=10
        node_group_config1 = autoscaler_pb2.NodeGroupConfig()
        node_group_config1.resources["CPU"] = 4
        node_group_config1.resources["memory"] = 1000
        node_group_config1.max_count = 10
        cluster_config.node_group_configs.append(node_group_config1)

        # Node type 2: 8 CPU, 2 GPU, 2000 memory, max_count=0 (should be skipped)
        node_group_config2 = autoscaler_pb2.NodeGroupConfig()
        node_group_config2.resources["CPU"] = 8
        node_group_config2.resources["GPU"] = 2
        node_group_config2.resources["memory"] = 2000
        node_group_config2.max_count = 0  # This should be skipped
        cluster_config.node_group_configs.append(node_group_config2)

        # Only the first node type should be discovered
        expected = {
            _NodeResourceSpec.of(cpu=4, gpu=0, mem=1000): 0,
        }

        with patch("ray.nodes", return_value=node_table):
            with patch(
                "ray._private.state.state.get_cluster_config",
                return_value=cluster_config,
            ):
                result = _get_node_resource_spec_and_count()
                assert result == expected

    def test_get_node_resource_spec_and_count_missing_all_resources(self):
        """Regression test for nodes with empty resources (ie missing CPU, GPU, and memory keys entirely)."""

        # Simulate a node with no standard resources defined
        node_empty_resources = {