From 54cbb4396fb64f1e4c6f1e182f15517b960f8220 Mon Sep 17 00:00:00 2001
From: Yuhong Guo <guoyuhong1985@outlook.com>
Date: Sat, 2 Feb 2019 12:53:36 +0800
Subject: [PATCH] Prepare socket file when start ray (#3925)

---
 python/ray/node.py | 22 +++++++++++++++++++++-
 test/runtest.py    |  8 ++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/python/ray/node.py b/python/ray/node.py
index df5ff652a..7f6e99bad 100644
--- a/python/ray/node.py
+++ b/python/ray/node.py
@@ -16,7 +16,8 @@ from ray.tempfile_services import (
     get_logs_dir_path, get_object_store_socket_name, get_raylet_socket_name,
     new_log_monitor_log_file, new_monitor_log_file,
     new_raylet_monitor_log_file, new_plasma_store_log_file,
-    new_raylet_log_file, new_webui_log_file, set_temp_root)
+    new_raylet_log_file, new_webui_log_file, set_temp_root,
+    try_to_create_directory)
 
 # Logger for this module. It should be configured at the entry point
 # into the program using Ray. Ray configures it by default automatically
@@ -102,6 +103,23 @@ class Node(object):
         """Get the node's raylet socket name."""
         return self._raylet_socket_name
 
+    def prepare_socket_file(self, socket_path):
+        """Prepare the socket file for raylet and plasma.
+
+        This method helps to prepare a socket file.
+        1. Make the directory if the directory does not exist.
+        2. If the socket file exists, raise exception.
+
+        Args:
+            socket_path (string): the socket file to prepare.
+        """
+        if not os.path.exists(socket_path):
+            path = os.path.dirname(socket_path)
+            if not os.path.isdir(path):
+                try_to_create_directory(path)
+        else:
+            raise Exception("Socket file {} exists!".format(socket_path))
+
     def start_redis(self):
         """Start the Redis servers."""
         assert self._redis_address is None
@@ -155,6 +173,7 @@ class Node(object):
         self._plasma_store_socket_name = (
             self._ray_params.plasma_store_socket_name
             or get_object_store_socket_name())
+        self.prepare_socket_file(self._plasma_store_socket_name)
         stdout_file, stderr_file = (new_plasma_store_log_file(
             self._ray_params.redirect_output))
         process_info = ray.services.start_plasma_store(
@@ -186,6 +205,7 @@ class Node(object):
         # If the user specified a socket name, use it.
         self._raylet_socket_name = (self._ray_params.raylet_socket_name
                                     or get_raylet_socket_name())
+        self.prepare_socket_file(self._raylet_socket_name)
         stdout_file, stderr_file = new_raylet_log_file(
             redirect_output=self._ray_params.redirect_worker_output)
         process_info = ray.services.start_raylet(
diff --git a/test/runtest.py b/test/runtest.py
index e081b596d..4bf1f1b33 100644
--- a/test/runtest.py
+++ b/test/runtest.py
@@ -2607,3 +2607,11 @@ def test_pandas_parquet_serialization():
     pd.DataFrame({"col1": [0, 1], "col2": [0, 1]}).to_parquet(filename)
     # Clean up
     shutil.rmtree(tempdir)
+
+
+def test_socket_dir_not_existing(shutdown_only):
+    random_name = ray.ObjectID(_random_string()).hex()
+    temp_raylet_socket_dir = "/tmp/ray/tests/{}".format(random_name)
+    temp_raylet_socket_name = os.path.join(temp_raylet_socket_dir,
+                                           "raylet_socket")
+    ray.init(num_cpus=1, raylet_socket_name=temp_raylet_socket_name)