From 9392cdbf749206279252894fe4e188d7e0867cd7 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Fri, 20 Mar 2020 12:43:57 -0700 Subject: [PATCH] [rllib] Add high-performance external application connector (#7641) --- .travis.yml | 4 +- doc/source/rllib-env.rst | 90 +++++- doc/source/rllib-external.svg | 1 + doc/source/rllib-toc.rst | 12 +- doc/source/rllib-training.rst | 12 +- doc/source/rllib.rst | 7 + rllib/BUILD | 16 + rllib/agents/dqn/dqn.py | 2 +- rllib/env/__init__.py | 15 +- rllib/env/policy_client.py | 294 ++++++++++++++++++ rllib/env/policy_server_input.py | 191 ++++++++++++ rllib/env/serving_env.py | 4 - rllib/evaluation/rollout_worker.py | 13 + rllib/examples/serving/cartpole_client.py | 16 +- rllib/examples/serving/cartpole_server.py | 100 +++--- rllib/examples/serving/test.sh | 12 - .../examples/serving/test_local_inference.sh | 23 ++ .../examples/serving/test_remote_inference.sh | 24 ++ rllib/utils/policy_client.py | 7 +- rllib/utils/policy_server.py | 37 +-- 20 files changed, 757 insertions(+), 123 deletions(-) create mode 100644 doc/source/rllib-external.svg create mode 100644 rllib/env/policy_client.py create mode 100644 rllib/env/policy_server_input.py delete mode 100644 rllib/env/serving_env.py delete mode 100755 rllib/examples/serving/test.sh create mode 100755 rllib/examples/serving/test_local_inference.sh create mode 100755 rllib/examples/serving/test_remote_inference.sh diff --git a/.travis.yml b/.travis.yml index b731b8f36..5d2676572 100644 --- a/.travis.yml +++ b/.travis.yml @@ -253,8 +253,8 @@ matrix: - if [ $RAY_CI_RLLIB_FULL_AFFECTED != "1" ]; then exit; fi - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_A,examples_B --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_C --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... - - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_E,examples_M,examples_P --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... - - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_R,examples_T --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... + - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_E,examples_L,examples_M,examples_P --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... + - ./ci/keep_alive bazel test --build_tests_only --test_tag_filters=examples_R,examples_S,examples_T --spawn_strategy=local --flaky_test_attempts=3 --nocache_test_results --test_verbose_timeout_warnings --progress_report_interval=100 --show_progress_rate_limit=100 --show_timestamps --test_output=errors rllib/... # RLlib: tests_dir: Everything in rllib/tests/ directory (A-I). - os: linux diff --git a/doc/source/rllib-env.rst b/doc/source/rllib-env.rst index 1ed611e7a..c7bbd4dd6 100644 --- a/doc/source/rllib-env.rst +++ b/doc/source/rllib-env.rst @@ -298,30 +298,96 @@ In this setup, the appropriate rewards for training lower-level agents must be p See this file for a runnable example: `hierarchical_training.py `__. -Interfacing with External Agents +External Agents and Applications -------------------------------- In many situations, it does not make sense for an environment to be "stepped" by RLlib. For example, if a policy is to be used in a web serving system, then it is more natural for an agent to query a service that serves policy decisions, and for that service to learn from experience over time. This case also naturally arises with **external simulators** that run independently outside the control of RLlib, but may still want to leverage RLlib for training. RLlib provides the `ExternalEnv `__ class for this purpose. Unlike other envs, ExternalEnv has its own thread of control. At any point, agents on that thread can query the current policy for decisions via ``self.get_action()`` and reports rewards via ``self.log_returns()``. This can be done for multiple concurrent episodes as well. -ExternalEnv can be used to implement a simple REST policy `server `__ that learns over time using RLlib. In this example RLlib runs with ``num_workers=0`` to avoid port allocation issues, but in principle this could be scaled by increasing ``num_workers``. - Logging off-policy actions ~~~~~~~~~~~~~~~~~~~~~~~~~~ -ExternalEnv also provides a ``self.log_action()`` call to support off-policy actions. This allows the client to make independent decisions, e.g., to compare two different policies, and for RLlib to still learn from those off-policy actions. Note that this requires the algorithm used to support learning from off-policy decisions (e.g., DQN). - -Data ingest -~~~~~~~~~~~ - -The ``log_action`` API of ExternalEnv can be used to ingest data from offline logs. The pattern would be as follows: First, some policy is followed to produce experience data which is stored in some offline storage system. Then, RLlib creates a number of workers that use a ExternalEnv to read the logs in parallel and ingest the experiences. After a round of training completes, the new policy can be deployed to collect more experiences. - -Note that envs can read from different partitions of the logs based on the ``worker_index`` attribute of the `env context `__ passed into the environment constructor. +ExternalEnv provides a ``self.log_action()`` call to support off-policy actions. This allows the client to make independent decisions, e.g., to compare two different policies, and for RLlib to still learn from those off-policy actions. Note that this requires the algorithm used to support learning from off-policy decisions (e.g., DQN). .. seealso:: - `Offline Datasets `__ provide higher-level interfaces for working with offline experience datasets. + `Offline Datasets `__ provide higher-level interfaces for working with off-policy experience datasets. + +External Application Clients +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For applications that are running entirely outside the Ray cluster (i.e., cannot be packaged into a Python environment of any form), RLlib provides the ``PolicyServerInput`` application connector, which can be connected to over the network using ``PolicyClient`` instances. + +You can configure any Trainer to launch a policy server with the following config: + +.. code-block:: python + + trainer_config = { + # An environment class is still required, but it doesn't need to be runnable. + # You only need to define its action and observation space attributes. + "env": YOUR_ENV_STUB, + + # Use the policy server to generate experiences. + "input": ( + lambda ioctx: PolicyServerInput(ioctx, SERVER_ADDRESS, SERVER_PORT) + ), + # Use the existing trainer process to run the server. + "num_workers": 0, + # Disable OPE, since the rollouts are coming from online clients. + "input_evaluation": [], + } + +Clients can then connect in either *local* or *remote* inference mode. In local inference mode, copies of the policy are downloaded from the server and cached on the client for a configurable period of time. This allows actions to be computed by the client without requiring a network round trip each time. In remote inference mode, each computed action requires a network call to the server. + +Example: + +.. code-block:: python + + client = PolicyClient("http://localhost:9900", inference_mode="local") + episode_id = client.start_episode() + ... + action = client.get_action(episode_id, cur_obs) + ... + client.end_episode(episode_id, last_obs) + +To understand the difference between standard envs, external envs, and connecting with a ``PolicyClient``, refer to the following figure: + +.. https://docs.google.com/drawings/d/1hJvT9bVGHVrGTbnCZK29BYQIcYNRbZ4Dr6FOPMJDjUs/edit +.. image:: rllib-external.svg + +Try it yourself by launching a `cartpole_server.py `__, and connecting to it with any number of clients (`cartpole_client.py `__): + +.. code-block:: bash + + # Start the server by running: + >>> python rllib/examples/serving/cartpole_server.py --run=PPO + -- + -- Starting policy server at localhost:9900 + -- + + # To connect from a client with inference_mode="remote". + >>> python rllib/examples/serving/cartpole_client.py --inference-mode=remote + Total reward: 10.0 + Total reward: 58.0 + ... + Total reward: 200.0 + ... + + # To connect from a client with inference_mode="local" (faster). + >>> python rllib/examples/serving/cartpole_client.py --inference-mode=local + Querying server for new policy weights. + Generating new batch of experiences. + Total reward: 13.0 + Total reward: 11.0 + ... + Sending batch of 1000 steps back to server. + Querying server for new policy weights. + ... + Total reward: 200.0 + ... + +For the best performance, when possible we recommend using ``inference_mode="local"`` when possible. Advanced Integrations --------------------- diff --git a/doc/source/rllib-external.svg b/doc/source/rllib-external.svg new file mode 100644 index 000000000..18f7aa7f5 --- /dev/null +++ b/doc/source/rllib-external.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/doc/source/rllib-toc.rst b/doc/source/rllib-toc.rst index 354118325..3f9949543 100644 --- a/doc/source/rllib-toc.rst +++ b/doc/source/rllib-toc.rst @@ -53,7 +53,7 @@ Training APIs - `Stack Traces `__ -* `REST API `__ +* `External Application API `__ Environments ------------ @@ -62,7 +62,10 @@ Environments * `OpenAI Gym `__ * `Vectorized `__ * `Multi-Agent and Hierarchical `__ -* `Interfacing with External Agents `__ +* `External Agents and Applications `__ + + - `External Application Clients `__ + * `Advanced Integrations `__ Models, Preprocessors, and Action Distributions @@ -185,8 +188,9 @@ try setting ``OMP_NUM_THREADS=1``. Similarly, check configured system limits wit If you encounter out-of-memory errors, consider setting ``redis_max_memory`` and ``object_store_memory`` in ``ray.init()`` to reduce memory usage. For debugging unexpected hangs or performance problems, you can run ``ray stack`` to dump -the stack traces of all Ray workers on the current node, and ``ray timeline`` to dump -a timeline visualization of tasks to a file. +the stack traces of all Ray workers on the current node, ``ray timeline`` to dump +a timeline visualization of tasks to a file, and ``ray memory`` to list all object +references in the cluster. TensorFlow 2.0 ~~~~~~~~~~~~~~ diff --git a/doc/source/rllib-training.rst b/doc/source/rllib-training.rst index e3a9c135a..4556f55c4 100644 --- a/doc/source/rllib-training.rst +++ b/doc/source/rllib-training.rst @@ -913,15 +913,13 @@ Stack Traces You can use the ``ray stack`` command to dump the stack traces of all the Python workers on a single node. This can be useful for debugging unexpected hangs or performance issues. -REST API --------- +External Application API +------------------------ -In some cases (i.e., when interacting with an externally hosted simulator or production environment) it makes more sense to interact with RLlib as if were an independently running service, rather than RLlib hosting the simulations itself. This is possible via RLlib's external agents `interface `__. +In some cases (i.e., when interacting with an externally hosted simulator or production environment) it makes more sense to interact with RLlib as if it were an independently running service, rather than RLlib hosting the simulations itself. This is possible via RLlib's external applications interface `(full documentation) `__. -.. autoclass:: ray.rllib.utils.policy_client.PolicyClient +.. autoclass:: ray.rllib.env.policy_client.PolicyClient :members: -.. autoclass:: ray.rllib.utils.policy_server.PolicyServer +.. autoclass:: ray.rllib.env.policy_server_input.PolicyServerInput :members: - -For a full client / server example that you can run, see the example `client script `__ and also the corresponding `server script `__, here configured to serve a policy for the toy CartPole-v0 environment. diff --git a/doc/source/rllib.rst b/doc/source/rllib.rst index dc544f595..7232626c4 100644 --- a/doc/source/rllib.rst +++ b/doc/source/rllib.rst @@ -105,6 +105,13 @@ RLlib provides ways to customize almost all aspects of training, including the ` .. image:: rllib-components.svg +Application Support +~~~~~~~~~~~~~~~~~~~ + +Beyond environments defined in Python, RLlib supports batch training on `offline datasets `__, and also provides a variety of integration strategies for `external applications `__: + +.. image:: rllib-external.svg + To learn more, proceed to the `table of contents `__. .. |tensorflow| image:: tensorflow.png diff --git a/rllib/BUILD b/rllib/BUILD index c73de0a48..1e4752be9 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1336,6 +1336,22 @@ py_test( args = ["--num-cpus=4"] ) +sh_test( + name = "examples/serving/test_local_inference", + tags = ["examples", "examples_L", "exclusive"], + size = "medium", + srcs = ["examples/serving/test_local_inference.sh"], + data = glob(["examples/serving/*.py"]), +) + +sh_test( + name = "examples/serving/test_remote_inference", + tags = ["examples", "examples_R", "exclusive"], + size = "medium", + srcs = ["examples/serving/test_remote_inference.sh"], + data = glob(["examples/serving/*.py"]), +) + py_test( name = "examples/rock_paper_scissors_multiagent", main = "examples/rock_paper_scissors_multiagent.py", tags = ["examples", "examples_R"], diff --git a/rllib/agents/dqn/dqn.py b/rllib/agents/dqn/dqn.py index fe0ebb795..4774330cc 100644 --- a/rllib/agents/dqn/dqn.py +++ b/rllib/agents/dqn/dqn.py @@ -86,7 +86,7 @@ DEFAULT_CONFIG = with_common_config({ # Epsilon to add to the TD errors when updating priorities. "prioritized_replay_eps": 1e-6, # Whether to LZ4 compress observations - "compress_observations": True, + "compress_observations": False, # === Optimization === # Learning rate for adam optimizer diff --git a/rllib/env/__init__.py b/rllib/env/__init__.py index a454f0dd9..a9483cb7f 100644 --- a/rllib/env/__init__.py +++ b/rllib/env/__init__.py @@ -2,11 +2,20 @@ from ray.rllib.env.base_env import BaseEnv from ray.rllib.env.dm_env_wrapper import DMEnv from ray.rllib.env.multi_agent_env import MultiAgentEnv from ray.rllib.env.external_env import ExternalEnv -from ray.rllib.env.serving_env import ServingEnv +from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv from ray.rllib.env.vector_env import VectorEnv from ray.rllib.env.env_context import EnvContext +from ray.rllib.env.policy_client import PolicyClient +from ray.rllib.env.policy_server_input import PolicyServerInput __all__ = [ - "BaseEnv", "MultiAgentEnv", "ExternalEnv", "VectorEnv", "ServingEnv", - "EnvContext", "DMEnv" + "BaseEnv", + "MultiAgentEnv", + "ExternalEnv", + "ExternalMultiAgentEnv", + "VectorEnv", + "EnvContext", + "DMEnv", + "PolicyClient", + "PolicyServerInput", ] diff --git a/rllib/env/policy_client.py b/rllib/env/policy_client.py new file mode 100644 index 000000000..2fbace2d3 --- /dev/null +++ b/rllib/env/policy_client.py @@ -0,0 +1,294 @@ +"""REST client to interact with a policy server. + +This client supports both local and remote policy inference modes. Local +inference is faster but causes more compute to be done on the client. +""" + +import logging +import threading +import time + +import ray.cloudpickle as pickle +from ray.rllib.evaluation.rollout_worker import RolloutWorker +from ray.rllib.env import ExternalEnv, MultiAgentEnv, ExternalMultiAgentEnv +from ray.rllib.utils.annotations import PublicAPI + +logger = logging.getLogger(__name__) +logger.setLevel("INFO") # TODO(ekl) seems to be needed for cartpole_client.py + +try: + import requests # `requests` is not part of stdlib. +except ImportError: + requests = None + logger.warning( + "Couldn't import `requests` library. Be sure to install it on" + " the client side.") + + +@PublicAPI +class PolicyClient: + """REST client to interact with a RLlib policy server.""" + + # Commands for local inference mode. + GET_WORKER_ARGS = "GET_WORKER_ARGS" + GET_WEIGHTS = "GET_WEIGHTS" + REPORT_SAMPLES = "REPORT_SAMPLES" + + # Commands for remote inference mode. + START_EPISODE = "START_EPISODE" + GET_ACTION = "GET_ACTION" + LOG_ACTION = "LOG_ACTION" + LOG_RETURNS = "LOG_RETURNS" + END_EPISODE = "END_EPISODE" + + @PublicAPI + def __init__(self, address, inference_mode="local", update_interval=10.0): + """Create a PolicyClient instance. + + Args: + address (str): Server to connect to (e.g., "localhost:9090"). + inference_mode (str): Whether to use 'local' or 'remote' policy + inference for computing actions. + update_interval (float): If using 'local' inference mode, the + policy is refreshed after this many seconds have passed. + """ + self.address = address + if inference_mode == "local": + self.local = True + self._setup_local_rollout_worker(update_interval) + elif inference_mode == "remote": + self.local = False + else: + raise ValueError( + "inference_mode must be either 'local' or 'remote'") + + @PublicAPI + def start_episode(self, episode_id=None, training_enabled=True): + """Record the start of an episode. + + Arguments: + episode_id (str): Unique string id for the episode or None for + it to be auto-assigned. + training_enabled (bool): Whether to use experiences for this + episode to improve the policy. + + Returns: + episode_id (str): Unique string id for the episode. + """ + + if self.local: + self._update_local_policy() + return self.env.start_episode(episode_id, training_enabled) + + return self._send({ + "episode_id": episode_id, + "command": PolicyClient.START_EPISODE, + "training_enabled": training_enabled, + })["episode_id"] + + @PublicAPI + def get_action(self, episode_id, observation): + """Record an observation and get the on-policy action. + + Arguments: + episode_id (str): Episode id returned from start_episode(). + observation (obj): Current environment observation. + + Returns: + action (obj): Action from the env action space. + """ + + if self.local: + self._update_local_policy() + return self.env.get_action(episode_id, observation) + + return self._send({ + "command": PolicyClient.GET_ACTION, + "observation": observation, + "episode_id": episode_id, + })["action"] + + @PublicAPI + def log_action(self, episode_id, observation, action): + """Record an observation and (off-policy) action taken. + + Arguments: + episode_id (str): Episode id returned from start_episode(). + observation (obj): Current environment observation. + action (obj): Action for the observation. + """ + + if self.local: + self._update_local_policy() + return self.env.log_action(episode_id, observation, action) + + self._send({ + "command": PolicyClient.LOG_ACTION, + "observation": observation, + "action": action, + "episode_id": episode_id, + }) + + @PublicAPI + def log_returns(self, episode_id, reward, info=None): + """Record returns from the environment. + + The reward will be attributed to the previous action taken by the + episode. Rewards accumulate until the next action. If no reward is + logged before the next action, a reward of 0.0 is assumed. + + Arguments: + episode_id (str): Episode id returned from start_episode(). + reward (float): Reward from the environment. + """ + + if self.local: + self._update_local_policy() + return self.env.log_returns(episode_id, reward, info) + + self._send({ + "command": PolicyClient.LOG_RETURNS, + "reward": reward, + "info": info, + "episode_id": episode_id, + }) + + @PublicAPI + def end_episode(self, episode_id, observation): + """Record the end of an episode. + + Arguments: + episode_id (str): Episode id returned from start_episode(). + observation (obj): Current environment observation. + """ + + if self.local: + self._update_local_policy() + return self.env.end_episode(episode_id, observation) + + self._send({ + "command": PolicyClient.END_EPISODE, + "observation": observation, + "episode_id": episode_id, + }) + + def _send(self, data): + payload = pickle.dumps(data) + response = requests.post(self.address, data=payload) + if response.status_code != 200: + logger.error("Request failed {}: {}".format(response.text, data)) + response.raise_for_status() + parsed = pickle.loads(response.content) + return parsed + + def _setup_local_rollout_worker(self, update_interval): + self.update_interval = update_interval + self.last_updated = 0 + + logger.info("Querying server for rollout worker settings.") + kwargs = self._send({ + "command": PolicyClient.GET_WORKER_ARGS, + })["worker_args"] + + (self.rollout_worker, + self.inference_thread) = create_embedded_rollout_worker( + kwargs, self._send) + self.env = self.rollout_worker.env + + def _update_local_policy(self): + assert self.inference_thread.is_alive() + if time.time() - self.last_updated > self.update_interval: + logger.info("Querying server for new policy weights.") + resp = self._send({ + "command": PolicyClient.GET_WEIGHTS, + }) + weights = resp["weights"] + global_vars = resp["global_vars"] + logger.info( + "Updating rollout worker weights and global vars {}.".format( + global_vars)) + self.rollout_worker.set_weights(weights, global_vars) + self.last_updated = time.time() + + +class _LocalInferenceThread(threading.Thread): + """Thread that handles experience generation (worker.sample() loop).""" + + def __init__(self, rollout_worker, send_fn): + super().__init__() + self.daemon = True + self.rollout_worker = rollout_worker + self.send_fn = send_fn + + def run(self): + try: + while True: + logger.info("Generating new batch of experiences.") + samples = self.rollout_worker.sample() + metrics = self.rollout_worker.get_metrics() + logger.info("Sending batch of {} steps back to server.".format( + samples.count)) + self.send_fn({ + "command": PolicyClient.REPORT_SAMPLES, + "samples": samples, + "metrics": metrics, + }) + except Exception as e: + logger.info("Error: inference worker thread died!", e) + + +def auto_wrap_external(real_env_creator): + """Wrap an environment in the ExternalEnv interface if needed. + + Args: + real_env_creator (fn): Create an env given the env_config. + """ + + def wrapped_creator(env_config): + real_env = real_env_creator(env_config) + if not (isinstance(real_env, ExternalEnv) + or isinstance(real_env, ExternalMultiAgentEnv)): + logger.info( + "The env you specified is not a type of ExternalEnv. " + "Attempting to convert it automatically to ExternalEnv.") + + if isinstance(real_env, MultiAgentEnv): + external_cls = MultiAgentEnv + else: + external_cls = ExternalEnv + + class ExternalEnvWrapper(external_cls): + def __init__(self, real_env): + super().__init__(real_env.action_space, + real_env.observation_space) + + def run(self): + # Since we are calling methods on this class in the + # client, run doesn't need to do anything. + time.sleep(999999) + + return ExternalEnvWrapper(real_env) + + return wrapped_creator + + +def create_embedded_rollout_worker(kwargs, send_fn): + """Create a local rollout worker and a thread that samples from it. + + Arguments: + kwargs (dict): args for the RolloutWorker constructor. + send_fn (fn): function to send a JSON request to the server. + """ + + # Since the server acts as an input datasource, we have to reset the + # input config to the default, which runs env rollouts. + kwargs = kwargs.copy() + del kwargs["input_creator"] + logger.info("Creating rollout worker with kwargs={}".format(kwargs)) + real_env_creator = kwargs["env_creator"] + kwargs["env_creator"] = auto_wrap_external(real_env_creator) + + rollout_worker = RolloutWorker(**kwargs) + inference_thread = _LocalInferenceThread(rollout_worker, send_fn) + inference_thread.start() + return rollout_worker, inference_thread diff --git a/rllib/env/policy_server_input.py b/rllib/env/policy_server_input.py new file mode 100644 index 000000000..efb1ea996 --- /dev/null +++ b/rllib/env/policy_server_input.py @@ -0,0 +1,191 @@ +import logging +import six.moves.queue as queue +import threading +import traceback + +from http.server import SimpleHTTPRequestHandler, HTTPServer +from socketserver import ThreadingMixIn + +import ray.cloudpickle as pickle +from ray.rllib.offline.input_reader import InputReader +from ray.rllib.env.policy_client import PolicyClient, \ + create_embedded_rollout_worker +from ray.rllib.utils.annotations import override, PublicAPI + +logger = logging.getLogger(__name__) +logger.setLevel("INFO") # TODO(ekl) this is needed for cartpole_server.py + + +class PolicyServerInput(ThreadingMixIn, HTTPServer, InputReader): + """REST policy server that acts as an offline data source. + + This launches a multi-threaded server that listens on the specified host + and port to serve policy requests and forward experiences to RLlib. For + high performance experience collection, it implements InputReader. + + For an example, run `examples/cartpole_server.py` along + with `examples/cartpole_client.py --inference-mode=local|remote`. + + Examples: + >>> pg = PGTrainer( + ... env="CartPole-v0", config={ + ... "input": lambda ioctx: + ... PolicyServerInput(ioctx, addr, port), + ... "num_workers": 0, # Run just 1 server, in the trainer. + ... } + >>> while True: + pg.train() + + >>> client = PolicyClient("localhost:9900", inference_mode="local") + >>> eps_id = client.start_episode() + >>> action = client.get_action(eps_id, obs) + >>> ... + >>> client.log_returns(eps_id, reward) + >>> ... + >>> client.log_returns(eps_id, reward) + """ + + @PublicAPI + def __init__(self, ioctx, address, port): + """Create a PolicyServerInput. + + This class implements rllib.offline.InputReader, and can be used with + any Trainer by configuring + + {"num_workers": 0, + "input": lambda ioctx: PolicyServerInput(ioctx, addr, port)} + + Note that by setting num_workers: 0, the trainer will only create one + rollout worker / PolicyServerInput. Clients can connect to the launched + server using rllib.env.PolicyClient. + + Args: + ioctx (IOContext): IOContext provided by RLlib. + address (str): Server addr (e.g., "localhost"). + port (int): Server port (e.g., 9900). + """ + + self.rollout_worker = ioctx.worker + self.samples_queue = queue.Queue() + self.metrics_queue = queue.Queue() + + def get_metrics(): + completed = [] + while True: + try: + completed.append(self.metrics_queue.get_nowait()) + except queue.Empty: + break + return completed + + # Forwards client-reported rewards directly into the local rollout + # worker. This is a bit of a hack since it is patching the get_metrics + # function of the sampler. + self.rollout_worker.sampler.get_metrics = get_metrics + + handler = _make_handler(self.rollout_worker, self.samples_queue, + self.metrics_queue) + HTTPServer.__init__(self, (address, port), handler) + logger.info("") + logger.info("Starting connector server at {}:{}".format(address, port)) + logger.info("") + thread = threading.Thread(name="server", target=self.serve_forever) + thread.start() + + @override(InputReader) + def next(self): + return self.samples_queue.get() + + +def _make_handler(rollout_worker, samples_queue, metrics_queue): + # Only used in remote inference mode. We must create a new rollout worker + # then since the original worker doesn't have the env properly wrapped in + # an ExternalEnv interface. + child_rollout_worker = None + inference_thread = None + lock = threading.Lock() + + def setup_child_rollout_worker(): + nonlocal lock + nonlocal child_rollout_worker + nonlocal inference_thread + + with lock: + if child_rollout_worker is None: + (child_rollout_worker, + inference_thread) = create_embedded_rollout_worker( + rollout_worker.creation_args(), report_data) + child_rollout_worker.set_weights(rollout_worker.get_weights()) + + def report_data(data): + nonlocal child_rollout_worker + + batch = data["samples"] + batch.decompress_if_needed() + samples_queue.put(batch) + for rollout_metric in data["metrics"]: + metrics_queue.put(rollout_metric) + + if child_rollout_worker is not None: + child_rollout_worker.set_weights(rollout_worker.get_weights(), + rollout_worker.get_global_vars()) + + class Handler(SimpleHTTPRequestHandler): + def __init__(self, *a, **kw): + super().__init__(*a, **kw) + + def do_POST(self): + content_len = int(self.headers.get("Content-Length"), 0) + raw_body = self.rfile.read(content_len) + parsed_input = pickle.loads(raw_body) + try: + response = self.execute_command(parsed_input) + self.send_response(200) + self.end_headers() + self.wfile.write(pickle.dumps(response)) + except Exception: + self.send_error(500, traceback.format_exc()) + + def execute_command(self, args): + command = args["command"] + response = {} + # Local inference commands: + if command == PolicyClient.GET_WORKER_ARGS: + logger.info("Sending worker creation args to client.") + response["worker_args"] = rollout_worker.creation_args() + elif command == PolicyClient.GET_WEIGHTS: + logger.info("Sending worker weights to client.") + response["weights"] = rollout_worker.get_weights() + response["global_vars"] = rollout_worker.get_global_vars() + elif command == PolicyClient.REPORT_SAMPLES: + logger.info("Got sample batch of size {} from client.".format( + args["samples"].count)) + report_data(args) + # Remote inference commands: + elif command == PolicyClient.START_EPISODE: + setup_child_rollout_worker() + assert inference_thread.is_alive() + response["episode_id"] = ( + child_rollout_worker.env.start_episode( + args["episode_id"], args["training_enabled"])) + elif command == PolicyClient.GET_ACTION: + assert inference_thread.is_alive() + response["action"] = child_rollout_worker.env.get_action( + args["episode_id"], args["observation"]) + elif command == PolicyClient.LOG_ACTION: + assert inference_thread.is_alive() + child_rollout_worker.env.log_action( + args["episode_id"], args["observation"], args["action"]) + elif command == PolicyClient.LOG_RETURNS: + assert inference_thread.is_alive() + child_rollout_worker.env.log_returns( + args["episode_id"], args["reward"], args["info"]) + elif command == PolicyClient.END_EPISODE: + assert inference_thread.is_alive() + child_rollout_worker.env.end_episode(args["episode_id"], + args["observation"]) + else: + raise ValueError("Unknown command: {}".format(command)) + return response + + return Handler diff --git a/rllib/env/serving_env.py b/rllib/env/serving_env.py deleted file mode 100644 index b14f6800f..000000000 --- a/rllib/env/serving_env.py +++ /dev/null @@ -1,4 +0,0 @@ -from ray.rllib.env.external_env import ExternalEnv - -# renamed to ExternalEnv in 0.6 -ServingEnv = ExternalEnv diff --git a/rllib/evaluation/rollout_worker.py b/rllib/evaluation/rollout_worker.py index db3fc801a..e527220a8 100644 --- a/rllib/evaluation/rollout_worker.py +++ b/rllib/evaluation/rollout_worker.py @@ -240,6 +240,8 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker): to ensure each remote worker has unique exploration behavior. _fake_sampler (bool): Use a fake (inf speed) sampler for testing. """ + self._original_kwargs = locals().copy() + del self._original_kwargs["self"] global _global_worker _global_worker = self @@ -279,6 +281,7 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker): self.compress_observations = compress_observations self.preprocessing_enabled = True self.last_batch = None + self.global_vars = None self._fake_sampler = _fake_sampler self.env = _validate_env(env_creator(env_context)) @@ -769,6 +772,11 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker): @DeveloperAPI def set_global_vars(self, global_vars): self.foreach_policy(lambda p, _: p.on_global_var_update(global_vars)) + self.global_vars = global_vars + + @DeveloperAPI + def get_global_vars(self): + return self.global_vars @DeveloperAPI def export_policy_model(self, export_dir, policy_id=DEFAULT_POLICY_ID): @@ -786,6 +794,11 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker): def stop(self): self.async_env.stop() + @DeveloperAPI + def creation_args(self): + """Returns the args used to create this worker.""" + return self._original_kwargs + def _build_policy_map(self, policy_dict, policy_config): policy_map = {} preprocessors = {} diff --git a/rllib/examples/serving/cartpole_client.py b/rllib/examples/serving/cartpole_client.py index ce2448b7b..2026e73aa 100755 --- a/rllib/examples/serving/cartpole_client.py +++ b/rllib/examples/serving/cartpole_client.py @@ -1,18 +1,23 @@ -"""Example of querying a policy server. Copy this file for your use case. +#!/usr/bin/env python +"""Example of training with a policy server. Copy this file for your use case. To try this out, in two separate shells run: - $ python cartpole_server.py - $ python cartpole_client.py + $ python cartpole_server.py --run=[PPO|DQN] + $ python cartpole_client.py --inference-mode=local|remote + +Local inference mode offloads inference to the client for better performance. """ import argparse import gym -from ray.rllib.utils.policy_client import PolicyClient +from ray.rllib.env.policy_client import PolicyClient parser = argparse.ArgumentParser() parser.add_argument( "--no-train", action="store_true", help="Whether to disable training.") +parser.add_argument( + "--inference-mode", type=str, required=True, choices=["local", "remote"]) parser.add_argument( "--off-policy", action="store_true", @@ -26,7 +31,8 @@ parser.add_argument( if __name__ == "__main__": args = parser.parse_args() env = gym.make("CartPole-v0") - client = PolicyClient("http://localhost:9900") + client = PolicyClient( + "http://localhost:9900", inference_mode=args.inference_mode) eid = client.start_episode(training_enabled=not args.no_train) obs = env.reset() diff --git a/rllib/examples/serving/cartpole_server.py b/rllib/examples/serving/cartpole_server.py index f72d5149c..fdc77cea9 100755 --- a/rllib/examples/serving/cartpole_server.py +++ b/rllib/examples/serving/cartpole_server.py @@ -1,66 +1,84 @@ +#!/usr/bin/env python """Example of running a policy server. Copy this file for your use case. To try this out, in two separate shells run: $ python cartpole_server.py - $ python cartpole_client.py + $ python cartpole_client.py --inference-mode=local|remote """ +import argparse import os -from gym import spaces -import numpy as np import ray from ray.rllib.agents.dqn import DQNTrainer -from ray.rllib.env.external_env import ExternalEnv -from ray.rllib.utils.policy_server import PolicyServer +from ray.rllib.agents.ppo import PPOTrainer +from ray.rllib.env.policy_server_input import PolicyServerInput from ray.tune.logger import pretty_print -from ray.tune.registry import register_env SERVER_ADDRESS = "localhost" SERVER_PORT = 9900 -CHECKPOINT_FILE = "last_checkpoint.out" - - -class CartpoleServing(ExternalEnv): - def __init__(self): - ExternalEnv.__init__( - self, spaces.Discrete(2), - spaces.Box(low=-10, high=10, shape=(4, ), dtype=np.float32)) - - def run(self): - print("Starting policy server at {}:{}".format(SERVER_ADDRESS, - SERVER_PORT)) - server = PolicyServer(self, SERVER_ADDRESS, SERVER_PORT) - server.serve_forever() +CHECKPOINT_FILE = "last_checkpoint_{}.out" +parser = argparse.ArgumentParser() +parser.add_argument("--run", type=str, default="DQN") if __name__ == "__main__": + args = parser.parse_args() ray.init() - register_env("srv", lambda _: CartpoleServing()) - # We use DQN since it supports off-policy actions, but you can choose and - # configure any agent. - dqn = DQNTrainer( - env="srv", - config={ - # Use a single process to avoid needing to set up a load balancer - "num_workers": 0, - # Configure the agent to run short iterations for debugging - "exploration_fraction": 0.01, - "learning_starts": 100, - "timesteps_per_iteration": 200, - }) + env = "CartPole-v0" + connector_config = { + # Use the connector server to generate experiences. + "input": ( + lambda ioctx: PolicyServerInput( \ + ioctx, SERVER_ADDRESS, SERVER_PORT) + ), + # Use a single worker process to run the server. + "num_workers": 0, + # Disable OPE, since the rollouts are coming from online clients. + "input_evaluation": [], + } + + if args.run == "DQN": + # Example of using DQN (supports off-policy actions). + trainer = DQNTrainer( + env=env, + config=dict( + connector_config, **{ + "exploration_config": { + "type": "EpsilonGreedy", + "initial_epsilon": 1.0, + "final_epsilon": 0.02, + "epsilon_timesteps": 1000, + }, + "learning_starts": 100, + "timesteps_per_iteration": 200, + "log_level": "INFO", + })) + elif args.run == "PPO": + # Example of using PPO (does NOT support off-policy actions). + trainer = PPOTrainer( + env=env, + config=dict( + connector_config, **{ + "sample_batch_size": 1000, + "train_batch_size": 4000, + })) + else: + raise ValueError("--run must be DQN or PPO") + + checkpoint_path = CHECKPOINT_FILE.format(args.run) # Attempt to restore from checkpoint if possible. - if os.path.exists(CHECKPOINT_FILE): - checkpoint_path = open(CHECKPOINT_FILE).read() + if os.path.exists(checkpoint_path): + checkpoint_path = open(checkpoint_path).read() print("Restoring from checkpoint path", checkpoint_path) - dqn.restore(checkpoint_path) + trainer.restore(checkpoint_path) # Serving and training loop while True: - print(pretty_print(dqn.train())) - checkpoint_path = dqn.save() - print("Last checkpoint", checkpoint_path) - with open(CHECKPOINT_FILE, "w") as f: - f.write(checkpoint_path) + print(pretty_print(trainer.train())) + checkpoint = trainer.save() + print("Last checkpoint", checkpoint) + with open(checkpoint_path, "w") as f: + f.write(checkpoint) diff --git a/rllib/examples/serving/test.sh b/rllib/examples/serving/test.sh deleted file mode 100755 index d1dfa1e89..000000000 --- a/rllib/examples/serving/test.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/bin/bash - -pkill -f cartpole_server.py -(python cartpole_server.py 2>&1 | grep -v 200) & -pid=$! - -while ! curl localhost:9900; do - sleep 1 -done - -python cartpole_client.py --stop-at-reward=100 -kill $pid diff --git a/rllib/examples/serving/test_local_inference.sh b/rllib/examples/serving/test_local_inference.sh new file mode 100755 index 000000000..9d4c98cfc --- /dev/null +++ b/rllib/examples/serving/test_local_inference.sh @@ -0,0 +1,23 @@ +#!/bin/bash + +rm -f last_checkpoint.out +pkill -f cartpole_server.py +sleep 1 + +if [ -f cartpole_server.py ]; then + basedir="." +else + basedir="rllib/examples/serving" # In bazel. +fi + +(python $basedir/cartpole_server.py --run=PPO 2>&1 | grep -v 200) & +pid=$! + +echo "Waiting for server to start" +while ! curl localhost:9900; do + sleep 1 +done + +sleep 2 +python $basedir/cartpole_client.py --stop-at-reward=100 --inference-mode=local +kill $pid diff --git a/rllib/examples/serving/test_remote_inference.sh b/rllib/examples/serving/test_remote_inference.sh new file mode 100755 index 000000000..925c3887e --- /dev/null +++ b/rllib/examples/serving/test_remote_inference.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +rm -f last_checkpoint.out +pkill -f cartpole_server.py +sleep 1 + +if [ -f cartpole_server.py ]; then + basedir="." +else + basedir="rllib/examples/serving" # In bazel. +fi + +(python $basedir/cartpole_server.py --run=DQN 2>&1 | grep -v 200) & +pid=$! + +echo "Waiting for server to start" +while ! curl localhost:9900; do + sleep 1 +done + +sleep 2 +python $basedir/cartpole_client.py --stop-at-reward=100 --inference-mode=remote +kill $pid + diff --git a/rllib/utils/policy_client.py b/rllib/utils/policy_client.py index 6071c258c..1b542e9ae 100644 --- a/rllib/utils/policy_client.py +++ b/rllib/utils/policy_client.py @@ -1,7 +1,10 @@ +"""DEPRECATED: Please use rllib.env.PolicyClient instead.""" + import logging import pickle from ray.rllib.utils.annotations import PublicAPI +from ray.rllib.utils.deprecation import deprecation_warning logger = logging.getLogger(__name__) @@ -16,7 +19,7 @@ except ImportError: @PublicAPI class PolicyClient: - """REST client to interact with a RLlib policy server.""" + """DEPRECATED: Please use rllib.env.PolicyClient instead.""" START_EPISODE = "START_EPISODE" GET_ACTION = "GET_ACTION" @@ -26,6 +29,8 @@ class PolicyClient: @PublicAPI def __init__(self, address): + deprecation_warning( + "rllib.utils.PolicyServer", new="rllib.env.PolicyServerInput") self._address = address @PublicAPI diff --git a/rllib/utils/policy_server.py b/rllib/utils/policy_server.py index be69480e3..b163a7984 100644 --- a/rllib/utils/policy_server.py +++ b/rllib/utils/policy_server.py @@ -1,3 +1,5 @@ +"""DEPRECATED: Please use rllib.env.PolicyServerInput instead.""" + import pickle import traceback @@ -6,44 +8,17 @@ from socketserver import ThreadingMixIn from ray.rllib.utils.annotations import PublicAPI from ray.rllib.utils.policy_client import PolicyClient +from ray.rllib.utils.deprecation import deprecation_warning @PublicAPI class PolicyServer(ThreadingMixIn, HTTPServer): - """REST server than can be launched from a ExternalEnv. - - This launches a multi-threaded server that listens on the specified host - and port to serve policy requests and forward experiences to RLlib. - - Examples: - >>> class CartpoleServing(ExternalEnv): - def __init__(self): - ExternalEnv.__init__( - self, spaces.Discrete(2), - spaces.Box( - low=-10, - high=10, - shape=(4,), - dtype=np.float32)) - def run(self): - server = PolicyServer(self, "localhost", 8900) - server.serve_forever() - >>> register_env("srv", lambda _: CartpoleServing()) - >>> pg = PGTrainer(env="srv", config={"num_workers": 0}) - >>> while True: - pg.train() - - >>> client = PolicyClient("localhost:8900") - >>> eps_id = client.start_episode() - >>> action = client.get_action(eps_id, obs) - >>> ... - >>> client.log_returns(eps_id, reward) - >>> ... - >>> client.log_returns(eps_id, reward) - """ + """DEPRECATED: Please use rllib.env.PolicyServerInput instead.""" @PublicAPI def __init__(self, external_env, address, port): + deprecation_warning( + "rllib.utils.PolicyClient", new="rllib.env.PolicyClient") handler = _make_handler(external_env) HTTPServer.__init__(self, (address, port), handler)