1__author__ = "Johannes Köster"
2__copyright__ = "Copyright 2021, Johannes Köster"
3__email__ = "johannes.koester@uni-due.de"
4__license__ = "MIT"
5
6import json
7import os
8import threading
9
10from flask import Flask, render_template, request
11
12from snakemake.common import __version__
13
14LOCK = threading.Lock()
15
16app = Flask("snakemake", template_folder=os.path.dirname(__file__))
17# app.debug=True
18app.extensions = {
19    "dag": None,
20    "run_snakemake": None,
21    "progress": "",
22    "log": [],
23    "status": {"running": False},
24    "args": None,
25    "targets": [],
26    "rule_info": [],
27    "resources": [],
28}
29
30
31def register(run_snakemake, args):
32    app.extensions["run_snakemake"] = run_snakemake
33    app.extensions["args"] = dict(
34        targets=args.target,
35        cluster=args.cluster,
36        workdir=args.directory,
37        touch=args.touch,
38        forcetargets=args.force,
39        forceall=args.forceall,
40        forcerun=args.forcerun,
41        prioritytargets=args.prioritize,
42        stats=args.stats,
43        keepgoing=args.keep_going,
44        jobname=args.jobname,
45        immediate_submit=args.immediate_submit,
46        ignore_ambiguity=args.allow_ambiguity,
47        lock=not args.nolock,
48        force_incomplete=args.rerun_incomplete,
49        ignore_incomplete=args.ignore_incomplete,
50        jobscript=args.jobscript,
51        notemp=args.notemp,
52        latency_wait=args.latency_wait,
53    )
54
55    target_rules = []
56
57    def log_handler(msg):
58        if msg["level"] == "rule_info":
59            target_rules.append(msg["name"])
60
61    run_snakemake(list_target_rules=True, log_handler=[log_handler])
62    for target in args.target:
63        target_rules.remove(target)
64    app.extensions["targets"] = args.target + target_rules
65
66    resources = []
67
68    def log_handler(msg):
69        if msg["level"] == "info":
70            resources.append(msg["msg"])
71
72    run_snakemake(list_resources=True, log_handler=[log_handler])
73    app.extensions["resources"] = resources
74    app.extensions["snakefilepath"] = os.path.abspath(args.snakefile)
75
76
77def run_snakemake(**kwargs):
78    args = dict(app.extensions["args"])
79    args.update(kwargs)
80    app.extensions["run_snakemake"](**args)
81
82
83@app.route("/")
84def index():
85    args = app.extensions["args"]
86    return render_template(
87        "gui.html",
88        targets=app.extensions["targets"],
89        cores_label="Nodes" if args["cluster"] else "Cores",
90        resources=app.extensions["resources"],
91        snakefilepath=app.extensions["snakefilepath"],
92        version=__version__,
93        node_width=15,
94        node_padding=10,
95    )
96
97
98@app.route("/dag")
99def dag():
100    if app.extensions["dag"] is None:
101
102        def record(msg):
103            if msg["level"] == "d3dag":
104                app.extensions["dag"] = msg
105            elif msg["level"] in ("error", "info"):
106                app.extensions["log"].append(msg)
107
108        run_snakemake(printd3dag=True, log_handler=[record])
109    return json.dumps(app.extensions["dag"])
110
111
112@app.route("/log/<int:id>")
113def log(id):
114    log = app.extensions["log"][id:]
115    return json.dumps(log)
116
117
118@app.route("/progress")
119def progress():
120    return json.dumps(app.extensions["progress"])
121
122
123def _run(dryrun=False, cores=1):
124    def log_handler(msg):
125        level = msg["level"]
126        if level == "progress":
127            app.extensions["progress"] = msg
128        elif level in ("info", "error", "job_info", "job_finished", "job_error"):
129            app.extensions["log"].append(msg)
130
131    with LOCK:
132        app.extensions["status"]["running"] = True
133    run_snakemake(log_handler=[log_handler], dryrun=dryrun, cores=cores)
134    with LOCK:
135        app.extensions["status"]["running"] = False
136    return ""
137
138
139@app.route("/run/<int:cores>")
140def run(cores):
141    return _run(cores=cores)
142
143
144@app.route("/dryrun")
145def dryrun():
146    return _run(dryrun=True)
147
148
149@app.route("/status")
150def status():
151    with LOCK:
152        return json.dumps(app.extensions["status"])
153
154
155@app.route("/targets")
156def targets():
157    return json.dumps(app.extensions["targets"])
158
159
160@app.route("/get_args")
161def get_args():
162    return json.dumps(app.extensions["args"])
163
164
165@app.route("/set_args", methods=["POST"])
166def set_args():
167    app.extensions["args"].update(
168        {name: value for name, value in request.form.items() if not name.endswith("[]")}
169    )
170    targets = request.form.getlist("targets[]")
171    if targets != app.extensions["args"]["targets"]:
172        app.extensions["dag"] = None
173    app.extensions["args"]["targets"] = targets
174    return ""
175