Fix tempfile issues (#4605)

python/ray/tests/test_tempfile.py

@@ -9,12 +9,6 @@ import pytest
 import ray
 from ray.tests.cluster_utils import Cluster
 
-# Py2 compatibility
-try:
-    FileNotFoundError
-except NameError:
-    FileNotFoundError = OSError
-
 
 def test_conn_cluster():
     # plasma_store_socket_name
@@ -45,13 +39,25 @@ def test_conn_cluster():
 
 
 def test_tempdir():
+    shutil.rmtree("/tmp/ray", ignore_errors=True)
     ray.init(temp_dir="/tmp/i_am_a_temp_dir")
     assert os.path.exists(
         "/tmp/i_am_a_temp_dir"), "Specified temp dir not found."
+    assert not os.path.exists("/tmp/ray"), "Default temp dir should not exist."
     ray.shutdown()
     shutil.rmtree("/tmp/i_am_a_temp_dir", ignore_errors=True)
 
 
+def test_tempdir_commandline():
+    shutil.rmtree("/tmp/ray", ignore_errors=True)
+    os.system("ray start --head --temp-dir=/tmp/i_am_a_temp_dir2")
+    assert os.path.exists(
+        "/tmp/i_am_a_temp_dir2"), "Specified temp dir not found."
+    assert not os.path.exists("/tmp/ray"), "Default temp dir should not exist."
+    os.system("ray stop")
+    shutil.rmtree("/tmp/i_am_a_temp_dir2", ignore_errors=True)
+
+
 def test_raylet_socket_name():
     ray.init(raylet_socket_name="/tmp/i_am_a_temp_socket")
     assert os.path.exists(
@@ -59,7 +65,7 @@ def test_raylet_socket_name():
     ray.shutdown()
     try:
         os.remove("/tmp/i_am_a_temp_socket")
-    except FileNotFoundError:
+    except OSError:
         pass  # It could have been removed by Ray.
     cluster = Cluster(True)
     cluster.add_node(raylet_socket_name="/tmp/i_am_a_temp_socket_2")
@@ -68,7 +74,7 @@ def test_raylet_socket_name():
     cluster.shutdown()
     try:
         os.remove("/tmp/i_am_a_temp_socket_2")
-    except FileNotFoundError:
+    except OSError:
         pass  # It could have been removed by Ray.
 
 
@@ -79,7 +85,7 @@ def test_temp_plasma_store_socket():
     ray.shutdown()
     try:
         os.remove("/tmp/i_am_a_temp_socket")
-    except FileNotFoundError:
+    except OSError:
         pass  # It could have been removed by Ray.
     cluster = Cluster(True)
     cluster.add_node(plasma_store_socket_name="/tmp/i_am_a_temp_socket_2")
@@ -88,14 +94,14 @@ def test_temp_plasma_store_socket():
     cluster.shutdown()
     try:
         os.remove("/tmp/i_am_a_temp_socket_2")
-    except FileNotFoundError:
+    except OSError:
         pass  # It could have been removed by Ray.
 
 
 def test_raylet_tempfiles():
     ray.init(num_cpus=0)
     node = ray.worker._global_node
-    top_levels = set(os.listdir(node.get_temp_dir_path()))
+    top_levels = set(os.listdir(node.get_session_dir_path()))
     assert top_levels.issuperset({"sockets", "logs"})
     log_files = set(os.listdir(node.get_logs_dir_path()))
     assert log_files.issuperset({
@@ -110,7 +116,7 @@ def test_raylet_tempfiles():
 
     ray.init(num_cpus=2)
     node = ray.worker._global_node
-    top_levels = set(os.listdir(node.get_temp_dir_path()))
+    top_levels = set(os.listdir(node.get_session_dir_path()))
     assert top_levels.issuperset({"sockets", "logs"})
     time.sleep(3)  # wait workers to start
     log_files = set(os.listdir(node.get_logs_dir_path()))
@@ -128,3 +134,20 @@ def test_raylet_tempfiles():
     socket_files = set(os.listdir(node.get_sockets_dir_path()))
     assert socket_files == {"plasma_store", "raylet"}
     ray.shutdown()
+
+
+def test_tempdir_privilege():
+    os.chmod("/tmp/ray", 0o000)
+    ray.init(num_cpus=1)
+    session_dir = ray.worker._global_node.get_session_dir_path()
+    assert os.path.exists(session_dir), "Specified socket path not found."
+    ray.shutdown()
+
+
+def test_session_dir_uniqueness():
+    session_dirs = set()
+    for _ in range(3):
+        ray.init(num_cpus=1)
+        session_dirs.add(ray.worker._global_node.get_session_dir_path)
+        ray.shutdown()
+    assert len(session_dirs) == 3