chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,607 @@
|
||||
"""
|
||||
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__))
|
||||
@@ -0,0 +1,278 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from typing import Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
import aiohttp
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_PROMETHEUS_HOST = "http://localhost:9090"
|
||||
PROMETHEUS_HOST_ENV_VAR = "RAY_PROMETHEUS_HOST"
|
||||
RETRIES = 3
|
||||
|
||||
|
||||
class PrometheusQueryError(Exception):
|
||||
def __init__(self, status, message):
|
||||
self.message = (
|
||||
"Error fetching data from prometheus. "
|
||||
f"status: {status}, message: {message}"
|
||||
)
|
||||
super().__init__(self.message)
|
||||
|
||||
|
||||
class PrometheusClient:
|
||||
def __init__(self) -> None:
|
||||
self.http_session = aiohttp.ClientSession()
|
||||
self.prometheus_host = os.environ.get(
|
||||
PROMETHEUS_HOST_ENV_VAR, DEFAULT_PROMETHEUS_HOST
|
||||
)
|
||||
|
||||
async def query_prometheus(self, query_type, **kwargs):
|
||||
url = f"{self.prometheus_host}/api/v1/{query_type}?" + "&".join(
|
||||
[f"{k}={quote(str(v), safe='')}" for k, v in kwargs.items()]
|
||||
)
|
||||
query_str = kwargs.get("query", url)
|
||||
logger.debug(f"Running Prometheus query {url}")
|
||||
last_error = None
|
||||
for attempt in range(RETRIES):
|
||||
try:
|
||||
async with self.http_session.get(url) as resp:
|
||||
if resp.status == 200:
|
||||
prom_data = await resp.json()
|
||||
return prom_data["data"]["result"]
|
||||
body = (await resp.text())[:500]
|
||||
last_error = f"non-200 status {resp.status}: {body}"
|
||||
logger.warning(
|
||||
f"Prometheus query returned non-200 status {resp.status} "
|
||||
f"(attempt {attempt + 1}/{RETRIES}). Query: {query_str!r}. "
|
||||
f"Body: {body}"
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
last_error = "request timed out"
|
||||
logger.warning(
|
||||
f"Prometheus query timed out "
|
||||
f"(attempt {attempt + 1}/{RETRIES}). Query: {query_str!r}."
|
||||
)
|
||||
except aiohttp.ClientError as e:
|
||||
last_error = f"connection error: {e}"
|
||||
logger.warning(
|
||||
f"Prometheus query connection error "
|
||||
f"(attempt {attempt + 1}/{RETRIES}). Query: {query_str!r}. "
|
||||
f"Error: {e}"
|
||||
)
|
||||
if attempt < RETRIES - 1:
|
||||
await asyncio.sleep(1)
|
||||
logger.error(
|
||||
f"Prometheus query failed after {RETRIES} attempts and returned no data. "
|
||||
f"Query: {query_str!r}. Last error: {last_error}. "
|
||||
"This is a metrics-collection failure (Prometheus unreachable/erroring), "
|
||||
"NOT an empty result for a healthy metric."
|
||||
)
|
||||
return None
|
||||
|
||||
async def close(self):
|
||||
await self.http_session.close()
|
||||
|
||||
|
||||
# Metrics here mirror what we have in Grafana.
|
||||
async def _get_prometheus_metrics(
|
||||
start_time: float, end_time: float, session_name: Optional[str] = None
|
||||
) -> dict:
|
||||
client = PrometheusClient()
|
||||
kwargs = {
|
||||
"query_type": "query_range",
|
||||
"start": int(start_time),
|
||||
"end": int(end_time),
|
||||
"step": 15,
|
||||
}
|
||||
sf = f'{{SessionName="{session_name}"}}' if session_name else ""
|
||||
sf_spilled = (
|
||||
f'{{SessionName="{session_name}",State="Spilled"}}'
|
||||
if session_name
|
||||
else '{State="Spilled"}'
|
||||
)
|
||||
metrics = {
|
||||
"cpu_utilization": client.query_prometheus(
|
||||
query=f"ray_node_cpu_utilization{sf} * ray_node_cpu_count{sf} / 100",
|
||||
**kwargs,
|
||||
),
|
||||
"cpu_count": client.query_prometheus(query=f"ray_node_cpu_count{sf}", **kwargs),
|
||||
"gpu_utilization": client.query_prometheus(
|
||||
query=f"ray_node_gpus_utilization{sf} / 100", **kwargs
|
||||
),
|
||||
"gpu_count": client.query_prometheus(
|
||||
query=f"ray_node_gpus_available{sf}", **kwargs
|
||||
),
|
||||
"disk_usage": client.query_prometheus(
|
||||
query=f"ray_node_disk_usage{sf}", **kwargs
|
||||
),
|
||||
"disk_space": client.query_prometheus(
|
||||
query=f"sum(ray_node_disk_free{sf}) + sum(ray_node_disk_usage{sf})",
|
||||
**kwargs,
|
||||
),
|
||||
"memory_usage": client.query_prometheus(
|
||||
query=f"ray_node_mem_used{sf}", **kwargs
|
||||
),
|
||||
"total_memory": client.query_prometheus(
|
||||
query=f"ray_node_mem_total{sf}", **kwargs
|
||||
),
|
||||
"memory_usage_host": client.query_prometheus(
|
||||
query=f"ray_node_mem_used_host{sf}", **kwargs
|
||||
),
|
||||
"total_memory_host": client.query_prometheus(
|
||||
query=f"ray_node_mem_total_host{sf}", **kwargs
|
||||
),
|
||||
"memory_usage_cgroup": client.query_prometheus(
|
||||
query=f"ray_node_cgroup_mem_used{sf}", **kwargs
|
||||
),
|
||||
"total_memory_cgroup": client.query_prometheus(
|
||||
query=f"ray_node_cgroup_mem_total{sf}", **kwargs
|
||||
),
|
||||
"gpu_memory_usage": client.query_prometheus(
|
||||
query=f"ray_node_gram_used{sf} * 1024 * 1024", **kwargs
|
||||
),
|
||||
"gpu_total_memory": client.query_prometheus(
|
||||
query=(
|
||||
f"(sum(ray_node_gram_available{sf}) + sum(ray_node_gram_used{sf}))"
|
||||
" * 1024 * 1024"
|
||||
),
|
||||
**kwargs,
|
||||
),
|
||||
"network_receive_speed": client.query_prometheus(
|
||||
query=f"ray_node_network_receive_speed{sf}", **kwargs
|
||||
),
|
||||
"network_send_speed": client.query_prometheus(
|
||||
query=f"ray_node_network_send_speed{sf}", **kwargs
|
||||
),
|
||||
"cluster_active_nodes": client.query_prometheus(
|
||||
query=f"ray_cluster_active_nodes{sf}", **kwargs
|
||||
),
|
||||
"cluster_failed_nodes": client.query_prometheus(
|
||||
query=f"ray_cluster_failed_nodes{sf}", **kwargs
|
||||
),
|
||||
"cluster_pending_nodes": client.query_prometheus(
|
||||
query=f"ray_cluster_pending_nodes{sf}", **kwargs
|
||||
),
|
||||
"worker_oom_kills": client.query_prometheus(
|
||||
query=(
|
||||
f"sum(ray_memory_manager_worker_eviction_total{sf}) by (Type, Name)"
|
||||
),
|
||||
**kwargs,
|
||||
),
|
||||
"unexpected_worker_failures": client.query_prometheus(
|
||||
query=f"sum(ray_node_manager_unexpected_worker_failure_total{sf}) by (Type, Name)",
|
||||
**kwargs,
|
||||
),
|
||||
# `State="Spilled"` is the cumulative-bytes counter (the other States
|
||||
# are point-in-time / transient); `> 0` drops always-emitted 0 points.
|
||||
"spilled_bytes": client.query_prometheus(
|
||||
query=f"sum(ray_spill_manager_objects_bytes{sf_spilled}) > 0",
|
||||
**kwargs,
|
||||
),
|
||||
}
|
||||
metrics = {k: await v for k, v in metrics.items()}
|
||||
await client.close()
|
||||
|
||||
# Summarise the outcome so a glance at the logs tells whether the metrics
|
||||
# are trustworthy. `None` => the query failed to collect (infra/timeout);
|
||||
# `[]` => collected fine but no matching series (e.g. no OOMs happened);
|
||||
# truthy => collected data.
|
||||
failed = sorted(k for k, v in metrics.items() if v is None)
|
||||
empty = sorted(k for k, v in metrics.items() if v == [])
|
||||
with_data = sorted(k for k, v in metrics.items() if v)
|
||||
logger.info(
|
||||
f"Prometheus collection summary: {len(with_data)} metric(s) with data, "
|
||||
f"{len(empty)} empty, {len(failed)} failed to collect "
|
||||
f"(out of {len(metrics)} total)."
|
||||
)
|
||||
if failed:
|
||||
logger.error(
|
||||
f"{len(failed)} metric(s) FAILED to collect and will be null in the "
|
||||
f"output: {failed}. See the per-query warnings above for the cause "
|
||||
"(timeout / connection error / non-200). This indicates a "
|
||||
"metrics-collection/infra problem, not a real metric signal."
|
||||
)
|
||||
return metrics
|
||||
|
||||
|
||||
def get_prometheus_metrics(start_time: float, end_time: float) -> dict:
|
||||
session_name = None
|
||||
try:
|
||||
import ray
|
||||
|
||||
if not ray.is_initialized():
|
||||
ray.init("auto")
|
||||
session_name = ray.get_runtime_context().get_session_name()
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Couldn't obtain Ray session name for Prometheus query filtering. "
|
||||
f"Exception below:\n{traceback.format_exc()}"
|
||||
)
|
||||
try:
|
||||
return asyncio.run(_get_prometheus_metrics(start_time, end_time, session_name))
|
||||
except Exception:
|
||||
logger.error(
|
||||
"Couldn't obtain Prometheus metrics. "
|
||||
f"Exception below:\n{traceback.format_exc()}"
|
||||
)
|
||||
return {}
|
||||
|
||||
|
||||
def save_prometheus_metrics(
|
||||
start_time: float,
|
||||
end_time: Optional[float] = None,
|
||||
path: Optional[str] = None,
|
||||
use_ray: bool = False,
|
||||
) -> bool:
|
||||
path = path or os.environ.get("METRICS_OUTPUT_JSON", None)
|
||||
if path:
|
||||
if not end_time:
|
||||
end_time = time.time()
|
||||
if use_ray:
|
||||
import ray
|
||||
from ray.air.util.node import _force_on_current_node
|
||||
|
||||
addr = os.environ.get("RAY_ADDRESS", None)
|
||||
ray.init(addr)
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
def get_metrics():
|
||||
end_time = time.time()
|
||||
return get_prometheus_metrics(start_time, end_time)
|
||||
|
||||
remote_run = _force_on_current_node(get_metrics)
|
||||
ref = remote_run.remote()
|
||||
metrics = ray.get(ref, timeout=900)
|
||||
else:
|
||||
metrics = get_prometheus_metrics(start_time, end_time)
|
||||
with open(path, "w") as metrics_output_file:
|
||||
json.dump(metrics, metrics_output_file)
|
||||
return path
|
||||
return None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="[%(levelname)s %(asctime)s] %(filename)s: %(lineno)d %(message)s",
|
||||
)
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("start_time", type=float, help="Start time")
|
||||
parser.add_argument(
|
||||
"--path", default="", type=str, help="Where to save the metrics json"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_ray",
|
||||
default=False,
|
||||
action="store_true",
|
||||
help="Whether to run this script in a ray.remote call (for Ray Client)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
save_prometheus_metrics(args.start_time, path=args.path, use_ray=args.use_ray)
|
||||
@@ -0,0 +1,54 @@
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import ray
|
||||
|
||||
ray.init(address="auto")
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"num_nodes", type=int, help="Wait for this number of nodes (includes head)"
|
||||
)
|
||||
|
||||
parser.add_argument("max_time_s", type=int, help="Wait for this number of seconds")
|
||||
|
||||
parser.add_argument(
|
||||
"--feedback_interval_s",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Wait for this number of seconds",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
curr_nodes = 0
|
||||
start = time.time()
|
||||
next_feedback = start
|
||||
max_time = start + args.max_time_s
|
||||
|
||||
while not curr_nodes >= args.num_nodes:
|
||||
now = time.time()
|
||||
|
||||
if now >= max_time:
|
||||
raise RuntimeError(
|
||||
f"Maximum wait time reached, but only "
|
||||
f"{curr_nodes}/{args.num_nodes} nodes came up. Aborting."
|
||||
)
|
||||
|
||||
if now >= next_feedback:
|
||||
passed = now - start
|
||||
print(
|
||||
f"Waiting for more nodes to come up: "
|
||||
f"{curr_nodes}/{args.num_nodes} "
|
||||
f"({passed:.0f} seconds passed)"
|
||||
)
|
||||
next_feedback = now + args.feedback_interval_s
|
||||
|
||||
time.sleep(5)
|
||||
curr_nodes = sum(1 for node in ray.nodes() if node["Alive"])
|
||||
|
||||
passed = time.time() - start
|
||||
print(
|
||||
f"Cluster is up: {curr_nodes}/{args.num_nodes} nodes online after "
|
||||
f"{passed:.0f} seconds"
|
||||
)
|
||||
@@ -0,0 +1,434 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
import shutil
|
||||
import tempfile
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
|
||||
from ray_release.cloud_util import (
|
||||
convert_abfss_uri_to_https,
|
||||
generate_tmp_cloud_storage_path,
|
||||
upload_working_dir_to_azure,
|
||||
)
|
||||
from ray_release.cluster_manager.cluster_manager import ClusterManager
|
||||
from ray_release.command_runner.command_runner import CommandRunner
|
||||
from ray_release.exception import (
|
||||
FetchResultError,
|
||||
JobBrokenError,
|
||||
JobNoLogsError,
|
||||
JobOutOfRetriesError,
|
||||
LogsError,
|
||||
PrepareCommandError,
|
||||
PrepareCommandTimeout,
|
||||
TestCommandError,
|
||||
TestCommandTimeout,
|
||||
)
|
||||
from ray_release.file_manager.job_file_manager import JobFileManager
|
||||
from ray_release.job_manager.anyscale_job_manager import (
|
||||
JOB_SOFT_INFRA_ERROR,
|
||||
JOB_STATE_UNKNOWN,
|
||||
AnyscaleJobManager,
|
||||
)
|
||||
from ray_release.logger import logger
|
||||
from ray_release.reporter.artifacts import DEFAULT_ARTIFACTS_DIR
|
||||
from ray_release.util import (
|
||||
AZURE_CLOUD_STORAGE,
|
||||
AZURE_STORAGE_CONTAINER,
|
||||
S3_CLOUD_STORAGE,
|
||||
get_anyscale_sdk,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from anyscale.sdk.anyscale_client.sdk import AnyscaleSDK
|
||||
|
||||
TIMEOUT_RETURN_CODE = 124
|
||||
|
||||
|
||||
def _join_cloud_storage_paths(*paths: str):
|
||||
paths = list(paths)
|
||||
if len(paths) > 1:
|
||||
for i in range(1, len(paths)):
|
||||
while paths[i][0] == "/":
|
||||
paths[i] = paths[i][1:]
|
||||
joined_path = os.path.join(*paths)
|
||||
while joined_path[-1] == "/":
|
||||
joined_path = joined_path[:-1]
|
||||
return joined_path
|
||||
|
||||
|
||||
def _get_env_str(env: Dict[str, str]) -> str:
|
||||
if env:
|
||||
env_str = " ".join(f"{k}={v}" for k, v in env.items()) + " "
|
||||
else:
|
||||
env_str = ""
|
||||
return env_str
|
||||
|
||||
|
||||
class AnyscaleJobRunner(CommandRunner):
|
||||
# the directory for runners to dump files to (on buildkite runner instances).
|
||||
# Write to this directory. run_release_tests.sh will ensure that the content
|
||||
# shows up under buildkite job's "Artifacts" UI tab.
|
||||
_DEFAULT_ARTIFACTS_DIR = DEFAULT_ARTIFACTS_DIR
|
||||
|
||||
# the artifact file name put under s3 bucket root.
|
||||
# AnyscalejobWrapper will upload user generated artifact to this path
|
||||
# and AnyscaleJobRunner will then download from there.
|
||||
_USER_GENERATED_ARTIFACT = "user_generated_artifact"
|
||||
|
||||
# the path where result json will be written to on both head node
|
||||
# as well as the relative path where result json will be uploaded to on s3.
|
||||
_RESULT_OUTPUT_JSON = "/tmp/release_test_out.json"
|
||||
|
||||
# the path where output json will be written to on both head node
|
||||
# as well as the relative path where metrics json will be uploaded to on s3.
|
||||
_METRICS_OUTPUT_JSON = "/tmp/metrics_test_out.json"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cluster_manager: ClusterManager,
|
||||
file_manager: JobFileManager,
|
||||
working_dir: str,
|
||||
sdk: Optional["AnyscaleSDK"] = None,
|
||||
artifact_path: Optional[str] = None,
|
||||
):
|
||||
super().__init__(
|
||||
cluster_manager=cluster_manager,
|
||||
working_dir=working_dir,
|
||||
)
|
||||
self.file_manager = file_manager
|
||||
self.sdk = sdk or get_anyscale_sdk()
|
||||
self.job_manager = AnyscaleJobManager(cluster_manager)
|
||||
|
||||
self.last_command_scd_id = None
|
||||
self.path_in_bucket = _join_cloud_storage_paths(
|
||||
"working_dirs",
|
||||
self.cluster_manager.test.get_name().replace(" ", "_"),
|
||||
generate_tmp_cloud_storage_path(),
|
||||
)
|
||||
# The root cloud storage bucket path. result, metric, artifact files
|
||||
# will be uploaded to under it on cloud storage.
|
||||
cloud_storage_provider = os.environ.get(
|
||||
"ANYSCALE_CLOUD_STORAGE_PROVIDER",
|
||||
S3_CLOUD_STORAGE,
|
||||
)
|
||||
|
||||
if cloud_storage_provider == AZURE_CLOUD_STORAGE:
|
||||
# Azure ABFSS involves container and account name in the path
|
||||
# and in a specific format/order.
|
||||
self.upload_path = _join_cloud_storage_paths(
|
||||
f"{AZURE_CLOUD_STORAGE}://{AZURE_STORAGE_CONTAINER}@{self.file_manager.bucket}.dfs.core.windows.net",
|
||||
self.path_in_bucket,
|
||||
)
|
||||
else:
|
||||
self.upload_path = _join_cloud_storage_paths(
|
||||
f"{cloud_storage_provider}://{self.file_manager.bucket}",
|
||||
self.path_in_bucket,
|
||||
)
|
||||
self.output_json = "/tmp/output.json"
|
||||
self.prepare_commands = []
|
||||
self._wait_for_nodes_timeout = 0
|
||||
|
||||
self._results_uploaded = True
|
||||
self._metrics_uploaded = True
|
||||
|
||||
# artifact related
|
||||
# user provided path to where they write the artifact to.
|
||||
self._artifact_path = artifact_path
|
||||
self._artifact_uploaded = artifact_path is not None
|
||||
|
||||
def _copy_script_to_working_dir(self, script_name):
|
||||
script = os.path.join(os.path.dirname(__file__), f"_{script_name}")
|
||||
shutil.copy(script, script_name)
|
||||
|
||||
def prepare_remote_env(self):
|
||||
self._copy_script_to_working_dir("anyscale_job_wrapper.py")
|
||||
self._copy_script_to_working_dir("wait_cluster.py")
|
||||
self._copy_script_to_working_dir("prometheus_metrics.py")
|
||||
|
||||
def run_prepare_command(
|
||||
self, command: str, env: Optional[Dict] = None, timeout: float = 3600.0
|
||||
):
|
||||
self.prepare_commands.append((command, env, timeout))
|
||||
|
||||
def wait_for_nodes(self, num_nodes: int, timeout: float = 900):
|
||||
self._wait_for_nodes_timeout = timeout
|
||||
self.job_manager.cluster_startup_timeout += timeout
|
||||
|
||||
# Give 30 seconds more to account for communication
|
||||
self.run_prepare_command(
|
||||
f"python wait_cluster.py {num_nodes} {timeout}", timeout=timeout + 30
|
||||
)
|
||||
|
||||
def _handle_command_output(
|
||||
self, job_return_code: int, raise_on_timeout: bool = True
|
||||
):
|
||||
if job_return_code == JOB_SOFT_INFRA_ERROR:
|
||||
raise JobOutOfRetriesError(
|
||||
"Job returned non-success state: 'FAILED' "
|
||||
"(command has not been run or no logs could be obtained)."
|
||||
)
|
||||
|
||||
if job_return_code == JOB_STATE_UNKNOWN:
|
||||
raise JobBrokenError("Job state is 'UNKNOWN'.")
|
||||
|
||||
# First try to obtain the output.json from S3.
|
||||
# If that fails, try logs.
|
||||
try:
|
||||
output_json = self.fetch_output()
|
||||
except Exception:
|
||||
logger.exception("Exception when obtaining output from S3.")
|
||||
try:
|
||||
logs = self.get_last_logs()
|
||||
output_json = re.search(r"### JSON \|([^\|]*)\| ###", logs)
|
||||
output_json = json.loads(output_json.group(1))
|
||||
except Exception:
|
||||
output_json = None
|
||||
|
||||
workload_status_code = None
|
||||
if output_json:
|
||||
logger.info(f"Output: {output_json}")
|
||||
workload_status_code = output_json["return_code"]
|
||||
workload_time_taken = output_json["workload_time_taken"]
|
||||
prepare_return_codes = output_json["prepare_return_codes"]
|
||||
last_prepare_time_taken = output_json["last_prepare_time_taken"]
|
||||
|
||||
# If we know results/metrics were not uploaded, we can fail fast
|
||||
# fetching later.
|
||||
self._results_uploaded = output_json["uploaded_results"]
|
||||
self._metrics_uploaded = output_json["uploaded_metrics"]
|
||||
self._artifact_uploaded = output_json["uploaded_artifact"]
|
||||
|
||||
if prepare_return_codes and prepare_return_codes[-1] != 0:
|
||||
if prepare_return_codes[-1] == TIMEOUT_RETURN_CODE:
|
||||
raise PrepareCommandTimeout(
|
||||
"Prepare command timed out after "
|
||||
f"{last_prepare_time_taken} seconds."
|
||||
)
|
||||
raise PrepareCommandError(
|
||||
f"Prepare command '{self.prepare_commands[-1]}' returned "
|
||||
f"non-success status: {prepare_return_codes[-1]}."
|
||||
)
|
||||
else:
|
||||
raise JobNoLogsError("Could not obtain logs for the job.")
|
||||
|
||||
if workload_status_code == TIMEOUT_RETURN_CODE:
|
||||
if not raise_on_timeout:
|
||||
# Expected - treat as success.
|
||||
return
|
||||
|
||||
raise TestCommandTimeout(
|
||||
f"Command timed out after {workload_time_taken} seconds."
|
||||
)
|
||||
|
||||
if workload_status_code is None or workload_status_code != 0:
|
||||
raise TestCommandError(
|
||||
f"Command returned non-success status: {workload_status_code}."
|
||||
)
|
||||
|
||||
def _get_full_command_env(self, env: Optional[Dict[str, str]] = None):
|
||||
full_env = {
|
||||
"TEST_OUTPUT_JSON": self._RESULT_OUTPUT_JSON,
|
||||
"METRICS_OUTPUT_JSON": self._METRICS_OUTPUT_JSON,
|
||||
"USER_GENERATED_ARTIFACT": self._USER_GENERATED_ARTIFACT,
|
||||
"BUILDKITE_BRANCH": os.environ.get("BUILDKITE_BRANCH", ""),
|
||||
"PYTHONUNBUFFERED": "1",
|
||||
}
|
||||
if env:
|
||||
full_env.update(env)
|
||||
return full_env
|
||||
|
||||
def run_command(
|
||||
self,
|
||||
command: str,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
timeout: float = 3600.0,
|
||||
raise_on_timeout: bool = True,
|
||||
) -> float:
|
||||
prepare_command_strs = []
|
||||
prepare_command_timeouts = []
|
||||
# Convert the prepare commands, envs and timeouts into shell-compliant
|
||||
# strings that can be passed to the wrapper script
|
||||
for prepare_command, prepare_env, prepare_timeout in self.prepare_commands:
|
||||
prepare_env = self._get_full_command_env(prepare_env)
|
||||
env_str = _get_env_str(prepare_env)
|
||||
prepare_command_strs.append(f"{env_str} {prepare_command}")
|
||||
prepare_command_timeouts.append(prepare_timeout)
|
||||
|
||||
prepare_commands_shell = " ".join(
|
||||
shlex.quote(str(x)) for x in prepare_command_strs
|
||||
)
|
||||
prepare_commands_timeouts_shell = " ".join(
|
||||
shlex.quote(str(x)) for x in prepare_command_timeouts
|
||||
)
|
||||
|
||||
full_env = self._get_full_command_env(env)
|
||||
|
||||
no_raise_on_timeout_str = (
|
||||
" --test-no-raise-on-timeout" if not raise_on_timeout else ""
|
||||
)
|
||||
results_cloud_storage_uri = _join_cloud_storage_paths(
|
||||
self.upload_path, self._RESULT_OUTPUT_JSON
|
||||
)
|
||||
metrics_cloud_storage_uri = _join_cloud_storage_paths(
|
||||
self.upload_path, self._METRICS_OUTPUT_JSON
|
||||
)
|
||||
output_cloud_storage_uri = _join_cloud_storage_paths(
|
||||
self.upload_path, self.output_json
|
||||
)
|
||||
upload_cloud_storage_uri = self.upload_path
|
||||
# Convert ABFSS URI to HTTPS URI for Azure
|
||||
# since azcopy doesn't support ABFSS.
|
||||
# azcopy is used to fetch these artifacts on Buildkite
|
||||
# after job is done.
|
||||
if self.upload_path.startswith(AZURE_CLOUD_STORAGE):
|
||||
results_cloud_storage_uri = convert_abfss_uri_to_https(
|
||||
results_cloud_storage_uri
|
||||
)
|
||||
metrics_cloud_storage_uri = convert_abfss_uri_to_https(
|
||||
metrics_cloud_storage_uri
|
||||
)
|
||||
output_cloud_storage_uri = convert_abfss_uri_to_https(
|
||||
output_cloud_storage_uri
|
||||
)
|
||||
upload_cloud_storage_uri = convert_abfss_uri_to_https(
|
||||
upload_cloud_storage_uri
|
||||
)
|
||||
full_command = (
|
||||
f"python anyscale_job_wrapper.py '{command}' "
|
||||
f"--test-workload-timeout {timeout}{no_raise_on_timeout_str} "
|
||||
"--results-cloud-storage-uri "
|
||||
f"'{results_cloud_storage_uri}' "
|
||||
"--metrics-cloud-storage-uri "
|
||||
f"'"
|
||||
f"{metrics_cloud_storage_uri}' "
|
||||
"--output-cloud-storage-uri "
|
||||
f"'{output_cloud_storage_uri}' "
|
||||
f"--upload-cloud-storage-uri '{upload_cloud_storage_uri}' "
|
||||
f"--prepare-commands {prepare_commands_shell} "
|
||||
f"--prepare-commands-timeouts {prepare_commands_timeouts_shell} "
|
||||
)
|
||||
if self._artifact_path:
|
||||
full_command += f"--artifact-path '{self._artifact_path}' "
|
||||
|
||||
timeout = min(
|
||||
(self.cluster_manager.maximum_uptime_minutes - 1) * 60,
|
||||
# The timeout set here is just for the prepare commands + test workload
|
||||
# WITHOUT wait for nodes time included, as that is set separately.
|
||||
# Since wait for nodes is a part of prepare_commands, we manually
|
||||
# subtract the timeout for it here.
|
||||
# We also add 15 mins for upload & metrics collection.
|
||||
timeout
|
||||
+ sum(prepare_command_timeouts)
|
||||
- self._wait_for_nodes_timeout
|
||||
+ 900,
|
||||
)
|
||||
working_dir = "."
|
||||
# If running on Azure, upload working dir to Azure blob storage first
|
||||
if self.upload_path.startswith(AZURE_CLOUD_STORAGE):
|
||||
azure_file_path = upload_working_dir_to_azure(
|
||||
working_dir=os.getcwd(), azure_directory_uri=self.upload_path
|
||||
)
|
||||
working_dir = azure_file_path
|
||||
logger.info(f"Working dir uploaded to {working_dir}")
|
||||
|
||||
job_return_code, time_taken = self.job_manager.run_and_wait(
|
||||
full_command,
|
||||
full_env,
|
||||
working_dir=working_dir,
|
||||
upload_path=self.upload_path,
|
||||
timeout=int(timeout),
|
||||
)
|
||||
self._handle_command_output(job_return_code, raise_on_timeout=raise_on_timeout)
|
||||
|
||||
return time_taken
|
||||
|
||||
def get_last_logs_ex(self) -> Optional[str]:
|
||||
try:
|
||||
return self.job_manager.get_last_logs()
|
||||
except Exception as e:
|
||||
raise LogsError(f"Could not get last logs: {e}") from e
|
||||
|
||||
def _fetch_json(self, path: str) -> Dict[str, Any]:
|
||||
try:
|
||||
tmpfile = tempfile.mkstemp(suffix=".json")[1]
|
||||
logger.info(tmpfile)
|
||||
self.file_manager.download_from_cloud(
|
||||
path, tmpfile, delete_after_download=True
|
||||
)
|
||||
|
||||
with open(tmpfile, "rt") as f:
|
||||
content = f.read()
|
||||
|
||||
try:
|
||||
data = json.loads(content)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.info(f"Result content = {content}")
|
||||
raise e
|
||||
|
||||
os.unlink(tmpfile)
|
||||
return data
|
||||
except Exception as e:
|
||||
raise FetchResultError(f"Could not fetch results from session: {e}") from e
|
||||
|
||||
def fetch_results(self) -> Dict[str, Any]:
|
||||
if not self._results_uploaded:
|
||||
raise FetchResultError(
|
||||
"Could not fetch results from session as they were not uploaded."
|
||||
)
|
||||
return self._fetch_json(
|
||||
_join_cloud_storage_paths(self.path_in_bucket, self._RESULT_OUTPUT_JSON)
|
||||
)
|
||||
|
||||
def fetch_metrics(self) -> Dict[str, Any]:
|
||||
if not self._metrics_uploaded:
|
||||
raise FetchResultError(
|
||||
"Could not fetch metrics from session as they were not uploaded."
|
||||
)
|
||||
return self._fetch_json(
|
||||
_join_cloud_storage_paths(self.path_in_bucket, self._METRICS_OUTPUT_JSON)
|
||||
)
|
||||
|
||||
def fetch_artifact(self) -> None:
|
||||
"""Fetch artifact (file) from `self._artifact_path` on Anyscale cluster
|
||||
head node.
|
||||
|
||||
Note, an implementation detail here is that by the time this function is called,
|
||||
the artifact file is already present in s3 bucket by the name of
|
||||
`self._USER_GENERATED_ARTIFACT`. This is because, the uploading to s3 portion is
|
||||
done by `_anyscale_job_wrapper`.
|
||||
|
||||
The fetched artifact will be placed under `self._DEFAULT_ARTIFACTS_DIR`,
|
||||
which will ultimately show up in buildkite Artifacts UI tab.
|
||||
The fetched file will have the same filename and extension as the one
|
||||
on Anyscale cluster head node (same as `self._artifact_path`).
|
||||
"""
|
||||
if not self._artifact_uploaded:
|
||||
raise FetchResultError(
|
||||
"Could not fetch artifact from session as they "
|
||||
"were either not generated or not uploaded."
|
||||
)
|
||||
# first make sure that `self._DEFAULT_ARTIFACTS_DIR` exists.
|
||||
if not os.path.exists(self._DEFAULT_ARTIFACTS_DIR):
|
||||
os.makedirs(self._DEFAULT_ARTIFACTS_DIR, 0o755)
|
||||
|
||||
# we use the same artifact file name and extension specified by user
|
||||
# and put it under `self._DEFAULT_ARTIFACTS_DIR`.
|
||||
artifact_file_name = os.path.basename(self._artifact_path)
|
||||
self.file_manager.download_from_cloud(
|
||||
_join_cloud_storage_paths(
|
||||
self.path_in_bucket, self._USER_GENERATED_ARTIFACT
|
||||
),
|
||||
os.path.join(self._DEFAULT_ARTIFACTS_DIR, artifact_file_name),
|
||||
)
|
||||
|
||||
def fetch_output(self) -> Dict[str, Any]:
|
||||
return self._fetch_json(
|
||||
_join_cloud_storage_paths(self.path_in_bucket, self.output_json),
|
||||
)
|
||||
|
||||
def job_url(self) -> Optional[str]:
|
||||
return self.job_manager.job_url()
|
||||
|
||||
def job_id(self) -> Optional[str]:
|
||||
return self.job_manager.job_id()
|
||||
@@ -0,0 +1,85 @@
|
||||
import abc
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from click.exceptions import ClickException
|
||||
|
||||
from ray_release.cluster_manager.cluster_manager import ClusterManager
|
||||
from ray_release.logger import logger
|
||||
from ray_release.util import exponential_backoff_retry
|
||||
|
||||
|
||||
class CommandRunner(abc.ABC):
|
||||
"""This is run on Buildkite runners."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cluster_manager: ClusterManager,
|
||||
working_dir: str,
|
||||
artifact_path: Optional[str] = None,
|
||||
):
|
||||
self.cluster_manager = cluster_manager
|
||||
self.working_dir = working_dir
|
||||
|
||||
def prepare_remote_env(self):
|
||||
"""Prepare remote environment, e.g. upload files."""
|
||||
raise NotImplementedError
|
||||
|
||||
def wait_for_nodes(self, num_nodes: int, timeout: float = 900.0):
|
||||
"""Wait for cluster nodes to be up.
|
||||
|
||||
Args:
|
||||
num_nodes: Number of nodes to wait for.
|
||||
timeout: Timeout in seconds to wait for nodes before
|
||||
raising a ``PrepareCommandTimeoutError``.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
PrepareCommandTimeoutError
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def run_command(
|
||||
self,
|
||||
command: str,
|
||||
env: Optional[Dict] = None,
|
||||
timeout: float = 3600.0,
|
||||
raise_on_timeout: bool = True,
|
||||
) -> float:
|
||||
"""Run command."""
|
||||
raise NotImplementedError
|
||||
|
||||
def run_prepare_command(
|
||||
self, command: str, env: Optional[Dict] = None, timeout: float = 3600.0
|
||||
):
|
||||
"""Run prepare command.
|
||||
|
||||
Command runners may choose to run this differently than the
|
||||
test command.
|
||||
"""
|
||||
return exponential_backoff_retry(
|
||||
lambda: self.run_command(command, env, timeout),
|
||||
ClickException,
|
||||
initial_retry_delay_s=5,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
def get_last_logs(self) -> Optional[str]:
|
||||
try:
|
||||
return self.get_last_logs_ex()
|
||||
except Exception as e:
|
||||
logger.exception(f"Error fetching logs: {e}")
|
||||
return None
|
||||
|
||||
def get_last_logs_ex(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def fetch_results(self) -> Dict[str, Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
def fetch_metrics(self) -> Dict[str, Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
def fetch_artifact(self) -> None:
|
||||
raise NotImplementedError
|
||||
Reference in New Issue
Block a user