wehub-resource-sync
278 lines
9.1 KiB
Python
278 lines
9.1 KiB
Python
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)
|