def test_read_all():
from pykern import pkio
from pykern import pkjson
from pykern import pkunit
from pykern.pkunit import pkok, pkeq, pkre
from pykern.pkdebug import pkdp
from pykern.pkcli import rsmanifest
import re
with pkunit.save_chdir_work(is_pkunit_prefix=True) as d:
rsmanifest.add_code(
'code1',
version='1.1',
uri='http://x.com',
source_d='/tmp',
pyenv='py2',
)
v = pkjson.load_any(pkio.py_path(rsmanifest.USER_FILE)).version
pkjson.dump_pretty(
{
'version': v,
'image': {
'type': 'docker'
}
},
filename=rsmanifest.CONTAINER_FILE,
)
m = rsmanifest.read_all()
pkeq(v, m.version)
pkeq('docker', m.image.type)
pkeq('1.1', m.codes.py2.code1.version)
def _create_file():
from sklearn.preprocessing import LabelEncoder
# POSIT: Matches logic in package_data.template.ml.scale.py.jinja.read_data_and_encode_output_column()
data = simulation_db.read_json(
frame_args.run_dir.join(template_common.INPUT_BASE_NAME),
)
v = np.genfromtxt(
str(simulation_db.simulation_lib_dir(SIM_TYPE).join(
_filename(data.models.dataFile.file),
)),
delimiter=',',
skip_header=data.models.columnInfo.hasHeaderRow,
dtype=None,
encoding='utf-8',
)
o = data.models.columnInfo.inputOutput.index('output')
c = v[f'f{o}']
e = LabelEncoder().fit(c)
res = PKDict(
zip(
e.transform(e.classes_).astype(np.float).tolist(),
e.classes_,
),
)
pkjson.dump_pretty(
res,
filename=_OUTPUT_FILE.classificationOutputColEncodingFile,
)
return res
def test_read_all():
from pykern import pkio
from pykern import pkjson
from pykern import pkunit
from pykern.pkunit import pkok, pkeq, pkre
from pykern.pkdebug import pkdp
from pykern.pkcli import rsmanifest
import re
with pkunit.save_chdir_work(is_pkunit_prefix=True) as d:
rsmanifest.add_code(
'code1',
version='1.1',
uri='http://x.com',
source_d='/tmp',
pyenv='py2',
)
v = pkjson.load_any(pkio.py_path(rsmanifest.USER_FILE)).version
pkjson.dump_pretty(
{'version': v, 'image': {'type': 'docker'}},
filename=rsmanifest.CONTAINER_FILE,
)
m = rsmanifest.read_all()
pkeq(v, m.version)
pkeq('docker', m.image.type)
pkeq('1.1', m.codes.py2.code1.version)
def host_init(j2_ctx, host):
from rsconf import db
jf = db.secret_path(j2_ctx, _PASSWD_SECRET_JSON_F, visibility=db.VISIBILITY_GLOBAL)
if jf.check():
with jf.open() as f:
y = pkjson.load_any(f)
else:
y = pkcollections.Dict()
if not host in y:
y[host] = _passwd_entry(j2_ctx, host)
pkjson.dump_pretty(y, filename=jf)
return """install -m 600 /dev/stdin /root/.netrc <<'EOF'
machine {} login {} password {}
EOF
curl {} | install_server={} bash -s {}
# On {}: ssh {} true""".format(
_vhost(j2_ctx),
host,
y[host],
j2_ctx.rsconf_db.http_host,
j2_ctx.rsconf_db.http_host,
host,
j2_ctx.bkp.primary,
host,
)
def _pop_users_flattened(self):
jc = self.j2_ctx
z = jc.dovecot
pw_f = db.secret_path(
jc,
_PASSWORD_SECRET_JSON_F,
visibility=_PASSWORD_VISIBILITY,
)
pw_modified = False
if pw_f.check():
with pw_f.open() as f:
pw_db = pkjson.load_any(f)
else:
pw_db = pkcollections.Dict()
res = []
for u, v in z.pop_users.items():
if not isinstance(v, dict):
v = pkcollections.Dict(password=v, username=u)
if not v.username in pw_db:
pw_modified = True
pw_db[v.username] \
= '{' + z.passdb_scheme + '}' + _sha512_crypt(v.password)
i = z.base_users.user_spec(u)
i.pw_hash = pw_db[v.username]
i.username = v.username
i.home_d = db.user_home_path(jc, u)
res.append(i)
self._setup_procmail(i)
self.install_access(mode='700', owner=i.uid, group=i.gid)
self.install_directory(i.home_d.join(z.user_mail_d))
if pw_modified:
pkjson.dump_pretty(pw_db, filename=pw_f)
return sorted(res, key=lambda x: x.username)
async def start_report_job(self, run_dir, jhash, backend, cmd, tmp_dir):
# First make sure there's no-one else using the run_dir
current_jhash, current_status = self.run_dir_status(run_dir)
if current_status is runner_client.JobStatus.RUNNING:
# Something's running.
if current_jhash == jhash:
# It's already the requested job, so we have nothing to
# do. Throw away the tmp_dir and move on.
pkdlog(
'job is already running; skipping (run_dir={}, jhash={}, tmp_dir={})',
run_dir,
jhash,
tmp_dir,
)
pkio.unchecked_remove(tmp_dir)
return
else:
# It's some other job. Better kill it before doing
# anything else.
# XX TODO: should we check some kind of sequence number
# here? I don't know how those work.
pkdlog(
'stale job is still running; killing it (run_dir={}, jhash={})',
run_dir,
jhash,
)
await self.kill_all(run_dir)
# Okay, now we have the dir to ourselves. Set up the new run_dir:
assert run_dir not in self.report_jobs
pkio.unchecked_remove(run_dir)
tmp_dir.rename(run_dir)
# Start the job:
report_job = await _BACKENDS[backend].start_report_job(run_dir, cmd)
# And update our records so we know it's running:
job_info = _JobInfo(
run_dir,
jhash,
runner_client.JobStatus.RUNNING,
report_job,
)
self.report_jobs[run_dir] = job_info
pkjson.dump_pretty(
{
'version': 1,
'backend': backend,
'backend_info': report_job.backend_info,
},
filename=run_dir.join(_RUNNER_INFO_BASENAME),
)
# And finally, start a background task to watch over it.
self._nursery.start_soon(
self._supervise_report_job,
run_dir,
jhash,
job_info,
)
def __repr__(self):
"""Writes tmp_d/db.json and returns filename
This is a hack to support pkdc below.
"""
f = self.rsconf_db.tmp_d.join('db.json')
if self.rsconf_db.tmp_d:
pkjson.dump_pretty(self.__str(self), filename=f)
return str(f)
def write_to_log(stdout, stderr, filename):
p = pkio.py_path(self._local_user_dir).join(
'agent-sbatch', self.cfg.host)
pkio.mkdir_parent(p)
pkjson.dump_pretty(
PKDict(stdout=stdout, stderr=stderr),
p.join(
f'{datetime.datetime.now().strftime("%Y%m%d%H%M%S")}-{filename}.log'
),
)
def test_convert_srw_to_shadow():
from pykern import pkio, pkjson
from pykern.pkunit import pkeq
from sirepo.template.srw_shadow_converter import SRWShadowConverter
with pkunit.save_chdir_work():
for name in ('crl', 'gaussian', 'grating'):
srw = _read_json_from_data_dir(f'{name}-srw.json')
actual = SRWShadowConverter().srw_to_shadow(srw.models)
del actual['version']
pkjson.dump_pretty(actual, f'{name}-shadow.json')
expect = _read_json_from_data_dir(f'{name}-shadow.json')
pkeq(expect, actual)
def update_sim_from_par(sim_id, par_path):
sim_path = _sim_path_from_id(sim_id)
data = pkjson.load_any(pkio.read_text(sim_path))
parser = flash_parser.ParameterParser()
values = parser.parse(data, pkio.read_text(par_path))
# reset all model fields to default and override with new par values
for (m, fields) in parser.schema.model.items():
for f in fields:
if f in values:
data.models[m][f] = values[f]
else:
data.models[m][f] = fields[f][2]
pkjson.dump_pretty(data, sim_path)
def pkunit_setup():
"""Create rsmanifest files"""
from pykern import pkjson
pkjson.dump_pretty(
{
'version': FILE_VERSION,
'image': {
'type': 'pkunit',
},
},
filename=CONTAINER_FILE,
)
add_code('pkunit', '1.1', 'https://pykern.org', '/tmp')
def units(src_path):
res = []
p = pkio.py_path(src_path).join('source')
for d, _, _ in os.walk(p):
m = re.search(
r'^(?:{})(.*\/[A-Z0-9][A-Za-z0-9-_\.]+$)'.format(re.escape(
str(p))),
d,
)
if m:
s = m.group(1)[1:]
res.append([s, s])
res.sort()
pkjson.dump_pretty(res, filename='res.json')
def test_parse_madx_file():
from pykern import pkio, pkjson
from pykern.pkunit import pkeq
from sirepo.template import madx, madx_parser
with pkunit.save_chdir_work():
for name in ('particle_track', 'alba'):
actual = madx_parser.parse_file(pkio.read_text(
pkunit.data_dir().join(f'{name}.madx')))
del actual['version']
outfile = f'{name}.json'
pkjson.dump_pretty(actual, outfile)
expect = pkjson.load_any(pkunit.data_dir().join(outfile))
pkeq(expect, actual)
def host_init(j2_ctx, host):
from rsconf import db
jf = _sasl_password_path(j2_ctx)
if jf.check():
with jf.open() as f:
y = pkjson.load_any(f)
else:
y = pkcollections.Dict()
u = _SASL_PASSWORD_PREFIX + host
if not u in y:
y[u] = db.random_string()
pkjson.dump_pretty(y, filename=jf)
return u, y[u]
def _issues():
if not repo.has_issues:
return
base = bd + '.issues'
d = pkio.mkdir_parent(base)
for i in repo.issues(state='all'):
j = _trim_body(i)
j['comments'] = [_trim_body(c) for c in i.comments()]
p = i.pull_request()
if p:
j['review_comments'] = [
_trim_body(c) for c in p.review_comments()
]
pkjson.dump_pretty(j, filename=d.join(str(i.number) + '.json'))
_tar(base)
def test_importer(import_req):
from pykern import pkcollections
from pykern import pkjson
from pykern.pkunit import pkeq
from sirepo.template import zgoubi
import sirepo.sim_data
with pkunit.save_chdir_work() as w:
for fn in pkio.sorted_glob(pkunit.data_dir().join('*.dat')):
error = None
try:
data = zgoubi.import_file(import_req(fn), unit_test_mode=True)
sirepo.sim_data.get_class('zgoubi').fixup_old_data(data)
#TODO(pjm): easier way to convert nested dict to pkcollections.Dict?
data = pkcollections.json_load_any(pkjson.dump_pretty(data))
except Exception as e:
pkdlog(pkdexc())
error = e.message
if error:
actual = error
else:
actual = zgoubi.python_source_for_model(data)
outfile = fn.basename + '.txt'
pkio.write_text(outfile, actual)
e = pkunit.data_dir().join(outfile)
expect = pkio.read_text(e)
pkeq(expect, actual, 'diff {} {}', e, w.join(outfile))
def _create_zip(sim, want_python, out_dir):
"""Zip up the json file and its dependencies
Args:
sim (req): simulation
want_python (bool): include template's python source?
out_dir (py.path): where to write to
Returns:
py.path.Local: zip file name
"""
path = out_dir.join(sim.id + '.zip')
data = simulation_db.open_json_file(sim.type, sid=sim.id)
simulation_db.update_rsmanifest(data)
data.pkdel('report')
files = sim_data.get_class(data).lib_files_for_export(data)
if want_python:
files.append(_python(data))
with zipfile.ZipFile(
str(path),
mode='w',
compression=zipfile.ZIP_DEFLATED,
allowZip64=True,
) as z:
for f in files:
z.write(str(f), f.basename)
z.writestr(
simulation_db.SIMULATION_DATA_FILE,
pkjson.dump_pretty(data, pretty=True),
)
return path, data
def default_command(in_file):
"""Reads `in_file` passes to `msg.jobCmd`
Must be called in run_dir
Writes its output on stdout.
Args:
in_file (str): json parsed to msg
Returns:
str: json output of command, e.g. status msg
"""
try:
job.init()
f = pkio.py_path(in_file)
msg = pkjson.load_any(f)
#TODO(e-carlin): find common place to serialize/deserialize paths
msg.runDir = pkio.py_path(msg.runDir)
f.remove()
res = globals()['_do_' + msg.jobCmd](msg,
sirepo.template.import_module(
msg.simulationType))
if res is None:
return
r = PKDict(res).pksetdefault(state=job.COMPLETED)
except Exception as e:
r = PKDict(
state=job.ERROR,
error=e.sr_args.error
if isinstance(e, sirepo.util.UserAlert) else str(e),
stack=pkdexc(),
)
return pkjson.dump_pretty(r, pretty=False)
def _rsdockerspawner(self, z):
from rsconf.component import docker
z.tls_d = self.__run_d.join(_DOCKER_TLS_SUBDIR)
c = pkcollections.Dict(
port_base=z.setdefault('port_base', _DEFAULT_PORT_BASE),
tls_dir=str(z.tls_d),
)
seen = pkcollections.Dict(
hosts=pkcollections.Dict(),
users=pkcollections.Dict(),
)
# POSIT: notebook_dir in
# radiasoft/container-beamsim-jupyter/container-conf/build.sh
# parameterize anyway, because matches above
z.setdefault('volumes', pkcollections.Dict())
z.volumes.setdefault(
str(z.user_d.join('{username}')),
pkcollections.Dict(bind=str(z.home_d.join('jupyter'))),
)
self._rsdockerspawner_v3(z, seen)
c.user_groups = z.get('user_groups', PKDict())
c.volumes = z.volumes
c.pools = z.pools
z.rsdockerspawner_cfg = pkjson.dump_pretty(c)
z.docker_hosts = list(seen.hosts.keys())
def f():
if isinstance(content_or_path, pkconst.PY_PATH_LOCAL_TYPE):
return flask.send_file(str(content_or_path))
if content_type == 'application/json':
return flask.current_app.response_class(
pkjson.dump_pretty(content_or_path))
return flask.current_app.response_class(content_or_path)
def test_import():
from pykern import pkjson
from pykern.pkunit import pkeq
from sirepo.template import flash_parser
import re
def _parse_config(fn):
return flash_parser.ConfigParser().parse(pkio.read_text(fn))
def _parse_par(fn):
data_file = fn.basename.replace('-flash.par', '')
return flash_parser.ParameterParser().parse(
pkjson.load_any(
pkio.read_text(
pkunit.data_dir().join(f'{data_file}-sirepo-data.json'))),
pkio.read_text(fn),
)
with pkunit.save_chdir_work():
for fn in pkio.sorted_glob(pkunit.data_dir().join('*')):
if re.search(r'-Config$', fn.basename):
parser = _parse_config
elif re.search(r'flash.par$', fn.basename):
parser = _parse_par
else:
continue
try:
actual = pkjson.dump_pretty(parser(fn))
except Exception as e:
pkdlog(pkdexc())
actual = str(e)
outfile = f'{fn.basename}.out'
pkio.write_text(outfile, actual)
expect = pkio.read_text(pkunit.data_dir().join(outfile))
pkeq(expect, actual)
def add_code(name, version, uri, source_d, virtual_env=None, pyenv=None):
"""Add a new code to ~?rsmanifest.json
Args:
name (str): name of the package
version (str): commit or version
uri (str): repo, source link
source_d (str): directory containing
virtual_env (str): DEPRECATED
pyenv (str): pyenv version
"""
from pykern import pkcollections
from pykern import pkio
from pykern import pkjson
import datetime
import json
fn = pkio.py_path(USER_FILE)
try:
values = pkcollections.json_load_any(fn)
except Exception as e:
if not (pkio.exception_is_not_found(e) or isinstance(e, ValueError)):
raise
values = pkcollections.Dict(
version=FILE_VERSION,
codes=pkcollections.Dict({_NO_PYENV: pkcollections.Dict()}),
)
if pyenv:
assert not virtual_env, \
'only one of pyenv or virtual-env (DEPRECATED)'
elif virtual_env:
assert not pyenv, \
'only one of pyenv or virtual-env (DEPRECATED)'
pyenv = virtual_env
if not pyenv:
pyenv = _NO_PYENV
v = values.codes.get(pyenv) or pkcollections.Dict()
v[name.lower()] = pkcollections.Dict(
installed=datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%SZ'),
name=name,
source_d=source_d,
uri=uri,
version=version,
)
values.codes[pyenv] = v
pkjson.dump_pretty(values, filename=fn)
def run_extract_job(run_dir, jhash, subcmd, *args):
return _rpc({
'action': 'run_extract_job',
'run_dir': str(run_dir),
'jhash': jhash,
'subcmd': subcmd,
'arg': pkjson.dump_pretty(args),
})
def start_sbatch():
def get_host():
h = socket.gethostname()
if '.' not in h:
h = socket.getfqdn()
return h
def kill_agent(pid_file):
if get_host() == pid_file.host:
os.kill(pid_file.pid, signal.SIGKILL)
else:
try:
subprocess.run(
('ssh', pid_file.host, 'kill', '-KILL', str(pid_file.pid)),
capture_output=True,
text=True,
).check_returncode()
except subprocess.CalledProcessError as e:
if '({}) - No such process'.format(
pid_file.pid) not in e.stderr:
pkdlog('cmd={cmd} returncode={returncode} stderr={stderr}',
**vars(e))
f = None
try:
f = pkjson.load_any(pkio.py_path(_PID_FILE))
except Exception as e:
if not pkio.exception_is_not_found(e):
pkdlog('error={} stack={}', e, pkdexc())
try:
if f:
kill_agent(f)
except Exception as e:
pkdlog('error={} stack={}', e, pkdexc())
pkjson.dump_pretty(
PKDict(
host=get_host(),
pid=os.getpid(),
),
_PID_FILE,
)
try:
start()
finally:
#TODO(robnagler) https://github.com/radiasoft/sirepo/issues/2195
pkio.unchecked_remove(_PID_FILE)
def _write_parallel_status(msg, template, is_running):
sys.stdout.write(
pkjson.dump_pretty(
PKDict(
state=job.RUNNING,
parallelStatus=_background_percent_complete(
msg, template, is_running),
),
pretty=False,
) + '\n', )
def _rsmanifest():
from pykern import pkcollections
from pykern import pkjson
from pykern.pkcli import rsmanifest
import cpuinfo
import datetime
import os
import subprocess
m = rsmanifest.read_all()
m['sim'] = {
'run': {
'datetime': datetime.datetime.utcnow().isoformat(),
'cpu_info': cpuinfo.get_cpu_info(),
'pyenv': _pyenv_version(),
#TODO(robnagler) can't include because of auth/credential
# values in environment variables
#'environ': pkcollections.Dict(os.environ),
},
}
pkjson.dump_pretty(m, filename=rsmanifest.BASENAME)
def _rsmanifest():
from pykern import pkcollections
from pykern import pkjson
from pykern.pkcli import rsmanifest
import cpuinfo
import datetime
import os
import subprocess
m = rsmanifest.read_all()
m['sim'] = {
'run': {
'datetime': datetime.datetime.utcnow().isoformat(),
'cpu_info': cpuinfo.get_cpu_info(),
'pyenv': _pyenv_version(),
#TODO(robnagler) can't include because of auth/credential
# values in environment variables
#'environ': pkcollections.Dict(os.environ),
},
}
pkjson.dump_pretty(m, filename=rsmanifest.BASENAME)
def _json(gen, suffix):
base = bd + suffix
with open(base, 'wt') as f:
sep = '[\n'
for i in gen:
f.write(sep)
f.write(pkjson.dump_pretty(_trim_body(i)))
sep = ',\n'
if '[' in sep:
# Empty iteration
f.write(sep)
f.write(']\n')
_shell(['xz', base])
def host_init(hdb, host):
jf = db.secret_path(hdb,
_PASSWD_SECRET_JSON_F,
visibility=_PASSWD_VISIBILITY)
if jf.check():
with jf.open() as f:
y = pkjson.load_any(f)
else:
y = PKDict()
assert not host in y, \
'{}: host already exists'.format(host)
y[host] = db.random_string()
pkjson.dump_pretty(y, filename=jf)
pf = db.secret_path(hdb, _PASSWD_SECRET_F, visibility=_PASSWD_VISIBILITY)
with pf.open(mode='at') as f:
f.write(
'{}:{}\n'.format(
host,
pkcompat.from_bytes(
bcrypt.hashpw(
pkcompat.to_bytes(y[host]),
bcrypt.gensalt(5),
), ),
), )
def test_dump_pretty():
from pykern import pkjson
from pykern.pkunit import pkeq
class Other(object):
def __init__(self, x):
self.x = x
def __str__(self):
return str(self.x)
v = {'d': ['a', 'b'], 'c': Other('xyz')}
a = pkjson.dump_pretty(v)
pkeq(
'''{
"c": "xyz",
"d": [
"a",
"b"
]
}
''',
a,
)
def _gen_exception_reply_SRException(args):
r = args.routeName
p = args.params or PKDict()
try:
t = sirepo.http_request.sim_type(p.pkdel('sim_type'))
s = simulation_db.get_schema(sim_type=t)
except Exception as e:
pkdc('exception={} stack={}', e, pkdexc())
# sim_type is bad so don't cascade errors, just
# try to get the schema without the type
t = None
s = simulation_db.get_schema(sim_type=None)
# If default route or always redirect/reload
if r:
assert r in s.localRoutes, \
'route={} not found in schema for type={}'.format(r, t)
else:
r = sirepo.uri.default_local_route_name(s)
p = PKDict(reload_js=True)
if (
# must be first, to always delete reload_js
not p.pkdel('reload_js') and flask.request.method == 'POST'
and r not in _RELOAD_JS_ROUTES):
pkdc('POST response={} route={} params={}', SR_EXCEPTION_STATE, r, p)
return gen_json(
PKDict({
_STATE: SR_EXCEPTION_STATE,
SR_EXCEPTION_STATE: args,
}), )
pkdc('redirect to route={} params={} type={}', r, p, t)
return gen_redirect_for_local_route(
t,
route=r,
params=p,
sr_exception=pkjson.dump_pretty(args, pretty=False),
)
def _create_in_file(self):
f = self.run_dir.join(_IN_FILE.format(job.unique_key()), )
pkjson.dump_pretty(self.msg, filename=f, pretty=False)
return f
