1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23 """
24 This modules provides all of the data-related function for gridmap.
25
26 @author: Christian Widmer
27 @author: Cheng Soon Ong
28 @author: Dan Blanchard (dblanchard@ets.org)
29
30 @var MAX_TRIES: Maximum number of times to try to get the output of a job from
31 the Redis database before giving up and assuming the job died
32 before writing its output; can be overriden by setting the
33 GRID_MAP_MAX_TRIES environment variable.
34 @var SLEEP_TIME: Number of seconds to sleep between attempts to retrieve job
35 output from the Redis database; can be overriden by setting the
36 GRID_MAP_SLEEP_TIME environment variable.
37 """
38
39 from __future__ import absolute_import, print_function, unicode_literals
40
41 import bz2
42 import os
43 try:
44 import cPickle as pickle
45 except ImportError:
46 import pickle
47 import re
48 from time import sleep
49
50
51
52 MAX_TRIES = int(os.getenv('GRID_MAP_MAX_TRIES', '10'))
53 SLEEP_TIME = int(os.getenv('GRID_MAP_SLEEP_TIME', '3'))
54
55
57 '''
58 Replace all weird SAN paths with normal paths. This is really
59 ETS-specific, but shouldn't harm anyone else.
60 '''
61
62 path = re.sub(r'/\.automount/\w+/SAN/NLP/(\w+)-(dynamic|static)',
63 r'/home/nlp-\1/\2', path)
64 path = re.sub(r'/\.automount/[^/]+/SAN/Research/HomeResearch',
65 '/home/research', path)
66 return path
67
68
69 -def zsave_db(obj, redis_server, prefix, job_num):
70 """
71 Saves an object/function as bz2-compressed pickled data in a Redis database
72
73 @param obj: The object/function to store.
74 @type obj: C{object} or C{function}
75 @param redis_server: An open connection to the database
76 @type redis_server: C{StrictRedis}
77 @param prefix: The prefix to use for the key for this data.
78 @type prefix: C{basestring}
79 @param job_num: The ID of the job this data is for.
80 @type job_num: C{int}
81 """
82
83
84 pickled_data = bz2.compress(pickle.dumps(obj, pickle.HIGHEST_PROTOCOL), 9)
85
86
87 redis_server.set('{0}_{1}'.format(prefix, job_num), pickled_data)
88
89
90 -def zload_db(redis_server, prefix, job_num):
91 """
92 Loads bz2-compressed pickled object from a Redis database
93
94 @param redis_server: An open connection to the database
95 @type redis_server: C{StrictRedis}
96 @param prefix: The prefix to use for the key for this data.
97 @type prefix: C{basestring}
98 @param job_num: The ID of the job this data is for.
99 @type job_num: C{int}
100 """
101 attempt = 0
102 pickled_data = None
103 while pickled_data is None and attempt < MAX_TRIES:
104 pickled_data = redis_server.get('{0}_{1}'.format(prefix, job_num))
105 attempt += 1
106 sleep(SLEEP_TIME)
107 return pickle.loads(bz2.decompress(pickled_data))
108