gryf/coach
1
0
Fork 0
mirror of https://github.com/gryf/coach.git synced 2025-12-17 19:20:19 +01:00
Code Issues Projects Releases Wiki Activity

Getting only the model_checkpoint_path files

Browse Source
This commit is contained in:
Ajay Deshpande
2018-10-05 13:48:10 -07:00
committed by zach dwiel
parent 052bbc8f19
commit a2e57a44f1
1 changed files with 13 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
rl_coach/data_stores/s3_data_store.py
View File

@@ -3,6 +3,8 @@ from kubernetes import client as k8sclient
from minio import Minio
from minio.error import ResponseError
from configparser import ConfigParser, Error
from google.protobuf import text_format
from tensorflow.python.training.checkpoint_state_pb2 import CheckpointState
import os
import time
import io
@@ -88,9 +90,16 @@ class S3DataStore(DataStore):
filename = os.path.abspath(os.path.join(self.params.checkpoint_dir, "checkpoint"))
self.mc.fget_object(self.params.bucket_name, "checkpoint", filename)
objects = self.mc.list_objects_v2(self.params.bucket_name, recursive=True)
for obj in objects:
filename = os.path.abspath(os.path.join(self.params.checkpoint_dir, obj.object_name))
self.mc.fget_object(obj.bucket_name, obj.object_name, filename)
ckpt = CheckpointState()
if os.path.exists(filename):
contents = open(filename, 'r').read()
text_format.Merge(contents, ckpt)
rel_path = os.path.relpath(ckpt.model_checkpoint_path, self.params.checkpoint_dir)
objects = self.mc.list_objects_v2(self.params.bucket_name, prefix=rel_path, recursive=True)
for obj in objects:
filename = os.path.abspath(os.path.join(self.params.checkpoint_dir, obj.object_name))
self.mc.fget_object(obj.bucket_name, obj.object_name, filename)
except ResponseError as e:
print("Got exception: %s\n while loading from S3", e)