wehub-resource-sync
144 lines
2.9 KiB
Python
144 lines
2.9 KiB
Python
import logging
|
|
|
|
import boto3
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import simulate_s3_bucket
|
|
from ray.cluster_utils import Cluster
|
|
|
|
# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
|
|
from ray.tests.conftest import (
|
|
propagate_logs, # noqa
|
|
pytest_runtest_makereport, # noqa
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_4_cpus():
|
|
address_info = ray.init(num_cpus=4)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_1_cpu_1_gpu():
|
|
address_info = ray.init(num_cpus=1, num_gpus=1)
|
|
yield address_info
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_4_cpus_2_gpus():
|
|
address_info = ray.init(num_cpus=4, num_gpus=2)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def shutdown_only():
|
|
yield None
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_2_node_2_gpu():
|
|
cluster = Cluster()
|
|
for _ in range(2):
|
|
cluster.add_node(num_cpus=4, num_gpus=2)
|
|
|
|
ray.init(address=cluster.address)
|
|
|
|
yield
|
|
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_2_node_2_neuron_cores():
|
|
cluster = Cluster()
|
|
for _ in range(2):
|
|
cluster.add_node(num_cpus=4, resources={"neuron_cores": 2})
|
|
|
|
ray.init(address=cluster.address)
|
|
|
|
yield
|
|
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_2_cpus():
|
|
address_info = ray.init(num_cpus=2)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_4_node_4_cpu():
|
|
cluster = Cluster()
|
|
for _ in range(4):
|
|
cluster.add_node(num_cpus=4)
|
|
|
|
ray.init(address=cluster.address)
|
|
|
|
yield
|
|
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_2_node_4_gpu():
|
|
cluster = Cluster()
|
|
for _ in range(2):
|
|
cluster.add_node(num_cpus=2, num_gpus=4)
|
|
|
|
ray.init(address=cluster.address)
|
|
|
|
yield
|
|
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_2_node_2_cpu():
|
|
cluster = Cluster()
|
|
for _ in range(2):
|
|
cluster.add_node(num_cpus=2)
|
|
|
|
ray.init(address=cluster.address)
|
|
|
|
yield
|
|
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_s3_bucket_uri():
|
|
from ray.air._internal.uri_utils import URI
|
|
|
|
port = 5002
|
|
region = "us-west-2"
|
|
with simulate_s3_bucket(port=port, region=region) as s3_uri:
|
|
s3 = boto3.client(
|
|
"s3", region_name=region, endpoint_url=f"http://localhost:{port}"
|
|
)
|
|
# Bucket name will be autogenerated/unique per test
|
|
bucket_name = URI(s3_uri).name
|
|
s3.create_bucket(
|
|
Bucket=bucket_name,
|
|
CreateBucketConfiguration={"LocationConstraint": region},
|
|
)
|
|
# Disable server HTTP request logging
|
|
logging.getLogger("werkzeug").setLevel(logging.WARNING)
|
|
yield s3_uri
|
|
logging.getLogger("werkzeug").setLevel(logging.INFO)
|