def dumps(x):
""" Manage between cloudpickle and pickle
1. Try pickle
2. If it is short then check if it contains __main__
3. If it is long, then first check type, then check __main__
"""
try:
result = pickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
if len(result) < 1000:
if b'__main__' in result:
return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
else:
return result
else:
if isinstance(x, pickle_types) or b'__main__' not in result:
return result
else:
return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
except:
try:
return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
except Exception:
logger.info("Failed to serialize %s", x, exc_info=True)
raise
def test_mpi_objects():
# Neighbours
grid = Grid(shape=(4, 4, 4))
obj = grid.distributor._obj_neighborhood
pkl_obj = pickle.dumps(obj)
new_obj = pickle.loads(pkl_obj)
assert obj.name == new_obj.name
assert obj.pname == new_obj.pname
assert obj.pfields == new_obj.pfields
# Communicator
obj = grid.distributor._obj_comm
pkl_obj = pickle.dumps(obj)
new_obj = pickle.loads(pkl_obj)
assert obj.name == new_obj.name
assert obj.dtype == new_obj.dtype
# Status
obj = MPIStatusObject(name='status')
pkl_obj = pickle.dumps(obj)
new_obj = pickle.loads(pkl_obj)
assert obj.name == new_obj.name
assert obj.dtype == new_obj.dtype
# Request
obj = MPIRequestObject(name='request')
pkl_obj = pickle.dumps(obj)
new_obj = pickle.loads(pkl_obj)
assert obj.name == new_obj.name
assert obj.dtype == new_obj.dtype
def test_closed_file(self):
# Write & close
with open(self.tmpfilepath, 'w') as f:
f.write(self.teststring)
with pytest.raises(pickle.PicklingError) as excinfo:
cloudpickle.dumps(f)
assert "Cannot pickle closed files" in str(excinfo.value)
os.remove(self.tmpfilepath)
def test_itemgetter(self):
d = range(10)
getter = itemgetter(1)
getter2 = pickle.loads(cloudpickle.dumps(getter))
self.assertEqual(getter(d), getter2(d))
getter = itemgetter(0, 3)
getter2 = pickle.loads(cloudpickle.dumps(getter))
self.assertEqual(getter(d), getter2(d))
def test_queue_serde(zk):
queue = Queue(zk, '/satyr/serde')
queue.put(cp.dumps({'a': 1, 'b': 2}))
queue.put(cp.dumps({'c': 3}))
pickled_queue = cp.dumps(queue)
unpickled_queue = cp.loads(pickled_queue)
assert cp.loads(unpickled_queue.get()) == {'a': 1, 'b': 2}
assert cp.loads(unpickled_queue.get()) == {'c': 3}
def NOT_WORKING_test_tty(self):
# FIXME: Mocking 'file' is not trivial... and fails for now
from sys import version_info
if version_info.major == 2:
import __builtin__ as builtins # pylint:disable=import-error
else:
import builtins # pylint:disable=import-error
with patch.object(builtins, 'open', mock_open(), create=True):
with open('foo', 'w+') as handle:
cloudpickle.dumps(handle)
def test_locking_queue_serde(zk):
queue = LockingQueue(zk, '/satyr/serde_locking')
queue.put(cp.dumps({'a': 1, 'b': 2}))
queue.put(cp.dumps({'c': 3}))
pickled_queue = cp.dumps(queue)
unpickled_queue = cp.loads(pickled_queue)
assert cp.loads(unpickled_queue.get()) == {'a': 1, 'b': 2}
unpickled_queue.consume()
assert cp.loads(unpickled_queue.get()) == {'c': 3}
unpickled_queue.consume()
def test_internal_symbols():
s = dSymbol(name='s', dtype=np.float32)
pkl_s = pickle.dumps(s)
new_s = pickle.loads(pkl_s)
assert new_s.name == s.name
assert new_s.dtype is np.float32
s = Scalar(name='s', dtype=np.int32, is_const=True)
pkl_s = pickle.dumps(s)
new_s = pickle.loads(pkl_s)
assert new_s.name == s.name
assert new_s.dtype is np.int32
assert new_s.is_const is True
def send_spyder_msg(self, spyder_msg_type, content=None, data=None):
"""publish custom messages to the spyder frontend
Parameters
----------
spyder_msg_type: str
The spyder message type
content: dict
The (JSONable) content of the message
data: any
Any object that is serializable by cloudpickle (should be most
things). Will arrive as cloudpickled bytes in `.buffers[0]`.
"""
import cloudpickle
if content is None:
content = {}
content['spyder_msg_type'] = spyder_msg_type
self.session.send(
self.iopub_socket,
'spyder_msg',
content=content,
buffers=[cloudpickle.dumps(data, protocol=PICKLE_PROTOCOL)],
parent=self._parent_header,
)
def subprocess_pickle_echo(input_data):
"""Echo function with a child Python process
Pickle the input data into a buffer, send it to a subprocess via
stdin, expect the subprocess to unpickle, re-pickle that data back
and send it back to the parent process via stdout for final unpickling.
>>> subprocess_pickle_echo([1, 'a', None])
[1, 'a', None]
"""
pickled_input_data = dumps(input_data)
cmd = [sys.executable, __file__]
cwd = os.getcwd()
proc = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE, cwd=cwd)
try:
comm_kwargs = {}
if timeout_supported:
comm_kwargs['timeout'] = 5
out, err = proc.communicate(pickled_input_data, **comm_kwargs)
if proc.returncode != 0 or len(err):
message = "Subprocess returned %d: " % proc.returncode
message += err.decode('utf-8')
raise RuntimeError(message)
return loads(out)
except TimeoutExpired:
proc.kill()
out, err = proc.communicate()
message = u"\n".join([out.decode('utf-8'), err.decode('utf-8')])
raise RuntimeError(message)
def _normalize_function(func):
if isinstance(func, curry):
func = func._partial
if isinstance(func, Compose):
first = getattr(func, 'first', None)
funcs = reversed((first,) + func.funcs) if first else func.funcs
return tuple(normalize_function(f) for f in funcs)
elif isinstance(func, partial):
args = tuple(normalize_token(i) for i in func.args)
if func.keywords:
kws = tuple((k, normalize_token(v))
for k, v in sorted(func.keywords.items()))
else:
kws = None
return (normalize_function(func.func), args, kws)
else:
try:
result = pickle.dumps(func, protocol=0)
if b'__main__' not in result: # abort on dynamic functions
return result
except Exception:
pass
try:
import cloudpickle
return cloudpickle.dumps(func, protocol=0)
except Exception:
return str(func)
def _send_op(result, foo, chunk, op, index, target_ip, own_ip, port, timeout):
'''
Sends an operation over the network for a server to process, and
receives the result. Since we want each chunk to be sent in
parallel, this should be threaded
:param result: empty list passed by reference which will contain the result.
Necessary because threads don't allow standard return values
:param foo: function to use for map, filter, or reduce calls
:param chunk: chunk to perform operation on
:param op: string corresponding to operation to perform:
'map', 'filter', 'reduce'
:param index: chunk number to allow ordering of processed chunks
:param port: port of server
'''
try:
dict_sending = {'func': foo, 'chunk': chunk, 'op': op, 'index': index}
csts = threading.Thread(
target = _client_socket_thread_send,
args = (target_ip, port, pickle.dumps(dict_sending), timeout))
csts.start()
queue = Queue.Queue()
cstr = threading.Thread(
target = _client_socket_thread_receive,
args = (own_ip, port+1, queue, timeout))
cstr.start()
cstr.join(timeout = None)
response = pickle.loads(queue.get())
result[response['index']] = response['chunk']
except RuntimeError, socket.timeout:
return #do nothing on error, just end and the client will restart the sending protocol
def test_geometry():
shape = (50, 50, 50)
spacing = [10. for _ in shape]
nbpml = 10
nrec = 10
tn = 150.
# Create two-layer model from preset
model = demo_model(preset='layers-isotropic', vp_top=1., vp_bottom=2.,
spacing=spacing, shape=shape, nbpml=nbpml)
# Source and receiver geometries
src_coordinates = np.empty((1, len(spacing)))
src_coordinates[0, :] = np.array(model.domain_size) * .5
if len(shape) > 1:
src_coordinates[0, -1] = model.origin[-1] + 2 * spacing[-1]
rec_coordinates = np.empty((nrec, len(spacing)))
rec_coordinates[:, 0] = np.linspace(0., model.domain_size[0], num=nrec)
if len(shape) > 1:
rec_coordinates[:, 1] = np.array(model.domain_size)[1] * .5
rec_coordinates[:, -1] = model.origin[-1] + 2 * spacing[-1]
geometry = AcquisitionGeometry(model, rec_coordinates, src_coordinates,
t0=0.0, tn=tn, src_type='Ricker', f0=0.010)
pkl_geom = pickle.dumps(geometry)
new_geom = pickle.loads(pkl_geom)
assert np.all(new_geom.src_positions == geometry.src_positions)
assert np.all(new_geom.rec_positions == geometry.rec_positions)
assert new_geom.f0 == geometry.f0
assert np.all(new_geom.src_type == geometry.src_type)
assert np.all(new_geom.src.data == geometry.src.data)
assert new_geom.t0 == geometry.t0
assert new_geom.tn == geometry.tn
def _pack(input_data, protocol=None):
pickled_input_data = dumps(input_data, protocol=protocol)
# Under Windows + Python 2.7, subprocess / communicate truncate the data
# on some specific bytes. To avoid this issue, let's use the pure ASCII
# Base32 encoding to encapsulate the pickle message sent to the child
# process.
return base64.b32encode(pickled_input_data)
def dumps(x):
try:
return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
except Exception as e:
logger.info("Failed to serialize %s", x)
logger.exception(e)
raise
def test_plus_mode(self):
# Write, then seek to 0
with open(self.tmpfilepath, 'w+') as f:
f.write(self.teststring)
f.seek(0)
self.assertEquals(self.teststring, pickle.loads(cloudpickle.dumps(f)).read())
os.remove(self.tmpfilepath)
def test_w_mode(self):
with open(self.tmpfilepath, 'w') as f:
f.write(self.teststring)
f.seek(0)
self.assertRaises(pickle.PicklingError,
lambda: cloudpickle.dumps(f))
os.remove(self.tmpfilepath)
def test_extended_arg(self):
# Functions with more than 65535 global vars prefix some global
# variable references with the EXTENDED_ARG opcode.
nvars = 65537 + 258
names = ['g%d' % i for i in range(1, nvars)]
r = random.Random(42)
d = {name: r.randrange(100) for name in names}
# def f(x):
# x = g1, g2, ...
# return zlib.crc32(bytes(bytearray(x)))
code = """
import zlib
def f():
x = {tup}
return zlib.crc32(bytes(bytearray(x)))
""".format(tup=', '.join(names))
exec(textwrap.dedent(code), d, d)
f = d['f']
res = f()
data = cloudpickle.dumps([f, f], protocol=self.protocol)
d = f = None
f2, f3 = pickle.loads(data)
self.assertTrue(f2 is f3)
self.assertEqual(f2(), res)
def thunk(*args, **kwargs):
serialized_fn = base64.b64encode(cloudpickle.dumps(lambda: fn(*args, **kwargs)))
subprocess.check_call([
'mpiexec','-n', str(nproc),
sys.executable,
'-m', 'baselines.common.tests.test_with_mpi',
serialized_fn
], env=os.environ, timeout=timeout)
def _send_job(self, command, job):
pickled_job = cloudpickle.dumps(job)
base64_pickled_job = base64.b64encode(pickled_job).decode('utf-8')
base64_pickled_job_data = {'job': base64_pickled_job}
handle = JobHandle(self._conn, self._session_id,
self._executor)
handle._start(command, base64_pickled_job_data)
return handle
def test_closed_file(self):
# Write & close
with open(self.tmpfilepath, 'w') as f:
f.write(self.teststring)
# Cloudpickle returns an empty (& closed!) StringIO if the file was closed...
unpickled = pickle.loads(cloudpickle.dumps(f))
self.assertTrue(unpickled.closed)
os.remove(self.tmpfilepath)
def test_temp_file(self):
with tempfile.NamedTemporaryFile(mode='ab+') as fp:
fp.write(self.teststring.encode('UTF-8'))
fp.seek(0)
f = fp.file
# FIXME this doesn't work yet: cloudpickle.dumps(fp)
newfile = pickle.loads(cloudpickle.dumps(f))
self.assertEquals(self.teststring, newfile.read())
def test_queue_size(zk):
queue = Queue(zk, '/satyr/size')
assert queue.empty()
assert queue.qsize() == 0
queue.put(cp.dumps(range(5)))
assert queue.empty() is False
assert queue.qsize() == 1
def test_symbolics():
a = Symbol('a')
id = IntDiv(a, 3)
pkl_id = pickle.dumps(id)
new_id = pickle.loads(pkl_id)
assert id == new_id
ffp = FunctionFromPointer('foo', a, ['b', 'c'])
pkl_ffp = pickle.dumps(ffp)
new_ffp = pickle.loads(pkl_ffp)
assert ffp == new_ffp
li = ListInitializer(['a', 'b'])
pkl_li = pickle.dumps(li)
new_li = pickle.loads(pkl_li)
assert li == new_li
def call_func(payload, protocol):
"""Remote function call that uses cloudpickle to transport everthing"""
func, args, kwargs = loads(payload)
try:
result = func(*args, **kwargs)
except BaseException as e:
result = e
return dumps(result, protocol=protocol)
def test_r_mode(self):
# Write & close
with open(self.tmpfilepath, 'w') as f:
f.write(self.teststring)
# Open for reading
with open(self.tmpfilepath, 'r') as f:
self.assertEquals(self.teststring, pickle.loads(cloudpickle.dumps(f)).read())
os.remove(self.tmpfilepath)
def dumps(x):
try:
if isinstance(x, pickle_types):
return pickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
else:
return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
except Exception as e:
logger.info("Failed to serialize %s", x, exc_info=True)
raise
def test_timers():
"""Pickling for Timers used in Operators for C-level profiling."""
timer = Timer('timer', ['sec0', 'sec1'])
pkl_obj = pickle.dumps(timer)
new_obj = pickle.loads(pkl_obj)
assert new_obj.name == timer.name
assert new_obj.sections == timer.sections
assert new_obj.value._obj.sec0 == timer.value._obj.sec0 == 0.0
assert new_obj.value._obj.sec1 == timer.value._obj.sec1 == 0.0
def test_operator_parameters():
grid = Grid(shape=(3, 3, 3))
f = Function(name='f', grid=grid)
g = TimeFunction(name='g', grid=grid)
h = TimeFunction(name='h', grid=grid, save=10)
op = Operator(Eq(h.forward, h + g + f + 1))
for i in op.parameters:
pkl_i = pickle.dumps(i)
pickle.loads(pkl_i)
def save(self, filename):
"""
:param filename: file name to save to
:type name: str
Save optimizer state to disk
"""
with open(filename, "wb") as output_file:
output_file.write(cloudpickle.dumps(self.get_state()))
import time
from common.utils.serialize_utils import Serializer
from models.user.user_info import UserToken, UserService
import cloudpickle
import pickle
if __name__ == '__main__':
uk = UserToken(ip='0.0.0.0', port='9999', user_id=1, serviceId=1277)
us = UserService(user_token=uk)
# print(type(globals()['UserToken']))
# print(globals()['UserToken'])
# temp = Serializer.serialize(uk)
# print(temp)
# print(type(temp))
temp = cloudpickle.dumps(us)
print(temp)
print(type(temp))
recover = pickle.loads(temp)
print(recover)
print(type(recover))
print(type(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))))
def serialize_function(function):
return cloudpickle.dumps(function)
def _dumps(x):
return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
def run(self, arg=None, convertToStr=False, collect=True):
self.pipe.run(arg, convertToStr, collect)
selfBytes = cloudpickle.dumps(self.pipe)
results = self.r.execute_command('RG.PYEXECUTEREMOTE', selfBytes)
res, errs = results
return res, errs
def run_experiment(method_call=None,
batch_tasks=None,
exp_prefix='experiment',
exp_name=None,
log_dir=None,
script='garage.experiment.experiment_wrapper',
python_command='python',
dry=False,
env=None,
variant=None,
force_cpu=False,
pre_commands=None,
**kwargs):
"""Serialize the method call and run the experiment using the
specified mode.
Args:
method_call (callable): A method call.
batch_tasks (list[dict]): A batch of method calls.
exp_prefix (str): Name prefix for the experiment.
exp_name (str): Name of the experiment.
log_dir (str): Log directory for the experiment.
script (str): The name of the entrance point python script.
python_command (str): Python command to run the experiment.
dry (bool): Whether to do a dry-run, which only prints the
commands without executing them.
env (dict): Extra environment variables.
variant (dict): If provided, should be a dictionary of parameters.
force_cpu (bool): Whether to set all GPU devices invisible
to force use CPU.
pre_commands (str): Pre commands to run the experiment.
kwargs (dict): Additional parameters.
"""
# pylint: disable=missing-raises-doc,global-statement,too-many-branches
if method_call is None and batch_tasks is None:
raise Exception(
'Must provide at least either method_call or batch_tasks')
for task in (batch_tasks or [method_call]):
if not hasattr(task, '__call__'):
raise ValueError('batch_tasks should be callable')
# ensure variant exists
if variant is None:
variant = dict()
if batch_tasks is None:
batch_tasks = [
dict(kwargs,
pre_commands=pre_commands,
method_call=method_call,
exp_name=exp_name,
log_dir=log_dir,
env=env,
variant=variant)
]
global exp_count
if force_cpu:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
for task in batch_tasks:
call = task.pop('method_call')
data = base64.b64encode(cloudpickle.dumps(call)).decode('utf-8')
task['args_data'] = data
exp_count += 1
if task.get('exp_name', None) is None:
task['exp_name'] = '{}_{}_{:04n}'.format(exp_prefix, timestamp,
exp_count)
if task.get('log_dir', None) is None:
task['log_dir'] = (
'{log_dir}/local/{exp_prefix}/{exp_name}'.format(
log_dir=osp.join(os.getcwd(), 'data'),
exp_prefix=exp_prefix.replace('_', '-'),
exp_name=task['exp_name']))
if task.get('variant', None) is not None:
variant = task.pop('variant')
if 'exp_name' not in variant:
variant['exp_name'] = task['exp_name']
task['variant_data'] = base64.b64encode(
pickle.dumps(variant)).decode('utf-8')
elif 'variant' in task:
del task['variant']
task['env'] = task.get('env', dict()) or dict()
task['env']['GARAGE_FORCE_CPU'] = str(force_cpu)
for task in batch_tasks:
env = task.pop('env', None)
command = to_local_command(task,
python_command=python_command,
script=script)
print(command)
if dry:
return
try:
if env is None:
env = dict()
subprocess.run(command,
shell=True,
env=dict(os.environ, **env),
check=True)
except Exception as e:
print(e)
raise
def test_is_pickleable(self):
e = Executor()
post = cloudpickle.loads(cloudpickle.dumps(e))
assert isinstance(post, Executor)
def fnc_send(c):
if c.port == baseport:
x = 13
c.send(get_min_port_friend(c).name, cloudpickle.dumps(lambda: x**2))
if c.port == baseport + 1:
wait_for(c, 169, lambda x: x())
def call_experiment(exp_name, thunk, seed=0, num_cpu=1, data_dir=None,
datestamp=False, **kwargs):
"""
Run a function (thunk) with hyperparameters (kwargs), plus configuration.
This wraps a few pieces of functionality which are useful when you want
to run many experiments in sequence, including logger configuration and
splitting into multiple processes for MPI.
There's also a SpinningUp-specific convenience added into executing the
thunk: if ``env_name`` is one of the kwargs passed to call_experiment, it's
assumed that the thunk accepts an argument called ``env_fn``, and that
the ``env_fn`` should make a gym environment with the given ``env_name``.
The way the experiment is actually executed is slightly complicated: the
function is serialized to a string, and then ``run_entrypoint.py`` is
executed in a subprocess call with the serialized string as an argument.
``run_entrypoint.py`` unserializes the function call and executes it.
We choose to do it this way---instead of just calling the function
directly here---to avoid leaking state between successive experiments.
Args:
exp_name (string): Name for experiment.
thunk (callable): A python function.
seed (int): Seed for random number generators.
num_cpu (int): Number of MPI processes to split into. Also accepts
'auto', which will set up as many procs as there are cpus on
the machine.
data_dir (string): Used in configuring the logger, to decide where
to store experiment results. Note: if left as None, data_dir will
default to ``DEFAULT_DATA_DIR`` from ``spinup/user_config.py``.
**kwargs: All kwargs to pass to thunk.
"""
# Determine number of CPU cores to run on
num_cpu = psutil.cpu_count(logical=False) if num_cpu=='auto' else num_cpu
# Send random seed to thunk
kwargs['seed'] = seed
# Be friendly and print out your kwargs, so we all know what's up
print(colorize('Running experiment:\n', color='cyan', bold=True))
print(exp_name + '\n')
print(colorize('with kwargs:\n', color='cyan', bold=True))
kwargs_json = convert_json(kwargs)
print(json.dumps(kwargs_json, separators=(',',':\t'), indent=4, sort_keys=True))
print('\n')
# Set up logger output directory
if 'logger_kwargs' not in kwargs:
kwargs['logger_kwargs'] = setup_logger_kwargs(exp_name, seed, data_dir, datestamp)
else:
print('Note: Call experiment is not handling logger_kwargs.\n')
def thunk_plus():
# Make 'env_fn' from 'env_name'
if 'env_name' in kwargs:
import gym
env_name = kwargs['env_name']
kwargs['env_fn'] = lambda : gym.make(env_name)
del kwargs['env_name']
# Fork into multiple processes
mpi_fork(num_cpu)
# Run thunk
thunk(**kwargs)
# Prepare to launch a script to run the experiment
pickled_thunk = cloudpickle.dumps(thunk_plus)
encoded_thunk = base64.b64encode(zlib.compress(pickled_thunk)).decode('utf-8')
entrypoint = osp.join(osp.abspath(osp.dirname(__file__)),'run_entrypoint.py')
cmd = [sys.executable if sys.executable else 'python', entrypoint, encoded_thunk]
try:
subprocess.check_call(cmd, env=os.environ)
except CalledProcessError:
err_msg = '\n'*3 + '='*DIV_LINE_WIDTH + '\n' + dedent("""
There appears to have been an error in your experiment.
Check the traceback above to see what actually went wrong. The
traceback below, included for completeness (but probably not useful
for diagnosing the error), shows the stack leading up to the
experiment launch.
""") + '='*DIV_LINE_WIDTH + '\n'*3
print(err_msg)
raise
# Tell the user about where results are, and how to check them
logger_kwargs = kwargs['logger_kwargs']
plot_cmd = 'python -m spinup.run plot '+logger_kwargs['output_dir']
plot_cmd = colorize(plot_cmd, 'green')
test_cmd = 'python -m spinup.run test_policy '+logger_kwargs['output_dir']
test_cmd = colorize(test_cmd, 'green')
output_msg = '\n'*5 + '='*DIV_LINE_WIDTH +'\n' + dedent("""\
End of experiment.
Plot results from this run with:
%s
Watch the trained agent with:
%s
"""%(plot_cmd,test_cmd)) + '='*DIV_LINE_WIDTH + '\n'*5
print(output_msg)
def serialize(obj):
""" Should take a complex object and pickle it"""
pickled = clp.dumps(obj)
compressed = lz4.frame.compress(pickled)
return compressed
def work_on_population(redis: StrictRedis,
start_time: float,
max_runtime_s: float,
kill_handler: KillHandler):
"""
Here the actual sampling happens.
"""
# set timers
population_start_time = time()
cumulative_simulation_time = 0
# read from pipeline
pipeline = redis.pipeline()
# extract bytes
ssa_b, batch_size_b, all_accepted_b, n_req_b, n_acc_b \
= (pipeline.get(SSA).get(BATCH_SIZE)
.get(ALL_ACCEPTED).get(N_REQ).get(N_ACC).execute())
if ssa_b is None:
return
kill_handler.exit = False
if n_acc_b is None:
return
# convert from bytes
simulate_one, sample_factory = pickle.loads(ssa_b)
batch_size = int(batch_size_b.decode())
all_accepted = bool(int(all_accepted_b.decode()))
n_req = int(n_req_b.decode())
# notify sign up as worker
n_worker = redis.incr(N_WORKER)
logger.info(
f"Begin population, batch size {batch_size}. "
f"I am worker {n_worker}")
# counter for number of simulations
internal_counter = 0
# create empty sample
sample = sample_factory()
# loop until no more particles required
while int(redis.get(N_ACC).decode()) < n_req \
and (not all_accepted or int(redis.get(N_EVAL).decode()) < n_req):
if kill_handler.killed:
logger.info(
f"Worker {n_worker} received stop signal. "
f"Terminating in the middle of a population "
f"after {internal_counter} samples.")
# notify quit
redis.decr(N_WORKER)
sys.exit(0)
# check whether time's up
current_runtime = time() - start_time
if current_runtime > max_runtime_s:
logger.info(
f"Worker {n_worker} stops during population because "
f"runtime {current_runtime} exceeds "
f"max runtime {max_runtime_s}")
# notify quit
redis.decr(N_WORKER)
return
# increase global number of evaluations counter
particle_max_id = redis.incr(N_EVAL, batch_size)
# timer for current simulation until batch_size acceptances
this_sim_start = time()
# collect accepted particles
accepted_samples = []
# make batch_size attempts
for n_batched in range(batch_size):
# increase evaluation counter
internal_counter += 1
try:
# simulate
new_sim = simulate_one()
# append to current sample
sample.append(new_sim)
# check for acceptance
if new_sim.accepted:
# the order of the IDs is reversed, but this does not
# matter. Important is only that the IDs are specified
# before the simulation starts
# append to accepted list
accepted_samples.append(
cloudpickle.dumps(
(particle_max_id - n_batched, sample)))
# initialize new sample
sample = sample_factory()
except Exception as e:
logger.warning(f"Redis worker number {n_worker} failed. "
f"Error message is: {e}")
# initialize new sample to be sure
sample = sample_factory()
# update total simulation-specific time
cumulative_simulation_time += time() - this_sim_start
# push to pipeline if at least one sample got accepted
if len(accepted_samples) > 0:
# new pipeline
pipeline = redis.pipeline()
# update particles counter
pipeline.incr(N_ACC, len(accepted_samples))
# note: samples are appended 1-by-1
pipeline.rpush(QUEUE, *accepted_samples)
# execute all commands
pipeline.execute()
# end of sampling loop
# notify quit
redis.decr(N_WORKER)
kill_handler.exit = True
population_total_time = time() - population_start_time
logger.info(
f"Finished population, did {internal_counter} samples. "
f"Simulation time: {cumulative_simulation_time:.2f}s, "
f"total time {population_total_time:.2f}.")
def _get_java_python_function_operator(self, func: Union[Function,
FunctionWrapper],
type_info: TypeInformation,
func_name: str, func_type: int):
"""
Create a flink operator according to user provided function object, data types,
function name and function type.
:param func: a function object that implements the Function interface.
:param type_info: the data type of the function output data.
:param func_name: function name.
:param func_type: function type, supports MAP, FLAT_MAP, etc.
:return: A flink java operator which is responsible for execution user defined python
function.
"""
gateway = get_gateway()
import cloudpickle
serialized_func = cloudpickle.dumps(func)
j_input_types = self._j_data_stream.getTransformation().getOutputType()
if type_info is None:
output_type_info = PickledBytesTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO(
)
else:
if isinstance(type_info, list):
output_type_info = RowTypeInfo(type_info)
else:
output_type_info = type_info
DataStreamPythonFunction = gateway.jvm.org.apache.flink.datastream.runtime.functions \
.python.DataStreamPythonFunction
j_python_data_stream_scalar_function = DataStreamPythonFunction(
func_name, bytearray(serialized_func), _get_python_env())
DataStreamPythonFunctionInfo = gateway.jvm. \
org.apache.flink.datastream.runtime.functions.python \
.DataStreamPythonFunctionInfo
j_python_data_stream_function_info = DataStreamPythonFunctionInfo(
j_python_data_stream_scalar_function, func_type)
j_conf = gateway.jvm.org.apache.flink.configuration.Configuration()
# set max bundle size to 1 to force synchronize process for reduce function.
from pyflink.fn_execution.flink_fn_execution_pb2 import UserDefinedDataStreamFunction
if func_type == UserDefinedDataStreamFunction.REDUCE:
j_conf.setInteger(
gateway.jvm.org.apache.flink.python.PythonOptions.
MAX_BUNDLE_SIZE, 1)
DataStreamPythonReduceFunctionOperator = gateway.jvm.org.apache.flink.datastream \
.runtime.operators.python.DataStreamPythonReduceFunctionOperator
j_output_type_info = j_input_types.getTypeAt(1)
j_python_data_stream_function_operator = DataStreamPythonReduceFunctionOperator(
j_conf, j_input_types, j_output_type_info,
j_python_data_stream_function_info)
return j_python_data_stream_function_operator, j_output_type_info
else:
DataStreamPythonFunctionOperator = gateway.jvm.org.apache.flink.datastream.runtime \
.operators.python.DataStreamPythonStatelessFunctionOperator
j_python_data_stream_function_operator = DataStreamPythonFunctionOperator(
j_conf, j_input_types, output_type_info.get_java_type_info(),
j_python_data_stream_function_info)
return j_python_data_stream_function_operator, output_type_info.get_java_type_info(
)
def run_experiment_old(task,
exp_prefix='default',
seed=None,
variant=None,
time_it=True,
save_profile=False,
profile_file='time_log.prof',
mode='here',
exp_id=0,
unique_id=None,
prepend_date_to_exp_prefix=True,
use_gpu=False,
snapshot_mode='last',
snapshot_gap=1,
n_parallel=0,
base_log_dir=None,
**run_experiment_lite_kwargs):
"""
Run a task via the rllab interface, i.e. serialize it and then run it via
the run_experiment_lite script.
This will soon be deprecated.
:param task:
:param exp_prefix:
:param seed:
:param variant:
:param time_it: Add a "time" command to the python command?
:param save_profile: Create a cProfile log?
:param profile_file: Where to save the cProfile log.
:param mode: 'here' will run the code in line, without any serialization
Other options include 'local', 'local_docker', and 'ec2'. See
run_experiment_lite documentation to learn what those modes do.
:param exp_id: Experiment ID. Should be unique across all
experiments. Note that one experiment may correspond to multiple seeds.
:param unique_id: Unique ID should be unique across all runs--even different
seeds!
:param prepend_date_to_exp_prefix: If True, prefix "month-day_" to
exp_prefix
:param run_experiment_lite_kwargs: kwargs to be passed to
`run_experiment_lite`
:return:
"""
if seed is None:
seed = random.randint(0, 100000)
if variant is None:
variant = {}
if unique_id is None:
unique_id = str(uuid.uuid4())
if prepend_date_to_exp_prefix:
exp_prefix = time.strftime("%m-%d") + "_" + exp_prefix
variant['seed'] = str(seed)
variant['exp_id'] = str(exp_id)
variant['unique_id'] = str(unique_id)
logger.log("Variant:")
logger.log(json.dumps(ppp.dict_to_safe_json(variant), indent=2))
command_words = []
if time_it:
command_words.append('time')
command_words.append('python')
if save_profile:
command_words += ['-m cProfile -o', profile_file]
repo = git.Repo(os.getcwd())
diff_string = repo.git.diff(None)
commit_hash = repo.head.commit.hexsha
script_name = "tmp"
if mode == 'here':
log_dir, exp_name = create_log_dir(exp_prefix, exp_id, seed,
base_log_dir)
data = dict(
log_dir=log_dir,
exp_name=exp_name,
mode=mode,
variant=variant,
exp_id=exp_id,
exp_prefix=exp_prefix,
seed=seed,
use_gpu=use_gpu,
snapshot_mode=snapshot_mode,
snapshot_gap=snapshot_gap,
diff_string=diff_string,
commit_hash=commit_hash,
n_parallel=n_parallel,
base_log_dir=base_log_dir,
script_name=script_name,
)
save_experiment_data(data, log_dir)
if mode == 'here':
run_experiment_here(
task,
exp_prefix=exp_prefix,
variant=variant,
exp_id=exp_id,
seed=seed,
use_gpu=use_gpu,
snapshot_mode=snapshot_mode,
snapshot_gap=snapshot_gap,
code_diff=diff_string,
commit_hash=commit_hash,
script_name=script_name,
n_parallel=n_parallel,
base_log_dir=base_log_dir,
)
else:
if mode == "ec2" and use_gpu:
if not query_yes_no("EC2 is more expensive with GPUs. Confirm?"):
sys.exit(1)
code_diff = (base64.b64encode(
cloudpickle.dumps(diff_string)).decode("utf-8"))
run_experiment_lite(task,
snapshot_mode=snapshot_mode,
snapshot_gap=snapshot_gap,
exp_prefix=exp_prefix,
variant=variant,
seed=seed,
use_cloudpickle=True,
python_command=' '.join(command_words),
mode=mode,
use_gpu=use_gpu,
script="railrl/scripts/run_experiment_lite.py",
code_diff=code_diff,
commit_hash=commit_hash,
script_name=script_name,
n_parallel=n_parallel,
**run_experiment_lite_kwargs)
def __init__(
self,
func: ty.Callable,
audit_flags: AuditFlag = AuditFlag.NONE,
cache_dir=None,
cache_locations=None,
input_spec: ty.Optional[SpecInfo] = None,
messenger_args=None,
messengers=None,
name=None,
output_spec: ty.Optional[BaseSpec] = None,
**kwargs,
):
"""
Initialize this task.
Parameters
----------
func : :obj:`callable`
A Python executable function.
audit_flags : :obj:`pydra.utils.messenger.AuditFlag`
Auditing configuration
cache_dir : :obj:`os.pathlike`
Cache directory
cache_locations : :obj:`list` of :obj:`os.pathlike`
List of alternative cache locations.
input_spec : :obj:`pydra.engine.specs.SpecInfo`
Specification of inputs.
messenger_args :
TODO
messengers :
TODO
name : :obj:`str`
Name of this task.
output_spec : :obj:`pydra.engine.specs.BaseSpec`
Specification of inputs.
"""
if input_spec is None:
input_spec = SpecInfo(
name="Inputs",
fields=[(
val.name,
val.annotation,
dc.field(
default=val.default,
metadata={
"help_string":
f"{val.name} parameter from {func.__name__}"
},
),
) if val.default is not inspect.Signature.empty else (
val.name,
val.annotation,
dc.field(metadata={"help_string": val.name}),
) for val in inspect.signature(func).parameters.values()] +
[("_func", str, cp.dumps(func))],
bases=(BaseSpec, ),
)
else:
input_spec.fields.append(("_func", str, cp.dumps(func)))
self.input_spec = input_spec
if name is None:
name = func.__name__
super(FunctionTask, self).__init__(
name,
inputs=kwargs,
audit_flags=audit_flags,
messengers=messengers,
messenger_args=messenger_args,
cache_dir=cache_dir,
cache_locations=cache_locations,
)
if output_spec is None:
if "return" not in func.__annotations__:
output_spec = SpecInfo(name="Output",
fields=[("out", ty.Any)],
bases=(BaseSpec, ))
else:
return_info = func.__annotations__["return"]
if hasattr(return_info, "__name__") and hasattr(
return_info, "__annotations__"):
output_spec = SpecInfo(
name=return_info.__name__,
fields=list(return_info.__annotations__.items()),
bases=(BaseSpec, ),
)
# Objects like int, float, list, tuple, and dict do not have __name__ attribute.
elif hasattr(return_info, "__annotations__"):
output_spec = SpecInfo(
name="Output",
fields=list(return_info.__annotations__.items()),
bases=(BaseSpec, ),
)
elif isinstance(return_info, dict):
output_spec = SpecInfo(
name="Output",
fields=list(return_info.items()),
bases=(BaseSpec, ),
)
else:
if not isinstance(return_info, tuple):
return_info = (return_info, )
output_spec = SpecInfo(
name="Output",
fields=[("out{}".format(n + 1), t)
for n, t in enumerate(return_info)],
bases=(BaseSpec, ),
)
elif "return" in func.__annotations__:
raise NotImplementedError("Branch not implemented")
self.output_spec = output_spec
def _capture_function_code_using_cloudpickle(
func, modules_to_capture: List[str] = None) -> str:
import base64
import sys
import cloudpickle
import pickle
if modules_to_capture is None:
modules_to_capture = [func.__module__]
# Hack to force cloudpickle to capture the whole function instead of just referencing the code file. See https://github.com/cloudpipe/cloudpickle/blob/74d69d759185edaeeac7bdcb7015cfc0c652f204/cloudpickle/cloudpickle.py#L490
old_modules = {}
try: # Try is needed to restore the state if something goes wrong
for module_name in modules_to_capture:
if module_name in sys.modules:
old_modules[module_name] = sys.modules.pop(module_name)
func_pickle = base64.b64encode(
cloudpickle.dumps(func, pickle.DEFAULT_PROTOCOL))
finally:
sys.modules.update(old_modules)
function_loading_code = '''\
import sys
try:
import cloudpickle as _cloudpickle
except ImportError:
import subprocess
try:
print("cloudpickle is not installed. Installing it globally", file=sys.stderr)
subprocess.run([sys.executable, "-m", "pip", "install", "cloudpickle==1.1.1", "--quiet"], env={"PIP_DISABLE_PIP_VERSION_CHECK": "1"}, check=True)
print("Installed cloudpickle globally", file=sys.stderr)
except:
print("Failed to install cloudpickle globally. Installing for the current user.", file=sys.stderr)
subprocess.run([sys.executable, "-m", "pip", "install", "cloudpickle==1.1.1", "--user", "--quiet"], env={"PIP_DISABLE_PIP_VERSION_CHECK": "1"}, check=True)
print("Installed cloudpickle for the current user", file=sys.stderr)
# Enable loading from user-installed package directory. Python does not add it to sys.path if it was empty at start. Running pip does not refresh `sys.path`.
import site
sys.path.append(site.getusersitepackages())
import cloudpickle as _cloudpickle
print("cloudpickle loaded successfully after installing.", file=sys.stderr)
''' + '''
pickler_python_version = {pickler_python_version}
current_python_version = tuple(sys.version_info)
if (
current_python_version[0] != pickler_python_version[0] or
current_python_version[1] < pickler_python_version[1] or
current_python_version[0] == 3 and ((pickler_python_version[1] < 6) != (current_python_version[1] < 6))
):
raise RuntimeError("Incompatible python versions: " + str(current_python_version) + " instead of " + str(pickler_python_version))
if current_python_version != pickler_python_version:
print("Warning!: Different python versions. The code may crash! Current environment python version: " + str(current_python_version) + ". Component code python version: " + str(pickler_python_version), file=sys.stderr)
import base64
import pickle
{func_name} = pickle.loads(base64.b64decode({func_pickle}))
'''.format(
func_name=func.__name__,
func_pickle=repr(func_pickle),
pickler_python_version=repr(tuple(sys.version_info)),
)
return function_loading_code
def run_experiment_lite(
stub_method_call=None,
batch_tasks=None,
exp_prefix="experiment",
exp_name=None,
log_dir=None,
script="scripts/run_experiment_lite.py",
python_command="python",
mode="local",
dry=False,
docker_image=None,
aws_config=None,
env=None,
variant=None,
use_gpu=False,
sync_s3_pkl=False,
sync_log_on_termination=True,
confirm_remote=True,
terminate_machine=True,
periodic_sync=True,
periodic_sync_interval=15,
sync_all_data_node_to_s3=True,
use_cloudpickle=False,
**kwargs):
"""
Serialize the stubbed method call and run the experiment using the specified mode.
:param stub_method_call: A stubbed method call.
:param script: The name of the entrance point python script
:param mode: Where & how to run the experiment. Should be one of "local", "local_docker", "ec2",
and "lab_kube".
:param dry: Whether to do a dry-run, which only prints the commands without executing them.
:param exp_prefix: Name prefix for the experiments
:param docker_image: name of the docker image. Ignored if using local mode.
:param aws_config: configuration for AWS. Only used under EC2 mode
:param env: extra environment variables
:param kwargs: All other parameters will be passed directly to the entrance python script.
:param variant: If provided, should be a dictionary of parameters
:param use_gpu: Whether the launched task is running on GPU. This triggers a few configuration changes including
certain environment flags
:param sync_s3_pkl: Whether to sync pkl files during execution of the experiment (they will always be synced at
the end of the experiment)
:param confirm_remote: Whether to confirm before launching experiments remotely
:param terminate_machine: Whether to terminate machine after experiment finishes. Only used when using
mode="ec2". This is useful when one wants to debug after an experiment finishes abnormally.
:param periodic_sync: Whether to synchronize certain experiment files periodically during execution.
:param periodic_sync_interval: Time interval between each periodic sync, in seconds.
"""
assert stub_method_call is not None or batch_tasks is not None, "Must provide at least either stub_method_call or batch_tasks"
if batch_tasks is None:
batch_tasks = [
dict(
kwargs,
stub_method_call=stub_method_call,
exp_name=exp_name,
log_dir=log_dir,
env=env,
variant=variant,
use_cloudpickle=use_cloudpickle
)
]
global exp_count
global remote_confirmed
config.USE_GPU = use_gpu
# params_list = []
for task in batch_tasks:
call = task.pop("stub_method_call")
if use_cloudpickle:
import cloudpickle
data = base64.b64encode(cloudpickle.dumps(call)).decode("utf-8")
else:
data = base64.b64encode(pickle.dumps(call)).decode("utf-8")
task["args_data"] = data
exp_count += 1
params = dict(kwargs)
if task.get("exp_name", None) is None:
task["exp_name"] = "%s_%s_%04d" % (
exp_prefix, timestamp, exp_count)
if task.get("log_dir", None) is None:
task["log_dir"] = config.LOG_DIR + "/local/" + \
exp_prefix.replace("_", "-") + "/" + task["exp_name"]
if task.get("variant", None) is not None:
variant = task.pop("variant")
if "exp_name" not in variant:
variant["exp_name"] = task["exp_name"]
task["variant_data"] = base64.b64encode(pickle.dumps(variant)).decode("utf-8")
elif "variant" in task:
del task["variant"]
task["remote_log_dir"] = osp.join(
config.AWS_S3_PATH, exp_prefix.replace("_", "-"), task["exp_name"])
if mode not in ["local", "local_docker"] and not remote_confirmed and not dry and confirm_remote:
remote_confirmed = query_yes_no(
"Running in (non-dry) mode %s. Confirm?" % mode)
if not remote_confirmed:
sys.exit(1)
if mode == "local":
for task in batch_tasks:
del task["remote_log_dir"]
env = task.pop("env", None)
command = to_local_command(
task, python_command=python_command, script=osp.join(config.PROJECT_PATH, script), use_gpu=use_gpu)
print(command)
if dry:
return
try:
if env is None:
env = dict()
subprocess.call(
command, shell=True, env=dict(os.environ, **env))
except Exception as e:
print(e)
if isinstance(e, KeyboardInterrupt):
raise
elif mode == "local_docker":
if docker_image is None:
docker_image = config.DOCKER_IMAGE
for task in batch_tasks:
del task["remote_log_dir"]
env = task.pop("env", None)
command = to_docker_command(
task,
docker_image=docker_image,
script=script,
env=env,
use_gpu=use_gpu,
use_tty=True,
)
print(command)
if dry:
return
p = subprocess.Popen(command, shell=True)
try:
p.wait()
except KeyboardInterrupt:
try:
print("terminating")
p.terminate()
except OSError:
print("os error!")
pass
p.wait()
elif mode == "ec2":
if docker_image is None:
docker_image = config.DOCKER_IMAGE
s3_code_path = s3_sync_code(config, dry=dry)
launch_ec2(batch_tasks,
exp_prefix=exp_prefix,
docker_image=docker_image,
python_command=python_command,
script=script,
aws_config=aws_config,
dry=dry,
terminate_machine=terminate_machine,
use_gpu=use_gpu,
code_full_path=s3_code_path,
sync_s3_pkl=sync_s3_pkl,
sync_log_on_termination=sync_log_on_termination,
periodic_sync=periodic_sync,
periodic_sync_interval=periodic_sync_interval)
elif mode == "lab_kube":
# assert env is None
# first send code folder to s3
s3_code_path = s3_sync_code(config, dry=dry)
if docker_image is None:
docker_image = config.DOCKER_IMAGE
for task in batch_tasks:
# if 'env' in task:
# assert task.pop('env') is None
# TODO: dangerous when there are multiple tasks?
task["resources"] = params.pop(
"resources", config.KUBE_DEFAULT_RESOURCES)
task["node_selector"] = params.pop(
"node_selector", config.KUBE_DEFAULT_NODE_SELECTOR)
task["exp_prefix"] = exp_prefix
pod_dict = to_lab_kube_pod(
task, code_full_path=s3_code_path, docker_image=docker_image, script=script, is_gpu=use_gpu,
python_command=python_command,
sync_s3_pkl=sync_s3_pkl, periodic_sync=periodic_sync, periodic_sync_interval=periodic_sync_interval,
sync_all_data_node_to_s3=sync_all_data_node_to_s3,
terminate_machine=terminate_machine,
)
pod_str = json.dumps(pod_dict, indent=1)
if dry:
print(pod_str)
dir = "{pod_dir}/{exp_prefix}".format(
pod_dir=config.POD_DIR, exp_prefix=exp_prefix)
ensure_dir(dir)
fname = "{dir}/{exp_name}.json".format(
dir=dir,
exp_name=task["exp_name"]
)
with open(fname, "w") as fh:
fh.write(pod_str)
kubecmd = "kubectl create -f %s" % fname
print(kubecmd)
if dry:
return
retry_count = 0
wait_interval = 1
while retry_count <= 5:
try:
return_code = subprocess.call(kubecmd, shell=True)
if return_code == 0:
break
retry_count += 1
print("trying again...")
time.sleep(wait_interval)
except Exception as e:
if isinstance(e, KeyboardInterrupt):
raise
print(e)
else:
raise NotImplementedError
def _serialize(self, value, attr, obj, **kwargs):
return base64.b64encode(cloudpickle.dumps(value)).decode('ascii')
def __reduce__(self):
s = cloudpickle.dumps(self._obj)
return cloudpickle.loads, (s,)
def test_is_pickleable(self, executor):
post = cloudpickle.loads(cloudpickle.dumps(executor))
assert isinstance(post, DaskExecutor)
assert post.client is None
def __getstate__(self):
state = self.__dict__.copy()
state["input_spec"] = cp.dumps(state["input_spec"])
state["output_spec"] = cp.dumps(state["output_spec"])
state["inputs"] = dc.asdict(state["inputs"])
return state
def __getstate__(self):
import cloudpickle
return cloudpickle.dumps(self.x)
def learn(network, env, seed, total_timesteps=int(40e6), gamma=0.99, log_interval=100, nprocs=32, nsteps=20,
ent_coef=0.01, vf_coef=0.5, vf_fisher_coef=1.0, lr=0.25, max_grad_norm=0.5,
kfac_clip=0.001, save_interval=None, lrschedule='linear', load_path=None, is_async=True, **network_kwargs):
set_global_seeds(seed)
if network == 'cnn':
network_kwargs['one_dim_bias'] = True
policy = build_policy(env, network, **network_kwargs)
nenvs = env.num_envs
ob_space = env.observation_space
ac_space = env.action_space
make_model = lambda : Model(policy, ob_space, ac_space, nenvs, total_timesteps, nprocs=nprocs, nsteps
=nsteps, ent_coef=ent_coef, vf_coef=vf_coef, vf_fisher_coef=
vf_fisher_coef, lr=lr, max_grad_norm=max_grad_norm, kfac_clip=kfac_clip,
lrschedule=lrschedule, is_async=is_async)
if save_interval and logger.get_dir():
import cloudpickle
with open(osp.join(logger.get_dir(), 'make_model.pkl'), 'wb') as fh:
fh.write(cloudpickle.dumps(make_model))
model = make_model()
if load_path is not None:
model.load(load_path)
runner = Runner(env, model, nsteps=nsteps, gamma=gamma)
epinfobuf = deque(maxlen=100)
nbatch = nenvs*nsteps
tstart = time.time()
coord = tf.train.Coordinator()
if is_async:
enqueue_threads = model.q_runner.create_threads(model.sess, coord=coord, start=True)
else:
enqueue_threads = []
for update in range(1, total_timesteps//nbatch+1):
obs, states, rewards, masks, actions, values, epinfos = runner.run()
epinfobuf.extend(epinfos)
policy_loss, value_loss, policy_entropy = model.train(obs, states, rewards, masks, actions, values)
model.old_obs = obs
nseconds = time.time()-tstart
fps = int((update*nbatch)/nseconds)
if update % log_interval == 0 or update == 1:
ev = explained_variance(values, rewards)
logger.record_tabular("nupdates", update)
logger.record_tabular("total_timesteps", update*nbatch)
logger.record_tabular("fps", fps)
logger.record_tabular("policy_entropy", float(policy_entropy))
logger.record_tabular("policy_loss", float(policy_loss))
logger.record_tabular("value_loss", float(value_loss))
logger.record_tabular("explained_variance", float(ev))
logger.record_tabular("eprewmean", safemean([epinfo['r'] for epinfo in epinfobuf]))
logger.record_tabular("eplenmean", safemean([epinfo['l'] for epinfo in epinfobuf]))
logger.dump_tabular()
if save_interval and (update % save_interval == 0 or update == 1) and logger.get_dir():
savepath = osp.join(logger.get_dir(), 'checkpoint%.5i'%update)
print('Saving to', savepath)
model.save(savepath)
coord.request_stop()
coord.join(enqueue_threads)
return model
def get_encoded_class(classObj):
pickled = cloudpickle.dumps(classObj)
encoding = base64.b64encode(pickled).decode()
return encoding
def test_is_pickleable_after_start(self):
e = LocalDaskExecutor()
with e.start():
post = cloudpickle.loads(cloudpickle.dumps(e))
assert isinstance(post, LocalDaskExecutor)
assert post._pool is None
def test_is_pickleable_after_start(self):
e = Executor()
with e.start():
post = cloudpickle.loads(cloudpickle.dumps(e))
assert isinstance(post, Executor)
def data_send(c):
if c.port == baseport:
c.send(get_min_port_friend(c).name, cloudpickle.dumps(123))
if c.port == baseport + 1:
wait_for(c, 123)
def learn(*,
network,
env,
total_timesteps,
seed=None,
nsteps=2048,
ent_coef=0.0,
lr=3e-4,
vf_coef=0.5,
max_grad_norm=0.5,
gamma=0.99,
lam=0.95,
log_interval=10,
nminibatches=4,
noptepochs=4,
cliprange=0.2,
save_interval=0,
load_path=None,
**network_kwargs):
'''
Learn policy using PPO algorithm (https://arxiv.org/abs/1707.06347)
Parameters:
----------
network: policy network architecture. Either string (mlp, lstm, lnlstm, cnn_lstm, cnn, cnn_small, conv_only - see baselines.common/models.py for full list)
specifying the standard network architecture, or a function that takes tensorflow tensor as input and returns
tuple (output_tensor, extra_feed) where output tensor is the last network layer output, extra_feed is None for feed-forward
neural nets, and extra_feed is a dictionary describing how to feed state into the network for recurrent neural nets.
See baselines.common/policies.py/lstm for more details on using recurrent nets in policies
env: baselines.common.vec_env.VecEnv environment. Needs to be vectorized for parallel environment simulation.
The environments produced by gym.make can be wrapped using baselines.common.vec_env.DummyVecEnv class.
nsteps: int number of steps of the vectorized environment per update (i.e. batch size is nsteps * nenv where
nenv is number of environment copies simulated in parallel)
total_timesteps: int number of timesteps (i.e. number of actions taken in the environment)
ent_coef: float policy entropy coefficient in the optimization objective
lr: float or function learning rate, constant or a schedule function [0,1] -> R+ where 1 is beginning of the
training and 0 is the end of the training.
vf_coef: float value function loss coefficient in the optimization objective
max_grad_norm: float or None gradient norm clipping coefficient
gamma: float discounting factor
lam: float advantage estimation discounting factor (lambda in the paper)
log_interval: int number of timesteps between logging events
nminibatches: int number of training minibatches per update
noptepochs: int number of training epochs per update
cliprange: float or function clipping range, constant or schedule function [0,1] -> R+ where 1 is beginning of the training
and 0 is the end of the training
save_interval: int number of timesteps between saving events
load_path: str path to load the model from
**network_kwargs: keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
For instance, 'mlp' network architecture has arguments num_hidden and num_layers.
'''
set_global_seeds(seed)
if isinstance(lr, float): lr = constfn(lr)
else: assert callable(lr)
if isinstance(cliprange, float): cliprange = constfn(cliprange)
else: assert callable(cliprange)
total_timesteps = int(total_timesteps)
policy = build_policy(env, network, **network_kwargs)
nenvs = env.num_envs
ob_space = env.observation_space
ac_space = env.action_space
nbatch = nenvs * nsteps
nbatch_train = nbatch // nminibatches
make_model = lambda: Model(policy=policy,
ob_space=ob_space,
ac_space=ac_space,
nbatch_act=nenvs,
nbatch_train=nbatch_train,
nsteps=nsteps,
ent_coef=ent_coef,
vf_coef=vf_coef,
max_grad_norm=max_grad_norm)
if save_interval and logger.get_dir():
import cloudpickle
with open(osp.join(logger.get_dir(), 'make_model.pkl'), 'wb') as fh:
fh.write(cloudpickle.dumps(make_model))
model = make_model()
if load_path is not None:
model.load(load_path)
runner = Runner(env=env, model=model, nsteps=nsteps, gamma=gamma, lam=lam)
epinfobuf = deque(maxlen=100)
tfirststart = time.time()
nupdates = total_timesteps // nbatch
for update in range(1, nupdates + 1):
assert nbatch % nminibatches == 0
tstart = time.time()
frac = 1.0 - (update - 1.0) / nupdates
lrnow = lr(frac)
cliprangenow = cliprange(frac)
obs, returns, masks, actions, values, neglogpacs, states, epinfos = runner.run(
) #pylint: disable=E0632
epinfobuf.extend(epinfos)
mblossvals = []
if states is None: # nonrecurrent version
inds = np.arange(nbatch)
for _ in range(noptepochs):
np.random.shuffle(inds)
for start in range(0, nbatch, nbatch_train):
end = start + nbatch_train
mbinds = inds[start:end]
slices = (arr[mbinds]
for arr in (obs, returns, masks, actions, values,
neglogpacs))
mblossvals.append(model.train(lrnow, cliprangenow,
*slices))
else: # recurrent version
assert nenvs % nminibatches == 0
envsperbatch = nenvs // nminibatches
envinds = np.arange(nenvs)
flatinds = np.arange(nenvs * nsteps).reshape(nenvs, nsteps)
envsperbatch = nbatch_train // nsteps
for _ in range(noptepochs):
np.random.shuffle(envinds)
for start in range(0, nenvs, envsperbatch):
end = start + envsperbatch
mbenvinds = envinds[start:end]
mbflatinds = flatinds[mbenvinds].ravel()
slices = (arr[mbflatinds]
for arr in (obs, returns, masks, actions, values,
neglogpacs))
mbstates = states[mbenvinds]
mblossvals.append(
model.train(lrnow, cliprangenow, *slices, mbstates))
lossvals = np.mean(mblossvals, axis=0)
tnow = time.time()
fps = int(nbatch / (tnow - tstart))
if update % log_interval == 0 or update == 1:
ev = explained_variance(values, returns)
logger.logkv("serial_timesteps", update * nsteps)
logger.logkv("nupdates", update)
logger.logkv("total_timesteps", update * nbatch)
logger.logkv("fps", fps)
logger.logkv("explained_variance", float(ev))
logger.logkv('eprewmean',
safemean([epinfo['r'] for epinfo in epinfobuf]))
logger.logkv('eplenmean',
safemean([epinfo['l'] for epinfo in epinfobuf]))
logger.logkv('time_elapsed', tnow - tfirststart)
for (lossval, lossname) in zip(lossvals, model.loss_names):
logger.logkv(lossname, lossval)
if MPI.COMM_WORLD.Get_rank() == 0:
logger.dumpkvs()
if save_interval and (
update % save_interval == 0 or update
== 1) and logger.get_dir() and MPI.COMM_WORLD.Get_rank() == 0:
checkdir = osp.join(logger.get_dir(), 'checkpoints')
os.makedirs(checkdir, exist_ok=True)
savepath = osp.join(checkdir, '%.5i' % update)
print('Saving to', savepath)
model.save(savepath)
env.close()
return model
def serialize_function(function, pickle_protocol=None):
return cloudpickle.dumps(function, protocol=pickle_protocol)
def store(key, value, chunksize=950000):
serialized = cloudpickle.dumps(value, 2)
values = {}
for i in xrange(0, len(serialized), chunksize):
values['%s.%s' % (key, i // chunksize)] = serialized[i:i + chunksize]
return memcache.set_multi(values)
def test_empty_file(self):
# Empty file
open(self.tmpfilepath, 'w').close()
with open(self.tmpfilepath, 'r') as f:
self.assertEqual('', pickle.loads(cloudpickle.dumps(f)).read())
os.remove(self.tmpfilepath)
def __getstate__(self):
return cloudpickle.dumps(self.x)
