chore: import upstream snapshot with attribution

python/ray/air/tests/test_air_usage.py

@@ -0,0 +1,210 @@
+"""Unit tests for AIR telemetry."""
+
+import json
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pyarrow.fs
+import pytest
+from packaging.version import Version
+
+import ray
+from ray import train, tune
+from ray._common.usage.usage_lib import TagKey
+from ray.air._internal import usage as air_usage
+from ray.air._internal.usage import AirEntrypoint
+from ray.air.integrations import comet, mlflow, wandb
+from ray.train._internal.storage import StorageContext
+from ray.tune.callback import Callback
+from ray.tune.experiment.experiment import Experiment
+from ray.tune.logger import LoggerCallback
+from ray.tune.utils.callback import DEFAULT_CALLBACK_CLASSES
+
+
+def _mock_record_from_module(module, monkeypatch):
+    recorded = {}
+
+    def mock_record_extra_usage_tag(key: TagKey, value: str):
+        recorded[key] = value
+
+    monkeypatch.setattr(
+        module,
+        "record_extra_usage_tag",
+        mock_record_extra_usage_tag,
+    )
+    return recorded
+
+
+@pytest.fixture
+def mock_record(monkeypatch):
+    import ray.air._internal.usage
+
+    yield _mock_record_from_module(ray.air._internal.usage, monkeypatch=monkeypatch)
+
+
+def train_fn(config):
+    train.report({"score": 1})
+
+
+@pytest.fixture
+def tuner(tmp_path):
+    yield tune.Tuner(train_fn, run_config=tune.RunConfig(storage_path=str(tmp_path)))
+
+
+@pytest.fixture
+def trainer(tmp_path):
+    from ray.train.data_parallel_trainer import DataParallelTrainer
+
+    yield DataParallelTrainer(
+        train_loop_per_worker=train_fn,
+        scaling_config=train.ScalingConfig(num_workers=2),
+        run_config=train.RunConfig(storage_path=str(tmp_path)),
+    )
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "storage_path_filesystem_expected",
+    [
+        ("/tmp/test", None, "local"),
+        ("s3://", None, "s3"),
+        ("gs://test", None, "gcs"),
+        ("mock://test", None, "mock"),
+        ("test", pyarrow.fs.LocalFileSystem(), "custom"),
+    ],
+)
+def test_tag_storage_type(storage_path_filesystem_expected, mock_record, monkeypatch):
+    # Don't write anything to storage for the test.
+    monkeypatch.setattr(StorageContext, "_create_validation_file", lambda _: None)
+    monkeypatch.setattr(StorageContext, "_check_validation_file", lambda _: None)
+
+    storage_path, storage_filesystem, expected = storage_path_filesystem_expected
+
+    if Version(pyarrow.__version__) < Version("17.0.0") and storage_path.startswith(
+        "gs://"
+    ):
+        pytest.skip("GCS support requires pyarrow >= 17.0.0")
+
+    storage = StorageContext(
+        storage_path=storage_path,
+        experiment_dir_name="test",
+        storage_filesystem=storage_filesystem,
+    )
+    air_usage.tag_storage_type(storage)
+    assert mock_record[TagKey.AIR_STORAGE_CONFIGURATION] == expected
+
+
+class _CustomLoggerCallback(LoggerCallback):
+    pass
+
+
+class _CustomCallback(Callback):
+    pass
+
+
+_TEST_CALLBACKS = [
+    wandb.WandbLoggerCallback,
+    mlflow.MLflowLoggerCallback,
+    comet.CometLoggerCallback,
+    _CustomLoggerCallback,
+    _CustomLoggerCallback,
+    _CustomCallback,
+]
+
+
+def test_tag_setup_wandb(mock_record):
+    from ray.air.integrations.wandb import _setup_wandb
+
+    with patch.dict(os.environ, {wandb.WANDB_MODE_ENV_VAR: "disabled"}):
+        _setup_wandb(trial_id="a", trial_name="b", config={}, _wandb=MagicMock())
+    assert mock_record[TagKey.AIR_SETUP_WANDB_INTEGRATION_USED] == "1"
+
+
+def test_tag_setup_mlflow(mock_record, monkeypatch):
+    from ray.air.integrations.mlflow import setup_mlflow
+
+    monkeypatch.setattr(ray.air.integrations.mlflow, "_MLflowLoggerUtil", MagicMock())
+    setup_mlflow()
+    assert mock_record[TagKey.AIR_SETUP_MLFLOW_INTEGRATION_USED] == "1"
+
+
+@pytest.mark.parametrize(
+    "callback_classes_expected",
+    [
+        (None, None),
+        ([], None),
+        ([lambda: None], None),
+        (
+            DEFAULT_CALLBACK_CLASSES,
+            {cls.__name__: 1 for cls in DEFAULT_CALLBACK_CLASSES},
+        ),
+        (
+            _TEST_CALLBACKS,
+            {
+                "WandbLoggerCallback": 1,
+                "MLflowLoggerCallback": 1,
+                "CometLoggerCallback": 1,
+                "CustomLoggerCallback": 2,
+                "CustomCallback": 1,
+            },
+        ),
+    ],
+)
+def test_tag_callbacks(mock_record, callback_classes_expected):
+    callback_classes, expected = callback_classes_expected
+
+    callbacks = (
+        [callback_cls() for callback_cls in callback_classes]
+        if callback_classes
+        else None
+    )
+
+    air_usage.tag_callbacks(callbacks)
+
+    callback_usage_str = mock_record.pop(TagKey.AIR_CALLBACKS, None)
+    callback_counts = json.loads(callback_usage_str) if callback_usage_str else None
+    assert callback_counts == expected
+
+
+def test_tag_env_vars(ray_start_4_cpus, mock_record, tuner):
+    """Test that env vars are recorded properly, and arbitrary user environment
+    variables are ignored."""
+    env_vars_to_record = {
+        "TUNE_GLOBAL_CHECKPOINT_S": "20",
+        "TUNE_MAX_PENDING_TRIALS_PG": "1",
+    }
+    untracked_env_vars = {"RANDOM_USER_ENV_VAR": "asdf"}
+
+    with patch.dict(os.environ, {**env_vars_to_record, **untracked_env_vars}):
+        tuner.fit()
+
+    recorded_env_vars = json.loads(mock_record[TagKey.AIR_ENV_VARS])
+    assert sorted(env_vars_to_record) == sorted(recorded_env_vars)
+
+
+@pytest.mark.parametrize("entrypoint", list(AirEntrypoint))
+def test_tag_air_entrypoint(ray_start_4_cpus, mock_record, entrypoint, tuner, trainer):
+    if entrypoint == AirEntrypoint.TUNE_RUN:
+        tune.run(train_fn)
+    elif entrypoint == AirEntrypoint.TUNE_RUN_EXPERIMENTS:
+        experiment_spec = Experiment("experiment", train_fn)
+        tune.run_experiments(experiments=experiment_spec)
+    elif entrypoint == AirEntrypoint.TUNER:
+        tuner.fit()
+    elif entrypoint == AirEntrypoint.TRAINER:
+        trainer.fit()
+
+    assert mock_record[TagKey.AIR_ENTRYPOINT] == entrypoint.value
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))