Completed
Push — master ( 2a4f2c...e53b42 )
by Mathieu
54s
created

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