def submit():
parser = argparse.ArgumentParser(
description="Submit task to yascheduler daemon")
parser.add_argument('script')
args = parser.parse_args()
if not os.path.isfile(args.script):
raise ValueError("Script parameter is not a file name")
inputs = {}
with open(args.script) as f:
for l in f.readlines():
try:
k, v = l.split('=')
inputs[k.strip()] = v.strip()
except ValueError:
pass
config = ConfigParser()
config.read(CONFIG_FILE)
yac = Yascheduler(config)
task_id = yac.queue_submit_task(
inputs['LABEL'], {
'structure': open(inputs['STRUCT']).read(),
'input': open(inputs['INPUT']).read(),
'local_folder': os.getcwd()
})
print("Successfully submitted task: {}".format(task_id))
yac.connection.close()
def show_nodes():
config = ConfigParser()
config.read(CONFIG_FILE)
yac = Yascheduler(config)
yac.cursor.execute(
'SELECT ip, label, task_id FROM yascheduler_tasks WHERE status=%s;',
[yac.STATUS_RUNNING])
tasks_running = {row[0]: [row[1], row[2]] for row in yac.cursor.fetchall()}
yac.cursor.execute(
'SELECT ip, ncpus, enabled, cloud from yascheduler_nodes;')
for item in yac.cursor.fetchall():
print("ip=%s ncpus=%s enabled=%s occupied_by=%s (task_id=%s) %s" %
tuple([item[0], item[1] or 'MAX', item[2]] +
tasks_running.get(item[0], ['-', '-']) + [item[3] or '']))
def _init_db(install_path):
# database initialization
config = ConfigParser()
config.read(CONFIG_FILE)
yac = Yascheduler(config)
schema = open(os.path.join(install_path, 'data', 'schema.sql')).read()
try:
for line in schema.split(';'):
if not line:
continue
yac.cursor.execute(line)
yac.connection.commit()
except ProgrammingError as e:
if "already exists" in str(e.args[0]):
print("Database already initialized!")
raise
import numpy as np
from aiida_crystal_dft.io.f34 import Fort34
from yascheduler import CONFIG_FILE
from yascheduler.scheduler import Yascheduler
from mpds_client import MPDSDataRetrieval
from mpds_aiida.common import get_template, get_basis_sets, get_mpds_structures, get_input
ela = ['Li', 'Na', 'K', 'Rb', 'Be', 'Mg', 'Ca', 'Sr']
elb = ['F', 'Cl', 'Br', 'I', 'O', 'S', 'Se', 'Te']
config = ConfigParser()
config.read(CONFIG_FILE)
yac = Yascheduler(config)
client = MPDSDataRetrieval()
calc_setup = get_template()
bs_repo = get_basis_sets(calc_setup['basis_family'])
try:
how_many = int(sys.argv[1])
except (IndexError, ValueError):
how_many = False
counter = 0
random.shuffle(ela)
random.shuffle(elb)
work_folder = os.path.dirname(target)
setup_input = open(target).read()
try:
sys.argv[2]
except IndexError:
folder = None
print('**To save calc in a local repo**')
else:
folder = work_folder
print('**To save calc in an input folder**')
if os.path.exists(os.path.join(work_folder, 'fort.34')):
assert "EXTERNAL" in setup_input
struct_input = open(os.path.join(work_folder, 'fort.34')).read()
else:
assert "EXTERNAL" not in setup_input
struct_input = "UNUSED"
label = setup_input.splitlines()[0]
config = ConfigParser()
config.read(CONFIG_FILE)
yac = Yascheduler(config)
result = yac.queue_submit_task(
label, dict(structure=struct_input, input=setup_input,
local_folder=folder))
print(label)
print(result)
def check_status():
parser = argparse.ArgumentParser(
description="Submit task to yascheduler daemon")
parser.add_argument('-j',
'--jobs',
required=False,
default=None,
nargs='*')
parser.add_argument('-v',
'--view',
required=False,
default=None,
nargs='?',
type=bool,
const=True)
parser.add_argument('-o',
'--convergence',
required=False,
default=None,
nargs='?',
type=bool,
const=True,
help='needs -v option')
parser.add_argument('-i',
'--info',
required=False,
default=None,
nargs='?',
type=bool,
const=True)
parser.add_argument('-k',
'--kill',
required=False,
default=None,
nargs='?',
type=bool,
const=True)
args = parser.parse_args()
config = ConfigParser()
config.read(CONFIG_FILE)
yac = Yascheduler(config)
statuses = {
yac.STATUS_TO_DO: "QUEUED",
yac.STATUS_RUNNING: "RUNNING",
yac.STATUS_DONE: "FINISHED"
}
local_parsing_ready, local_calc_snippet = False, False
if args.jobs:
tasks = yac.queue_get_tasks(jobs=args.jobs)
else:
tasks = yac.queue_get_tasks(status=(yac.STATUS_RUNNING,
yac.STATUS_TO_DO))
if args.view or args.kill:
if not tasks:
print('NO MATCHING TASKS FOUND')
return
ssh_custom_key = {}
for filename in os.listdir(config.get('local', 'data_dir')):
if not filename.startswith('yakey') or not os.path.isfile(
os.path.join(config.get('local', 'data_dir'), filename)):
continue
key_path = os.path.join(config.get('local', 'data_dir'), filename)
pmk_key = RSAKey.from_private_key_file(key_path)
print('LOADED KEY %s' % key_path)
ssh_custom_key = {'pkey': pmk_key}
break
if args.convergence:
try:
from pycrystal import CRYSTOUT
from numpy import nan
local_parsing_ready = True
except:
pass
if args.view:
yac.cursor.execute(
'SELECT task_id, label, metadata, ip FROM yascheduler_tasks WHERE status=%s AND task_id IN (%s);'
% (yac.STATUS_RUNNING, ', '.join(
[str(task['task_id']) for task in tasks])))
for row in yac.cursor.fetchall():
print("." * 50 + "ID%s %s at %s@%s:%s" %
(row[0], row[1], config.get('remote', 'user'), row[3],
row[2]['remote_folder']))
ssh_conn = SSH_Connection(host=row[3],
user=config.get('remote', 'user'),
connect_kwargs=ssh_custom_key)
try:
result = ssh_conn.run('tail -n15 %s/OUTPUT' %
row[2]['remote_folder'],
hide=True)
except UnexpectedExit:
print('OUTDATED TASK, SKIPPING')
else:
print(result.stdout)
if local_parsing_ready:
local_calc_snippet = os.path.join(
config.get('local', 'data_dir'), 'local_calc_snippet.tmp')
try:
ssh_conn.get(row[2]['remote_folder'] + '/OUTPUT',
local_calc_snippet)
except IOError as err:
continue
calc = CRYSTOUT(local_calc_snippet)
output_lines = ''
if calc.info['convergence']:
output_lines += str(calc.info['convergence']) + "\n"
if calc.info['optgeom']:
for n in range(len(calc.info['optgeom'])):
try:
ncycles = calc.info['ncycles'][n]
except IndexError:
ncycles = "^"
output_lines += "{:8f}".format(calc.info['optgeom'][n][0] or nan) + " " + \
"{:8f}".format(calc.info['optgeom'][n][1] or nan) + " " + \
"{:8f}".format(calc.info['optgeom'][n][2] or nan) + " " + \
"{:8f}".format(calc.info['optgeom'][n][3] or nan) + " " + \
"E={:12f}".format(calc.info['optgeom'][n][4] or nan) + " eV" + " " + \
"(%s)" % ncycles + "\n"
print(output_lines)
elif args.kill:
if not args.jobs:
print('NO JOBS GIVEN')
return
yac.cursor.execute(
'SELECT ip FROM yascheduler_tasks WHERE status=%s AND task_id IN (%s);'
% (yac.STATUS_RUNNING, ', '.join(
[str(task['task_id']) for task in tasks])))
for row in yac.cursor.fetchall():
ssh_conn = SSH_Connection(host=row[0],
user=config.get('remote', 'user'),
connect_kwargs=ssh_custom_key)
try:
result = ssh_conn.run('pkill %s' % yac.RUNNING_MARKER,
hide=True)
except:
pass
elif args.info:
for task in tasks:
print('task_id={}\tstatus={}\tlabel={}\tip={}'.format(
task['task_id'], statuses[task['status']], task['label'],
task['ip'] or '-'))
else:
for task in tasks:
print('{} {}'.format(task['task_id'], statuses[task['status']]))
yac.connection.close()
if local_calc_snippet and os.path.exists(local_calc_snippet):
os.unlink(local_calc_snippet)
def manage_node():
parser = argparse.ArgumentParser(
description="Add nodes to yascheduler daemon")
parser.add_argument('host', help='IP[~ncpus]')
parser.add_argument('--remove-soft',
required=False,
default=None,
nargs='?',
type=bool,
const=True,
help='Remove IP delayed')
parser.add_argument('--remove-hard',
required=False,
default=None,
nargs='?',
type=bool,
const=True,
help='Remove IP immediate')
args = parser.parse_args()
config = ConfigParser()
config.read(CONFIG_FILE)
ncpus = None
if '~' in args.host:
args.host, ncpus = args.host.split('~')
ncpus = int(ncpus)
yac = Yascheduler(config)
already_there = has_node(config, args.host)
if already_there and not args.remove_hard and not args.remove_soft:
print('Host already in DB: {}'.format(args.host))
return False
if not already_there and (args.remove_hard or args.remove_soft):
print('Host NOT in DB: {}'.format(args.host))
return False
if args.remove_hard:
yac.cursor.execute(
'SELECT task_id from yascheduler_tasks WHERE ip=%s AND status=%s;',
[args.host, yac.STATUS_RUNNING])
result = yac.cursor.fetchall() or []
for item in result: # only one item is expected, but here we also account inconsistency case
yac.cursor.execute(
'UPDATE yascheduler_tasks SET status=%s WHERE task_id=%s;',
[yac.STATUS_DONE, item[0]])
yac.connection.commit()
print('An associated task %s at %s is now marked done!' %
(item[0], args.host))
remove_node(config, args.host)
print('Removed host from yascheduler: {}'.format(args.host))
return True
elif args.remove_soft:
yac.cursor.execute(
'SELECT task_id from yascheduler_tasks WHERE ip=%s AND status=%s;',
[args.host, yac.STATUS_RUNNING])
if yac.cursor.fetchall():
print('A task associated, prevent from assigning the new tasks')
yac.cursor.execute(
'UPDATE yascheduler_nodes SET enabled=FALSE WHERE ip=%s;',
[args.host])
yac.connection.commit()
print('Prevented from assigning the new tasks: {}'.format(
args.host))
return True
else:
print('No tasks associated, remove node immediately')
remove_node(config, args.host)
print('Removed host from yascheduler: {}'.format(args.host))
return True
if not yac.ssh_check_node(args.host) or not add_node(
config, args.host, ncpus):
print('Failed to add host to yascheduler: {}'.format(args.host))
return False
print('Added host to yascheduler: {}'.format(args.host))
return True
assert not error, error
assert 'disordered' not in ase_obj.info
try: symprec = float(sys.argv[2])
except: symprec = 3E-02 # NB needs tuning
print('symprec = %s' % symprec)
label = sys.argv[1].split(os.sep)[-1].split('.')[0] + \
" " + spglib.get_spacegroup(ase_obj, symprec=symprec)
ase_obj, error = refine(ase_obj, accuracy=symprec, conventional_cell=True)
assert not error, error
yaconfig = ConfigParser()
yaconfig.read(CONFIG_FILE)
yac = Yascheduler(yaconfig)
calc_setup = get_template()
bs_repo = get_basis_sets(calc_setup['basis_family'])
#bs_repo_tzvp = get_basis_sets('./tzvp_RE')
#bs_repo_other = get_basis_sets('./hand_made_bs')
#bs_repo = {**bs_repo_tzvp , **bs_repo_other}
elements = list(set(ase_obj.get_chemical_symbols()))
f34_input = Fort34([bs_repo[el] for el in elements])
struct_input = str(f34_input.from_ase(ase_obj))
setup_input = get_input(calc_setup['parameters']['crystal'], elements, bs_repo, label)
result = yac.queue_submit_task(label, dict(structure=struct_input, input=setup_input))
print(label)