def run_sp(func, *args, **kwargs):
"""
Run a function in a subprocess and return its value. This is for achieving subprocess
isolation, not parallelism. The subprocess is configured so things like logging work
correctly, and is initialized with a derived random seed.
"""
ctx = LKContext.INSTANCE
rq = ctx.SimpleQueue()
seed = derive_seed(none_on_old_numpy=True)
worker_args = (log_queue(), seed, rq, func, args, kwargs)
_log.debug('spawning subprocess to run %s', func)
proc = ctx.Process(target=_sp_worker, args=worker_args)
proc.start()
_log.debug('waiting for process %s to return', proc)
success, payload = rq.get()
_log.debug('received success=%s', success)
_log.debug('waiting for process %s to exit', proc)
proc.join()
if proc.exitcode:
_log.error('subprocess failed with code %d', proc.exitcode)
raise RuntimeError('subprocess failed with code ' + str(proc.exitcode))
if success:
return payload
else:
_log.error('subprocess raised exception: %s', payload)
raise ChildProcessError('error in child process', payload)
def _initialize_mp_worker(mkey, func, threads, log_queue, seed):
seed = derive_seed(mp.current_process().name,
base=seed,
none_on_old_numpy=True)
_initialize_worker(log_queue, seed)
global __work_model, __work_func
nnt_env = os.environ.get('NUMBA_NUM_THREADS', None)
if nnt_env is None or int(nnt_env) > threads:
_log.debug('configuring Numba thread count')
import numba
numba.config.NUMBA_NUM_THREADS = threads
try:
import mkl
_log.debug('configuring MKL thread count')
mkl.set_num_threads(threads)
except ImportError:
pass
__work_model = mkey
# deferred function unpickling to minimize imports before initialization
__work_func = pickle.loads(func)
_log.debug('worker %d ready (process %s)', os.getpid(),
mp.current_process())
def test_derive_seed_str():
random.init_rng(42, propagate=False)
s2 = random.derive_seed(b'wombat')
assert s2.entropy == 42
assert s2.spawn_key == (zlib.crc32(b'wombat'), )
def test_derive_seed_intkey():
random.init_rng(42, propagate=False)
s2 = random.derive_seed(10, 7)
assert s2.entropy == 42
assert s2.spawn_key == (10, 7)
def main(opts: STANOptions):
stan.init()
if opts.init_only:
return # we're done
keys = ['stan', opts.model]
if opts.joint:
inf_dir = data_dir / 'joint-inference'
mfile = mod_dir / f'joint-{opts.model}.stan'
else:
if opts.data:
inf_dir = data_dir / opts.data / 'inference'
mfile = mod_dir / f'{opts.model}.stan'
keys.append(opts.data)
model = stan.compile(mfile)
if opts.compile_only:
return # we're done
_log.info('sampling %s on %s (variant %s)', opts.model, opts.data,
opts.variant)
seed = derive_seed(*keys)
sample_args = stan.sample_options(seed, opts)
if opts.input:
input = inf_dir / opts.input
else:
input = inf_dir / (opts.basename + '-inputs.json')
output = inf_dir / opts.basename
sum_file = inf_dir / (opts.basename + '-summary.csv')
dia_file = inf_dir / (opts.basename + '-diag.txt')
draw_file = output / 'samples.zarr'
_log.info('running %d chains with %d warmup and %d sampling iterations',
sample_args['chains'], sample_args['iter_warmup'],
sample_args['iter_sampling'])
output.mkdir(exist_ok=True)
mcmc = model.sample(os.fspath(input),
output_dir=os.fspath(output),
**sample_args)
_log.info('starting sample compression')
comp = stan.start_compression(output)
with ThreadPoolExecutor() as pool:
_log.info('spawning summarization process')
summary = pool.submit(lambda m: m.summary(), mcmc)
_log.info('spawning diagnostic process')
diag = pool.submit(lambda m: m.diagnose(), mcmc)
_log.info('saving model samples')
stan.save_samples(draw_file, mcmc)
_log.info('waiting for summary results')
summary = summary.result()
_log.info('saving model summary')
summary.to_csv(sum_file, index=True)
_log.info('waiting for diagnostic results')
diag = diag.result()
_log.info('saving diagnostics')
dia_file.write_text(diag)
_log.info('finalizing sample compression')
stan.finish_compression(output, comp)
def stan_seed(*keys):
seed = derive_seed(*keys)
stan_seed = seed.generate_state(1)[0] & 0x80000000 # put in range of int
_log.info('using random seed %s (from %s)', stan_seed, seed)
def test_derive_seed_str():
random.init_rng(42, propagate=False)
s2 = random.derive_seed(b'wombat')
assert s2.entropy == 42