Skip to content

Commit

Permalink
Merge pull request #888 from sphuber/fix_318_verdi_work_list
Browse files Browse the repository at this point in the history
Improve verdi work list and verdi work report
  • Loading branch information
sphuber authored Nov 2, 2017
2 parents 8a0b83f + 6da647d commit 205fffa
Showing 1 changed file with 126 additions and 81 deletions.
207 changes: 126 additions & 81 deletions aiida/cmdline/commands/work.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,65 +42,103 @@ def cli(self, *args):


@work.command('list', context_settings=CONTEXT_SETTINGS)
@click.option('-p', '--past-days', type=int, default=1,
help="add a filter to show only workflows created in the past N"
" days")
@click.option('-a', '--all', 'all_nodes', is_flag=True, help='Return all nodes. Overrides the -l flag')
@click.option('-l', '--limit', type=int, default=None,
help="Limit to this many results")
def do_list(past_days, all_nodes, limit):
@click.option(
'-p', '--past-days', type=int, default=1,
help='add a filter to show only work calculations created in the past N days'
)
@click.option(
'-a', '--all-states', 'all_states', is_flag=True,
help='return all work calculations regardless of their sealed state'
)
@click.option(
'-l', '--limit', type=int, default=None,
help='limit the final result set to this size'
)
def do_list(past_days, all_states, limit):
"""
Return a list of running workflows on screen
"""
from aiida.common.utils import str_timedelta
from aiida.backends.utils import load_dbenv, is_dbenv_loaded
if not is_dbenv_loaded():
load_dbenv()
import aiida.utils.timezone as timezone
from aiida.orm.mixins import Sealable
_SEALED_ATTRIBUTE_KEY = 'attributes.{}'.format(Sealable.SEALED_KEY)

now = timezone.now()
from aiida.common.utils import str_timedelta
from aiida.utils import timezone
from aiida.orm.mixins import Sealable
from aiida.orm.calculation.work import WorkCalculation

if all_nodes:
past_days = None
_SEALED_ATTRIBUTE_KEY = 'attributes.{}'.format(Sealable.SEALED_KEY)
_ABORTED_ATTRIBUTE_KEY = 'attributes.{}'.format(WorkCalculation.ABORTED_KEY)
_FAILED_ATTRIBUTE_KEY = 'attributes.{}'.format(WorkCalculation.FAILED_KEY)
_FINISHED_ATTRIBUTE_KEY = 'attributes.{}'.format(WorkCalculation.FINISHED_KEY)

table = []
for res in _build_query(limit=limit, past_days=past_days, order_by={'ctime': 'desc'}):
calc = res['calculation']

calculation = res['calculation']

creation_time = str_timedelta(
timezone.delta(calc['ctime'], now), negative_to_zero=True,
max_num_fields=1)
timezone.delta(calculation['ctime'], timezone.now()), negative_to_zero=True,
max_num_fields=1
)

if _SEALED_ATTRIBUTE_KEY in calculation and calculation[_SEALED_ATTRIBUTE_KEY]:
sealed = True
else:
sealed = False

if _FINISHED_ATTRIBUTE_KEY in calculation and calculation[_FINISHED_ATTRIBUTE_KEY]:
state = 'Finished'
elif _FAILED_ATTRIBUTE_KEY in calculation and calculation[_FAILED_ATTRIBUTE_KEY]:
state = 'Failed'
elif _ABORTED_ATTRIBUTE_KEY in calculation and calculation[_ABORTED_ATTRIBUTE_KEY]:
state = 'Aborted'
elif sealed:
# If it is not in terminal state but sealed, we have an inconsistent state
state = 'Unknown'
else:
state = 'Running'

# By default we only display unsealed entries, unless all_states flag is set
if sealed and not all_states:
continue

table.append([
calc['id'],
calculation['id'],
creation_time,
calc['attributes._process_label'],
str(calc[_SEALED_ATTRIBUTE_KEY])
state,
str(sealed),
calculation['attributes._process_label']
])

# Revert table:
# in this way, I order by 'desc', so I start by the most recent, but then
# I print this as the las one (like 'verdi calculation list' does)
# This is useful when 'limit' is set to not None
# Since we sorted by descending creation time, we revert the list to print the most
# recent entries last
table = table[::-1]
print(tabulate(table, headers=["PID", "Creation time", "ProcessLabel", "Sealed"]))

print(tabulate(table, headers=['PK', 'Creation', 'State', 'Sealed', 'ProcessLabel']))


@work.command('report', context_settings=CONTEXT_SETTINGS)
@click.argument('pk', nargs=1, type=int)
@click.option('-i', '--indent-size', type=int, default=2)
@click.option('-l', '--levelname',
type=click.Choice(LOG_LEVELS),
default='REPORT',
help='Filter the results by name of the log level'
)
@click.option('-o', '--order-by',
type=click.Choice(['id', 'time', 'levelname']),
default='time',
help='Order the results by column'
)
def report(pk, levelname, order_by, indent_size):
@click.argument(
'pk', nargs=1, type=int
)
@click.option(
'-i', '--indent-size', type=int, default=2,
help='Set the number of spaces to indent each level by'
)
@click.option(
'-l', '--levelname', type=click.Choice(LOG_LEVELS), default='REPORT',
help='Filter the results by name of the log level'
)
@click.option(
'-o', '--order-by', type=click.Choice(['id', 'time', 'levelname']), default='time',
help='Order the results by column'
)
@click.option(
'-m', '--max-depth', 'max_depth', type=int, default=None,
help='Limit the number of levels to be printed'
)
def report(pk, levelname, order_by, indent_size, max_depth):
"""
Return a list of recorded log messages for the WorkChain with pk=PK
"""
Expand Down Expand Up @@ -157,7 +195,12 @@ def print_subtree(tree, prepend=""):

workchain_tree = get_subtree(pk)

reports = list(itertools.chain(*[get_report_messages(pk, depth, levelname) for pk, depth in workchain_tree]))
if max_depth:
report_list = [get_report_messages(pk, depth, levelname) for pk, depth in workchain_tree if depth < max_depth]
else:
report_list = [get_report_messages(pk, depth, levelname) for pk, depth in workchain_tree]

reports = list(itertools.chain(*report_list))
reports.sort(key=lambda r: r[0].time)

if reports is None or len(reports) == 0:
Expand Down Expand Up @@ -224,48 +267,6 @@ def checkpoint(pks):
except ValueError:
print("Unable to show checkpoint for calculation '{}'".format(pk))


def _build_query(order_by=None, limit=None, past_days=None):
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.calculation.work import WorkCalculation
import aiida.utils.timezone as timezone
import datetime
from aiida.orm.mixins import Sealable
_SEALED_ATTRIBUTE_KEY = 'attributes.{}'.format(Sealable.SEALED_KEY)

# The things that we want to get out
calculation_projections = \
['id', 'ctime', 'attributes._process_label', _SEALED_ATTRIBUTE_KEY]

now = timezone.now()

# The things to filter by
calculation_filters = {}

if past_days is not None:
n_days_ago = now - datetime.timedelta(days=past_days)
calculation_filters['ctime'] = {'>': n_days_ago}

qb = QueryBuilder()

# Build the quiery
qb.append(
cls=WorkCalculation,
filters=calculation_filters,
project=calculation_projections,
tag='calculation'
)

# ORDER
if order_by is not None:
qb.order_by({'calculation': order_by})

# LIMIT
if limit is not None:
qb.limit(limit)

return qb.iterdict()

@work.command('kill', context_settings=CONTEXT_SETTINGS)
@click.argument('pks', nargs=-1, type=int)
def kill(pks):
Expand Down Expand Up @@ -293,3 +294,47 @@ def kill(pks):
click.echo('Abort!')
else:
click.echo('No pks of valid running workchains given.')


def _build_query(order_by=None, limit=None, past_days=None):
import datetime
from aiida.utils import timezone
from aiida.orm.mixins import Sealable
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.calculation.work import WorkCalculation

_SEALED_ATTRIBUTE_KEY = 'attributes.{}'.format(Sealable.SEALED_KEY)
_ABORTED_ATTRIBUTE_KEY = 'attributes.{}'.format(WorkCalculation.ABORTED_KEY)
_FAILED_ATTRIBUTE_KEY = 'attributes.{}'.format(WorkCalculation.FAILED_KEY)
_FINISHED_ATTRIBUTE_KEY = 'attributes.{}'.format(WorkCalculation.FINISHED_KEY)

calculation_projections = [
'id', 'ctime', 'attributes._process_label',
_SEALED_ATTRIBUTE_KEY, _ABORTED_ATTRIBUTE_KEY, _FAILED_ATTRIBUTE_KEY, _FINISHED_ATTRIBUTE_KEY
]

# Define filters
calculation_filters = {}

if past_days is not None:
n_days_ago = timezone.now() - datetime.timedelta(days=past_days)
calculation_filters['ctime'] = {'>': n_days_ago}

# Build the query
qb = QueryBuilder()
qb.append(
cls=WorkCalculation,
filters=calculation_filters,
project=calculation_projections,
tag='calculation'
)

# Ordering of queryset
if order_by is not None:
qb.order_by({'calculation': order_by})

# Limiting the queryset
if limit is not None:
qb.limit(limit)

return qb.iterdict()

0 comments on commit 205fffa

Please sign in to comment.