mirror of
https://github.com/wassname/ray.git
synced 2026-07-14 11:17:54 +08:00
[autoscaler] Use RLock in addition to FileLock
This commit is contained in:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user