From 8e9f2c923fb5ddccdfb6b87311cf3d9af3b2afa9 Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Sun, 10 Feb 2019 19:16:43 -0800
Subject: [PATCH] [autoscaler] Use RLock in addition to FileLock

---
 python/ray/autoscaler/local/node_provider.py | 81 +++++++++++---------
 1 file changed, 43 insertions(+), 38 deletions(-)

diff --git a/python/ray/autoscaler/local/node_provider.py b/python/ray/autoscaler/local/node_provider.py
index 291f2189c..3e8d81939 100644
--- a/python/ray/autoscaler/local/node_provider.py
+++ b/python/ray/autoscaler/local/node_provider.py
@@ -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):