Address comments pt2

python/ray/_private/test_utils.py

@@ -9,12 +9,12 @@
 import sys
 import time
 import timeit
-import socket
 import math
 import traceback
 from typing import Optional, Any, List, Dict
 from contextlib import redirect_stdout, redirect_stderr
 import yaml
+import socket
 import pytest
 import tempfile
 

python/ray/_private/tls_utils.py

@@ -59,7 +59,7 @@ def generate_self_signed_tls_certs():
 
 
 def add_port_to_grpc_server(server, address):
- if os.environ.get("RAY_USE_TLS", "0") == "1":
+ if os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"):
  server_cert_chain, private_key, ca_cert = load_certs_from_env()
  credentials = grpc.ssl_server_credentials(
  [(private_key, server_cert_chain)],
@@ -71,7 +71,14 @@ def add_port_to_grpc_server(server, address):
 
 
 def load_certs_from_env():
- if os.environ.get("RAY_USE_TLS", "0") == "1":
+ if os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"):
+ if ("RAY_TLS_SERVER_CERT" not in os.environ) or \
+ ("RAY_TLS_SERVER_KEY" not in os.environ):
+ raise RuntimeError(
+ "If the environment variable RAY_USE_TLS is set to true"
+ "then both RAY_TLS_SERVER_CERT and RAY_TLS_SERVER_KEY must "
+ "also be set.")
+
  with open(os.environ["RAY_TLS_SERVER_CERT"], "rb") as f:
  server_cert_chain = f.read()
  with open(os.environ["RAY_TLS_SERVER_KEY"], "rb") as f:

python/ray/_private/utils.py

@@ -1116,7 +1116,7 @@ def init_grpc_channel(address: str,
  options: Optional[Sequence[Tuple[str, Any]]] = None,
  asynchronous: bool = False):
  grpc_module = aiogrpc if asynchronous else grpc
- if os.environ.get("RAY_USE_TLS", "0") == "1":
+ if os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"):
  server_cert_chain, private_key, ca_cert = load_certs_from_env()
  credentials = grpc.ssl_channel_credentials(
  certificate_chain=server_cert_chain,

python/ray/tests/conftest.py

@@ -89,17 +89,15 @@ def ray_start_regular_shared(request):
 
 
 @pytest.fixture(
- scope="module",
- params=[{
+ scope="module", params=[{
  "local_mode": True
  }, {
  "local_mode": False
  }])
 def ray_start_shared_local_modes(request):
  param = getattr(request, "param", {})
- use_tls = param.pop("use_tls", False)
  with _ray_start(**param) as res:
-  yield res
+ yield res
 
 
 @pytest.fixture

python/ray/tests/test_client_builder.py

@@ -55,7 +55,7 @@ def test_namespace():
  put in the same namespace.
 
  This test checks that:
- RayConfig::instance().RAY_USE_TLS() * When two drivers don't specify a namespace, they are placed in different
+ * When two drivers don't specify a namespace, they are placed in different
  anonymous namespaces.
  * When two drivers specify a namespace, they collide.
  * The namespace name (as provided by the runtime context) is correct.

python/ray/tests/test_tls_auth.py

@@ -2,18 +2,37 @@
 import logging
 import os
 import sys
+import subprocess
 
 import pytest
 
-import ray
+from ray._private.test_utils import run_string_as_driver
 
 logger = logging.getLogger(__name__)
 
 
 @pytest.mark.parametrize("use_tls", [True], indirect=True)
-def test_put_get_with_tls(shutdown_only, use_tls):
- ray.init(num_cpus=0)
+def test_init_with_tls(use_tls):
+ # Run as a new process to pick up environment variables set
+ # in the use_tls fixture
+ run_string_as_driver(
+ """
+import ray
+try:
+ ray.init()
+finally:
+ ray.shutdown()
+ """,
+ env=os.environ)
 
+
+@pytest.mark.parametrize("use_tls", [True], indirect=True)
+def test_put_get_with_tls(use_tls):
+ run_string_as_driver(
+ """
+import ray
+ray.init()
+try:
  for i in range(100):
  value_before = i * 10**6
  object_ref = ray.put(value_before)
@@ -37,45 +56,82 @@ def test_put_get_with_tls(shutdown_only, use_tls):
  object_ref = ray.put(value_before)
  value_after = ray.get(object_ref)
  assert value_before == value_after
+finally:
+ ray.shutdown()
+ """,
+ env=os.environ)
 
 
-@pytest.mark.parametrize("use_tls", [True], indirect=True)
-def test_submit_with_tls(shutdown_only, use_tls):
- ray.init(num_cpus=2, num_gpus=1, resources={"Custom": 1})
+@pytest.mark.parametrize("use_tls", [True], indirect=True, scope="module")
+def test_submit_with_tls(use_tls):
+ run_string_as_driver(
+ """
+import ray
+ray.init(num_cpus=2, num_gpus=1, resources={"Custom": 1})
 
- @ray.remote
- def f(n):
-  return list(range(n))
+@ray.remote
+def f(n):
+ return list(range(n))
 
- id1, id2, id3 = f._remote(args=[3], num_returns=3)
- assert ray.get([id1, id2, id3]) == [0, 1, 2]
+id1, id2, id3 = f._remote(args=[3], num_returns=3)
+assert ray.get([id1, id2, id3]) == [0, 1, 2]
 
- @ray.remote
- class Actor:
-  def __init__(self, x, y=0):
-  self.x = x
-  self.y = y
+@ray.remote
+class Actor:
+ def __init__(self, x, y=0):
+ self.x = x
+ self.y = y
 
-  def method(self, a, b=0):
-  return self.x, self.y, a, b
+ def method(self, a, b=0):
+ return self.x, self.y, a, b
 
- a = Actor._remote(
-  args=[0], kwargs={"y": 1}, num_gpus=1, resources={"Custom": 1})
+a = Actor._remote(
+ args=[0], kwargs={"y": 1}, num_gpus=1, resources={"Custom": 1})
 
- id1, id2, id3, id4 = a.method._remote(
- args=["test"], kwargs={"b": 2}, num_returns=4)
- assert ray.get([id1, id2, id3, id4]) == [0, 1, "test", 2]
+id1, id2, id3, id4 = a.method._remote(
+ args=["test"], kwargs={"b": 2}, num_returns=4)
+assert ray.get([id1, id2, id3, id4]) == [0, 1, "test", 2]
+ """,
+ env=os.environ)
 
 
 @pytest.mark.skipif(
  sys.platform == "darwin",
  reason=("Cryptography doesn't install in Mac build pipeline"))
 @pytest.mark.parametrize("use_tls", [True], indirect=True)
-def test_client_connect_to_tls_server(use_tls, init_and_serve):
- from ray.util.client import ray as ray_client
- os.environ["RAY_USE_TLS"] = "0"
- with pytest.raises(ConnectionError):
- ray_client.connect("localhost:50051")
-
- os.environ["RAY_USE_TLS"] = "1"
- ray_client.connect("localhost:50051")
+def test_client_connect_to_tls_server(use_tls, call_ray_start):
+ for k, v in os.environ.items():
+ if k.startswith("RAY_"):
+ print("export {}={}".format(k, v))
+
+ tls_env = os.environ.copy(
+ ) # use_tls fixture sets TLS environment variables
+ without_tls_env = {}
+
+ # Attempt to connect without TLS
+ try:
+ out = run_string_as_driver(
+ """
+from ray.util.client import ray as ray_client
+ray_client.connect("localhost:10001")
+ """,
+ env=without_tls_env)
+ except subprocess.CalledProcessError as e:
+ assert "ConnectionError" in e.output.decode("utf-8")
+
+ # Attempt to connect with TLS
+ out = run_string_as_driver(
+ """
+import ray
+from ray.util.client import ray as ray_client
+ray_client.connect("localhost:10001")
+print(ray.is_initialized())
+ """,
+ env=tls_env)
+ assert out == "True\n"
+
+
+if __name__ == "__main__":
+ import pytest
+ import sys
+ sys.exit(pytest.main(["-v", __file__]))

python/ray/util/client/worker.py

@@ -101,7 +101,8 @@ def __init__(
  self.server = None
  self._conn_state = grpc.ChannelConnectivity.IDLE
  self._converted: Dict[str, ClientStub] = {}
- self._secure = secure or os.environ.get("RAY_USE_TLS", "0") == "1"
+ self._secure = secure or os.environ.get("RAY_USE_TLS",
+ "0").lower() in ("1", "true")
  self._conn_str = conn_str
  self._connection_retries = connection_retries
 
@@ -160,7 +161,7 @@ def _connect_channel(self, reconnecting=False) -> None:
  if self._secure:
  if self._credentials is not None:
  credentials = self._credentials
- elif os.environ.get("RAY_USE_TLS", "0") == "1":
+ elif os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"):
  server_cert_chain, private_key, ca_cert = ray._private.utils \
  .load_certs_from_env()
  credentials = grpc.ssl_channel_credentials(
@@ -362,8 +363,8 @@ def get(self, vals, *, timeout: Optional[float] = None) -> Any:
  logger.debug("Internal retry for get {}".format(to_get))
  if len(to_get) != len(res):
  raise Exception(
- "Mismatched number of items in request ({}) and response ({})".
- format(len(to_get), len(res)))
+ "Mismatched number of items in request ({}) and response ({})"
+ .format(len(to_get), len(res)))
  if isinstance(vals, ClientObjectRef):
  res = res[0]
  return res

src/ray/common/ray_config_def.h

@@ -490,9 +490,20 @@ RAY_CONFIG(bool, scheduler_avoid_gpu_nodes, true)
 RAY_CONFIG(bool, runtime_env_skip_local_gc, false)
 
 /// Whether or not use TLS.
-RAY_CONFIG(int64_t, USE_TLS, 0)
+RAY_CONFIG(bool, USE_TLS,
+ std::getenv("RAY_USE_TLS") != nullptr &&
+ (std::getenv("RAY_USE_TLS") == std::string("true") ||
+ std::getenv("RAY_USE_TLS") == std::string("1")))
 
 /// Location of TLS credentials
-RAY_CONFIG(std::string, TLS_SERVER_CERT, "")
-RAY_CONFIG(std::string, TLS_SERVER_KEY, "")
-RAY_CONFIG(std::string, TLS_CA_CERT, "")
+RAY_CONFIG(std::string, TLS_SERVER_CERT,
+ std::getenv("RAY_TLS_SERVER_CERT") != nullptr
+ ? std::getenv("RAY_TLS_SERVER_CERT")
+ : "")
+RAY_CONFIG(std::string, TLS_SERVER_KEY,
+ std::getenv("RAY_TLS_SERVER_KEY") != nullptr
+ ? std::getenv("RAY_TLS_SERVER_KEY")
+ : "")
+RAY_CONFIG(std::string, TLS_CA_CERT,
+ std::getenv("RAY_TLS_CA_CERT") != nullptr ? std::getenv("RAY_TLS_CA_CERT")
+ : "")

src/ray/rpc/common.cc

@@ -12,11 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "ray/rpc/common.h"
+
 #include <fstream>
 #include <sstream>
 
-#include "ray/rpc/common.h"
-
 namespace ray::rpc {
 
 std::string ReadCert(const std::string &cert_filepath) {
@@ -26,4 +26,4 @@ std::string ReadCert(const std::string &cert_filepath) {
  return buffer.str();
 };
 
-} // namespace rpc::ray
+} // namespace ray::rpc

src/ray/rpc/common.h

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <string>
+
 namespace ray::rpc {
 
 // Utility to read cert file from a particular location

src/ray/rpc/grpc_client.h

@@ -45,7 +45,7 @@ template <class GrpcService>
 class GrpcClient {
  public:
  GrpcClient(const std::string &address, const int port, ClientCallManager &call_manager,
- bool use_tls = true)
+ bool use_tls = false)
  : client_call_manager_(call_manager), use_tls_(use_tls) {
  grpc::ChannelArguments argument;
  // Disable http proxy since it disrupts local connections. TODO(ekl) we should make
@@ -107,7 +107,8 @@ class GrpcClient {
  const std::string &address, int port) {
  std::shared_ptr<grpc::Channel> channel;
  if (::RayConfig::instance().USE_TLS()) {
- std::string server_cert_file = std::string(::RayConfig::instance().TLS_SERVER_CERT());
+ std::string server_cert_file =
+ std::string(::RayConfig::instance().TLS_SERVER_CERT());
  std::string server_key_file = std::string(::RayConfig::instance().TLS_SERVER_KEY());
  std::string root_cert_file = std::string(::RayConfig::instance().TLS_CA_CERT());
  std::string server_cert_chain = ReadCert(server_cert_file);
@@ -127,7 +128,6 @@ class GrpcClient {
  }
  return channel;
  };
-
 };
 
 } // namespace rpc