wehub-resource-sync
475 lines
15 KiB
Python
475 lines
15 KiB
Python
"""Integration tests for token-based authentication in Ray (part 1).
|
|
|
|
Cluster lifecycle, connection, and end-to-end operation tests.
|
|
"""
|
|
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray._private.test_utils import client_test_enabled
|
|
|
|
try:
|
|
from ray._raylet import AuthenticationTokenLoader
|
|
|
|
_RAYLET_AVAILABLE = True
|
|
except ImportError:
|
|
_RAYLET_AVAILABLE = False
|
|
AuthenticationTokenLoader = None
|
|
|
|
from ray._private.authentication_test_utils import (
|
|
clear_auth_token_sources,
|
|
reset_auth_token_state,
|
|
set_auth_mode,
|
|
set_env_auth_token,
|
|
)
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not _RAYLET_AVAILABLE,
|
|
reason="Authentication tests require ray._raylet (not available in minimal installs)",
|
|
)
|
|
|
|
|
|
def _run_ray_start_and_verify_status(
|
|
args: list, env: dict, expect_success: bool = True, timeout: int = 30
|
|
) -> subprocess.CompletedProcess:
|
|
"""Helper to run ray start command with proper error handling."""
|
|
result = subprocess.run(
|
|
["ray", "start"] + args,
|
|
env={"RAY_ENABLE_WINDOWS_OR_OSX_CLUSTER": "1", **env},
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=timeout,
|
|
)
|
|
|
|
if expect_success:
|
|
assert result.returncode == 0, (
|
|
f"ray start should have succeeded. "
|
|
f"stdout: {result.stdout}, stderr: {result.stderr}"
|
|
)
|
|
else:
|
|
assert result.returncode != 0, (
|
|
f"ray start should have failed but succeeded. "
|
|
f"stdout: {result.stdout}, stderr: {result.stderr}"
|
|
)
|
|
# Check that error message mentions token
|
|
error_output = result.stdout + result.stderr
|
|
assert (
|
|
"authentication token" in error_output.lower()
|
|
or "token" in error_output.lower()
|
|
), f"Error message should mention token. Got: {error_output}"
|
|
|
|
return result
|
|
|
|
|
|
def _cleanup_ray_start(env: Optional[dict] = None):
|
|
"""Helper to clean up ray start processes."""
|
|
# Ensure any ray.init() connection is closed first
|
|
if ray.is_initialized():
|
|
ray.shutdown()
|
|
|
|
# Stop with a longer timeout
|
|
subprocess.run(
|
|
["ray", "stop", "--force"],
|
|
env=env,
|
|
capture_output=True,
|
|
timeout=60, # Increased timeout for flaky cleanup
|
|
check=False, # Don't raise on non-zero exit
|
|
)
|
|
|
|
# Wait for ray processes to actually stop
|
|
def ray_stopped():
|
|
result = subprocess.run(
|
|
["ray", "status"],
|
|
capture_output=True,
|
|
check=False,
|
|
)
|
|
# ray status returns non-zero when no cluster is running
|
|
return result.returncode != 0
|
|
|
|
try:
|
|
wait_for_condition(ray_stopped, timeout=10, retry_interval_ms=500)
|
|
except Exception:
|
|
# Best effort - don't fail the test if we can't verify it stopped
|
|
pass
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _auto_clean_token_sources(clean_token_sources):
|
|
"""Opt in to the shared clean_token_sources fixture for every test."""
|
|
yield
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
client_test_enabled(),
|
|
reason="This test is for starting a new local cluster, not compatible with client mode",
|
|
)
|
|
def test_local_cluster_generates_token():
|
|
"""Test ray.init() generates token for local cluster when auth_mode=token is set."""
|
|
# Ensure no token exists
|
|
default_token_path = Path.home() / ".ray" / "auth_token"
|
|
assert (
|
|
not default_token_path.exists()
|
|
), f"Token file already exists at {default_token_path}"
|
|
|
|
# Enable token auth via environment variable
|
|
set_auth_mode("token")
|
|
reset_auth_token_state()
|
|
|
|
# Initialize Ray with token auth
|
|
ray.init()
|
|
|
|
try:
|
|
# Verify token file was created
|
|
assert default_token_path.exists(), (
|
|
f"Token file was not created at {default_token_path}. "
|
|
f"HOME={os.environ.get('HOME')}, "
|
|
f"Files in {default_token_path.parent}: {list(default_token_path.parent.iterdir()) if default_token_path.parent.exists() else 'directory does not exist'}"
|
|
)
|
|
token = default_token_path.read_text().strip()
|
|
assert len(token) == 64
|
|
assert all(c in "0123456789abcdef" for c in token)
|
|
|
|
# Verify cluster is working
|
|
assert ray.is_initialized()
|
|
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
|
|
def test_connect_without_token_raises_error(setup_cluster_with_token_auth):
|
|
"""Test ray.init(address=...) without token fails when auth_mode=token is set."""
|
|
cluster_info = setup_cluster_with_token_auth
|
|
cluster = cluster_info["cluster"]
|
|
|
|
# Disconnect the current driver session and drop token state before retrying.
|
|
ray.shutdown()
|
|
set_auth_mode("disabled")
|
|
clear_auth_token_sources(remove_default=True)
|
|
reset_auth_token_state()
|
|
|
|
# Ensure no token exists
|
|
token_loader = AuthenticationTokenLoader.instance()
|
|
assert not token_loader.has_token()
|
|
|
|
# Try to connect to the cluster without a token - should raise RuntimeError
|
|
with pytest.raises(ConnectionError):
|
|
ray.init(address=cluster.address)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"token,expected_status",
|
|
[
|
|
(None, 401), # No token -> Unauthorized
|
|
("wrong_token", 403), # Wrong token -> Forbidden
|
|
],
|
|
ids=["no_token", "wrong_token"],
|
|
)
|
|
def test_state_api_auth_failure(token, expected_status, setup_cluster_with_token_auth):
|
|
"""Test that state API calls fail with missing or incorrect token."""
|
|
import requests
|
|
|
|
cluster_info = setup_cluster_with_token_auth
|
|
dashboard_url = cluster_info["dashboard_url"]
|
|
|
|
# Make direct HTTP request to state API endpoint
|
|
headers = {}
|
|
if token is not None:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
|
|
response = requests.get(f"{dashboard_url}/api/v0/actors", headers=headers)
|
|
|
|
assert response.status_code == expected_status, (
|
|
f"State API should return {expected_status}, got {response.status_code}: "
|
|
f"{response.text}"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("tokens_match", [True, False])
|
|
def test_cluster_token_authentication(tokens_match, setup_cluster_with_token_auth):
|
|
"""Test cluster authentication with matching and non-matching tokens."""
|
|
cluster_info = setup_cluster_with_token_auth
|
|
cluster = cluster_info["cluster"]
|
|
cluster_token = cluster_info["token"]
|
|
|
|
# Reconfigure the driver token state to simulate fresh connections.
|
|
ray.shutdown()
|
|
set_auth_mode("token")
|
|
|
|
if tokens_match:
|
|
client_token = cluster_token # Same token - should succeed
|
|
else:
|
|
client_token = "b" * 64 # Different token - should fail
|
|
|
|
set_env_auth_token(client_token)
|
|
reset_auth_token_state()
|
|
|
|
if tokens_match:
|
|
# Should succeed - test gRPC calls work
|
|
ray.init(address=cluster.address)
|
|
|
|
obj_ref = ray.put("test_data")
|
|
result = ray.get(obj_ref)
|
|
assert result == "test_data"
|
|
|
|
@ray.remote
|
|
def test_func():
|
|
return "success"
|
|
|
|
result = ray.get(test_func.remote())
|
|
assert result == "success"
|
|
|
|
ray.shutdown()
|
|
|
|
else:
|
|
# Should fail - connection or gRPC calls should fail
|
|
with pytest.raises((ConnectionError, RuntimeError)):
|
|
ray.init(address=cluster.address)
|
|
try:
|
|
ray.put("test")
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
client_test_enabled(),
|
|
reason="Uses subprocess ray start, not compatible with client mode",
|
|
)
|
|
@pytest.mark.parametrize("is_head", [True, False])
|
|
def test_ray_start_without_token_raises_error(is_head, request):
|
|
"""Test that ray start fails when auth_mode=token but no token exists."""
|
|
# Set up environment with token auth enabled but no token
|
|
env = os.environ.copy()
|
|
env["RAY_AUTH_MODE"] = "token"
|
|
env.pop("RAY_AUTH_TOKEN", None)
|
|
env.pop("RAY_AUTH_TOKEN_PATH", None)
|
|
|
|
# Ensure no default token file exists (already cleaned by fixture)
|
|
default_token_path = Path.home() / ".ray" / "auth_token"
|
|
assert not default_token_path.exists()
|
|
|
|
# When specifying an address, we need a head node to connect to
|
|
cluster_info = None
|
|
if not is_head:
|
|
cluster_info = request.getfixturevalue("setup_cluster_with_token_auth")
|
|
cluster = cluster_info["cluster"]
|
|
ray.shutdown()
|
|
|
|
# Prepare arguments
|
|
if is_head:
|
|
args = ["--head", "--port=0"]
|
|
else:
|
|
args = [f"--address={cluster.address}"]
|
|
|
|
# Try to start node - should fail
|
|
_run_ray_start_and_verify_status(args, env, expect_success=False)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
client_test_enabled(),
|
|
reason="Uses subprocess ray start, not compatible with client mode",
|
|
)
|
|
def test_ray_start_head_with_token_succeeds():
|
|
"""Test that ray start --head succeeds when token auth is enabled with a valid token."""
|
|
# Set up environment with token auth and a valid token
|
|
test_token = "a" * 64
|
|
env = os.environ.copy()
|
|
env["RAY_AUTH_TOKEN"] = test_token
|
|
env["RAY_AUTH_MODE"] = "token"
|
|
|
|
try:
|
|
# Start head node - should succeed
|
|
_run_ray_start_and_verify_status(
|
|
["--head", "--port=0"], env, expect_success=True
|
|
)
|
|
|
|
# Verify we can connect to the cluster with ray.init()
|
|
set_env_auth_token(test_token)
|
|
set_auth_mode("token")
|
|
reset_auth_token_state()
|
|
|
|
# Wait for cluster to be ready
|
|
def cluster_ready():
|
|
try:
|
|
ray.init(address="auto")
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
wait_for_condition(cluster_ready, timeout=10)
|
|
assert ray.is_initialized()
|
|
|
|
# Test basic operations work
|
|
@ray.remote
|
|
def test_func():
|
|
return "success"
|
|
|
|
result = ray.get(test_func.remote())
|
|
assert result == "success"
|
|
|
|
finally:
|
|
# Cleanup handles ray.shutdown() internally
|
|
_cleanup_ray_start(env)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
client_test_enabled(),
|
|
reason="Uses subprocess ray start, not compatible with client mode",
|
|
)
|
|
@pytest.mark.parametrize("token_match", ["correct", "incorrect"])
|
|
def test_ray_start_address_with_token(token_match, setup_cluster_with_token_auth):
|
|
"""Test ray start --address=... with correct or incorrect token."""
|
|
cluster_info = setup_cluster_with_token_auth
|
|
cluster = cluster_info["cluster"]
|
|
cluster_token = cluster_info["token"]
|
|
|
|
# Reset the driver connection to reuse the fixture-backed cluster.
|
|
ray.shutdown()
|
|
set_auth_mode("token")
|
|
|
|
# Set up environment for worker
|
|
env = os.environ.copy()
|
|
env["RAY_AUTH_MODE"] = "token"
|
|
|
|
if token_match == "correct":
|
|
env["RAY_AUTH_TOKEN"] = cluster_token
|
|
expect_success = True
|
|
else:
|
|
env["RAY_AUTH_TOKEN"] = "b" * 64
|
|
expect_success = False
|
|
|
|
# Start worker node
|
|
_run_ray_start_and_verify_status(
|
|
[f"--address={cluster.address}", "--num-cpus=1"],
|
|
env,
|
|
expect_success=expect_success,
|
|
)
|
|
|
|
if token_match == "correct":
|
|
try:
|
|
# Connect and verify the cluster has 2 nodes (head + worker)
|
|
set_env_auth_token(cluster_token)
|
|
reset_auth_token_state()
|
|
ray.init(address=cluster.address)
|
|
|
|
def worker_joined():
|
|
return len(ray.nodes()) >= 2
|
|
|
|
wait_for_condition(worker_joined, timeout=10)
|
|
|
|
nodes = ray.nodes()
|
|
assert (
|
|
len(nodes) >= 2
|
|
), f"Expected at least 2 nodes, got {len(nodes)}: {nodes}"
|
|
|
|
finally:
|
|
if ray.is_initialized():
|
|
ray.shutdown()
|
|
_cleanup_ray_start(env)
|
|
|
|
|
|
def test_e2e_operations_with_token_auth(setup_cluster_with_token_auth):
|
|
"""Test that e2e operations work with token authentication enabled.
|
|
|
|
This verifies that with token auth enabled:
|
|
1. Tasks execute successfully
|
|
2. Actors can be created and called
|
|
3. State API works (list_nodes, list_actors, list_tasks)
|
|
4. Job submission works
|
|
"""
|
|
cluster_info = setup_cluster_with_token_auth
|
|
|
|
# Test 1: Submit a simple task
|
|
@ray.remote
|
|
def simple_task(x):
|
|
return x + 1
|
|
|
|
result = ray.get(simple_task.remote(41))
|
|
assert result == 42, f"Task should return 42, got {result}"
|
|
|
|
# Test 2: Create and use an actor
|
|
@ray.remote
|
|
class SimpleActor:
|
|
def __init__(self):
|
|
self.value = 0
|
|
|
|
def increment(self):
|
|
self.value += 1
|
|
return self.value
|
|
|
|
actor = SimpleActor.remote()
|
|
result = ray.get(actor.increment.remote())
|
|
assert result == 1, f"Actor method should return 1, got {result}"
|
|
|
|
# Test 3: State API operations (uses HTTP with auth headers)
|
|
from ray.util.state import list_actors, list_nodes, list_tasks
|
|
|
|
# List nodes - should include at least the head node
|
|
wait_for_condition(lambda: len(list_nodes()) >= 1)
|
|
|
|
# List actors - should include our SimpleActor
|
|
def check_actors():
|
|
actors = list_actors()
|
|
if len(actors) < 1:
|
|
return False
|
|
return "SimpleActor" in actors[0].class_name
|
|
|
|
wait_for_condition(check_actors)
|
|
|
|
# List tasks - should include completed tasks
|
|
wait_for_condition(lambda: len(list_tasks()) >= 1)
|
|
|
|
# Test 4: Submit a job and wait for completion
|
|
from ray.job_submission import JobSubmissionClient
|
|
|
|
# Create job submission client (uses HTTP with auth headers)
|
|
client = JobSubmissionClient(address=cluster_info["dashboard_url"])
|
|
|
|
# Submit a simple job
|
|
job_id = client.submit_job(
|
|
entrypoint="echo 'Hello from job'",
|
|
)
|
|
|
|
# Wait for job to complete
|
|
def job_finished():
|
|
status = client.get_job_status(job_id)
|
|
return status in ["SUCCEEDED", "FAILED", "STOPPED"]
|
|
|
|
wait_for_condition(job_finished, timeout=30)
|
|
|
|
final_status = client.get_job_status(job_id)
|
|
assert (
|
|
final_status == "SUCCEEDED"
|
|
), f"Job should succeed, got status: {final_status}"
|
|
|
|
|
|
def test_logs_api_with_token_auth(setup_cluster_with_token_auth):
|
|
"""Test that log APIs work with token authentication enabled."""
|
|
from ray.util.state import get_log, list_logs
|
|
|
|
# Get node ID for log queries
|
|
node_id = ray.nodes()[0]["NodeID"]
|
|
|
|
# Test list_logs() with valid auth
|
|
logs = list_logs(node_id=node_id)
|
|
assert isinstance(logs, dict), f"list_logs should return a dict, got {type(logs)}"
|
|
|
|
# Test get_log() with valid auth (fetch raylet.out which will always exist)
|
|
chunks_received = 0
|
|
for chunk in get_log(filename="raylet.out", node_id=node_id, tail=10):
|
|
assert isinstance(chunk, str), f"get_log chunk should be str, got {type(chunk)}"
|
|
chunks_received += 1
|
|
break
|
|
|
|
assert chunks_received > 0, "Should have received at least one log chunk"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-vv", __file__]))
|