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
31 from __future__ import absolute_import, print_function, unicode_literals
32
33 import bz2
34 import os
35 try:
36 import cPickle as pickle
37 except ImportError:
38 import pickle
39 import re
40 from time import sleep
41
42
43
44 MAX_TRIES = int(os.getenv('GRID_MAP_MAX_TRIES', '10'))
45 SLEEP_TIME = int(os.getenv('GRID_MAP_SLEEP_TIME', '3'))
46
47
49 '''
50 Replace all weird SAN paths with normal paths. This is really
51 ETS-specific, but shouldn't harm anyone else.
52 '''
53
54 path = re.sub(r'/\.automount/\w+/SAN/NLP/(\w+)-(dynamic|static)',
55 r'/home/nlp-\1/\2', path)
56 path = re.sub(r'/\.automount/[^/]+/SAN/Research/HomeResearch',
57 '/home/research', path)
58 return path
59
60
61 -def zsave_db(obj, redis_server, prefix, job_num):
62 """
63 Saves an object/function as bz2-compressed pickled data in a Redis database
64
65 @param obj: The object/function to store.
66 @type obj: C{object} or C{function}
67 @param redis_server: An open connection to the database
68 @type redis_server: C{StrictRedis}
69 @param prefix: The prefix to use for the key for this data.
70 @type prefix: C{basestring}
71 @param job_num: The ID of the job this data is for.
72 @type job_num: C{int}
73 """
74
75
76 pickled_data = bz2.compress(pickle.dumps(obj, pickle.HIGHEST_PROTOCOL), 9)
77
78
79 redis_server.set('{0}_{1}'.format(prefix, job_num), pickled_data)
80
81
82 -def zload_db(redis_server, prefix, job_num):
83 """
84 Loads bz2-compressed pickled object from a Redis database
85
86 @param redis_server: An open connection to the database
87 @type redis_server: C{StrictRedis}
88 @param prefix: The prefix to use for the key for this data.
89 @type prefix: C{basestring}
90 @param job_num: The ID of the job this data is for.
91 @type job_num: C{int}
92 """
93 attempt = 0
94 pickled_data = None
95 while pickled_data is None and attempt < MAX_TRIES:
96 pickled_data = redis_server.get('{0}_{1}'.format(prefix, job_num))
97 attempt += 1
98 sleep(SLEEP_TIME)
99 return pickle.loads(bz2.decompress(pickled_data))
100