From ec9357b486bd43c3ce943ec970fcbf08e144d382 Mon Sep 17 00:00:00 2001
From: Thomas Desrosiers <681004+thomasdesr@users.noreply.github.com>
Date: Tue, 5 May 2020 20:22:03 -0400
Subject: [PATCH] [autoscaler] Fix filesystem permission race conditions
 (#8327)

---
 python/ray/autoscaler/aws/config.py |  8 ++++++--
 python/ray/autoscaler/gcp/config.py | 11 +++++++++--
 python/ray/autoscaler/updater.py    | 11 ++---------
 3 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/python/ray/autoscaler/aws/config.py b/python/ray/autoscaler/aws/config.py
index 77b061a7a..6af48d7fc 100644
--- a/python/ray/autoscaler/aws/config.py
+++ b/python/ray/autoscaler/aws/config.py
@@ -1,4 +1,5 @@
 from distutils.version import StrictVersion
+from functools import partial
 import json
 import os
 import time
@@ -160,9 +161,12 @@ def _configure_key_pair(config):
             logger.info("_configure_key_pair: "
                         "Creating new key pair {}".format(key_name))
             key = ec2.create_key_pair(KeyName=key_name)
-            with open(key_path, "w") as f:
+
+            # We need to make sure to _create_ the file with the right
+            # permissions. In order to do that we need to change the default
+            # os.open behavior to include the mode we want.
+            with open(key_path, "w", opener=partial(os.open, mode=0o600)) as f:
                 f.write(key.key_material)
-            os.chmod(key_path, 0o600)
             break
 
     if not key:
diff --git a/python/ray/autoscaler/gcp/config.py b/python/ray/autoscaler/gcp/config.py
index 4071d48dc..ff54d1286 100644
--- a/python/ray/autoscaler/gcp/config.py
+++ b/python/ray/autoscaler/gcp/config.py
@@ -1,3 +1,4 @@
+from functools import partial
 import os
 import logging
 import time
@@ -236,9 +237,15 @@ def _configure_key_pair(config):
 
             _create_project_ssh_key_pair(project, public_key, ssh_user)
 
-            with open(private_key_path, "w") as f:
+            # We need to make sure to _create_ the file with the right
+            # permissions. In order to do that we need to change the default
+            # os.open behavior to include the mode we want.
+            with open(
+                    private_key_path,
+                    "w",
+                    opener=partial(os.open, mode=0o600),
+            ) as f:
                 f.write(private_key)
-            os.chmod(private_key_path, 0o600)
 
             with open(public_key_path, "w") as f:
                 f.write(public_key)
diff --git a/python/ray/autoscaler/updater.py b/python/ray/autoscaler/updater.py
index 8fcf03405..8c08fedd4 100644
--- a/python/ray/autoscaler/updater.py
+++ b/python/ray/autoscaler/updater.py
@@ -223,17 +223,10 @@ class SSHCommandRunner:
         #   the ControlPath directory exists, allowing SSH to maintain
         #   persistent sessions later on.
         try:
-            self.process_runner.check_call(
-                ["mkdir", "-p", self.ssh_control_path])
-        except subprocess.CalledProcessError as e:
+            os.makedirs(self.ssh_control_path, mode=0o700, exist_ok=True)
+        except OSError as e:
             logger.warning(e)
 
-        try:
-            self.process_runner.check_call(
-                ["chmod", "0700", self.ssh_control_path])
-        except subprocess.CalledProcessError as e:
-            logger.warning(self.log_prefix + str(e))
-
     def run(self,
             cmd,
             timeout=120,