def reset(self):
utils.send_message(self.sock,
gym_uds_pb2.Request(type=gym_uds_pb2.Request.RESET))
state_pb = utils.recv_message(self.sock, gym_uds_pb2.State)
observation = np.asarray(state_pb.observation.data).reshape(
state_pb.observation.shape)
return observation
def request_chunk_from_storage(storage_ip, chunk_path):
sock = socket.socket()
sock.connect((storage_ip, 9005))
send_message(sock, chunk_path)
response = recv_message(sock)
sock.close()
return response
def clients_commands(conn):
while True:
request = recv_message(conn)
print request
sys.stdout.flush()
req = request.split()
client_pwd = req[0]
command = req[1]
args = req[2:]
response = "nothing"
if command == "quit":
conn.close()
break
elif command == "pwd":
response = pwd(client_pwd)
elif command == "ls":
response = ls(client_pwd, args)
elif command == "cd":
response = cd(client_pwd, args)
elif command == "mkdir":
response = mkdir(client_pwd, args)
elif command == "cp":
response = cp(client_pwd, list(connected_storages), args)
elif command == "rm":
response = rm(client_pwd, args)
elif command == "stat":
response = stat(client_pwd, args)
elif command == "cat":
response = cat(client_pwd, connected_storages, args)
elif command == "init":
response = init(client_pwd, args)
send_message(conn, response)
def run(self):
while True:
request = utils.recv_message(self.sock, gym_uds_pb2.Request)
if request.type == gym_uds_pb2.Request.DONE: break
elif request.type == gym_uds_pb2.Request.RESET: self.reset()
elif request.type == gym_uds_pb2.Request.STEP: self.step()
elif request.type == gym_uds_pb2.Request.SAMPLE: self.sample()
def naming_commands(conn):
request = recv_message(conn)
filepath = request
chunk_name = os.path.normpath(fake_root + filepath)
chunk_data = ''
with open(chunk_name, 'r') as chunk:
chunk_data += chunk.read()
send_message(conn, chunk_data)
def step(self, action):
utils.send_message(self.sock,
gym_uds_pb2.Request(type=gym_uds_pb2.Request.STEP))
utils.send_message(self.sock, gym_uds_pb2.Action(value=action))
state_pb = utils.recv_message(self.sock, gym_uds_pb2.State)
observation = np.asarray(state_pb.observation.data).reshape(
state_pb.observation.shape)
return observation, state_pb.reward, state_pb.done
def send_file_to_storage(storage_ip, chunk_data):
sock = socket.socket()
sock.connect((storage_ip, 9004))
send_message(sock, chunk_data)
logger.info('Command cp response: Sent %s to the storage %s' %
(os.path.basename(chunk_data.splitlines()[0]), storage_ip))
response = recv_message(sock)
logger.info('Command cp response: Storage %s response: %s' %
(storage_ip, response))
sock.close()
def clients_commands(conn):
request = recv_message(conn)
file = request.splitlines()[0]
chunk_name = os.path.normpath(fake_root + os.path.normpath(file))
filepath = os.path.dirname(file)
mkdir(filepath)
with open(chunk_name, "wb") as out_file:
for line in request.splitlines()[1:]:
out_file.write(line + os.linesep)
print 'Added ', chunk_name
response = 'OK'
send_message(conn, response)
def step(self):
action = utils.recv_message(self.sock, gym_uds_pb2.Action)
observation, reward, done, _ = self.env.step(action.value)
assert type(observation) is np.ndarray
observation_pb = gym_uds_pb2.Observation(data=observation.ravel(),
shape=observation.shape)
utils.send_message(
self.sock,
gym_uds_pb2.State(observation=observation_pb,
reward=reward,
done=done))
def send_heartbeat(connected_storages):
while True:
lost = []
for s in connected_storages:
try:
sock = socket.socket()
sock.connect((s, 9003))
sock.settimeout(10)
data = recv_message(sock)
except:
duplicate_lost_data(s, connected_storages)
lost.append(s)
sock.close()
for s in lost:
connected_storages.remove(s)
sleep(5)
def sample(self):
utils.send_message(
self.sock, gym_uds_pb2.Request(type=gym_uds_pb2.Request.SAMPLE))
action_pb = utils.recv_message(self.sock, gym_uds_pb2.Action)
return action_pb.value
