"""
Module for easy access to aiida-vasp calculations.
"""
from functools import wraps
import click
from aiida.cmdline.params.arguments import (
CALCULATION,
PROCESS,
WORKFLOW,
)
from aiida.cmdline.utils.echo import (
echo_critical,
echo_error,
echo_info,
echo_success,
)
from aiida_vasp.utils.aiida_utils import cmp_load_verdi_data
VERDI_DATA = cmp_load_verdi_data()
@VERDI_DATA.group('vasp-tools')
def tools():
"""Tool for aiida-vasp related data"""
@tools.command('export')
@PROCESS('process')
@click.argument('folder')
@click.option(
'--include-potcar',
default=False,
is_flag=True,
help='Whether to include POTCAR in the export folder',
)
@click.option(
'--decompress',
default=False,
is_flag=True,
help='Wether to decompress the contents',
)
def export(process, folder, decompress, include_potcar):
"""Export a VASP calculation, works for both `VaspCalculation` or `VaspWorkChain`"""
from aiida_vasp.utils.export import export_vasp
export_vasp(process, folder, decompress, include_potcar)
[docs]
def select_calcjob_from_work(func):
"""Select calcjob from work"""
@wraps(func)
def wrapper(*args, **kwargs):
calcjob = kwargs['calcjob']
index = kwargs.pop('index', 0)
from aiida import orm
if isinstance(calcjob, orm.WorkChainNode):
valid = [process for process in calcjob.called_descendants if not process.is_finished]
if len(valid) > 0:
echo_info(
f'More than one calculations are still running: {", ".join(str(process.pk) for process in valid)}.'
f' Selected {valid[index].pk}'
)
calcjob = valid[index]
else:
calcjob = valid[0]
kwargs['calcjob'] = calcjob
kwargs['index'] = index
func_ = click.option(
'--index',
'-i',
help='The index of the calculation to cat if multiple calculations are still running',
default=0,
)(func)
return func_(*args, **kwargs)
return wrapper
@tools.command('remotecat')
@PROCESS('calcjob')
@click.argument('fname')
@click.option('--save-to', '-s', help='Name of the file to save to')
@select_calcjob_from_work
def remotecat(calcjob, fname, save_to):
"""
Print the conetent of a remote file to STDOUT
This command for printing the content of a remote file to STDOUT.
Useful for analysing running calculations.
"""
import os
import sys
import tempfile
from shutil import copyfileobj
rfolder = calcjob.outputs.remote_folder
if save_to is None:
fd, temppath = tempfile.mkstemp()
else:
temppath = save_to
rfolder.getfile(fname, temppath)
with open(temppath, 'rb') as fhandle:
copyfileobj(fhandle, sys.stdout.buffer)
if save_to is None:
os.close(fd)
os.remove(temppath)
@tools.command('remotepull')
@PROCESS('calcjob')
@click.argument('dest')
@click.option(
'--max-size',
'-m',
help='Maximum size of the files to be retrieved - this is passed to rsync',
)
@select_calcjob_from_work
def remotepull(calcjob, dest, max_size):
"""
Pull a calculation folder from the remote
This command for pull a calculation folder to a local folder.
`rsync` is used for doing the heavy lifting.
"""
import subprocess
rfolder = calcjob.outputs.remote_folder
cmd_args = ['rsync', '-av']
if max_size:
cmd_args.extend(['--max-size', max_size])
cmd_args.append(f'{rfolder.computer.hostname}:{rfolder.get_remote_path()}/')
if not dest.endswith('/'):
dest = dest + '/'
cmd_args.append(dest)
echo_info('Running commands: {}'.format(' '.join(cmd_args)))
completed = subprocess.run(cmd_args, check=False)
if completed.returncode != 0:
echo_error('Failled to pull data using rsync')
else:
echo_success(f'Remote folder pulled to {dest}')
@tools.command('remotetail')
@CALCULATION('calcjob')
@click.argument('fname')
def remotetail(calcjob, fname):
"""
Follow a file on the remote computer
This command will launch a ssh session dedicated for following a file
using the `tail -f` command
"""
import os
from aiida.common.exceptions import NotExistent
try:
transport = calcjob.get_transport()
except NotExistent as exception:
echo_critical(repr(exception))
remote_workdir = calcjob.get_remote_workdir()
if not remote_workdir:
echo_critical('no remote work directory for this calcjob, maybe the daemon did not submit it yet')
command = tailf_command(transport, remote_workdir, fname)
os.system(command)
@tools.command('relaxcat')
@WORKFLOW('workflow')
@click.argument('fname')
def relaxcat(workflow, fname):
"""Cat the output of the last calculation of a finished workflow"""
from aiida.cmdline.commands.cmd_calcjob import calcjob_outputcat
from aiida.orm import CalcJobNode, QueryBuilder, WorkChainNode
q = QueryBuilder()
q.append(WorkChainNode, filters={'id': workflow.id})
q.append(WorkChainNode)
q.append(CalcJobNode, tag='calc', project=['*', 'ctime'])
q.order_by({'calc': {'ctime': 'desc'}})
calc, ctime = q.first()
click.Context(calcjob_outputcat).invoke(calcjob_outputcat, calcjob=calc, path=fname)
[docs]
def tailf_command(transport, remotedir, fname):
"""
Specific gotocomputer string to connect to a given remote computer via
ssh and directly go to the calculation folder and then do tail -f of the target file.
"""
from aiida.common.escaping import escape_for_bash
further_params = []
if 'username' in transport._connect_args:
further_params.append('-l {}'.format(escape_for_bash(transport._connect_args['username'])))
if transport._connect_args.get('port'):
further_params.append('-p {}'.format(transport._connect_args['port']))
if transport._connect_args.get('key_filename'):
further_params.append('-i {}'.format(escape_for_bash(transport._connect_args['key_filename'])))
further_params_str = ' '.join(further_params)
connect_string = (
""" "if [ -d {escaped_remotedir} ] ;"""
""" then cd {escaped_remotedir} ; {bash_command} -c 'tail -f {fname}' ; else echo ' ** The directory' ; """
"""echo ' ** {remotedir}' ; echo ' ** seems to have been deleted, I logout...' ; fi" """.format(
bash_command=transport._bash_command_str,
escaped_remotedir=f"'{remotedir}'",
remotedir=remotedir,
fname=fname,
)
)
cmd = 'ssh -t {machine} {further_params} {connect_string}'.format(
further_params=further_params_str,
machine=transport._machine,
connect_string=connect_string,
)
return cmd