chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_telemetry.py

@@ -0,0 +1,129 @@
+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__]))