[autoscaler] Use RLock in addition to FileLock

This commit is contained in:
Eric Liang
2019-02-10 19:16:43 -08:00
committed by GitHub
parent 5fb1efd60d
commit 8e9f2c923f
+43 -38
View File
@@ -3,6 +3,7 @@ from __future__ import division
from __future__ import print_function
from filelock import FileLock
from threading import RLock
import json
import os
import socket
@@ -19,58 +20,62 @@ filelock_logger.setLevel(logging.WARNING)
class ClusterState(object):
def __init__(self, lock_path, save_path, provider_config):
self.lock = RLock()
self.file_lock = FileLock(lock_path)
self.save_path = save_path
with self.file_lock:
if os.path.exists(self.save_path):
workers = json.loads(open(self.save_path).read())
else:
workers = {}
logger.info("ClusterState: "
"Loaded cluster state: {}".format(workers))
for worker_ip in provider_config["worker_ips"]:
if worker_ip not in workers:
workers[worker_ip] = {
with self.lock:
with self.file_lock:
if os.path.exists(self.save_path):
workers = json.loads(open(self.save_path).read())
else:
workers = {}
logger.info("ClusterState: "
"Loaded cluster state: {}".format(workers))
for worker_ip in provider_config["worker_ips"]:
if worker_ip not in workers:
workers[worker_ip] = {
"tags": {
TAG_RAY_NODE_TYPE: "worker"
},
"state": "terminated",
}
else:
assert workers[worker_ip]["tags"][
TAG_RAY_NODE_TYPE] == "worker"
if provider_config["head_ip"] not in workers:
workers[provider_config["head_ip"]] = {
"tags": {
TAG_RAY_NODE_TYPE: "worker"
TAG_RAY_NODE_TYPE: "head"
},
"state": "terminated",
}
else:
assert workers[worker_ip]["tags"][
TAG_RAY_NODE_TYPE] == "worker"
if provider_config["head_ip"] not in workers:
workers[provider_config["head_ip"]] = {
"tags": {
TAG_RAY_NODE_TYPE: "head"
},
"state": "terminated",
}
else:
assert workers[provider_config["head_ip"]]["tags"][
TAG_RAY_NODE_TYPE] == "head"
assert len(workers) == len(provider_config["worker_ips"]) + 1
with open(self.save_path, "w") as f:
logger.info("ClusterState: "
"Writing cluster state: {}".format(workers))
f.write(json.dumps(workers))
assert workers[provider_config["head_ip"]]["tags"][
TAG_RAY_NODE_TYPE] == "head"
assert len(workers) == len(provider_config["worker_ips"]) + 1
with open(self.save_path, "w") as f:
logger.info("ClusterState: "
"Writing cluster state: {}".format(workers))
f.write(json.dumps(workers))
def get(self):
with self.file_lock:
workers = json.loads(open(self.save_path).read())
return workers
with self.lock:
with self.file_lock:
workers = json.loads(open(self.save_path).read())
return workers
def put(self, worker_id, info):
assert "tags" in info
assert "state" in info
with self.file_lock:
workers = self.get()
workers[worker_id] = info
with open(self.save_path, "w") as f:
logger.info("ClusterState: "
"Writing cluster state: {}".format(workers))
f.write(json.dumps(workers))
with self.lock:
with self.file_lock:
workers = self.get()
workers[worker_id] = info
with open(self.save_path, "w") as f:
logger.info("ClusterState: "
"Writing cluster state: {}".format(workers))
f.write(json.dumps(workers))
class LocalNodeProvider(NodeProvider):