data.py 2.96 KB
Newer Older
1
2
3
#! /usr/bin/env python3

import json
4
5
import argparse
import db.aggregator as aggregator
6
import conf.config as conf
7
8
9
import subprocess
import re
import os
10
from format import format
11
from rcm import rcm
12

Azat Khuziyakhmetov's avatar
Azat Khuziyakhmetov committed
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def merge_and_out(job_id, aggr, rcm, type, out_dir=None):
    formatted = format.format(aggr, type)

    formatted["recommendations"] = rcm

    if out_dir is None:
        print(json.dumps(formatted))
    else:
        filename = "{:s}/{:s}.{:s}.json".format(out_dir, job_id, type)
        with open(filename, 'w') as outfile:
            json.dump(formatted, outfile)

        print("{:s} data was exported in {:s}".format(type, filename))

27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def check_user(jobid):
    euid = os.geteuid()
    uid = os.getuid()

    if euid == 0 or uid == 0:
        return

    if conf.BATCH_SYSTEM == "SLURM":
        job_uid_comm = conf.job_uid_comm["slurm"].format(jobid = jobid)
    elif conf.BATCH_SYSTEM == "LSF":
        job_uid_comm = conf.job_uid_comm["lsf"].format(jobid = jobid)
    else:
        print("Cannot check the user ID, no batch system specified")
        exit(1)

    result = subprocess.run(job_uid_comm, stdout=subprocess.PIPE, shell=True,
                            executable='/bin/bash')
    out = result.stdout.decode("utf-8")
    m = re.search("[0-9]+", out)

    if m is None or m.group(0) is None:
        print("Cannot parse UID. {:s}".format(out))
        exit(1)

    job_uid = m.group(0)

    if int(job_uid) != int(uid):
        print("Access denied. UIDs of the user and the job do not match.")
        exit(1)

    return

59
60
def main():
    parser = argparse.ArgumentParser(description="""
Azat Khuziyakhmetov's avatar
Azat Khuziyakhmetov committed
61
        Gets the job information required for generating text or PDF reports
62
        and outputs it in JSON format.
63
    """, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
64

65
    parser.add_argument("-t", "--type", help="type of the output",
Azat Khuziyakhmetov's avatar
Azat Khuziyakhmetov committed
66
67
68
69
                        choices=['text', 'pdf', 'all'], default="text")

    parser.add_argument("-o", "--output-dir", help="output directory",
                        const="./out", required=False, nargs='?')
70

71
72
73
74
    parser.add_argument("JOBID",
                        help="job ID used in the batch system")

    args = parser.parse_args()
75
76
77
78
79
80

    job_id = args.JOBID

    if conf.SECUSER:
        check_user(job_id)

Azat Khuziyakhmetov's avatar
Azat Khuziyakhmetov committed
81
82
83
84
85
86
    # Errors in arguments
    if args.output_dir is None:
        if args.type == "all":
            print("Cannot print both data in STDOUT")
            exit()
    # End of errors in arguments
87

88
89
90
91
    aggr = aggregator.get_aggregator(job_id, "pdf")

    recommendations = rcm.get_recommendations(aggr)

Azat Khuziyakhmetov's avatar
Azat Khuziyakhmetov committed
92
93
94
95
96
97
98
99
    if args.output_dir is None:
        merge_and_out(job_id, aggr, recommendations, args.type)
    else:
        if args.type == "all":
            merge_and_out(job_id, aggr, recommendations, "text", args.output_dir)
            merge_and_out(job_id, aggr, recommendations, "pdf", args.output_dir)
        else:
            merge_and_out(job_id, aggr, recommendations, args.type, args.output_dir)
100
101
102
103
104

    return 0

if __name__ == "__main__":
    main()