wehub-resource-sync
608 lines
20 KiB
Python
608 lines
20 KiB
Python
"""
|
|
This script provides extra functionality for Anyscale Jobs tests.
|
|
It will be ran on the cluster.
|
|
|
|
We need to reimplement some utility functions here as it will not
|
|
have access to the ray_release package.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import multiprocessing
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import List, Optional, Tuple
|
|
from urllib.parse import urlparse
|
|
|
|
AZURE_STORAGE_ACCOUNT = "rayreleasetests"
|
|
OUTPUT_JSON_FILENAME = "output.json"
|
|
AWS_CP_TIMEOUT = 300
|
|
TIMEOUT_RETURN_CODE = 124 # same as bash timeout
|
|
|
|
# Prometheus metric type for idle worker evictions. We expect Ray to kill idle workers
|
|
# under memory pressure, so we exclude them from the OOM check.
|
|
IDLE_WORKER_EVICTION_METRIC_TYPE = "MemoryManager.IdleWorkerEviction.Total"
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
handler = logging.StreamHandler(stream=sys.stderr)
|
|
formatter = logging.Formatter(
|
|
fmt="[%(levelname)s %(asctime)s] %(filename)s: %(lineno)d %(message)s"
|
|
)
|
|
handler.setFormatter(formatter)
|
|
logger.addHandler(handler)
|
|
|
|
|
|
def exponential_backoff_retry(
|
|
f, retry_exceptions, initial_retry_delay_s, max_retries
|
|
) -> None:
|
|
retry_cnt = 0
|
|
retry_delay_s = initial_retry_delay_s
|
|
while True:
|
|
try:
|
|
return f()
|
|
except retry_exceptions as e:
|
|
retry_cnt += 1
|
|
if retry_cnt > max_retries:
|
|
raise
|
|
logger.warning(
|
|
f"Retry function call failed due to {e} "
|
|
f"in {retry_delay_s} seconds..."
|
|
)
|
|
time.sleep(retry_delay_s)
|
|
retry_delay_s *= 2
|
|
|
|
|
|
def run_storage_cp(source: str, target: str):
|
|
if not source or not target:
|
|
return False
|
|
|
|
if not Path(source).exists():
|
|
logger.warning(f"Couldn't upload to cloud storage: '{source}' does not exist.")
|
|
return False
|
|
|
|
storage_service = urlparse(target).scheme
|
|
if target.startswith(f"https://{AZURE_STORAGE_ACCOUNT}.dfs.core.windows.net"):
|
|
storage_service = "azure_blob"
|
|
cp_cmd_args = []
|
|
if storage_service == "s3":
|
|
cp_cmd_args = [
|
|
"aws",
|
|
"s3",
|
|
"cp",
|
|
source,
|
|
target,
|
|
"--acl",
|
|
"bucket-owner-full-control",
|
|
]
|
|
elif storage_service == "gs":
|
|
cp_cmd_args = [
|
|
"gcloud",
|
|
"storage",
|
|
"cp",
|
|
source,
|
|
target,
|
|
]
|
|
elif storage_service == "azure_blob":
|
|
subprocess.run(["azcopy", "login", "--identity"], check=True)
|
|
cp_cmd_args = [
|
|
"azcopy",
|
|
"copy",
|
|
source,
|
|
target,
|
|
]
|
|
else:
|
|
raise Exception(f"Not supporting storage service: {storage_service}")
|
|
|
|
try:
|
|
exponential_backoff_retry(
|
|
lambda: subprocess.run(
|
|
cp_cmd_args,
|
|
timeout=AWS_CP_TIMEOUT,
|
|
check=True,
|
|
),
|
|
subprocess.SubprocessError,
|
|
initial_retry_delay_s=10,
|
|
max_retries=3,
|
|
)
|
|
return True
|
|
except subprocess.SubprocessError:
|
|
logger.exception("Couldn't upload to cloud storage.")
|
|
return False
|
|
|
|
|
|
def collect_metrics(start_time: float, time_taken: float) -> bool:
|
|
if "METRICS_OUTPUT_JSON" not in os.environ:
|
|
return False
|
|
|
|
# Timeout is the time the test took divided by 200
|
|
# (~7 minutes for a 24h test) but no less than 90s
|
|
# and no more than 900s
|
|
metrics_timeout = max(90, min(time_taken / 200, 900))
|
|
logger.info(f"Collecting Prometheus metrics (timeout: {metrics_timeout:.0f}s).")
|
|
try:
|
|
subprocess.run(
|
|
[
|
|
"python",
|
|
"prometheus_metrics.py",
|
|
str(start_time),
|
|
"--path",
|
|
os.environ["METRICS_OUTPUT_JSON"],
|
|
],
|
|
timeout=metrics_timeout,
|
|
check=True,
|
|
)
|
|
logger.info("Metrics collection subprocess finished successfully.")
|
|
return True
|
|
# TimeoutExpired and CalledProcessError are SubprocessError subclasses, so
|
|
# they must be caught first to differentiate them in the logs.
|
|
except subprocess.TimeoutExpired:
|
|
logger.error(
|
|
f"Metrics collection TIMED OUT after {metrics_timeout:.0f}s. The metrics "
|
|
"file may be missing or incomplete. This is a metrics-collection timeout, "
|
|
"distinct from an actual metric/OOM/spill issue."
|
|
)
|
|
return False
|
|
except subprocess.CalledProcessError as e:
|
|
logger.error(
|
|
f"Metrics collection subprocess exited with non-zero return code "
|
|
f"{e.returncode}. See the prometheus_metrics.py output above for the "
|
|
"specific failure."
|
|
)
|
|
return False
|
|
except subprocess.SubprocessError:
|
|
logger.exception("Couldn't collect metrics due to an unexpected error.")
|
|
return False
|
|
|
|
|
|
# Has to be here so it can be pickled
|
|
def _run_bash_command_subprocess(command: str, timeout: float):
|
|
"""Ran in a multiprocessing process."""
|
|
try:
|
|
subprocess.run(command, check=True, timeout=timeout)
|
|
return_code = 0
|
|
except subprocess.TimeoutExpired:
|
|
return_code = TIMEOUT_RETURN_CODE
|
|
except subprocess.CalledProcessError as e:
|
|
return_code = e.returncode
|
|
print(f"Subprocess return code: {return_code}", file=sys.stderr)
|
|
# Exit so the return code is propagated to the outer process
|
|
sys.exit(return_code)
|
|
|
|
|
|
def run_bash_command(workload: str, timeout: float):
|
|
timeout = timeout if timeout > 0 else None
|
|
cwd = Path.cwd()
|
|
workload_path = cwd / "workload.sh"
|
|
workload_path = workload_path.resolve()
|
|
with open(workload_path, "w") as fp:
|
|
fp.write(workload)
|
|
|
|
command = ["bash", "-x", str(workload_path)]
|
|
logger.info(f"Running command {workload}")
|
|
|
|
# Pop job's runtime env to allow workload's runtime env to take precedence
|
|
# TODO: Confirm this is safe
|
|
os.environ.pop("RAY_JOB_CONFIG_JSON_ENV_VAR", None)
|
|
|
|
# We use multiprocessing with 'spawn' context to avoid
|
|
# forking (as happens when using subprocess directly).
|
|
# Forking messes up Ray interactions and causes deadlocks.
|
|
return_code = None
|
|
try:
|
|
ctx = multiprocessing.get_context("spawn")
|
|
p = ctx.Process(target=_run_bash_command_subprocess, args=(command, timeout))
|
|
p.start()
|
|
logger.info(f"Starting process {p.pid}.")
|
|
# Add a little extra to the timeout as _run_bash_command_subprocess
|
|
# also has a timeout internally and it's cleaner to use that
|
|
p.join(timeout=timeout + 10)
|
|
except multiprocessing.TimeoutError:
|
|
return_code = TIMEOUT_RETURN_CODE
|
|
except multiprocessing.ProcessError:
|
|
pass
|
|
finally:
|
|
if p.is_alive():
|
|
logger.warning(f"Terminating process {p.pid} forcefully.")
|
|
p.terminate()
|
|
if return_code is None:
|
|
return_code = p.exitcode
|
|
os.remove(str(workload_path))
|
|
logger.info(f"Process {p.pid} exited with return code {return_code}.")
|
|
assert return_code is not None
|
|
return return_code
|
|
|
|
|
|
def run_prepare_commands(
|
|
prepare_commands: List[str], prepare_commands_timeouts: List[float]
|
|
) -> Tuple[bool, List[int], float]:
|
|
"""Run prepare commands. All commands must pass. Fails fast."""
|
|
prepare_return_codes = []
|
|
prepare_passed = True
|
|
prepare_time_taken = None
|
|
|
|
if not prepare_commands:
|
|
return prepare_passed, prepare_return_codes, prepare_time_taken
|
|
|
|
logger.info("### Starting prepare commands ###")
|
|
for prepare_command, timeout in zip(prepare_commands, prepare_commands_timeouts):
|
|
command_start_time = time.monotonic()
|
|
prepare_return_codes.append(run_bash_command(prepare_command, timeout))
|
|
prepare_time_taken = time.monotonic() - command_start_time
|
|
return_code = prepare_return_codes[-1]
|
|
|
|
if return_code == 0:
|
|
continue
|
|
|
|
timed_out = return_code == TIMEOUT_RETURN_CODE
|
|
if timed_out:
|
|
logger.error(
|
|
"Prepare command timed out. " f"Time taken: {prepare_time_taken}"
|
|
)
|
|
else:
|
|
logger.info(
|
|
f"Prepare command finished with return code {return_code}. "
|
|
f"Time taken: {prepare_time_taken}"
|
|
)
|
|
logger.error("Prepare command failed.")
|
|
prepare_passed = False
|
|
break
|
|
|
|
return prepare_passed, prepare_return_codes, prepare_time_taken
|
|
|
|
|
|
def _load_metrics_for_check(check_name: str, env_var: str) -> Optional[dict]:
|
|
"""Load the Prometheus metrics file for a failure check.
|
|
|
|
Returns the parsed metrics dict, or ``None`` when the metrics could not be
|
|
obtained at all (file missing, unreadable, or an empty ``{}`` written
|
|
because every Prometheus query failed). In every ``None`` case this is a
|
|
metrics-collection/infra failure rather than an actual metric signal, and
|
|
the caller should treat it as such.
|
|
"""
|
|
metrics_path = os.environ.get("METRICS_OUTPUT_JSON", None)
|
|
if not (metrics_path and Path(metrics_path).exists()):
|
|
logger.error(
|
|
f"{check_name}: {env_var} is set to 1, but no metrics file was found "
|
|
f"at path: {metrics_path}. Metrics collection failed entirely."
|
|
)
|
|
return None
|
|
try:
|
|
with open(metrics_path, "r") as f:
|
|
metrics = json.load(f)
|
|
except (OSError, json.JSONDecodeError) as e:
|
|
logger.error(f"{check_name}: could not read metrics file {metrics_path}: {e}")
|
|
return None
|
|
if not isinstance(metrics, dict) or not metrics:
|
|
logger.error(
|
|
f"{check_name}: metrics file at {metrics_path} is empty. "
|
|
"See the prometheus_metrics.py output above for the cause."
|
|
)
|
|
return None
|
|
return metrics
|
|
|
|
|
|
def _metric_unavailable(check_name: str, metrics: dict, key: str) -> bool:
|
|
"""Return True if ``key`` could not be collected (missing or null).
|
|
|
|
Distinguishes a metrics-collection/infra failure (logged here) from an
|
|
actual metric signal, which the caller inspects when this returns False.
|
|
A ``None`` value means the Prometheus query failed; an empty list ``[]``
|
|
means the query succeeded but matched no series (i.e. a healthy result).
|
|
"""
|
|
if key not in metrics:
|
|
logger.error(
|
|
f"{check_name}: '{key}' is missing from the metrics file, likely a collection issue."
|
|
)
|
|
return True
|
|
if metrics[key] is None:
|
|
logger.error(
|
|
f"{check_name}: '{key}' is None, likely the Prometheus query failed "
|
|
"(timeout / connection error / non-200)"
|
|
)
|
|
return True
|
|
return False
|
|
|
|
|
|
def run_oom_check():
|
|
metrics = _load_metrics_for_check("OOM check", "RAYTEST_FAIL_ON_WORKER_OOM")
|
|
if metrics is None:
|
|
return 1
|
|
|
|
return_code = 0
|
|
if _metric_unavailable("OOM check", metrics, "worker_oom_kills"):
|
|
return_code = 1
|
|
else:
|
|
worker_oom_kills = _filter_idle_worker_kills(metrics["worker_oom_kills"])
|
|
if worker_oom_kills:
|
|
logger.error(
|
|
f"Test failed: OOM worker kills detected. Details: {worker_oom_kills}"
|
|
)
|
|
return_code = 1
|
|
|
|
if _metric_unavailable("OOM check", metrics, "unexpected_worker_failures"):
|
|
return_code = 1
|
|
elif metrics["unexpected_worker_failures"]:
|
|
logger.error(
|
|
"Test failed: Unexpected worker failures detected "
|
|
"(potential kernel OOM kills or SIGKILLs not captured by Ray's memory monitor). "
|
|
f"Details: {metrics['unexpected_worker_failures']}"
|
|
)
|
|
return_code = 1
|
|
return return_code
|
|
|
|
|
|
def _filter_idle_worker_kills(worker_oom_kills: list) -> list:
|
|
"""Drop idle-worker evictions from the worker OOM kill series.
|
|
|
|
Idle-worker evictions are expected behavior, so we exclude them and only keep task
|
|
and actor kills.
|
|
"""
|
|
return [
|
|
series
|
|
for series in worker_oom_kills
|
|
if series.get("metric", {}).get("Type") != IDLE_WORKER_EVICTION_METRIC_TYPE
|
|
]
|
|
|
|
|
|
def run_spilling_check():
|
|
metrics = _load_metrics_for_check("Spilling check", "RAYTEST_FAIL_ON_SPILLING")
|
|
if metrics is None:
|
|
return 1
|
|
|
|
return_code = 0
|
|
if _metric_unavailable("Spilling check", metrics, "spilled_bytes"):
|
|
return_code = 1
|
|
elif metrics["spilled_bytes"]:
|
|
logger.error(
|
|
"Test failed: unexpected object-store spilling detected. "
|
|
f"Details: {metrics['spilled_bytes']}"
|
|
)
|
|
return_code = 1
|
|
return return_code
|
|
|
|
|
|
def run_dead_node_check():
|
|
# Connect to the cluster and check for dead nodes
|
|
import ray
|
|
from ray.core.generated import common_pb2
|
|
|
|
return_code = 0
|
|
try:
|
|
ray.init(address="auto") # Connect to the local cluster
|
|
unexpected_termination = common_pb2.NodeDeathInfo.Reason.Value(
|
|
"UNEXPECTED_TERMINATION"
|
|
)
|
|
unspecified = common_pb2.NodeDeathInfo.Reason.Value("UNSPECIFIED")
|
|
dead_nodes = [
|
|
node["NodeID"]
|
|
for node in ray.nodes()
|
|
if not node["Alive"]
|
|
and node.get("DeathReason") in [unexpected_termination, unspecified]
|
|
]
|
|
if dead_nodes:
|
|
logger.error(f"Dead nodes found, node IDs: {dead_nodes}")
|
|
return_code = 1
|
|
except Exception as e:
|
|
logger.error(f"Error during dead node check: {e}")
|
|
return_code = 1
|
|
finally:
|
|
ray.shutdown() # Disconnect from the cluster
|
|
return return_code
|
|
|
|
|
|
def main(
|
|
test_workload: str,
|
|
test_workload_timeout: float,
|
|
test_no_raise_on_timeout: bool,
|
|
results_cloud_storage_uri: Optional[str],
|
|
metrics_cloud_storage_uri: Optional[str],
|
|
output_cloud_storage_uri: Optional[str],
|
|
upload_cloud_storage_uri: Optional[str],
|
|
artifact_path: Optional[str],
|
|
prepare_commands: List[str],
|
|
prepare_commands_timeouts: List[str],
|
|
):
|
|
"""
|
|
This function provides extra functionality for an Anyscale Job.
|
|
|
|
1. Runs prepare commands and handles their timeouts
|
|
2. Runs the actual test workload and handles its timeout
|
|
3. Uploads test results.json
|
|
4. Gathers prometheus metrics
|
|
5. Uploads prometheus metrics.json
|
|
6. Uploads output.json
|
|
"""
|
|
logger.info("### Starting ###")
|
|
start_time = time.monotonic()
|
|
|
|
if len(prepare_commands) != len(prepare_commands_timeouts):
|
|
raise ValueError(
|
|
"`prepare_commands` and `prepare_commands_timeouts` must "
|
|
"have the same length."
|
|
)
|
|
|
|
# Run prepare commands. All prepare commands must pass.
|
|
(
|
|
prepare_passed,
|
|
prepare_return_codes,
|
|
last_prepare_time_taken,
|
|
) = run_prepare_commands(prepare_commands, prepare_commands_timeouts)
|
|
|
|
uploaded_results = False
|
|
collected_metrics = False
|
|
uploaded_metrics = False
|
|
uploaded_artifact = artifact_path is not None
|
|
workload_time_taken = None
|
|
|
|
# If all prepare commands passed, run actual test workload.
|
|
if prepare_passed:
|
|
logger.info("### Starting entrypoint ###")
|
|
command_start_time = time.monotonic()
|
|
workload_start_unix_time = time.time()
|
|
return_code = run_bash_command(test_workload, test_workload_timeout)
|
|
workload_time_taken = time.monotonic() - command_start_time
|
|
|
|
timed_out = return_code == TIMEOUT_RETURN_CODE
|
|
if timed_out:
|
|
msg = f"Timed out. Time taken: {workload_time_taken}"
|
|
if test_no_raise_on_timeout:
|
|
logger.info(msg)
|
|
else:
|
|
logger.error(msg)
|
|
else:
|
|
logger.info(
|
|
f"Finished with return code {return_code}. "
|
|
f"Time taken: {workload_time_taken}"
|
|
)
|
|
|
|
test_fail_on_dead_nodes = os.environ.get("RAYTEST_FAIL_ON_DEAD_NODES") == "1"
|
|
|
|
if return_code == 0 and test_fail_on_dead_nodes:
|
|
return_code = run_dead_node_check()
|
|
|
|
# Upload results.json
|
|
uploaded_results = run_storage_cp(
|
|
os.environ.get("TEST_OUTPUT_JSON", None), results_cloud_storage_uri
|
|
)
|
|
|
|
# Collect prometheus metrics
|
|
collected_metrics = collect_metrics(
|
|
workload_start_unix_time, workload_time_taken
|
|
)
|
|
if collected_metrics:
|
|
# Upload prometheus metrics
|
|
uploaded_metrics = run_storage_cp(
|
|
os.environ.get("METRICS_OUTPUT_JSON", None), metrics_cloud_storage_uri
|
|
)
|
|
|
|
test_fail_on_worker_oom = os.environ.get("RAYTEST_FAIL_ON_WORKER_OOM") == "1"
|
|
|
|
# Fail if any OOM kills occurred
|
|
if return_code == 0 and test_fail_on_worker_oom:
|
|
return_code = run_oom_check()
|
|
|
|
test_fail_on_spilling = os.environ.get("RAYTEST_FAIL_ON_SPILLING") == "1"
|
|
|
|
# Fail if any object-store spilling occurred
|
|
if return_code == 0 and test_fail_on_spilling:
|
|
return_code = run_spilling_check()
|
|
|
|
uploaded_artifact = run_storage_cp(
|
|
artifact_path,
|
|
os.path.join(
|
|
upload_cloud_storage_uri, os.environ["USER_GENERATED_ARTIFACT"]
|
|
)
|
|
if "USER_GENERATED_ARTIFACT" in os.environ
|
|
else None,
|
|
)
|
|
|
|
else:
|
|
return_code = None
|
|
|
|
total_time_taken = time.monotonic() - start_time
|
|
output_json = {
|
|
"return_code": return_code,
|
|
"prepare_return_codes": prepare_return_codes,
|
|
"last_prepare_time_taken": last_prepare_time_taken,
|
|
"workload_time_taken": workload_time_taken,
|
|
"total_time_taken": total_time_taken,
|
|
"uploaded_results": uploaded_results,
|
|
"collected_metrics": collected_metrics,
|
|
"uploaded_metrics": uploaded_metrics,
|
|
"uploaded_artifact": uploaded_artifact,
|
|
}
|
|
output_json = json.dumps(
|
|
output_json, ensure_ascii=True, sort_keys=True, separators=(",", ":")
|
|
)
|
|
|
|
output_json_file = (Path.cwd() / OUTPUT_JSON_FILENAME).resolve()
|
|
with open(output_json_file, "w") as fp:
|
|
fp.write(output_json)
|
|
|
|
# Upload output.json
|
|
run_storage_cp(str(output_json_file), output_cloud_storage_uri)
|
|
|
|
logger.info("### Finished ###")
|
|
# This will be read by the AnyscaleJobRunner on the buildkite runner
|
|
# if output.json cannot be obtained from cloud storage
|
|
logger.info(f"### JSON |{output_json}| ###")
|
|
|
|
# Flush buffers
|
|
logging.shutdown()
|
|
print("", flush=True)
|
|
print("", file=sys.stderr, flush=True)
|
|
|
|
if return_code == TIMEOUT_RETURN_CODE and test_no_raise_on_timeout:
|
|
return_code = 0
|
|
elif return_code is None:
|
|
return_code = 1
|
|
|
|
time.sleep(1)
|
|
return return_code
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"test_workload", type=str, help="test workload, eg. python workloads/script.py"
|
|
)
|
|
parser.add_argument(
|
|
"--test-workload-timeout",
|
|
default=3600,
|
|
type=float,
|
|
help="test workload timeout (set to <0 for infinite)",
|
|
)
|
|
parser.add_argument(
|
|
"--test-no-raise-on-timeout",
|
|
action="store_true",
|
|
help="don't fail on timeout",
|
|
)
|
|
parser.add_argument(
|
|
"--results-cloud-storage-uri",
|
|
type=str,
|
|
help="bucket address to upload results.json to",
|
|
required=False,
|
|
)
|
|
parser.add_argument(
|
|
"--metrics-cloud-storage-uri",
|
|
type=str,
|
|
help="bucket address to upload metrics.json to",
|
|
required=False,
|
|
)
|
|
parser.add_argument(
|
|
"--output-cloud-storage-uri",
|
|
type=str,
|
|
help="bucket address to upload output.json to",
|
|
required=False,
|
|
)
|
|
parser.add_argument(
|
|
"--upload-cloud-storage-uri",
|
|
type=str,
|
|
help="root cloud-storage bucket address to upload stuff",
|
|
required=False,
|
|
)
|
|
parser.add_argument(
|
|
"--artifact-path",
|
|
type=str,
|
|
help="user provided artifact path (on head node), must be a single file path",
|
|
required=False,
|
|
)
|
|
parser.add_argument(
|
|
"--prepare-commands", type=str, nargs="*", help="prepare commands to run"
|
|
)
|
|
parser.add_argument(
|
|
"--prepare-commands-timeouts",
|
|
default=3600,
|
|
type=float,
|
|
nargs="*",
|
|
help="timeout for prepare commands (set to <0 for infinite)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
sys.exit(main(**args.__dict__))
|