chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,484 @@
|
||||
"""Utilities for e2e tests of KubeRay/Ray integration.
|
||||
For consistency, all K8s interactions use kubectl through subprocess calls.
|
||||
"""
|
||||
import atexit
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
import pathlib
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from typing import Any, Dict, Generator, List, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SCRIPTS_DIR = pathlib.Path(__file__).resolve().parent / "scripts"
|
||||
TEST_CR_PATH = (
|
||||
pathlib.Path(__file__).resolve().parent / "setup" / "raycluster_test.yaml"
|
||||
)
|
||||
TEST_CLUSTER_NAME = "raycluster-test"
|
||||
|
||||
# Parent directory of Ray repository
|
||||
RAY_PARENT = str(pathlib.Path(__file__).resolve().parents[5])
|
||||
|
||||
RAYCLUSTERS_QUALIFIED = "rayclusters.ray.io"
|
||||
|
||||
LOG_FORMAT = "[%(levelname)s %(asctime)s] %(filename)s: %(lineno)d %(message)s"
|
||||
|
||||
|
||||
def switch_to_ray_parent_dir():
|
||||
# Switch to parent of Ray repo, because that's what the doc examples do.
|
||||
logger.info("Switching to parent of Ray directory.")
|
||||
os.chdir(RAY_PARENT)
|
||||
|
||||
|
||||
def setup_kuberay_operator():
|
||||
"""Set up KubeRay operator and Ray autoscaler RBAC."""
|
||||
switch_to_ray_parent_dir()
|
||||
logger.info("Cloning KubeRay and setting up KubeRay configuration.")
|
||||
# For faster run-time when triggering the test locally, don't run the init
|
||||
# script if it has already been run.
|
||||
subprocess.check_call(
|
||||
[
|
||||
"bash",
|
||||
"-c",
|
||||
(
|
||||
"ls ray/python/ray/autoscaler/kuberay/config ||"
|
||||
" ./ray/python/ray/autoscaler/kuberay/init-config.sh"
|
||||
),
|
||||
]
|
||||
)
|
||||
logger.info("Creating KubeRay operator.")
|
||||
subprocess.check_call(
|
||||
[
|
||||
"kubectl",
|
||||
"create",
|
||||
"-k",
|
||||
"ray/python/ray/autoscaler/kuberay/config/default",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def teardown_kuberay_operator():
|
||||
logger.info("Switching to parent of Ray directory.")
|
||||
os.chdir(RAY_PARENT)
|
||||
|
||||
logger.info("Deleting operator.")
|
||||
subprocess.check_call(
|
||||
[
|
||||
"kubectl",
|
||||
"delete",
|
||||
"--ignore-not-found",
|
||||
"-k",
|
||||
"ray/python/ray/autoscaler/kuberay/config/default",
|
||||
]
|
||||
)
|
||||
|
||||
logger.info("Double-checking no pods left over in namespace ray-system.")
|
||||
wait_for_pods(goal_num_pods=0, namespace="ray-system")
|
||||
|
||||
|
||||
def wait_for_raycluster_crd(tries=60, backoff_s=5):
|
||||
"""CRD creation can take a bit of time after the client request.
|
||||
This function waits until the crd with the provided name is registered.
|
||||
"""
|
||||
switch_to_ray_parent_dir()
|
||||
logger.info("Making sure RayCluster CRD has been registered.")
|
||||
for i in range(tries):
|
||||
get_crd_output = subprocess.check_output(["kubectl", "get", "crd"]).decode()
|
||||
if RAYCLUSTERS_QUALIFIED in get_crd_output:
|
||||
logger.info("Confirmed existence of RayCluster CRD.")
|
||||
break
|
||||
elif i < tries - 1:
|
||||
logger.info("Still waiting to register RayCluster CRD.")
|
||||
time.sleep(backoff_s)
|
||||
else:
|
||||
raise Exception("Failed to register RayCluster CRD.")
|
||||
|
||||
# Create a test RayCluster CR to make sure that the CRD is fully registered.
|
||||
for i in range(tries):
|
||||
try:
|
||||
subprocess.check_call(["kubectl", "apply", "-f", TEST_CR_PATH])
|
||||
break
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.info("Can't create RayCluster CR.")
|
||||
if i < tries - 1:
|
||||
logger.info("Retrying.")
|
||||
time.sleep(backoff_s)
|
||||
else:
|
||||
logger.info("Giving up.")
|
||||
raise e from None
|
||||
|
||||
# Confirm the test RayCluster exists.
|
||||
out = subprocess.check_output(["kubectl", "get", RAYCLUSTERS_QUALIFIED]).decode()
|
||||
assert TEST_CLUSTER_NAME in out, out
|
||||
|
||||
# Delete the test RayCluster.
|
||||
subprocess.check_call(["kubectl", "delete", "-f", TEST_CR_PATH])
|
||||
# Make sure the associated resources are gone before proceeding.
|
||||
wait_for_pods(goal_num_pods=0, namespace="default")
|
||||
|
||||
|
||||
def wait_for_pods(goal_num_pods: int, namespace: str, tries=60, backoff_s=5) -> None:
|
||||
"""Wait for the number of pods in the `namespace` to be exactly `num_pods`.
|
||||
|
||||
Raise an exception after exceeding `tries` attempts with `backoff_s` second waits.
|
||||
"""
|
||||
for i in range(tries):
|
||||
cur_num_pods = _get_num_pods(namespace)
|
||||
if cur_num_pods == goal_num_pods:
|
||||
logger.info(f"Confirmed {goal_num_pods} pod(s) in namespace {namespace}.")
|
||||
return
|
||||
elif i < tries - 1:
|
||||
logger.info(
|
||||
f"The number of pods in namespace {namespace} is {cur_num_pods}."
|
||||
f" Waiting until the number of pods is {goal_num_pods}."
|
||||
f"{get_pod_names(namespace)}"
|
||||
)
|
||||
time.sleep(backoff_s)
|
||||
else:
|
||||
raise Exception(
|
||||
f"Failed to scale to {goal_num_pods} pod(s) in namespace {namespace}."
|
||||
)
|
||||
|
||||
|
||||
def _get_num_pods(namespace: str) -> int:
|
||||
return len(get_pod_names(namespace))
|
||||
|
||||
|
||||
def get_pod_names(namespace: str) -> List[str]:
|
||||
"""Get the list of pod names in the namespace."""
|
||||
get_pods_output = (
|
||||
subprocess.check_output(
|
||||
[
|
||||
"kubectl",
|
||||
"-n",
|
||||
namespace,
|
||||
"get",
|
||||
"pods",
|
||||
"-o",
|
||||
"custom-columns=POD:metadata.name",
|
||||
"--no-headers",
|
||||
]
|
||||
)
|
||||
.decode()
|
||||
.strip()
|
||||
)
|
||||
|
||||
# If there aren't any pods, the output is any empty string.
|
||||
if not get_pods_output:
|
||||
return []
|
||||
else:
|
||||
return get_pods_output.split("\n")
|
||||
|
||||
|
||||
def wait_for_pod_to_start(
|
||||
pod_name_filter: str, namespace: str, tries=60, backoff_s=5
|
||||
) -> None:
|
||||
"""Waits for a pod to have Running status.phase.
|
||||
|
||||
More precisely, waits until there is a pod with name containing `pod_name_filter`
|
||||
and the pod has Running status.phase."""
|
||||
for i in range(tries):
|
||||
pod = get_pod(pod_name_filter=pod_name_filter, namespace=namespace)
|
||||
if not pod:
|
||||
# We didn't get a matching pod.
|
||||
continue
|
||||
pod_status = (
|
||||
subprocess.check_output(
|
||||
[
|
||||
"kubectl",
|
||||
"-n",
|
||||
namespace,
|
||||
"get",
|
||||
"pod",
|
||||
pod,
|
||||
"-o",
|
||||
"custom-columns=POD:status.phase",
|
||||
"--no-headers",
|
||||
]
|
||||
)
|
||||
.decode()
|
||||
.strip()
|
||||
)
|
||||
# "not found" is part of the kubectl output if the pod's not there.
|
||||
if "not found" in pod_status:
|
||||
raise Exception(f"Pod {pod} not found.")
|
||||
elif pod_status == "Running":
|
||||
logger.info(f"Confirmed pod {pod} is Running.")
|
||||
return
|
||||
elif i < tries - 1:
|
||||
logger.info(
|
||||
f"Pod {pod} has status {pod_status}. Waiting for the pod to enter "
|
||||
"Running status."
|
||||
)
|
||||
time.sleep(backoff_s)
|
||||
else:
|
||||
raise Exception(f"Timed out waiting for pod {pod} to enter Running status.")
|
||||
|
||||
|
||||
def wait_for_ray_health(
|
||||
pod_name_filter: str,
|
||||
namespace: str,
|
||||
tries=60,
|
||||
backoff_s=5,
|
||||
ray_container="ray-head",
|
||||
) -> None:
|
||||
"""Waits until a Ray pod passes `ray health-check`.
|
||||
|
||||
More precisely, waits until a Ray pod whose name includes the string
|
||||
`pod_name_filter` passes `ray health-check`.
|
||||
(Ensures Ray has completely started in the pod.)
|
||||
|
||||
Use case: Wait until there is a Ray head pod with Ray running on it.
|
||||
"""
|
||||
for i in range(tries):
|
||||
try:
|
||||
pod = get_pod(pod_name_filter=pod_name_filter, namespace="default")
|
||||
assert pod, f"Couldn't find a pod matching {pod_name_filter}."
|
||||
# `ray health-check` yields 0 exit status iff it succeeds
|
||||
kubectl_exec(
|
||||
["ray", "health-check"], pod, namespace, container=ray_container
|
||||
)
|
||||
logger.info(f"ray health check passes for pod {pod}")
|
||||
return
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.info(f"Failed ray health check for pod {pod}.")
|
||||
if i < tries - 1:
|
||||
logger.info("Trying again.")
|
||||
time.sleep(backoff_s)
|
||||
else:
|
||||
logger.info("Giving up.")
|
||||
raise e from None
|
||||
|
||||
|
||||
def get_pod(pod_name_filter: str, namespace: str) -> Optional[str]:
|
||||
"""Gets pods in the `namespace`.
|
||||
|
||||
Returns the first pod that has `pod_name_filter` as a
|
||||
substring of its name. Returns None if there are no matches.
|
||||
"""
|
||||
pod_names = get_pod_names(namespace)
|
||||
matches = [pod_name for pod_name in pod_names if pod_name_filter in pod_name]
|
||||
if not matches:
|
||||
logger.warning(f"No match for `{pod_name_filter}` in namespace `{namespace}`.")
|
||||
return None
|
||||
return matches[0]
|
||||
|
||||
|
||||
def kubectl_exec(
|
||||
command: List[str],
|
||||
pod: str,
|
||||
namespace: str,
|
||||
container: Optional[str] = None,
|
||||
) -> str:
|
||||
"""kubectl exec the `command` in the given `pod` in the given `namespace`.
|
||||
If a `container` is specified, will specify that container for kubectl.
|
||||
|
||||
Prints and return kubectl's output as a string.
|
||||
"""
|
||||
container_option = ["-c", container] if container else []
|
||||
kubectl_exec_command = (
|
||||
["kubectl", "exec", "-it", pod] + container_option + ["--"] + command
|
||||
)
|
||||
# Print for debugging convenience.
|
||||
try:
|
||||
out = subprocess.check_output(kubectl_exec_command).decode().strip()
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.error(f"Error running command {kubectl_exec_command}.")
|
||||
logger.error(f"Output: {e.output.decode()}")
|
||||
raise e from None
|
||||
print(out)
|
||||
return out
|
||||
|
||||
|
||||
def kubectl_logs(
|
||||
pod: str,
|
||||
namespace: str,
|
||||
container: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Wrapper for kubectl logs.
|
||||
|
||||
Returns the logs as a string.
|
||||
"""
|
||||
container_option = ["-c", container] if container else []
|
||||
kubectl_logs_command = ["kubectl", "logs", pod] + container_option
|
||||
out = subprocess.check_output(kubectl_logs_command).decode().strip()
|
||||
return out
|
||||
|
||||
|
||||
def kubectl_exec_python_script(
|
||||
script_name: str,
|
||||
pod: str,
|
||||
namespace: str,
|
||||
container: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Runs a python script in a container via `kubectl exec`.
|
||||
Scripts live in `tests/kuberay/scripts`.
|
||||
|
||||
Prints and return kubectl's output as a string.
|
||||
"""
|
||||
script_path = SCRIPTS_DIR / script_name
|
||||
with open(script_path) as script_file:
|
||||
script_string = script_file.read()
|
||||
return kubectl_exec(["python", "-c", script_string], pod, namespace, container)
|
||||
|
||||
|
||||
def get_raycluster(raycluster: str, namespace: str) -> Dict[str, Any]:
|
||||
"""Gets the Ray CR with name `raycluster` in namespace `namespace`.
|
||||
|
||||
Returns the CR as a nested Dict.
|
||||
"""
|
||||
get_raycluster_output = (
|
||||
subprocess.check_output(
|
||||
["kubectl", "-n", namespace, "get", "raycluster", raycluster, "-o", "yaml"]
|
||||
)
|
||||
.decode()
|
||||
.strip()
|
||||
)
|
||||
return yaml.safe_load(get_raycluster_output)
|
||||
|
||||
|
||||
def _get_service_port(service: str, namespace: str, target_port: int) -> int:
|
||||
"""Given a K8s service and a port targetted by the service, returns the
|
||||
corresponding port exposed by the service.
|
||||
|
||||
Args:
|
||||
service: Name of a K8s service.
|
||||
namespace: Namespace to which the service belongs.
|
||||
target_port: Port targeted by the service.
|
||||
|
||||
Returns:
|
||||
service_port: The port exposed by the service.
|
||||
"""
|
||||
service_str = (
|
||||
subprocess.check_output(
|
||||
["kubectl", "-n", namespace, "get", "service", service, "-o", "yaml"]
|
||||
)
|
||||
.decode()
|
||||
.strip()
|
||||
)
|
||||
service_dict = yaml.safe_load(service_str)
|
||||
service_ports: List = service_dict["spec"]["ports"]
|
||||
matching_ports = [
|
||||
port for port in service_ports if port["targetPort"] == target_port
|
||||
]
|
||||
assert matching_ports
|
||||
service_port = matching_ports[0]["port"]
|
||||
return service_port
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _kubectl_port_forward(
|
||||
service: str, namespace: str, target_port: int, local_port: Optional[int] = None
|
||||
) -> Generator[int, None, None]:
|
||||
"""Context manager which creates a kubectl port-forward process targeting a
|
||||
K8s service.
|
||||
|
||||
Terminates the port-forwarding process upon exit.
|
||||
|
||||
Args:
|
||||
service: Name of a K8s service.
|
||||
namespace: Namespace to which the service belongs.
|
||||
target_port: The port targeted by the service.
|
||||
local_port: Forward from this port. Optional. By default, uses the port exposed
|
||||
by the service.
|
||||
|
||||
Yields:
|
||||
int: The local port. The service can then be accessed at
|
||||
127.0.0.1:<local_port>.
|
||||
"""
|
||||
# First, figure out which port the service exposes for the given target port.
|
||||
service_port = _get_service_port(service, namespace, target_port)
|
||||
if not local_port:
|
||||
local_port = service_port
|
||||
|
||||
process = subprocess.Popen(
|
||||
[
|
||||
"kubectl",
|
||||
"-n",
|
||||
namespace,
|
||||
"port-forward",
|
||||
f"service/{service}",
|
||||
f"{local_port}:{service_port}",
|
||||
]
|
||||
)
|
||||
|
||||
def terminate_process():
|
||||
process.terminate()
|
||||
# Wait 10 seconds for the process to terminate.
|
||||
# This cleans up the zombie entry from the process table.
|
||||
# 10 seconds is a deliberately excessive amount of time to wait.
|
||||
process.wait(timeout=10)
|
||||
|
||||
# Ensure clean-up in case of interrupt.
|
||||
atexit.register(terminate_process)
|
||||
# terminate_process is ok to execute multiple times.
|
||||
|
||||
try:
|
||||
yield local_port
|
||||
finally:
|
||||
terminate_process()
|
||||
|
||||
|
||||
def kubectl_patch(
|
||||
kind: str,
|
||||
name: str,
|
||||
namespace: str,
|
||||
patch: Dict[str, Any],
|
||||
patch_type: str = "strategic",
|
||||
):
|
||||
"""Wrapper for kubectl patch.
|
||||
|
||||
Args:
|
||||
kind: Kind of the K8s resource (e.g. pod)
|
||||
name: Name of the K8s resource.
|
||||
namespace: Namespace of the K8s resource.
|
||||
patch: The patch to apply, as a dict.
|
||||
patch_type: json, merge, or strategic
|
||||
"""
|
||||
with tempfile.NamedTemporaryFile("w") as patch_file:
|
||||
yaml.dump(patch, patch_file)
|
||||
patch_file.flush()
|
||||
subprocess.check_call(
|
||||
[
|
||||
"kubectl",
|
||||
"-n",
|
||||
f"{namespace}",
|
||||
"patch",
|
||||
f"{kind}",
|
||||
f"{name}",
|
||||
"--patch-file",
|
||||
f"{patch_file.name}",
|
||||
"--type",
|
||||
f"{patch_type}",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def kubectl_delete(kind: str, name: str, namespace: str, wait: bool = True):
|
||||
"""Wrapper for kubectl delete.
|
||||
|
||||
Args:
|
||||
kind: Kind of the K8s resource (e.g. pod)
|
||||
name: Name of the K8s resource.
|
||||
namespace: Namespace of the K8s resource.
|
||||
wait: Whether to pass ``--wait=true`` so ``kubectl`` blocks until the
|
||||
resource is fully removed.
|
||||
"""
|
||||
wait_str = "true" if wait else "false"
|
||||
subprocess.check_output(
|
||||
[
|
||||
"kubectl",
|
||||
"-n",
|
||||
f"{namespace}",
|
||||
"delete",
|
||||
f"{kind}",
|
||||
f"{name}",
|
||||
f"--wait={wait_str}",
|
||||
]
|
||||
)
|
||||
Reference in New Issue
Block a user