ray/rllib/tests/test_env_with_subprocess.py

"""Tests that envs clean up after themselves on agent exit."""
import os
import subprocess
import tempfile

import ray
from ray.tune import run_experiments
from ray.tune.registry import register_env
from ray.rllib.examples.env.env_with_subprocess import EnvWithSubprocess
from ray._private.test_utils import wait_for_condition


def leaked_processes():
    """Returns whether any subprocesses were leaked."""
    result = subprocess.check_output(
        "ps aux | grep '{}' | grep -v grep || true".format(
            EnvWithSubprocess.UNIQUE_CMD
        ),
        shell=True,
    )
    return result


if __name__ == "__main__":

    # Create 4 temp files, which the Env has to clean up after it's done.
    _, tmp1 = tempfile.mkstemp("test_env_with_subprocess")
    _, tmp2 = tempfile.mkstemp("test_env_with_subprocess")
    _, tmp3 = tempfile.mkstemp("test_env_with_subprocess")
    _, tmp4 = tempfile.mkstemp("test_env_with_subprocess")
    register_env("subproc", lambda c: EnvWithSubprocess(c))

    ray.init()
    # Check whether everything is ok.
    assert os.path.exists(tmp1)
    assert os.path.exists(tmp2)
    assert os.path.exists(tmp3)
    assert os.path.exists(tmp4)
    assert not leaked_processes()

    run_experiments(
        {
            "demo": {
                "run": "PG",
                "env": "subproc",
                "num_samples": 1,
                "config": {
                    "num_workers": 1,
                    "env_config": {
                        "tmp_file1": tmp1,
                        "tmp_file2": tmp2,
                        "tmp_file3": tmp3,
                        "tmp_file4": tmp4,
                    },
                    "framework": "tf",
                },
                "stop": {"training_iteration": 1},
            },
        }
    )

    ray.shutdown()

    # Check whether processes are still running.
    wait_for_condition(lambda: not leaked_processes(), timeout=30)

    print("OK")