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 _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