chore: import upstream snapshot with attribution

python/ray/air/_internal/usage.py

@@ -0,0 +1,277 @@
+import collections
+import json
+import os
+from enum import Enum
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
+
+from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
+
+if TYPE_CHECKING:
+    from ray.train._internal.storage import StorageContext
+    from ray.train.trainer import BaseTrainer
+    from ray.tune import Callback
+    from ray.tune.schedulers import TrialScheduler
+    from ray.tune.search import BasicVariantGenerator, Searcher
+
+
+AIR_TRAINERS = {
+    "HorovodTrainer",
+    "LightGBMTrainer",
+    "TensorflowTrainer",
+    "TorchTrainer",
+    "XGBoostTrainer",
+}
+
+TRAIN_V2_TRAINERS = {
+    "DataParallelTrainer",
+    "JaxTrainer",
+    "LightGBMTrainer",
+    "TensorflowTrainer",
+    "TorchTrainer",
+    "XGBoostTrainer",
+}
+
+# searchers implemented by Ray Tune.
+TUNE_SEARCHERS = {
+    "AxSearch",
+    "BayesOptSearch",
+    "TuneBOHB",
+    "HEBOSearch",
+    "HyperOptSearch",
+    "NevergradSearch",
+    "OptunaSearch",
+    "ZOOptSearch",
+}
+
+# These are just wrappers around real searchers.
+# We don't want to double tag in this case, otherwise, the real tag
+# will be overwritten.
+TUNE_SEARCHER_WRAPPERS = {
+    "ConcurrencyLimiter",
+    "Repeater",
+}
+
+TUNE_SCHEDULERS = {
+    "FIFOScheduler",
+    "AsyncHyperBandScheduler",
+    "MedianStoppingRule",
+    "HyperBandScheduler",
+    "HyperBandForBOHB",
+    "PopulationBasedTraining",
+    "PopulationBasedTrainingReplay",
+    "PB2",
+    "ResourceChangingScheduler",
+}
+
+
+class AirEntrypoint(Enum):
+    TUNER = "Tuner.fit"
+    TRAINER = "Trainer.fit"
+    TUNE_RUN = "tune.run"
+    TUNE_RUN_EXPERIMENTS = "tune.run_experiments"
+
+
+def _find_class_name(obj: object, allowed_module_path_prefix: str, whitelist: Set[str]):
+    """Find the class name of the object. If the object is not
+    under `allowed_module_path_prefix` or if its class is not in the whitelist,
+    return "Custom".
+
+    Args:
+        obj: The object under inspection.
+        allowed_module_path_prefix: If the `obj`'s class is not under
+            the `allowed_module_path_prefix`, its class name will be anonymized.
+        whitelist: If the `obj`'s class is not in the `whitelist`,
+            it will be anonymized.
+    Returns:
+        The class name to be tagged with telemetry.
+    """
+    module_path = obj.__module__
+    cls_name = obj.__class__.__name__
+    if module_path.startswith(allowed_module_path_prefix) and cls_name in whitelist:
+        return cls_name
+    else:
+        return "Custom"
+
+
+def tag_air_trainer(trainer: "BaseTrainer"):
+    from ray.train.trainer import BaseTrainer
+
+    assert isinstance(trainer, BaseTrainer)
+    trainer_name = _find_class_name(trainer, "ray.train", AIR_TRAINERS)
+    record_extra_usage_tag(TagKey.AIR_TRAINER, trainer_name)
+
+
+def tag_train_v2_trainer(trainer):
+    from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
+
+    assert isinstance(trainer, DataParallelTrainer)
+    trainer_name = _find_class_name(trainer, "ray.train", TRAIN_V2_TRAINERS)
+    record_extra_usage_tag(TagKey.TRAIN_TRAINER, trainer_name)
+
+
+def tag_searcher(searcher: Union["BasicVariantGenerator", "Searcher"]):
+    from ray.tune.search import BasicVariantGenerator, Searcher
+
+    if isinstance(searcher, BasicVariantGenerator):
+        # Note this could be highly inflated as all train flows are treated
+        # as using BasicVariantGenerator.
+        record_extra_usage_tag(TagKey.TUNE_SEARCHER, "BasicVariantGenerator")
+    elif isinstance(searcher, Searcher):
+        searcher_name = _find_class_name(
+            searcher, "ray.tune.search", TUNE_SEARCHERS.union(TUNE_SEARCHER_WRAPPERS)
+        )
+        if searcher_name in TUNE_SEARCHER_WRAPPERS:
+            # ignore to avoid double tagging with wrapper name.
+            return
+        record_extra_usage_tag(TagKey.TUNE_SEARCHER, searcher_name)
+    else:
+        assert False, (
+            "Not expecting a non-BasicVariantGenerator, "
+            "non-Searcher type passed in for `tag_searcher`."
+        )
+
+
+def tag_scheduler(scheduler: "TrialScheduler"):
+    from ray.tune.schedulers import TrialScheduler
+
+    assert isinstance(scheduler, TrialScheduler)
+    scheduler_name = _find_class_name(scheduler, "ray.tune.schedulers", TUNE_SCHEDULERS)
+    record_extra_usage_tag(TagKey.TUNE_SCHEDULER, scheduler_name)
+
+
+def tag_setup_wandb():
+    record_extra_usage_tag(TagKey.AIR_SETUP_WANDB_INTEGRATION_USED, "1")
+
+
+def tag_setup_mlflow():
+    record_extra_usage_tag(TagKey.AIR_SETUP_MLFLOW_INTEGRATION_USED, "1")
+
+
+def _count_callbacks(callbacks: Optional[List["Callback"]]) -> Dict[str, int]:
+    """Creates a map of callback class name -> count given a list of callbacks."""
+    from ray.air.integrations.comet import CometLoggerCallback
+    from ray.air.integrations.mlflow import MLflowLoggerCallback
+    from ray.air.integrations.wandb import WandbLoggerCallback
+    from ray.tune import Callback
+    from ray.tune.logger import LoggerCallback
+    from ray.tune.logger.aim import AimLoggerCallback
+    from ray.tune.utils.callback import DEFAULT_CALLBACK_CLASSES
+
+    built_in_callbacks = (
+        WandbLoggerCallback,
+        MLflowLoggerCallback,
+        CometLoggerCallback,
+        AimLoggerCallback,
+    ) + DEFAULT_CALLBACK_CLASSES
+
+    callback_names = [callback_cls.__name__ for callback_cls in built_in_callbacks]
+    callback_counts = collections.defaultdict(int)
+
+    callbacks = callbacks or []
+    for callback in callbacks:
+        if not isinstance(callback, Callback):
+            # This will error later, but don't include this as custom usage.
+            continue
+
+        callback_name = callback.__class__.__name__
+
+        if callback_name in callback_names:
+            callback_counts[callback_name] += 1
+        elif isinstance(callback, LoggerCallback):
+            callback_counts["CustomLoggerCallback"] += 1
+        else:
+            callback_counts["CustomCallback"] += 1
+
+    return callback_counts
+
+
+def tag_callbacks(callbacks: Optional[List["Callback"]]) -> bool:
+    """Records built-in callback usage via a JSON str representing a
+    dictionary mapping callback class name -> counts.
+
+    User-defined callbacks will increment the count under the `CustomLoggerCallback`
+    or `CustomCallback` key depending on which of the provided interfaces they subclass.
+    NOTE: This will NOT track the name of the user-defined callback,
+    nor its implementation.
+
+    This will NOT report telemetry if no callbacks are provided by the user.
+
+    Args:
+        callbacks: List of callbacks supplied by the user. May be ``None``.
+
+    Returns:
+        bool: True if usage was recorded, False otherwise.
+    """
+    if not callbacks:
+        # User didn't pass in any callbacks -> no usage recorded.
+        return False
+
+    callback_counts = _count_callbacks(callbacks)
+
+    if callback_counts:
+        callback_counts_str = json.dumps(callback_counts)
+        record_extra_usage_tag(TagKey.AIR_CALLBACKS, callback_counts_str)
+
+
+def tag_storage_type(storage: "StorageContext"):
+    """Records the storage configuration of an experiment.
+
+    The storage configuration is set by `RunConfig(storage_path, storage_filesystem)`.
+
+    The possible storage types (defined by `pyarrow.fs.FileSystem.type_name`) are:
+    - 'local' = pyarrow.fs.LocalFileSystem. This includes NFS usage.
+    - 'mock' = pyarrow.fs._MockFileSystem. This is used for testing.
+    - ('s3', 'gcs', 'abfs', 'hdfs'): Various remote storage schemes
+        with default implementations in pyarrow.
+    - 'custom' = All other storage schemes, which includes ALL cases where a
+        custom `storage_filesystem` is provided.
+    - 'other' = catches any other cases not explicitly handled above.
+    """
+    whitelist = {"local", "mock", "s3", "gcs", "abfs", "hdfs"}
+
+    if storage.custom_fs_provided:
+        storage_config_tag = "custom"
+    elif storage.storage_filesystem.type_name in whitelist:
+        storage_config_tag = storage.storage_filesystem.type_name
+    else:
+        storage_config_tag = "other"
+
+    record_extra_usage_tag(TagKey.AIR_STORAGE_CONFIGURATION, storage_config_tag)
+
+
+def tag_ray_air_env_vars() -> bool:
+    """Records usage of environment variables exposed by the Ray AIR libraries.
+
+    NOTE: This does not track the values of the environment variables, nor
+    does this track environment variables not explicitly included in the
+    `all_ray_air_env_vars` allow-list.
+
+    Returns:
+        bool: True if at least one environment var is supplied by the user.
+    """
+    from ray.air.constants import AIR_ENV_VARS
+    from ray.train.constants import TRAIN_ENV_VARS
+    from ray.tune.constants import TUNE_ENV_VARS
+
+    all_ray_air_env_vars = sorted(
+        set().union(AIR_ENV_VARS, TUNE_ENV_VARS, TRAIN_ENV_VARS)
+    )
+
+    user_supplied_env_vars = []
+
+    for env_var in all_ray_air_env_vars:
+        if env_var in os.environ:
+            user_supplied_env_vars.append(env_var)
+
+    if user_supplied_env_vars:
+        env_vars_str = json.dumps(user_supplied_env_vars)
+        record_extra_usage_tag(TagKey.AIR_ENV_VARS, env_vars_str)
+        return True
+
+    return False
+
+
+def tag_air_entrypoint(entrypoint: AirEntrypoint) -> None:
+    """Records the entrypoint to an AIR training run."""
+    assert entrypoint in AirEntrypoint
+    record_extra_usage_tag(TagKey.AIR_ENTRYPOINT, entrypoint.value)