wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-06-28 18:45:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

[tune] Properly handle closing files in Trainable (#4232)

Browse Source 
Fixes #3965.

Using the with keyword/block will close to file immediately after the block ends
This commit is contained in:
Adi Zimmerman
2019-03-03 14:23:05 -08:00
committed by Richard Liaw
parent 3483282254
commit 9551f2a92e
1 changed files with 13 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/tune/trainable.py
+13 -10
View File
@@ -245,14 +245,15 @@ class Trainable(object):
raise ValueError(
"`_save` must return a dict or string type: {}".format(
str(type(checkpoint))))
pickle.dump({
"experiment_id": self._experiment_id,
"iteration": self._iteration,
"timesteps_total": self._timesteps_total,
"time_total": self._time_total,
"episodes_total": self._episodes_total,
"saved_as_dict": saved_as_dict
}, open(checkpoint_path + ".tune_metadata", "wb"))
with open(checkpoint_path + ".tune_metadata", "wb") as f:
pickle.dump({
"experiment_id": self._experiment_id,
"iteration": self._iteration,
"timesteps_total": self._timesteps_total,
"time_total": self._time_total,
"episodes_total": self._episodes_total,
"saved_as_dict": saved_as_dict
}, f)
return checkpoint_path
def save_to_object(self):
@@ -271,7 +272,8 @@ class Trainable(object):
for path in os.listdir(base_dir):
path = os.path.join(base_dir, path)
if path.startswith(checkpoint_prefix):
data[os.path.basename(path)] = open(path, "rb").read()
with open(path, "rb") as f:
data[os.path.basename(path)] = f.read()
out = io.BytesIO()
data_dict = pickle.dumps({
@@ -294,7 +296,8 @@ class Trainable(object):
This method restores additional metadata saved with the checkpoint.
"""
metadata = pickle.load(open(checkpoint_path + ".tune_metadata", "rb"))
with open(checkpoint_path + ".tune_metadata", "rb") as f:
metadata = pickle.load(f)
self._experiment_id = metadata["experiment_id"]
self._iteration = metadata["iteration"]
self._timesteps_total = metadata["timesteps_total"]