wehub-resource-sync
130 lines
4.0 KiB
Python
130 lines
4.0 KiB
Python
import sys
|
|
import tempfile
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
import ray._common.usage.usage_lib as ray_usage_lib
|
|
from ray._common.test_utils import TelemetryCallsite, check_library_usage_telemetry
|
|
from ray.train import Checkpoint
|
|
from ray.train.v2.api.config import ScalingConfig
|
|
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
|
|
from ray.train.v2.api.report_config import CheckpointUploadMode
|
|
from ray.train.v2.api.validation_config import ValidationConfig
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_record(monkeypatch):
|
|
import ray._common.usage.usage_lib
|
|
import ray.air._internal.usage
|
|
|
|
recorded = {}
|
|
|
|
def mock_record_extra_usage_tag(key: ray_usage_lib.TagKey, value: str):
|
|
recorded[key] = value
|
|
|
|
monkeypatch.setattr(
|
|
ray.air._internal.usage,
|
|
"record_extra_usage_tag",
|
|
mock_record_extra_usage_tag,
|
|
)
|
|
monkeypatch.setattr(
|
|
ray._common.usage.usage_lib,
|
|
"record_extra_usage_tag",
|
|
mock_record_extra_usage_tag,
|
|
)
|
|
yield recorded
|
|
|
|
|
|
@pytest.fixture
|
|
def reset_usage_lib():
|
|
yield
|
|
ray.shutdown()
|
|
ray_usage_lib.reset_global_state()
|
|
|
|
|
|
@pytest.mark.parametrize("callsite", list(TelemetryCallsite))
|
|
def test_not_used_on_import(reset_usage_lib, callsite: TelemetryCallsite):
|
|
def _import_ray_train():
|
|
from ray import train # noqa: F401
|
|
|
|
check_library_usage_telemetry(
|
|
_import_ray_train, callsite=callsite, expected_library_usages=[set(), {"core"}]
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("callsite", list(TelemetryCallsite))
|
|
def test_used_on_trainer_fit(reset_usage_lib, callsite: TelemetryCallsite):
|
|
def _call_trainer_fit():
|
|
def train_fn():
|
|
tmpdir = tempfile.mkdtemp()
|
|
ray.train.report(
|
|
{},
|
|
checkpoint=Checkpoint.from_directory(tmpdir),
|
|
checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
|
|
validation=True,
|
|
)
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_fn, validation_config=ValidationConfig(fn=lambda x: {})
|
|
)
|
|
trainer.fit()
|
|
|
|
check_library_usage_telemetry(
|
|
_call_trainer_fit,
|
|
callsite=callsite,
|
|
expected_library_usages=[{"train"}, {"core", "train"}],
|
|
expected_extra_usage_tags={
|
|
"train_trainer": "DataParallelTrainer",
|
|
"train_checkpoint_mode": CheckpointUploadMode.ASYNC.value,
|
|
"train_asynchronous_validation": "1",
|
|
},
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info.major == 3 and sys.version_info.minor >= 12,
|
|
reason="Python 3.12+ does not have Tensorflow installed on CI due to dependency conflicts.",
|
|
)
|
|
def test_tag_train_entrypoint(mock_record):
|
|
"""Test that Train v2 entrypoints are recorded correctly."""
|
|
from ray.train.v2.lightgbm.lightgbm_trainer import LightGBMTrainer
|
|
from ray.train.v2.tensorflow.tensorflow_trainer import TensorflowTrainer
|
|
from ray.train.v2.torch.torch_trainer import TorchTrainer
|
|
from ray.train.v2.xgboost.xgboost_trainer import XGBoostTrainer
|
|
|
|
trainer_classes = [
|
|
TorchTrainer,
|
|
TensorflowTrainer,
|
|
XGBoostTrainer,
|
|
LightGBMTrainer,
|
|
]
|
|
for trainer_cls in trainer_classes:
|
|
trainer = trainer_cls(
|
|
lambda: None,
|
|
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
|
)
|
|
assert (
|
|
mock_record[ray_usage_lib.TagKey.TRAIN_TRAINER]
|
|
== trainer.__class__.__name__
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"scaling_config, elasticity_enabled",
|
|
[
|
|
(ScalingConfig(num_workers=(1, 2)), True),
|
|
(ScalingConfig(num_workers=2), False),
|
|
],
|
|
)
|
|
def test_tag_train_elasticity(mock_record, scaling_config, elasticity_enabled):
|
|
DataParallelTrainer(lambda: None, scaling_config=scaling_config)
|
|
if elasticity_enabled:
|
|
assert mock_record[ray_usage_lib.TagKey.TRAIN_ELASTICITY_ENABLED] == "1"
|
|
else:
|
|
assert ray_usage_lib.TagKey.TRAIN_ELASTICITY_ENABLED not in mock_record
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|