chore: import upstream snapshot with attribution

python/ray/tests/kuberay/scripts/check_cpu_and_memory.py

@@ -0,0 +1,17 @@
+import ray
+
+
+def main():
+    """This script runs in a container with 1 CPU limit and 1G memory limit.
+    Validate that Ray reads the correct limits.
+    """
+    cpu_limit = ray._private.utils.get_num_cpus()
+    mem_limit_gb = round(ray._common.utils.get_system_memory() / 10**9, 2)
+    assert cpu_limit == 1, cpu_limit
+    assert mem_limit_gb == 2.00, mem_limit_gb
+    print(f"Confirmed cpu limit {cpu_limit}.")
+    print(f"Confirmed memory limit {mem_limit_gb} gigabyte.")
+
+
+if __name__ == "__main__":
+    main()

python/ray/tests/kuberay/scripts/gpu_actor_placement.py

@@ -0,0 +1,18 @@
+import ray
+
+
+def main():
+    """Requests placement of a GPU actor."""
+
+    @ray.remote(num_gpus=1, num_cpus=1)
+    class GPUActor:
+        def where_am_i(self):
+            assert len(ray.get_gpu_ids()) == 1
+            return "on-a-gpu-node"
+
+    GPUActor.options(name="gpu_actor", lifetime="detached").remote()
+
+
+if __name__ == "__main__":
+    ray.init("auto", namespace="gpu-test")
+    main()

python/ray/tests/kuberay/scripts/gpu_actor_validation.py

@@ -0,0 +1,14 @@
+import ray
+
+
+def main():
+    """Confirms placement of a GPU actor."""
+    gpu_actor = ray.get_actor("gpu_actor")
+    actor_response = ray.get(gpu_actor.where_am_i.remote())
+    return actor_response
+
+
+if __name__ == "__main__":
+    ray.init("auto", namespace="gpu-test")
+    out = main()
+    print(out)

python/ray/tests/kuberay/scripts/non_terminated_nodes_count.py

@@ -0,0 +1,26 @@
+import ray
+from ray.autoscaler._private.kuberay.autoscaling_config import _generate_provider_config
+from ray.autoscaler._private.providers import _get_node_provider
+
+
+@ray.remote
+def count_non_terminated_nodes() -> int:
+    """Get the count of non terminated nodes for the Ray cluster raycluster-autoscaler
+    in namespace default.
+    """
+    provider_config = _generate_provider_config(ray_cluster_namespace="default")
+    kuberay_node_provider = _get_node_provider(
+        provider_config=provider_config, cluster_name="raycluster-autoscaler"
+    )
+    nodes = kuberay_node_provider.non_terminated_nodes({})
+    return len(nodes)
+
+
+def main() -> int:
+    return ray.get(count_non_terminated_nodes.remote())
+
+
+if __name__ == "__main__":
+    ray.init("auto")
+    out = main()
+    print(out)

python/ray/tests/kuberay/scripts/scale_down.py

@@ -0,0 +1,37 @@
+import ray
+from ray._common import test_utils
+
+
+def main():
+    """Removes CPU request, removes GPU actor.
+    Waits for autoscaler scale-down events to get emitted to stdout.
+
+    The worker idle timeout is set to 10 seconds and the autoscaler's update interval is
+    5 seconds, so it should be enough to wait 15 seconds.
+    """
+
+    # Before scale-down.
+    cluster_resources = ray.cluster_resources()
+    assert cluster_resources.get("CPU", 0) > 0, cluster_resources
+    assert cluster_resources.get("GPU", 0) > 0, cluster_resources
+
+    # Remove resource demands
+    ray.autoscaler.sdk.request_resources(num_cpus=0)
+    gpu_actor = ray.get_actor("gpu_actor")
+    ray.kill(gpu_actor)
+
+    # Wait for scale-down to happen.
+    def verify():
+        cluster_resources = ray.cluster_resources()
+        # From head node
+        assert cluster_resources.get("CPU", 0) == 1, cluster_resources
+        assert cluster_resources.get("GPU", 0) == 0, cluster_resources
+
+        return True
+
+    test_utils.wait_for_condition(verify, timeout=60, retry_interval_ms=2000)
+
+
+if __name__ == "__main__":
+    ray.init("auto", namespace="gpu-test")
+    main()

python/ray/tests/kuberay/scripts/scale_up.py

@@ -0,0 +1,27 @@
+import ray
+from ray._common.test_utils import wait_for_condition
+
+
+def main():
+    """Submits CPU request"""
+    ray.autoscaler.sdk.request_resources(num_cpus=2)
+    from ray.autoscaler.v2.sdk import get_cluster_status
+    from ray.autoscaler.v2.utils import ClusterStatusFormatter
+
+    gcs_address = ray.get_runtime_context().gcs_address
+
+    def verify():
+        cluster_resources = ray.cluster_resources()
+
+        cluster_status = get_cluster_status(gcs_address)
+        print(ClusterStatusFormatter.format(cluster_status, verbose=True))
+        assert cluster_resources.get("CPU", 0) == 2, cluster_resources
+
+        return True
+
+    wait_for_condition(verify, timeout=60, retry_interval_ms=2000)
+
+
+if __name__ == "__main__":
+    ray.init("auto")
+    main()

python/ray/tests/kuberay/scripts/scale_up_custom.py

@@ -0,0 +1,32 @@
+import time
+
+import ray
+
+
+def main():
+    """Submits custom resource request.
+
+    Also, validates runtime env data submitted with the Ray Job that executes
+    this script.
+    """
+
+    # Workers and head are annotated as having 5 "Custom2" capacity each,
+    # so this should trigger upscaling of two workers.
+    # (One of the bundles will be "placed" on the head.)
+    ray.autoscaler.sdk.request_resources(
+        bundles=[{"Custom2": 3}, {"Custom2": 3}, {"Custom2": 3}]
+    )
+
+    while (
+        ray.cluster_resources().get("Custom2", 0) < 3
+        and ray.cluster_resources().get("Custom2", 0) < 6
+    ):
+        time.sleep(0.1)
+
+    # Output something to validate the job logs.
+    print("Submitted custom scale request!")
+
+
+if __name__ == "__main__":
+    ray.init("auto")
+    main()