wehub-resource-sync
149 lines
4.8 KiB
Python
149 lines
4.8 KiB
Python
# @OldAPIStack
|
|
|
|
#!/usr/bin/env python
|
|
# Runs one or more memory leak tests.
|
|
#
|
|
# Example usage:
|
|
# $ python run_memory_leak_tests.py memory-leak-test-ppo.yaml
|
|
#
|
|
# When using in BAZEL (with py_test), e.g. see in ray/rllib/BUILD:
|
|
# py_test(
|
|
# name = "memory_leak_ppo",
|
|
# main = "tests/test_memory_leak.py",
|
|
# tags = ["memory_leak_tests"],
|
|
# size = "medium", # 5min timeout
|
|
# srcs = ["tests/test_memory_leak.py"],
|
|
# data = glob(["examples/algorithms/ppo/*.yaml"]),
|
|
# # Pass `BAZEL` option and the path to look for yaml files.
|
|
# args = ["BAZEL", "examples/algorithms/ppo/memory-leak-test-ppo.yaml"]
|
|
# )
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
import ray
|
|
from ray._common.deprecation import deprecation_warning
|
|
from ray.rllib.common import SupportedFileType
|
|
from ray.rllib.train import load_experiments_from_file
|
|
from ray.rllib.utils.debug.memory import check_memory_leaks
|
|
from ray.tune.registry import get_trainable_cls
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--framework",
|
|
required=False,
|
|
choices=["jax", "tf2", "tf", "torch", None],
|
|
default=None,
|
|
help="The deep learning framework to use.",
|
|
)
|
|
parser.add_argument(
|
|
"--dir",
|
|
type=str,
|
|
required=True,
|
|
help="The directory or file in which to find all tests.",
|
|
)
|
|
parser.add_argument(
|
|
"--local-mode",
|
|
action="store_true",
|
|
help=argparse.SUPPRESS, # Deprecated.
|
|
)
|
|
parser.add_argument(
|
|
"--to-check",
|
|
nargs="+",
|
|
default=["env", "policy", "rollout_worker", "learner"],
|
|
help="List of 'env', 'policy', 'rollout_worker', 'model', 'learner'.",
|
|
)
|
|
|
|
# Obsoleted arg, use --dir instead.
|
|
parser.add_argument("--yaml-dir", type=str, default="")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parser.parse_args()
|
|
|
|
if args.yaml_dir != "":
|
|
deprecation_warning(old="--yaml-dir", new="--dir", error=True)
|
|
|
|
# Bazel regression test mode: Get path to look for yaml files.
|
|
# Get the path or single file to use.
|
|
rllib_dir = Path(__file__).parent.parent.parent
|
|
print("rllib dir={}".format(rllib_dir))
|
|
|
|
abs_path = os.path.join(rllib_dir, args.dir)
|
|
# Single file given.
|
|
if os.path.isfile(abs_path):
|
|
files = [abs_path]
|
|
# Path given -> Get all py/yaml files in there via rglob.
|
|
elif os.path.isdir(abs_path):
|
|
files = []
|
|
for type_ in ["yaml", "yml", "py"]:
|
|
files += list(rllib_dir.rglob(args.dir + f"/*.{type_}"))
|
|
files = sorted(map(lambda path: str(path.absolute()), files), reverse=True)
|
|
# Given path/file does not exist.
|
|
else:
|
|
raise ValueError(f"--dir ({args.dir}) not found!")
|
|
|
|
print("Will run the following memory-leak tests:")
|
|
for file in files:
|
|
print("->", file)
|
|
|
|
# Loop through all collected files.
|
|
for file in files:
|
|
# For python files, need to make sure, we only deliver the module name into the
|
|
# `load_experiments_from_file` function (everything from "/ray/rllib" on).
|
|
if file.endswith(".py"):
|
|
if file.endswith("__init__.py"): # weird CI learning test (BAZEL) case
|
|
continue
|
|
experiments = load_experiments_from_file(file, SupportedFileType.python)
|
|
else:
|
|
experiments = load_experiments_from_file(file, SupportedFileType.yaml)
|
|
|
|
assert (
|
|
len(experiments) == 1
|
|
), "Error, can only run a single experiment per yaml file!"
|
|
|
|
experiment = list(experiments.values())[0]
|
|
|
|
# Add framework option to exp configs.
|
|
if args.framework:
|
|
experiment["config"]["framework"] = args.framework
|
|
# Create env on local_worker for memory leak testing just the env.
|
|
experiment["config"]["create_env_on_driver"] = True
|
|
# experiment["config"]["callbacks"] = MemoryTrackingCallbacks
|
|
|
|
# Move "env" specifier into config.
|
|
experiment["config"]["env"] = experiment["env"]
|
|
experiment.pop("env", None)
|
|
|
|
# Print out the actual config.
|
|
print("== Test config ==")
|
|
print(yaml.dump(experiment))
|
|
|
|
if args.local_mode:
|
|
raise ValueError("`--local-mode` is no longer supported.")
|
|
|
|
# Construct the Algorithm instance based on the given config.
|
|
leaking = True
|
|
try:
|
|
ray.init(num_cpus=5)
|
|
if isinstance(experiment["run"], str):
|
|
algo_cls = get_trainable_cls(experiment["run"])
|
|
else:
|
|
algo_cls = get_trainable_cls(experiment["run"].__name__)
|
|
algo = algo_cls(experiment["config"])
|
|
results = check_memory_leaks(algo, to_check=set(args.to_check))
|
|
if not results:
|
|
leaking = False
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
if not leaking:
|
|
print("Memory leak test PASSED")
|
|
else:
|
|
print("Memory leak test FAILED. Exiting with Error.")
|
|
sys.exit(1)
|