126 lines
2.7 KiB
Python
Executable File
126 lines
2.7 KiB
Python
Executable File
#!/usr/bin/python
|
|
|
|
|
|
import click
|
|
|
|
from solar.orchestration import graph
|
|
from solar.orchestration import tasks
|
|
|
|
import networkx as nx
|
|
import subprocess
|
|
|
|
@click.group()
|
|
def orchestration():
|
|
"""
|
|
\b
|
|
./cli.py create orch/examples/multi.yaml
|
|
<id>
|
|
./cli.py execute <id>
|
|
./cli.py report <id>
|
|
<task> -> <status>
|
|
./cli.py restart <id> --reset
|
|
"""
|
|
|
|
|
|
@click.command()
|
|
@click.argument('plan', type=click.File('rb'))
|
|
def create(plan):
|
|
click.echo(graph.create_plan(plan.read()))
|
|
|
|
|
|
@click.command()
|
|
@click.argument('uid')
|
|
@click.argument('plan', type=click.File('rb'))
|
|
def update(uid, plan):
|
|
graph.update_plan(uid, plan.read())
|
|
|
|
@click.command()
|
|
@click.argument('uid')
|
|
def report(uid):
|
|
colors = {
|
|
'PENDING': 'blue',
|
|
'ERROR': 'red',
|
|
'SUCCESS': 'green',
|
|
'INPROGRESS': 'yellow'}
|
|
|
|
report = graph.report_topo(uid)
|
|
for item in report:
|
|
msg = '{} -> {}'.format(item[0], item[1])
|
|
if item[2]:
|
|
msg += ' :: {}'.format(item[2])
|
|
click.echo(click.style(msg, fg=colors[item[1]]))
|
|
|
|
|
|
@click.command()
|
|
@click.argument('uid')
|
|
@click.option('--start', default=None)
|
|
@click.option('--end', default=None)
|
|
def execute(uid, start, end):
|
|
tasks.schedule_start.apply_async(
|
|
args=[uid],
|
|
kwargs={'start': start, 'end': end},
|
|
queue='scheduler')
|
|
|
|
|
|
@click.command()
|
|
@click.argument('uid')
|
|
def restart(uid):
|
|
graph.reset(uid)
|
|
tasks.schedule_start.apply_async(args=[uid], queue='scheduler')
|
|
|
|
|
|
@click.command()
|
|
@click.argument('uid')
|
|
def reset(uid):
|
|
graph.reset(uid)
|
|
|
|
|
|
@click.command()
|
|
@click.argument('uid')
|
|
def stop(uid):
|
|
# TODO(dshulyak) how to do "hard" stop?
|
|
# using revoke(terminate=True) will lead to inability to restart execution
|
|
# research possibility of customizations of
|
|
# app.control and Panel.register in celery
|
|
graph.soft_stop(uid)
|
|
|
|
|
|
@click.command()
|
|
@click.argument('uid')
|
|
def retry(uid):
|
|
graph.reset(uid, ['ERROR'])
|
|
tasks.schedule_start.apply_async(args=[uid], queue='scheduler')
|
|
|
|
|
|
@click.command()
|
|
@click.argument('uid')
|
|
def dg(uid):
|
|
plan = graph.get_graph(uid)
|
|
|
|
colors = {
|
|
'PENDING': 'blue',
|
|
'ERROR': 'red',
|
|
'SUCCESS': 'green',
|
|
'INPROGRESS': 'yellow'}
|
|
|
|
for n in plan:
|
|
color = colors[plan.node[n]['status']]
|
|
plan.node[n]['color'] = color
|
|
nx.write_dot(plan, 'graph.dot')
|
|
subprocess.call(['dot', '-Tpng', 'graph.dot', '-o', 'graph.png'])
|
|
|
|
|
|
orchestration.add_command(create)
|
|
orchestration.add_command(update)
|
|
orchestration.add_command(report)
|
|
orchestration.add_command(execute)
|
|
orchestration.add_command(restart)
|
|
orchestration.add_command(reset)
|
|
orchestration.add_command(stop)
|
|
orchestration.add_command(retry)
|
|
orchestration.add_command(dg)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
orchestration()
|