-
Notifications
You must be signed in to change notification settings - Fork 173
/
LocalContext.py
210 lines (184 loc) · 7.47 KB
/
LocalContext.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
import os,shutil,uuid,hashlib
import subprocess as sp
from glob import glob
from dpgen import dlog
class LocalSession (object) :
def __init__ (self, jdata) :
self.work_path = os.path.abspath(jdata['work_path'])
assert(os.path.exists(self.work_path))
def get_work_root(self) :
return self.work_path
class SPRetObj(object) :
def __init__ (self,
ret) :
self.data = ret
def read(self) :
return self.data
def readlines(self) :
lines = self.data.decode('utf-8').splitlines()
ret = []
for aa in lines:
ret.append(aa+'\n')
return ret
def _check_file_path(fname) :
dirname = os.path.dirname(fname)
if dirname != "":
os.makedirs(dirname, exist_ok=True)
def _identical_files(fname0, fname1) :
with open(fname0) as fp:
code0 = hashlib.sha1(fp.read().encode('utf-8')).hexdigest()
with open(fname1) as fp:
code1 = hashlib.sha1(fp.read().encode('utf-8')).hexdigest()
return code0 == code1
class LocalContext(object) :
def __init__ (self,
local_root,
work_profile,
job_uuid = None) :
"""
work_profile:
local_root:
"""
assert(type(local_root) == str)
self.local_root = os.path.abspath(local_root)
if job_uuid:
self.job_uuid=job_uuid
else:
self.job_uuid = str(uuid.uuid4())
self.remote_root = os.path.join(work_profile.get_work_root(), self.job_uuid)
dlog.debug("local_root is %s"% local_root)
dlog.debug("remote_root is %s"% self.remote_root)
os.makedirs(self.remote_root, exist_ok = True)
def get_job_root(self) :
return self.remote_root
def upload(self,
job_dirs,
local_up_files,
dereference = True) :
cwd = os.getcwd()
for ii in job_dirs :
local_job = os.path.join(self.local_root, ii)
remote_job = os.path.join(self.remote_root, ii)
os.makedirs(remote_job, exist_ok = True)
os.chdir(remote_job)
for jj in local_up_files :
if not os.path.exists(os.path.join(local_job, jj)):
os.chdir(cwd)
raise RuntimeError('cannot find upload file ' + os.path.join(local_job, jj))
if os.path.exists(os.path.join(remote_job, jj)) :
os.remove(os.path.join(remote_job, jj))
_check_file_path(jj)
os.symlink(os.path.join(local_job, jj),
os.path.join(remote_job, jj))
os.chdir(cwd)
def download(self,
job_dirs,
remote_down_files,
check_exists = False,
mark_failure = True,
back_error=False) :
cwd = os.getcwd()
for ii in job_dirs :
local_job = os.path.join(self.local_root, ii)
remote_job = os.path.join(self.remote_root, ii)
flist = remote_down_files
if back_error :
os.chdir(remote_job)
flist += glob('error*')
os.chdir(cwd)
for jj in flist :
rfile = os.path.join(remote_job, jj)
lfile = os.path.join(local_job, jj)
if not os.path.realpath(rfile) == os.path.realpath(lfile) :
if (not os.path.exists(rfile)) and (not os.path.exists(lfile)):
if check_exists :
if mark_failure:
with open(os.path.join(self.local_root, ii, 'tag_failure_download_%s' % jj), 'w') as fp: pass
else :
pass
else :
raise RuntimeError('do not find download file ' + rfile)
elif (not os.path.exists(rfile)) and (os.path.exists(lfile)) :
# already downloaded
pass
elif (os.path.exists(rfile)) and (not os.path.exists(lfile)) :
# trivial case, download happily
# If the file to be downloaded is a softlink, `cp` should be performed instead of `mv`.
# Otherwise, `lfile` is still a file linked to some original file,
# and when this file's removed, `lfile` will be invalid.
if os.path.islink(rfile):
shutil.copyfile(rfile,lfile)
else:
shutil.move(rfile, lfile)
elif (os.path.exists(rfile)) and (os.path.exists(lfile)) :
# both exists, replace!
dlog.info('find existing %s, replacing by %s' % (lfile, rfile))
if os.path.isdir(lfile):
shutil.rmtree(lfile, ignore_errors=True)
elif os.path.isfile(lfile) or os.path.islink(lfile):
os.remove(lfile)
shutil.move(rfile, lfile)
else :
raise RuntimeError('should not reach here!')
else :
# no nothing in the case of linked files
pass
os.chdir(cwd)
def block_checkcall(self,
cmd) :
cwd = os.getcwd()
os.chdir(self.remote_root)
proc = sp.Popen(cmd, shell=True, stdout = sp.PIPE, stderr = sp.PIPE)
o, e = proc.communicate()
stdout = SPRetObj(o)
stderr = SPRetObj(e)
code = proc.returncode
if code != 0:
os.chdir(cwd)
raise RuntimeError("Get error code %d in locally calling %s with job: %s ", (code, cmd, self.job_uuid))
os.chdir(cwd)
return None, stdout, stderr
def block_call(self, cmd) :
cwd = os.getcwd()
os.chdir(self.remote_root)
proc = sp.Popen(cmd, shell=True, stdout = sp.PIPE, stderr = sp.PIPE)
o, e = proc.communicate()
stdout = SPRetObj(o)
stderr = SPRetObj(e)
code = proc.returncode
os.chdir(cwd)
return code, None, stdout, stderr
def clean(self) :
shutil.rmtree(self.remote_root, ignore_errors=True)
def write_file(self, fname, write_str):
with open(os.path.join(self.remote_root, fname), 'w') as fp :
fp.write(write_str)
def read_file(self, fname):
with open(os.path.join(self.remote_root, fname), 'r') as fp:
ret = fp.read()
return ret
def check_file_exists(self, fname):
return os.path.isfile(os.path.join(self.remote_root, fname))
def call(self, cmd) :
cwd = os.getcwd()
os.chdir(self.remote_root)
proc = sp.Popen(cmd, shell=True, stdout = sp.PIPE, stderr = sp.PIPE)
os.chdir(cwd)
return proc
def kill(self, proc):
proc.kill()
def check_finish(self, proc):
return (proc.poll() != None)
def get_return(self, proc):
ret = proc.poll()
if ret is None:
return None, None, None
else :
try:
o, e = proc.communicate()
stdout = SPRetObj(o)
stderr = SPRetObj(e)
except ValueError:
stdout = None
stderr = None
return ret, stdout, stderr