Module bbrl.agents.remote
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 copy
import time
import torch
import torch.multiprocessing as mp
from bbrl.agents.agent import Agent
from bbrl.workspace import Workspace, _SplitSharedWorkspace
# Unused???
def f(agent, in_queue, out_queue, seed, verbose):
"""The function that is executed in a single process"""
out_queue.put("ok")
running = True
old_workspace = None
if verbose:
print("Seeding remote agent with ", seed)
agent.seed(seed)
while running:
command = in_queue.get()
if command[0] == "go_new_workspace":
_, workspace, kwargs = command
old_workspace = workspace
agent(workspace, **kwargs)
out_queue.put("ok")
elif command[0] == "go_reuse_workspace":
_, _, kwargs = command
agent(old_workspace, **kwargs)
out_queue.put("ok")
elif command[0] == "exit":
out_queue.put("ok")
return
elif command[0] == "eval_mode":
agent.eval()
out_queue.put("ok")
elif command[0] == "train_mode":
agent.train()
out_queue.put("ok")
class RemoteAgent(Agent):
"""It corresponds to an agent that is executed in another process
Args:
Agent ([bbrl.Agent]): the agent ot execute in another process
"""
def __init__(self, agent, name=None, verbose=False):
super().__init__(name=name, verbose=verbose)
self.agent = agent
self._is_running = False
self.process = None
self.last_workspace = None
self.train_mode = True
def get_by_name(self, n):
if self._name == n:
return [self] + self.agent.get_by_name(n)
else:
return self.agent.get_by_name(n)
def forward(self, **kwargs):
raise NotImplementedError
def _create_process(self):
if self.verbose:
print("[RemoteAgent] starting process...")
self.i_queue = mp.Queue()
self.o_queue = mp.Queue()
self.i_queue.cancel_join_thread()
self.o_queue.cancel_join_thread()
self.process = mp.Process(
target=f, args=(self.agent, self.i_queue, self.o_queue, self._seed)
)
self.process.daemon = False
self.process.start()
def __call__(self, workspace, **kwargs):
with torch.no_grad():
assert (
workspace.is_shared
), "You must use a shared workspace when using a Remote Agent"
if self.process is None:
self._create_process()
self.train(self.train_mode)
if not workspace == self.last_workspace:
self.i_queue.put(("go_new_workspace", workspace, kwargs))
self.last_workspace = workspace
r = self.o_queue.get()
assert r == "ok"
else:
self.i_queue.put(("go_reuse_workspace", workspace, kwargs))
r = self.o_queue.get()
assert r == "ok"
def _asynchronous_call(self, workspace, **kwargs):
"""Non-blocking forward. To use together with `is_running`"""
with torch.no_grad():
self._is_running = True
assert (
workspace.is_shared
), "You must use a shared workspace when using a Remote Agent"
if self.process is None:
self._create_process()
if not workspace == self.last_workspace:
self.i_queue.put(("go_new_workspace", workspace, kwargs))
self.last_workspace = workspace
else:
self.i_queue.put(("go_reuse_workspace", workspace, kwargs))
def train(self, f=True):
self.train_mode = f
if self.process is None:
return
if f:
self.i_queue.put(("train_mode",))
a = self.o_queue.get()
assert a == "ok"
else:
self.eval()
def eval(self):
self.train_mode = False
if self.process is None:
return
self.i_queue.put(("eval_mode",))
a = self.o_queue.get()
assert a == "ok"
def seed(self, _seed):
self._seed = _seed
def _running_queue(self):
return self.o_queue
def is_running(self):
if self._is_running:
try:
r = self.o_queue.get(False)
assert r == "ok"
self._is_running = False
except:
pass
return self._is_running
def close(self):
if self.process is None:
return
if self.verbose:
print("[RemoteAgent] closing process")
self.i_queue.put(("exit",))
self.o_queue.get()
time.sleep(0.1)
self.process.terminate()
self.process.join()
self.i_queue.close()
self.o_queue.close()
time.sleep(0.1)
del self.i_queue
del self.o_queue
self.process = None
def __del__(self):
self.close()
class NRemoteAgent(Agent):
"""Multiple agents executed in different processes. Use the `NRemoteAgent.create` function to create such an agent"""
def __init__(self, agents, batch_dims):
super().__init__()
self.agents = agents
self.batch_dims = batch_dims
def get_by_name(self, name):
r = []
if self._name == name:
r = [self]
for a in self.agents:
r = r + a.get_by_name(name)
return r
@staticmethod
def create(agent, num_processes=0, time_size=None, **extra_kwargs):
"""Returns a NRemote agent with num_processes copies of agent in different processes
Also returns the specific workspace to use with such an agent
Args:
agent ([bbrl.Agent]): The agent to execute in multiple processes
num_processes (int, optional): Number of processes to create. If 0, then no processes are created (for debugging). Defaults to 0.
time_size ([type], optional): If specified, it forces the created Workspace to have this particular time_size. Defaults to None.
Returns:
[bbrl.Agent,bbrl.SharedWorkspace]: The NRemoteAgent and the corresponding workspace
"""
agent.seed(0)
if num_processes == 0:
workspace = Workspace()
_agent = copy.deepcopy(agent)
agent(workspace, **extra_kwargs)
shared_workspace = workspace._convert_to_shared_workspace(
n_repeat=1, time_size=time_size
)
return _agent, shared_workspace
workspace = Workspace()
agents = [copy.deepcopy(agent) for t in range(num_processes)]
agent(workspace, **extra_kwargs)
b = workspace.batch_size()
batch_dims = [(k * b, k * b + b) for k, a in enumerate(agents)]
shared_workspace = workspace._convert_to_shared_workspace(
n_repeat=num_processes, time_size=time_size
)
agents = [RemoteAgent(a) for a in agents]
return NRemoteAgent(agents, batch_dims), shared_workspace
def __call__(self, workspace, **kwargs):
assert workspace.is_shared
for k in range(len(self.agents)):
_workspace = _SplitSharedWorkspace(workspace, self.batch_dims[k])
self.agents[k]._asynchronous_call(_workspace, **kwargs)
for a in self.agents:
ok = a._running_queue().get()
assert ok == "ok"
def seed(self, seed, inc=1):
s = seed
for a in self.agents:
a.seed(s)
s += inc
def _asynchronous_call(self, workspace, **kwargs):
assert workspace.is_shared
for k in range(len(self.agents)):
_workspace = _SplitSharedWorkspace(workspace, self.batch_dims[k])
self.agents[k]._asynchronous_call(_workspace, **kwargs)
def is_running(self):
for a in self.agents:
if a.is_running():
return True
return False
def train(self, f=True):
for a in self.agents:
a.train(f)
def eval(self):
for a in self.agents:
a.eval()
def close(self):
for a in self.agents:
a.close()
Functions
def f(agent, in_queue, out_queue, seed, verbose)
-
The function that is executed in a single process
Expand source code
def f(agent, in_queue, out_queue, seed, verbose): """The function that is executed in a single process""" out_queue.put("ok") running = True old_workspace = None if verbose: print("Seeding remote agent with ", seed) agent.seed(seed) while running: command = in_queue.get() if command[0] == "go_new_workspace": _, workspace, kwargs = command old_workspace = workspace agent(workspace, **kwargs) out_queue.put("ok") elif command[0] == "go_reuse_workspace": _, _, kwargs = command agent(old_workspace, **kwargs) out_queue.put("ok") elif command[0] == "exit": out_queue.put("ok") return elif command[0] == "eval_mode": agent.eval() out_queue.put("ok") elif command[0] == "train_mode": agent.train() out_queue.put("ok")
Classes
class NRemoteAgent (agents, batch_dims)
-
Multiple agents executed in different processes. Use the
NRemoteAgent.create()
function to create such an agentTo 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 NRemoteAgent(Agent): """Multiple agents executed in different processes. Use the `NRemoteAgent.create` function to create such an agent""" def __init__(self, agents, batch_dims): super().__init__() self.agents = agents self.batch_dims = batch_dims def get_by_name(self, name): r = [] if self._name == name: r = [self] for a in self.agents: r = r + a.get_by_name(name) return r @staticmethod def create(agent, num_processes=0, time_size=None, **extra_kwargs): """Returns a NRemote agent with num_processes copies of agent in different processes Also returns the specific workspace to use with such an agent Args: agent ([bbrl.Agent]): The agent to execute in multiple processes num_processes (int, optional): Number of processes to create. If 0, then no processes are created (for debugging). Defaults to 0. time_size ([type], optional): If specified, it forces the created Workspace to have this particular time_size. Defaults to None. Returns: [bbrl.Agent,bbrl.SharedWorkspace]: The NRemoteAgent and the corresponding workspace """ agent.seed(0) if num_processes == 0: workspace = Workspace() _agent = copy.deepcopy(agent) agent(workspace, **extra_kwargs) shared_workspace = workspace._convert_to_shared_workspace( n_repeat=1, time_size=time_size ) return _agent, shared_workspace workspace = Workspace() agents = [copy.deepcopy(agent) for t in range(num_processes)] agent(workspace, **extra_kwargs) b = workspace.batch_size() batch_dims = [(k * b, k * b + b) for k, a in enumerate(agents)] shared_workspace = workspace._convert_to_shared_workspace( n_repeat=num_processes, time_size=time_size ) agents = [RemoteAgent(a) for a in agents] return NRemoteAgent(agents, batch_dims), shared_workspace def __call__(self, workspace, **kwargs): assert workspace.is_shared for k in range(len(self.agents)): _workspace = _SplitSharedWorkspace(workspace, self.batch_dims[k]) self.agents[k]._asynchronous_call(_workspace, **kwargs) for a in self.agents: ok = a._running_queue().get() assert ok == "ok" def seed(self, seed, inc=1): s = seed for a in self.agents: a.seed(s) s += inc def _asynchronous_call(self, workspace, **kwargs): assert workspace.is_shared for k in range(len(self.agents)): _workspace = _SplitSharedWorkspace(workspace, self.batch_dims[k]) self.agents[k]._asynchronous_call(_workspace, **kwargs) def is_running(self): for a in self.agents: if a.is_running(): return True return False def train(self, f=True): for a in self.agents: a.train(f) def eval(self): for a in self.agents: a.eval() def close(self): for a in self.agents: a.close()
Ancestors
- Agent
- torch.nn.modules.module.Module
Class variables
var dump_patches : bool
var training : bool
Static methods
def create(agent, num_processes=0, time_size=None, **extra_kwargs)
-
Returns a NRemote agent with num_processes copies of agent in different processes Also returns the specific workspace to use with such an agent
Args
agent
:[bbrl.Agent]
- The agent to execute in multiple processes
num_processes
:int
, optional- Number of processes to create. If 0, then no processes are created (for debugging). Defaults to 0.
time_size
:[type]
, optional- If specified, it forces the created Workspace to have this particular time_size. Defaults to None.
Returns
[bbrl.Agent,bbrl.SharedWorkspace]
- The NRemoteAgent and the corresponding workspace
Expand source code
@staticmethod def create(agent, num_processes=0, time_size=None, **extra_kwargs): """Returns a NRemote agent with num_processes copies of agent in different processes Also returns the specific workspace to use with such an agent Args: agent ([bbrl.Agent]): The agent to execute in multiple processes num_processes (int, optional): Number of processes to create. If 0, then no processes are created (for debugging). Defaults to 0. time_size ([type], optional): If specified, it forces the created Workspace to have this particular time_size. Defaults to None. Returns: [bbrl.Agent,bbrl.SharedWorkspace]: The NRemoteAgent and the corresponding workspace """ agent.seed(0) if num_processes == 0: workspace = Workspace() _agent = copy.deepcopy(agent) agent(workspace, **extra_kwargs) shared_workspace = workspace._convert_to_shared_workspace( n_repeat=1, time_size=time_size ) return _agent, shared_workspace workspace = Workspace() agents = [copy.deepcopy(agent) for t in range(num_processes)] agent(workspace, **extra_kwargs) b = workspace.batch_size() batch_dims = [(k * b, k * b + b) for k, a in enumerate(agents)] shared_workspace = workspace._convert_to_shared_workspace( n_repeat=num_processes, time_size=time_size ) agents = [RemoteAgent(a) for a in agents] return NRemoteAgent(agents, batch_dims), shared_workspace
Methods
def close(self)
-
Expand source code
def close(self): for a in self.agents: a.close()
def eval(self)
-
Sets the module in evaluation mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:
Dropout
, :class:BatchNorm
, etc.This is equivalent with :meth:
self.train(False) <torch.nn.Module.train>
.See :ref:
locally-disable-grad-doc
for a comparison between.eval()
and several similar mechanisms that may be confused with it.Returns
Module
- self
Expand source code
def eval(self): for a in self.agents: a.eval()
def train(self, f=True)
-
Sets the module in training mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:
Dropout
, :class:BatchNorm
, etc.Args
mode
:bool
- whether to set training mode (
True
) or evaluation mode (False
). Default:True
.
Returns
Module
- self
Expand source code
def train(self, f=True): for a in self.agents: a.train(f)
Inherited members
class RemoteAgent (agent, name=None, verbose=False)
-
It corresponds to an agent that is executed in another process
Args
Agent
:[bbrl.Agent]
- the agent ot execute in another process
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 RemoteAgent(Agent): """It corresponds to an agent that is executed in another process Args: Agent ([bbrl.Agent]): the agent ot execute in another process """ def __init__(self, agent, name=None, verbose=False): super().__init__(name=name, verbose=verbose) self.agent = agent self._is_running = False self.process = None self.last_workspace = None self.train_mode = True def get_by_name(self, n): if self._name == n: return [self] + self.agent.get_by_name(n) else: return self.agent.get_by_name(n) def forward(self, **kwargs): raise NotImplementedError def _create_process(self): if self.verbose: print("[RemoteAgent] starting process...") self.i_queue = mp.Queue() self.o_queue = mp.Queue() self.i_queue.cancel_join_thread() self.o_queue.cancel_join_thread() self.process = mp.Process( target=f, args=(self.agent, self.i_queue, self.o_queue, self._seed) ) self.process.daemon = False self.process.start() def __call__(self, workspace, **kwargs): with torch.no_grad(): assert ( workspace.is_shared ), "You must use a shared workspace when using a Remote Agent" if self.process is None: self._create_process() self.train(self.train_mode) if not workspace == self.last_workspace: self.i_queue.put(("go_new_workspace", workspace, kwargs)) self.last_workspace = workspace r = self.o_queue.get() assert r == "ok" else: self.i_queue.put(("go_reuse_workspace", workspace, kwargs)) r = self.o_queue.get() assert r == "ok" def _asynchronous_call(self, workspace, **kwargs): """Non-blocking forward. To use together with `is_running`""" with torch.no_grad(): self._is_running = True assert ( workspace.is_shared ), "You must use a shared workspace when using a Remote Agent" if self.process is None: self._create_process() if not workspace == self.last_workspace: self.i_queue.put(("go_new_workspace", workspace, kwargs)) self.last_workspace = workspace else: self.i_queue.put(("go_reuse_workspace", workspace, kwargs)) def train(self, f=True): self.train_mode = f if self.process is None: return if f: self.i_queue.put(("train_mode",)) a = self.o_queue.get() assert a == "ok" else: self.eval() def eval(self): self.train_mode = False if self.process is None: return self.i_queue.put(("eval_mode",)) a = self.o_queue.get() assert a == "ok" def seed(self, _seed): self._seed = _seed def _running_queue(self): return self.o_queue def is_running(self): if self._is_running: try: r = self.o_queue.get(False) assert r == "ok" self._is_running = False except: pass return self._is_running def close(self): if self.process is None: return if self.verbose: print("[RemoteAgent] closing process") self.i_queue.put(("exit",)) self.o_queue.get() time.sleep(0.1) self.process.terminate() self.process.join() self.i_queue.close() self.o_queue.close() time.sleep(0.1) 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)
-
Expand source code
def close(self): if self.process is None: return if self.verbose: print("[RemoteAgent] closing process") self.i_queue.put(("exit",)) self.o_queue.get() time.sleep(0.1) self.process.terminate() self.process.join() self.i_queue.close() self.o_queue.close() time.sleep(0.1) del self.i_queue del self.o_queue self.process = None
def eval(self)
-
Sets the module in evaluation mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:
Dropout
, :class:BatchNorm
, etc.This is equivalent with :meth:
self.train(False) <torch.nn.Module.train>
.See :ref:
locally-disable-grad-doc
for a comparison between.eval()
and several similar mechanisms that may be confused with it.Returns
Module
- self
Expand source code
def eval(self): self.train_mode = False if self.process is None: return self.i_queue.put(("eval_mode",)) a = self.o_queue.get() assert a == "ok"
def train(self, f=True)
-
Sets the module in training mode.
This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:
Dropout
, :class:BatchNorm
, etc.Args
mode
:bool
- whether to set training mode (
True
) or evaluation mode (False
). Default:True
.
Returns
Module
- self
Expand source code
def train(self, f=True): self.train_mode = f if self.process is None: return if f: self.i_queue.put(("train_mode",)) a = self.o_queue.get() assert a == "ok" else: self.eval()
Inherited members