Module bbrl.agents.asynchronous

Expand source code
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#

import torch
import torch.multiprocessing as mp

from bbrl.agents.agent import Agent
from bbrl.workspace import Workspace


def f(agent, in_queue, out_queue):
    while True:
        args = in_queue.get()
        if args == "exit":
            out_queue.put("ok")
            return
        workspace = Workspace()
        with torch.no_grad():
            agent(workspace, **args)
        out_queue.put("ok")
        for k in workspace.keys():
            out_queue.put((k, workspace.get_full(k)))
        out_queue.put("ok")


class AsynchronousAgent(Agent):
    """Implements an agent that is executed aynchronously in another process, and that returns its own workspace

    Usage is:
    * agent(workspace)
    * while agent.is_running():
    *     .....
    * workspace=agent.get_workspace()
    """

    def __init__(self, agent, verbose=False):
        super().__init__(verbose=verbose)
        """ Create the AsynchronousAgent

        Args:
            agent ([bbrl.Agent]): The agent to execute in another process
        """
        self._is_running = False
        self.process = None
        self._workspace = None
        self.agent = agent

    def __call__(self, **kwargs):
        """Executes the agent in non-blocking mode. A new workspace is created by the agent."""
        assert not self._is_running
        if self.process is None:
            self.o_queue = mp.Queue()
            self.o_queue.cancel_join_thread()
            self.i_queue = mp.Queue()
            self.i_queue.cancel_join_thread()
            self.process = mp.Process(
                target=f, args=(self.agent, self.i_queue, self.o_queue)
            )
            self.process.daemon = False
            self.process.start()
        self._is_running = True
        self.i_queue.put(kwargs)

    def is_running(self):
        """Is the agent still running ?

        Returns:
            [bool]: True is the agent is running
        """
        if self._is_running:
            try:
                r = self.o_queue.get(False)
                assert r == "ok"
                self._is_running = False
                r = self.o_queue.get()
                workspace = Workspace()
                while r != "ok":
                    key, val = r
                    workspace.set_full(key, val)
                    r = self.o_queue.get()
                self._workspace = workspace.to("cpu")
            except:
                pass  # TODO: shouldn't we assert False?
        return self._is_running

    def get_workspace(self):
        """Returns the built workspace is the agent has stopped its execution

        Returns:
            [bbrl.Workspace]: The built workspace
        """
        if self.is_running():
            return None
        return self._workspace

    def close(self):
        """Close the agent and kills the corresponding process"""
        if self.process is None:
            return

        if self.verbose:
            print("[AsynchronousAgent] closing process")
        self.i_queue.put("exit")
        self.o_queue.get()
        self.process.terminate()
        self.process.join()
        self.i_queue.close()
        self.o_queue.close()
        del self.i_queue
        del self.o_queue
        self.process = None

    def __del__(self):
        self.close()

Functions

def f(agent, in_queue, out_queue)
Expand source code
def f(agent, in_queue, out_queue):
    while True:
        args = in_queue.get()
        if args == "exit":
            out_queue.put("ok")
            return
        workspace = Workspace()
        with torch.no_grad():
            agent(workspace, **args)
        out_queue.put("ok")
        for k in workspace.keys():
            out_queue.put((k, workspace.get_full(k)))
        out_queue.put("ok")

Classes

class AsynchronousAgent (agent, verbose=False)

Implements an agent that is executed aynchronously in another process, and that returns its own workspace

Usage is: * agent(workspace) * while agent.is_running(): * ..... * workspace=agent.get_workspace()

To create a new Agent

Args

name : [type], optional
An agent can have a name that will allow to perform operations

on agents that are composed into more complex agents.

Expand source code
class AsynchronousAgent(Agent):
    """Implements an agent that is executed aynchronously in another process, and that returns its own workspace

    Usage is:
    * agent(workspace)
    * while agent.is_running():
    *     .....
    * workspace=agent.get_workspace()
    """

    def __init__(self, agent, verbose=False):
        super().__init__(verbose=verbose)
        """ Create the AsynchronousAgent

        Args:
            agent ([bbrl.Agent]): The agent to execute in another process
        """
        self._is_running = False
        self.process = None
        self._workspace = None
        self.agent = agent

    def __call__(self, **kwargs):
        """Executes the agent in non-blocking mode. A new workspace is created by the agent."""
        assert not self._is_running
        if self.process is None:
            self.o_queue = mp.Queue()
            self.o_queue.cancel_join_thread()
            self.i_queue = mp.Queue()
            self.i_queue.cancel_join_thread()
            self.process = mp.Process(
                target=f, args=(self.agent, self.i_queue, self.o_queue)
            )
            self.process.daemon = False
            self.process.start()
        self._is_running = True
        self.i_queue.put(kwargs)

    def is_running(self):
        """Is the agent still running ?

        Returns:
            [bool]: True is the agent is running
        """
        if self._is_running:
            try:
                r = self.o_queue.get(False)
                assert r == "ok"
                self._is_running = False
                r = self.o_queue.get()
                workspace = Workspace()
                while r != "ok":
                    key, val = r
                    workspace.set_full(key, val)
                    r = self.o_queue.get()
                self._workspace = workspace.to("cpu")
            except:
                pass  # TODO: shouldn't we assert False?
        return self._is_running

    def get_workspace(self):
        """Returns the built workspace is the agent has stopped its execution

        Returns:
            [bbrl.Workspace]: The built workspace
        """
        if self.is_running():
            return None
        return self._workspace

    def close(self):
        """Close the agent and kills the corresponding process"""
        if self.process is None:
            return

        if self.verbose:
            print("[AsynchronousAgent] closing process")
        self.i_queue.put("exit")
        self.o_queue.get()
        self.process.terminate()
        self.process.join()
        self.i_queue.close()
        self.o_queue.close()
        del self.i_queue
        del self.o_queue
        self.process = None

    def __del__(self):
        self.close()

Ancestors

  • Agent
  • torch.nn.modules.module.Module

Class variables

var dump_patches : bool
var training : bool

Methods

def close(self)

Close the agent and kills the corresponding process

Expand source code
def close(self):
    """Close the agent and kills the corresponding process"""
    if self.process is None:
        return

    if self.verbose:
        print("[AsynchronousAgent] closing process")
    self.i_queue.put("exit")
    self.o_queue.get()
    self.process.terminate()
    self.process.join()
    self.i_queue.close()
    self.o_queue.close()
    del self.i_queue
    del self.o_queue
    self.process = None
def get_workspace(self)

Returns the built workspace is the agent has stopped its execution

Returns

[bbrl.Workspace]
The built workspace
Expand source code
def get_workspace(self):
    """Returns the built workspace is the agent has stopped its execution

    Returns:
        [bbrl.Workspace]: The built workspace
    """
    if self.is_running():
        return None
    return self._workspace
def is_running(self)

Is the agent still running ?

Returns

[bool]
True is the agent is running
Expand source code
def is_running(self):
    """Is the agent still running ?

    Returns:
        [bool]: True is the agent is running
    """
    if self._is_running:
        try:
            r = self.o_queue.get(False)
            assert r == "ok"
            self._is_running = False
            r = self.o_queue.get()
            workspace = Workspace()
            while r != "ok":
                key, val = r
                workspace.set_full(key, val)
                r = self.o_queue.get()
            self._workspace = workspace.to("cpu")
        except:
            pass  # TODO: shouldn't we assert False?
    return self._is_running

Inherited members