Completed
Push — master ( 0b9edb...ab2b1a )
by Mathieu
10s
created

get_job_folders()   B

Complexity

Conditions 5

Size

Total Lines 16

Duplication

Lines 0
Ratio 0 %
Metric Value
cc 5
dl 0
loc 16
rs 8.5454
1
from __future__ import absolute_import
2
3
import os
4
import re
5
import itertools
6
import time as t
7
from os.path import join as pjoin
8
9
import smartdispatch
10
from smartdispatch import utils
11
from smartdispatch.filelock import open_with_lock
12
from smartdispatch.argument_template import argument_templates
13
14
UID_TAG = "{UID}"
15
16
17
def generate_name_from_command(command, max_length_arg=None, max_length=None):
18
    ''' Generates name from a given command.
19
20
    Generate a name by replacing spaces in command with dashes and
21
    by trimming lengthty (as defined by max_length_arg) arguments.
22
23
    Parameters
24
    ----------
25
    command : str
26
        command from which to generate the name
27
    max_length_arg : int
28
        arguments longer than this will be trimmed keeping last characters (Default: inf)
29
    max_length : int
30
        trim name if longer than this keeping last characters (Default: inf)
31
32
    Returns
33
    -------
34
    name : str
35
        slugified name
36
    '''
37
    if max_length_arg is not None:
38
        max_length_arg = min(-max_length_arg, max_length_arg)
39
40
    name = t.strftime("%Y-%m-%d_%H-%M-%S_")
41
    name += '_'.join([utils.slugify(argvalue)[max_length_arg:] for argvalue in command.split()])
42
    return name[:max_length]
43
44
45
def get_commands_from_file(fileobj):
46
    ''' Reads commands from `fileobj`.
47
48
    Parameters
49
    ----------
50
    fileobj : file
51
        opened file where to read commands from
52
53
    Returns
54
    -------
55
    commands : list of str
56
        commands read from the file
57
    '''
58
    return fileobj.read().strip().split('\n')
59
60
61
def unfold_command(command):
62
    ''' Unfolds a command into a list of unfolded commands.
63
64
    Unfolding is performed for every folded arguments (see *Arguments templates*)
65
    found in `command`. Then, resulting commands are generated using the product
66
    of every unfolded arguments.
67
68
    Parameters
69
    ----------
70
    command : list of str
71
        command to unfold
72
73
    Returns
74
    -------
75
    commands : list of str
76
        commands obtained after unfolding `command`
77
78
    Arguments template
79
    ------------------
80
    *list*: "[item1 item2 ... itemN]"
81
    *range*: "[start:end]" or "[start:end:step]"
82
    '''
83
    text = utils.encode_escaped_characters(command)
84
85
    # Build the master regex with all argument's regex
86
    regex = "(" + "|".join(["(?P<{0}>{1})".format(name, arg.regex) for name, arg in argument_templates.items()]) + ")"
87
88
    pos = 0
89
    arguments = []
90
    for match in re.finditer(regex, text):
91
        # Add already unfolded argument
92
        arguments.append([text[pos:match.start()]])
93
94
        # Unfold argument
95
        argument_template_name, matched_text = next((k, v) for k, v in match.groupdict().items() if v is not None)
96
        arguments.append(argument_templates[argument_template_name].unfold(matched_text))
97
        pos = match.end()
98
99
    arguments.append([text[pos:]])  # Add remaining unfolded arguments
100
    arguments = [map(utils.decode_escaped_characters, argvalues) for argvalues in arguments]
101
    return ["".join(argvalues) for argvalues in itertools.product(*arguments)]
102
103
104
def replace_uid_tag(commands):
105
    return [command.replace("{UID}", utils.generate_uid_from_string(command)) for command in commands]
106
107
108
def get_available_queues(cluster_name=utils.detect_cluster()):
109
    """ Fetches all available queues on the current cluster """
110
    if cluster_name is None:
111
        return {}
112
113
    smartdispatch_dir, _ = os.path.split(smartdispatch.__file__)
114
    config_dir = pjoin(smartdispatch_dir, 'config')
115
116
    config_filename = cluster_name + ".json"
117
    config_filepath = pjoin(config_dir, config_filename)
118
119
    if not os.path.isfile(config_filepath):
120
        return {}  # Unknown cluster
121
122
    queues_infos = utils.load_dict_from_json_file(config_filepath)
123
    return queues_infos
124
125
126
def get_job_folders(path, jobname, create_if_needed=False):
127
    """ Get all folder paths for a specific job (creating them if needed). """
128
    path_job = pjoin(path, jobname)
129
    path_job_logs = pjoin(path_job, 'logs')
130
    path_job_commands = pjoin(path_job, 'commands')
131
132
    if not os.path.isdir(path_job_commands):
133
        os.makedirs(path_job_commands)
134
    if not os.path.isdir(path_job_logs):
135
        os.makedirs(path_job_logs)
136
    if not os.path.isdir(pjoin(path_job_logs, "worker")):
137
        os.makedirs(pjoin(path_job_logs, "worker"))
138
    if not os.path.isdir(pjoin(path_job_logs, "job")):
139
        os.makedirs(pjoin(path_job_logs, "job"))
140
141
    return path_job, path_job_logs, path_job_commands
142
143
144
def log_command_line(path_job, command_line):
145
    """ Logs a command line in a job folder.
146
147
    The command line is append to a file named 'command_line.log' that resides
148
    in the given job folder. The current date and time is also added along
149
    each command line logged.
150
151
    Notes
152
    -----
153
    Commands save in log file might differ from sys.argv since we want to make sure
154
    we can paste the command line as-is in the terminal. This means that the quotes
155
    symbole " and the square brackets will be escaped.
156
    """
157
    with open_with_lock(pjoin(path_job, "command_line.log"), 'a') as command_line_log:
158
        command_line_log.write(t.strftime("## %Y-%m-%d %H:%M:%S ##\n"))
159
        command_line = command_line.replace('"', r'\"')  # Make sure we can paste the command line as-is
160
        command_line = re.sub(r'(\[)([^\[\]]*\\ [^\[\]]*)(\])', r'"\1\2\3"', command_line)  # Make sure we can paste the command line as-is
161
        command_line_log.write(command_line + "\n\n")
162