wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/air/integrations/__init__.py
View File
python/ray/air/integrations/comet.py
+260
View File
@@ -0,0 +1,260 @@
import os
from pathlib import Path
from typing import Dict, List
import pyarrow.fs
from ray.tune.experiment import Trial
from ray.tune.logger import LoggerCallback
from ray.tune.utils import flatten_dict
def _import_comet():
"""Try importing comet_ml.
Used to check if comet_ml is installed and, otherwise, pass an informative
error message.
"""
if "COMET_DISABLE_AUTO_LOGGING" not in os.environ:
os.environ["COMET_DISABLE_AUTO_LOGGING"] = "1"
try:
import comet_ml # noqa: F401
except ImportError:
raise RuntimeError("pip install 'comet-ml' to use CometLoggerCallback")
return comet_ml
class CometLoggerCallback(LoggerCallback):
"""CometLoggerCallback for logging Tune results to Comet.
Comet (https://comet.ml/site/) is a tool to manage and optimize the
entire ML lifecycle, from experiment tracking, model optimization
and dataset versioning to model production monitoring.
This Ray Tune ``LoggerCallback`` sends metrics and parameters to
Comet for tracking.
In order to use the CometLoggerCallback you must first install Comet
via ``pip install comet_ml``
Then set the following environment variables
``export COMET_API_KEY=<Your API Key>``
Alternatively, you can also pass in your API Key as an argument to the
CometLoggerCallback constructor.
``CometLoggerCallback(api_key=<Your API Key>)``
Args:
online: Whether to make use of an Online or
Offline Experiment. Defaults to True.
tags: Tags to add to the logged Experiment.
Defaults to None.
save_checkpoints: If ``True``, model checkpoints will be saved to
Comet ML as artifacts. Defaults to ``False``.
**experiment_kwargs: Other keyword arguments will be passed to the
constructor for comet_ml.Experiment (or OfflineExperiment if
online=False).
Please consult the Comet ML documentation for more information on the
Experiment and OfflineExperiment classes: https://comet.ml/site/
Example:
.. code-block:: python
from ray.air.integrations.comet import CometLoggerCallback
tune.run(
train,
config=config
callbacks=[CometLoggerCallback(
True,
['tag1', 'tag2'],
workspace='my_workspace',
project_name='my_project_name'
)]
)
"""
# Do not enable these auto log options unless overridden
_exclude_autolog = [
"auto_output_logging",
"log_git_metadata",
"log_git_patch",
"log_env_cpu",
"log_env_gpu",
]
# Do not log these metrics.
_exclude_results = ["done", "should_checkpoint"]
# These values should be logged as system info instead of metrics.
_system_results = ["node_ip", "hostname", "pid", "date"]
# These values should be logged as "Other" instead of as metrics.
_other_results = ["trial_id", "experiment_id", "experiment_tag"]
_episode_results = ["hist_stats/episode_reward", "hist_stats/episode_lengths"]
def __init__(
self,
online: bool = True,
tags: List[str] = None,
save_checkpoints: bool = False,
**experiment_kwargs,
):
_import_comet()
self.online = online
self.tags = tags
self.save_checkpoints = save_checkpoints
self.experiment_kwargs = experiment_kwargs
# Disable the specific autologging features that cause throttling.
self._configure_experiment_defaults()
# Mapping from trial to experiment object.
self._trial_experiments = {}
self._to_exclude = self._exclude_results.copy()
self._to_system = self._system_results.copy()
self._to_other = self._other_results.copy()
self._to_episodes = self._episode_results.copy()
def _configure_experiment_defaults(self):
"""Disable the specific autologging features that cause throttling."""
for option in self._exclude_autolog:
if not self.experiment_kwargs.get(option):
self.experiment_kwargs[option] = False
def _check_key_name(self, key: str, item: str) -> bool:
"""
Check if key argument is equal to item argument or starts with item and
a forward slash. Used for parsing trial result dictionary into ignored
keys, system metrics, episode logs, etc.
"""
return key.startswith(item + "/") or key == item
def log_trial_start(self, trial: "Trial"):
"""
Initialize an Experiment (or OfflineExperiment if self.online=False)
and start logging to Comet.
Args:
trial: Trial object.
"""
_import_comet() # is this necessary?
from comet_ml import Experiment, OfflineExperiment
from comet_ml.config import set_global_experiment
if trial not in self._trial_experiments:
experiment_cls = Experiment if self.online else OfflineExperiment
experiment = experiment_cls(**self.experiment_kwargs)
self._trial_experiments[trial] = experiment
# Set global experiment to None to allow for multiple experiments.
set_global_experiment(None)
else:
experiment = self._trial_experiments[trial]
experiment.set_name(str(trial))
experiment.add_tags(self.tags)
experiment.log_other("Created from", "Ray")
config = trial.config.copy()
config.pop("callbacks", None)
experiment.log_parameters(config)
def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
"""
Log the current result of a Trial upon each iteration.
"""
if trial not in self._trial_experiments:
self.log_trial_start(trial)
experiment = self._trial_experiments[trial]
step = result["training_iteration"]
config_update = result.pop("config", {}).copy()
config_update.pop("callbacks", None) # Remove callbacks
for k, v in config_update.items():
if isinstance(v, dict):
experiment.log_parameters(flatten_dict({k: v}, "/"), step=step)
else:
experiment.log_parameter(k, v, step=step)
other_logs = {}
metric_logs = {}
system_logs = {}
episode_logs = {}
flat_result = flatten_dict(result, delimiter="/")
for k, v in flat_result.items():
if any(self._check_key_name(k, item) for item in self._to_exclude):
continue
if any(self._check_key_name(k, item) for item in self._to_other):
other_logs[k] = v
elif any(self._check_key_name(k, item) for item in self._to_system):
system_logs[k] = v
elif any(self._check_key_name(k, item) for item in self._to_episodes):
episode_logs[k] = v
else:
metric_logs[k] = v
experiment.log_others(other_logs)
experiment.log_metrics(metric_logs, step=step)
for k, v in system_logs.items():
experiment.log_system_info(k, v)
for k, v in episode_logs.items():
experiment.log_curve(k, x=range(len(v)), y=v, step=step)
def log_trial_save(self, trial: "Trial"):
comet_ml = _import_comet()
if self.save_checkpoints and trial.checkpoint:
experiment = self._trial_experiments[trial]
artifact = comet_ml.Artifact(
name=f"checkpoint_{(str(trial))}", artifact_type="model"
)
checkpoint_root = None
if isinstance(trial.checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
checkpoint_root = trial.checkpoint.path
# Todo: For other filesystems, we may want to use
# artifact.add_remote() instead. However, this requires a full
# URI. We can add this once we have a way to retrieve it.
# Walk through checkpoint directory and add all files to artifact
if checkpoint_root:
for root, dirs, files in os.walk(checkpoint_root):
rel_root = os.path.relpath(root, checkpoint_root)
for file in files:
local_file = Path(checkpoint_root, rel_root, file).as_posix()
logical_path = Path(rel_root, file).as_posix()
# Strip leading `./`
if logical_path.startswith("./"):
logical_path = logical_path[2:]
artifact.add(local_file, logical_path=logical_path)
experiment.log_artifact(artifact)
def log_trial_end(self, trial: "Trial", failed: bool = False):
self._trial_experiments[trial].end()
del self._trial_experiments[trial]
def __del__(self):
for trial, experiment in self._trial_experiments.items():
experiment.end()
self._trial_experiments = {}
python/ray/air/integrations/keras.py
+185
View File
@@ -0,0 +1,185 @@
import shutil
from typing import Dict, List, Optional, Union
from tensorflow.keras.callbacks import Callback as KerasCallback
import ray
from ray.train.tensorflow import TensorflowCheckpoint
from ray.util.annotations import PublicAPI
class _Callback(KerasCallback):
"""Base class for Air's Keras callbacks."""
_allowed = [
"epoch_begin",
"epoch_end",
"train_batch_begin",
"train_batch_end",
"test_batch_begin",
"test_batch_end",
"predict_batch_begin",
"predict_batch_end",
"train_begin",
"train_end",
"test_begin",
"test_end",
"predict_begin",
"predict_end",
]
def __init__(self, on: Union[str, List[str]] = "validation_end"):
super(_Callback, self).__init__()
if not isinstance(on, list):
on = [on]
if any(w not in self._allowed for w in on):
raise ValueError(
"Invalid trigger time selected: {}. Must be one of {}".format(
on, self._allowed
)
)
self._on = on
def _handle(self, logs: Dict, when: str):
raise NotImplementedError
def on_epoch_begin(self, epoch, logs=None):
if "epoch_begin" in self._on:
self._handle(logs, "epoch_begin")
def on_epoch_end(self, epoch, logs=None):
if "epoch_end" in self._on:
self._handle(logs, "epoch_end")
def on_train_batch_begin(self, batch, logs=None):
if "train_batch_begin" in self._on:
self._handle(logs, "train_batch_begin")
def on_train_batch_end(self, batch, logs=None):
if "train_batch_end" in self._on:
self._handle(logs, "train_batch_end")
def on_test_batch_begin(self, batch, logs=None):
if "test_batch_begin" in self._on:
self._handle(logs, "test_batch_begin")
def on_test_batch_end(self, batch, logs=None):
if "test_batch_end" in self._on:
self._handle(logs, "test_batch_end")
def on_predict_batch_begin(self, batch, logs=None):
if "predict_batch_begin" in self._on:
self._handle(logs, "predict_batch_begin")
def on_predict_batch_end(self, batch, logs=None):
if "predict_batch_end" in self._on:
self._handle(logs, "predict_batch_end")
def on_train_begin(self, logs=None):
if "train_begin" in self._on:
self._handle(logs, "train_begin")
def on_train_end(self, logs=None):
if "train_end" in self._on:
self._handle(logs, "train_end")
def on_test_begin(self, logs=None):
if "test_begin" in self._on:
self._handle(logs, "test_begin")
def on_test_end(self, logs=None):
if "test_end" in self._on:
self._handle(logs, "test_end")
def on_predict_begin(self, logs=None):
if "predict_begin" in self._on:
self._handle(logs, "predict_begin")
def on_predict_end(self, logs=None):
if "predict_end" in self._on:
self._handle(logs, "predict_end")
@PublicAPI(stability="alpha")
class ReportCheckpointCallback(_Callback):
"""Keras callback for Ray Train reporting and checkpointing.
.. note::
Metrics are always reported with checkpoints, even if the event isn't specified
in ``report_metrics_on``.
Example:
.. code-block:: python
############# Using it in TrainSession ###############
from ray.air.integrations.keras import ReportCheckpointCallback
def train_loop_per_worker():
strategy = tf.distribute.MultiWorkerMirroredStrategy()
with strategy.scope():
model = build_model()
model.fit(dataset_shard, callbacks=[ReportCheckpointCallback()])
Args:
checkpoint_on: When to save checkpoints. Must be one of the Keras event hooks
(less the ``on_``), e.g. "train_start" or "predict_end". Defaults to
"epoch_end".
report_metrics_on: When to report metrics. Must be one of
the Keras event hooks (less the ``on_``), e.g.
"train_start" or "predict_end". Defaults to "epoch_end".
metrics: Metrics to report. If this is a list, each item describes
the metric key reported to Keras, and it's reported under the
same name. If this is a dict, each key is the name reported
and the respective value is the metric key reported to Keras.
If this is None, all Keras logs are reported.
"""
def __init__(
self,
checkpoint_on: Union[str, List[str]] = "epoch_end",
report_metrics_on: Union[str, List[str]] = "epoch_end",
metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
):
if isinstance(checkpoint_on, str):
checkpoint_on = [checkpoint_on]
if isinstance(report_metrics_on, str):
report_metrics_on = [report_metrics_on]
on = list(set(checkpoint_on + report_metrics_on))
super().__init__(on=on)
self._checkpoint_on: List[str] = checkpoint_on
self._report_metrics_on: List[str] = report_metrics_on
self._metrics = metrics
def _handle(self, logs: Dict, when: str):
assert when in self._checkpoint_on or when in self._report_metrics_on
metrics = self._get_reported_metrics(logs)
should_checkpoint = when in self._checkpoint_on
if should_checkpoint:
checkpoint = TensorflowCheckpoint.from_model(self.model)
ray.train.report(metrics, checkpoint=checkpoint)
# Clean up temporary checkpoint
shutil.rmtree(checkpoint.path, ignore_errors=True)
else:
ray.train.report(metrics, checkpoint=None)
def _get_reported_metrics(self, logs: Dict) -> Dict:
assert isinstance(self._metrics, (type(None), str, list, dict))
if self._metrics is None:
reported_metrics = logs
elif isinstance(self._metrics, str):
reported_metrics = {self._metrics: logs[self._metrics]}
elif isinstance(self._metrics, list):
reported_metrics = {metric: logs[metric] for metric in self._metrics}
elif isinstance(self._metrics, dict):
reported_metrics = {
key: logs[metric] for key, metric in self._metrics.items()
}
assert isinstance(reported_metrics, dict)
return reported_metrics
python/ray/air/integrations/mlflow.py
+343
View File
@@ -0,0 +1,343 @@
import logging
from types import ModuleType
from typing import Dict, Optional, Union
import ray
from ray.air._internal import usage as air_usage
from ray.air._internal.mlflow import _MLflowLoggerUtil
from ray.air.constants import TRAINING_ITERATION
from ray.tune.experiment import Trial
from ray.tune.logger import LoggerCallback
from ray.tune.result import TIMESTEPS_TOTAL
from ray.tune.trainable.trainable_fn_utils import _in_tune_session
from ray.util.annotations import PublicAPI
try:
import mlflow
except ImportError:
mlflow = None
logger = logging.getLogger(__name__)
class _NoopModule:
def __getattr__(self, item):
return _NoopModule()
def __call__(self, *args, **kwargs):
return None
@PublicAPI(stability="alpha")
def setup_mlflow(
config: Optional[Dict] = None,
tracking_uri: Optional[str] = None,
registry_uri: Optional[str] = None,
experiment_id: Optional[str] = None,
experiment_name: Optional[str] = None,
tracking_token: Optional[str] = None,
artifact_location: Optional[str] = None,
run_name: Optional[str] = None,
create_experiment_if_not_exists: bool = False,
tags: Optional[Dict] = None,
rank_zero_only: bool = True,
) -> Union[ModuleType, _NoopModule]:
"""Set up a MLflow session.
This function can be used to initialize an MLflow session in a
(distributed) training or tuning run. The session will be created on the trainable.
By default, the MLflow experiment ID is the Ray trial ID and the
MLlflow experiment name is the Ray trial name. These settings can be overwritten by
passing the respective keyword arguments.
The ``config`` dict is automatically logged as the run parameters (excluding the
mlflow settings).
In distributed training with Ray Train, only the zero-rank worker will initialize
mlflow. All other workers will return a noop client, so that logging is not
duplicated in a distributed run. This can be disabled by passing
``rank_zero_only=False``, which will then initialize mlflow in every training
worker. Note: for Ray Tune, there's no concept of worker ranks, so the `rank_zero_only` is ignored.
This function will return the ``mlflow`` module or a noop module for
non-rank zero workers ``if rank_zero_only=True``. By using
``mlflow = setup_mlflow(config)`` you can ensure that only the rank zero worker
calls the mlflow API.
Args:
config: Configuration dict to be logged to mlflow as parameters.
tracking_uri: The tracking URI for MLflow tracking. If using
Tune in a multi-node setting, make sure to use a remote server for
tracking.
registry_uri: The registry URI for the MLflow model registry.
experiment_id: The id of an already created MLflow experiment.
All logs from all trials in ``tune.Tuner()`` will be reported to this
experiment. If this is not provided or the experiment with this
id does not exist, you must provide an``experiment_name``. This
parameter takes precedence over ``experiment_name``.
experiment_name: The name of an already existing MLflow
experiment. All logs from all trials in ``tune.Tuner()`` will be
reported to this experiment. If this is not provided, you must
provide a valid ``experiment_id``.
tracking_token: A token to use for HTTP authentication when
logging to a remote tracking server. This is useful when you
want to log to a Databricks server, for example. This value will
be used to set the MLFLOW_TRACKING_TOKEN environment variable on
all the remote training processes.
artifact_location: The location to store run artifacts.
If not provided, MLFlow picks an appropriate default.
Ignored if experiment already exists.
run_name: Name of the new MLflow run that will be created.
If not set, will default to the ``experiment_name``.
create_experiment_if_not_exists: Whether to create an
experiment with the provided name if it does not already
exist. Defaults to False.
tags: Tags to set for the new run.
rank_zero_only: If True, will return an initialized session only for the
rank 0 worker in distributed training. If False, will initialize a
session for all workers. Defaults to True.
Example:
Per default, you can just call ``setup_mlflow`` and continue to use
MLflow like you would normally do:
.. code-block:: python
from ray.air.integrations.mlflow import setup_mlflow
def training_loop(config):
mlflow = setup_mlflow(config)
# ...
mlflow.log_metric(key="loss", val=0.123, step=0)
In distributed data parallel training, you can utilize the return value of
``setup_mlflow``. This will make sure it is only invoked on the first worker
in distributed training runs.
.. code-block:: python
from ray.air.integrations.mlflow import setup_mlflow
def training_loop(config):
mlflow = setup_mlflow(config)
# ...
mlflow.log_metric(key="loss", val=0.123, step=0)
You can also use MlFlow's autologging feature if using a training
framework like Pytorch Lightning, XGBoost, etc. More information can be
found here
(https://mlflow.org/docs/latest/tracking.html#automatic-logging).
.. code-block:: python
from ray.air.integrations.mlflow import setup_mlflow
def train_fn(config):
mlflow = setup_mlflow(config)
mlflow.autolog()
xgboost_results = xgb.train(config, ...)
Returns:
The ``mlflow`` module, or a noop module for non-rank-zero workers when
``rank_zero_only`` is True.
"""
if not mlflow:
raise RuntimeError(
"mlflow was not found - please install with `pip install mlflow`"
)
default_trial_id = None
default_trial_name = None
try:
if _in_tune_session():
context: ray.tune.TuneContext = ray.tune.get_context()
default_trial_id = context.get_trial_id()
default_trial_name = context.get_trial_name()
else:
context: ray.train.TrainContext = ray.train.get_context()
if rank_zero_only and context.get_world_rank() != 0:
return _NoopModule()
except RuntimeError:
default_trial_id = None
default_trial_name = None
_config = config.copy() if config else {}
experiment_id = experiment_id or default_trial_id
experiment_name = experiment_name or default_trial_name
# Setup mlflow
mlflow_util = _MLflowLoggerUtil()
mlflow_util.setup_mlflow(
tracking_uri=tracking_uri,
registry_uri=registry_uri,
experiment_id=experiment_id,
experiment_name=experiment_name,
tracking_token=tracking_token,
artifact_location=artifact_location,
create_experiment_if_not_exists=create_experiment_if_not_exists,
)
mlflow_util.start_run(
run_name=run_name or experiment_name,
tags=tags,
set_active=True,
)
mlflow_util.log_params(_config)
# Record `setup_mlflow` usage when everything has setup successfully.
air_usage.tag_setup_mlflow()
return mlflow_util._mlflow
class MLflowLoggerCallback(LoggerCallback):
"""MLflow Logger to automatically log Tune results and config to MLflow.
MLflow (https://mlflow.org) Tracking is an open source library for
recording and querying experiments. This Ray Tune ``LoggerCallback``
sends information (config parameters, training results & metrics,
and artifacts) to MLflow for automatic experiment tracking.
Keep in mind that the callback will open an MLflow session on the driver and
not on the trainable. Therefore, it is not possible to call MLflow functions
like ``mlflow.log_figure()`` inside the trainable as there is no MLflow session
on the trainable. For more fine grained control, use
:func:`ray.air.integrations.mlflow.setup_mlflow`.
Args:
tracking_uri: The tracking URI for where to manage experiments
and runs. This can either be a local file path or a remote server.
This arg gets passed directly to mlflow
initialization. When using Tune in a multi-node setting, make sure
to set this to a remote server and not a local file path.
registry_uri: The registry URI that gets passed directly to
mlflow initialization.
experiment_name: The experiment name to use for this Tune run.
If the experiment with the name already exists with MLflow,
it will be reused. If not, a new experiment will be created with
that name.
tags: An optional dictionary of string keys and values to set
as tags on the run
tracking_token: Tracking token used to authenticate with MLflow.
save_artifact: If set to True, automatically save the entire
contents of the Tune local_dir as an artifact to the
corresponding run in MlFlow.
log_params_on_trial_end: If set to True, log parameters to MLflow
at the end of the trial instead of at the beginning
Example:
.. code-block:: python
from ray.air.integrations.mlflow import MLflowLoggerCallback
tags = { "user_name" : "John",
"git_commit_hash" : "abc123"}
tune.run(
train_fn,
config={
# define search space here
"parameter_1": tune.choice([1, 2, 3]),
"parameter_2": tune.choice([4, 5, 6]),
},
callbacks=[MLflowLoggerCallback(
experiment_name="experiment1",
tags=tags,
save_artifact=True,
log_params_on_trial_end=True)])
"""
def __init__(
self,
tracking_uri: Optional[str] = None,
*,
registry_uri: Optional[str] = None,
experiment_name: Optional[str] = None,
tags: Optional[Dict] = None,
tracking_token: Optional[str] = None,
save_artifact: bool = False,
log_params_on_trial_end: bool = False,
):
self.tracking_uri = tracking_uri
self.registry_uri = registry_uri
self.experiment_name = experiment_name
self.tags = tags
self.tracking_token = tracking_token
self.should_save_artifact = save_artifact
self.log_params_on_trial_end = log_params_on_trial_end
self.mlflow_util = _MLflowLoggerUtil()
if ray.util.client.ray.is_connected():
logger.warning(
"When using MLflowLoggerCallback with Ray Client, "
"it is recommended to use a remote tracking "
"server. If you are using a MLflow tracking server "
"backed by the local filesystem, then it must be "
"setup on the server side and not on the client "
"side."
)
def setup(self, *args, **kwargs):
# Setup the mlflow logging util.
self.mlflow_util.setup_mlflow(
tracking_uri=self.tracking_uri,
registry_uri=self.registry_uri,
experiment_name=self.experiment_name,
tracking_token=self.tracking_token,
)
if self.tags is None:
# Create empty dictionary for tags if not given explicitly
self.tags = {}
self._trial_runs = {}
def log_trial_start(self, trial: "Trial"):
# Create run if not already exists.
if trial not in self._trial_runs:
# Set trial name in tags
tags = self.tags.copy()
tags["trial_name"] = str(trial)
run = self.mlflow_util.start_run(tags=tags, run_name=str(trial))
self._trial_runs[trial] = run.info.run_id
run_id = self._trial_runs[trial]
# Log the config parameters.
config = trial.config
if not self.log_params_on_trial_end:
self.mlflow_util.log_params(run_id=run_id, params_to_log=config)
def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
step = result.get(TIMESTEPS_TOTAL) or result[TRAINING_ITERATION]
run_id = self._trial_runs[trial]
self.mlflow_util.log_metrics(run_id=run_id, metrics_to_log=result, step=step)
def log_trial_end(self, trial: "Trial", failed: bool = False):
run_id = self._trial_runs[trial]
# Log the artifact if set_artifact is set to True.
if self.should_save_artifact:
self.mlflow_util.save_artifacts(run_id=run_id, dir=trial.local_path)
# Stop the run once trial finishes.
status = "FINISHED" if not failed else "FAILED"
# Log the config parameters.
config = trial.config
if self.log_params_on_trial_end:
self.mlflow_util.log_params(run_id=run_id, params_to_log=config)
self.mlflow_util.end_run(run_id=run_id, status=status)
python/ray/air/integrations/wandb.py
+810
View File
@@ -0,0 +1,810 @@
import enum
import os
import pickle
import urllib
import warnings
from numbers import Number
from types import ModuleType
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
import numpy as np
import pyarrow.fs
import ray
from ray import logger
from ray._common.utils import load_class
from ray.air._internal import usage as air_usage
from ray.air.constants import TRAINING_ITERATION
from ray.air.util.node import _force_on_current_node
from ray.train._internal.session import get_session
from ray.train._internal.syncer import DEFAULT_SYNC_TIMEOUT
from ray.tune.experiment import Trial
from ray.tune.logger import LoggerCallback
from ray.tune.utils import flatten_dict
from ray.util import PublicAPI
from ray.util.queue import Queue
try:
import wandb
from wandb.sdk.data_types.base_types.wb_value import WBValue
from wandb.sdk.data_types.image import Image
from wandb.sdk.data_types.video import Video
from wandb.sdk.lib.disabled import RunDisabled
from wandb.util import json_dumps_safer
from wandb.wandb_run import Run
except ImportError:
wandb = json_dumps_safer = Run = RunDisabled = WBValue = None
WANDB_ENV_VAR = "WANDB_API_KEY"
WANDB_PROJECT_ENV_VAR = "WANDB_PROJECT_NAME"
WANDB_GROUP_ENV_VAR = "WANDB_GROUP_NAME"
WANDB_MODE_ENV_VAR = "WANDB_MODE"
# Hook that is invoked before wandb.init in the setup method of WandbLoggerCallback
# to populate the API key if it isn't already set when initializing the callback.
# It doesn't take in any arguments and returns the W&B API key.
# Example: "your.module.wandb_setup_api_key_hook".
WANDB_SETUP_API_KEY_HOOK = "WANDB_SETUP_API_KEY_HOOK"
# Hook that is invoked before wandb.init in the setup method of WandbLoggerCallback
# to populate environment variables to specify the location
# (project and group) of the W&B run.
# It doesn't take in any arguments and doesn't return anything, but it does populate
# WANDB_PROJECT_NAME and WANDB_GROUP_NAME.
# Example: "your.module.wandb_populate_run_location_hook".
WANDB_POPULATE_RUN_LOCATION_HOOK = "WANDB_POPULATE_RUN_LOCATION_HOOK"
# Hook that is invoked after running wandb.init in WandbLoggerCallback
# to process information about the W&B run.
# It takes in a W&B run object and doesn't return anything.
# Example: "your.module.wandb_process_run_info_hook".
WANDB_PROCESS_RUN_INFO_HOOK = "WANDB_PROCESS_RUN_INFO_HOOK"
@PublicAPI(stability="alpha")
def setup_wandb(
config: Optional[Dict] = None,
api_key: Optional[str] = None,
api_key_file: Optional[str] = None,
rank_zero_only: bool = True,
**kwargs,
) -> Union[Run, RunDisabled]:
"""Set up a Weights & Biases session.
This function can be used to initialize a Weights & Biases session in a
(distributed) training or tuning run.
By default, the run ID is the trial ID, the run name is the trial name, and
the run group is the experiment name. These settings can be overwritten by
passing the respective arguments as ``kwargs``, which will be passed to
``wandb.init()``.
In distributed training with Ray Train, only the zero-rank worker will initialize
wandb. All other workers will return a disabled run object, so that logging is not
duplicated in a distributed run. This can be disabled by passing
``rank_zero_only=False``, which will then initialize wandb in every training
worker.
The ``config`` argument will be passed to Weights and Biases and will be logged
as the run configuration.
If no API key or key file are passed, wandb will try to authenticate
using locally stored credentials, created for instance by running ``wandb login``.
Keyword arguments passed to ``setup_wandb()`` will be passed to
``wandb.init()`` and take precedence over any potential default settings.
Args:
config: Configuration dict to be logged to Weights and Biases. Can contain
arguments for ``wandb.init()`` as well as authentication information.
api_key: API key to use for authentication with Weights and Biases.
api_key_file: File pointing to API key for with Weights and Biases.
rank_zero_only: If True, will return an initialized session only for the
rank 0 worker in distributed training. If False, will initialize a
session for all workers.
**kwargs: Passed to ``wandb.init()``.
Example:
.. code-block:: python
from ray.air.integrations.wandb import setup_wandb
def training_loop(config):
wandb = setup_wandb(config)
# ...
wandb.log({"loss": 0.123})
Returns:
The initialized wandb run, or a disabled run for non-rank-zero workers
when ``rank_zero_only`` is True.
"""
if not wandb:
raise RuntimeError(
"Wandb was not found - please install with `pip install wandb`"
)
default_trial_id = None
default_trial_name = None
default_experiment_name = None
# Do a try-catch here if we are not in a train session
session = get_session()
if rank_zero_only:
# Check if we are in a train session and if we are not the rank 0 worker
if session and session.world_rank is not None and session.world_rank != 0:
return RunDisabled()
if session:
default_trial_id = session.trial_id
default_trial_name = session.trial_name
default_experiment_name = session.experiment_name
# Default init kwargs
wandb_init_kwargs = {
"trial_id": kwargs.get("trial_id") or default_trial_id,
"trial_name": kwargs.get("trial_name") or default_trial_name,
"group": kwargs.get("group") or default_experiment_name,
}
# Passed kwargs take precedence over default kwargs
wandb_init_kwargs.update(kwargs)
return _setup_wandb(
config=config, api_key=api_key, api_key_file=api_key_file, **wandb_init_kwargs
)
def _setup_wandb(
trial_id: str,
trial_name: str,
config: Optional[Dict] = None,
api_key: Optional[str] = None,
api_key_file: Optional[str] = None,
_wandb: Optional[ModuleType] = None,
**kwargs,
) -> Union[Run, RunDisabled]:
_config = config.copy() if config else {}
# If key file is specified, set
if api_key_file:
api_key_file = os.path.expanduser(api_key_file)
_set_api_key(api_key_file, api_key)
project = _get_wandb_project(kwargs.pop("project", None))
group = kwargs.pop("group", os.environ.get(WANDB_GROUP_ENV_VAR))
# Remove unpickleable items.
_config = _clean_log(_config)
wandb_init_kwargs = dict(
id=trial_id,
name=trial_name,
resume=True,
reinit=True,
allow_val_change=True,
config=_config,
project=project,
group=group,
)
# Update config (e.g. set any other parameters in the call to wandb.init)
wandb_init_kwargs.update(**kwargs)
# On windows, we can't fork
if os.name == "nt":
os.environ["WANDB_START_METHOD"] = "thread"
else:
os.environ["WANDB_START_METHOD"] = "fork"
_wandb = _wandb or wandb
run = _wandb.init(**wandb_init_kwargs)
_run_wandb_process_run_info_hook(run)
# Record `setup_wandb` usage when everything has setup successfully.
air_usage.tag_setup_wandb()
return run
def _is_allowed_type(obj):
"""Return True if type is allowed for logging to wandb"""
if isinstance(obj, np.ndarray) and obj.size == 1:
return isinstance(obj.item(), Number)
if isinstance(obj, Sequence) and len(obj) > 0:
return isinstance(obj[0], (Image, Video, WBValue))
return isinstance(obj, (Number, WBValue))
def _clean_log(
obj: Any,
*,
video_kwargs: Optional[Dict[str, Any]] = None,
image_kwargs: Optional[Dict[str, Any]] = None,
):
# Fixes https://github.com/ray-project/ray/issues/10631
if video_kwargs is None:
video_kwargs = {}
if image_kwargs is None:
image_kwargs = {}
if isinstance(obj, dict):
return {
k: _clean_log(v, video_kwargs=video_kwargs, image_kwargs=image_kwargs)
for k, v in obj.items()
}
elif isinstance(obj, (list, set)):
return [
_clean_log(v, video_kwargs=video_kwargs, image_kwargs=image_kwargs)
for v in obj
]
elif isinstance(obj, tuple):
return tuple(
_clean_log(v, video_kwargs=video_kwargs, image_kwargs=image_kwargs)
for v in obj
)
elif isinstance(obj, np.ndarray) and obj.ndim == 3:
# Must be single image (H, W, C).
return Image(obj, **image_kwargs)
elif isinstance(obj, np.ndarray) and obj.ndim == 4:
# Must be batch of images (N >= 1, H, W, C).
return (
_clean_log(
[Image(v, **image_kwargs) for v in obj],
video_kwargs=video_kwargs,
image_kwargs=image_kwargs,
)
if obj.shape[0] > 1
else Image(obj[0], **image_kwargs)
)
elif isinstance(obj, np.ndarray) and obj.ndim == 5:
# Must be batch of videos (N >= 1, T, C, W, H).
return (
_clean_log(
[Video(v, **video_kwargs) for v in obj],
video_kwargs=video_kwargs,
image_kwargs=image_kwargs,
)
if obj.shape[0] > 1
else Video(obj[0], **video_kwargs)
)
elif _is_allowed_type(obj):
return obj
# Else
try:
# This is what wandb uses internally. If we cannot dump
# an object using this method, wandb will raise an exception.
json_dumps_safer(obj)
# This is probably unnecessary, but left here to be extra sure.
pickle.dumps(obj)
return obj
except Exception:
# give up, similar to _SafeFallBackEncoder
fallback = str(obj)
# Try to convert to int
try:
fallback = int(fallback)
return fallback
except ValueError:
pass
# Try to convert to float
try:
fallback = float(fallback)
return fallback
except ValueError:
pass
# Else, return string
return fallback
def _get_wandb_project(project: Optional[str] = None) -> Optional[str]:
"""Get W&B project from environment variable or external hook if not passed
as and argument."""
if (
not project
and not os.environ.get(WANDB_PROJECT_ENV_VAR)
and os.environ.get(WANDB_POPULATE_RUN_LOCATION_HOOK)
):
# Try to populate WANDB_PROJECT_ENV_VAR and WANDB_GROUP_ENV_VAR
# from external hook
try:
load_class(os.environ[WANDB_POPULATE_RUN_LOCATION_HOOK])()
except Exception as e:
logger.exception(
f"Error executing {WANDB_POPULATE_RUN_LOCATION_HOOK} to "
f"populate {WANDB_PROJECT_ENV_VAR} and {WANDB_GROUP_ENV_VAR}: {e}",
exc_info=e,
)
if not project and os.environ.get(WANDB_PROJECT_ENV_VAR):
# Try to get project and group from environment variables if not
# passed through WandbLoggerCallback.
project = os.environ.get(WANDB_PROJECT_ENV_VAR)
return project
def _set_api_key(api_key_file: Optional[str] = None, api_key: Optional[str] = None):
"""Set WandB API key from `wandb_config`. Will pop the
`api_key_file` and `api_key` keys from `wandb_config` parameter.
The order of fetching the API key is:
1) From `api_key` or `api_key_file` arguments
2) From WANDB_API_KEY environment variables
3) User already logged in to W&B (wandb.api.api_key set)
4) From external hook WANDB_SETUP_API_KEY_HOOK
"""
if os.environ.get(WANDB_MODE_ENV_VAR) in {"offline", "disabled"}:
return
if api_key_file:
if api_key:
raise ValueError("Both WandB `api_key_file` and `api_key` set.")
with open(api_key_file, "rt") as fp:
api_key = fp.readline().strip()
if not api_key and not os.environ.get(WANDB_ENV_VAR):
# Check if user is already logged into wandb.
try:
wandb.ensure_configured()
if wandb.api.api_key:
logger.info("Already logged into W&B.")
return
except AttributeError:
pass
# Try to get API key from external hook
if WANDB_SETUP_API_KEY_HOOK in os.environ:
try:
api_key = load_class(os.environ[WANDB_SETUP_API_KEY_HOOK])()
except Exception as e:
logger.exception(
f"Error executing {WANDB_SETUP_API_KEY_HOOK} to setup API key: {e}",
exc_info=e,
)
if api_key:
os.environ[WANDB_ENV_VAR] = api_key
elif not os.environ.get(WANDB_ENV_VAR):
raise ValueError(
"No WandB API key found. Either set the {} environment "
"variable, pass `api_key` or `api_key_file` to the"
"`WandbLoggerCallback` class as arguments, "
"or run `wandb login` from the command line".format(WANDB_ENV_VAR)
)
def _run_wandb_process_run_info_hook(run: Any) -> None:
"""Run external hook to process information about wandb run"""
if WANDB_PROCESS_RUN_INFO_HOOK in os.environ:
try:
load_class(os.environ[WANDB_PROCESS_RUN_INFO_HOOK])(run)
except Exception as e:
logger.exception(
f"Error calling {WANDB_PROCESS_RUN_INFO_HOOK}: {e}", exc_info=e
)
class _QueueItem(enum.Enum):
END = enum.auto()
RESULT = enum.auto()
CHECKPOINT = enum.auto()
class _WandbLoggingActor:
"""
Wandb assumes that each trial's information should be logged from a
separate process. We use Ray actors as forking multiprocessing
processes is not supported by Ray and spawn processes run into pickling
problems.
We use a queue for the driver to communicate with the logging process.
The queue accepts the following items:
- If it's a dict, it is assumed to be a result and will be logged using
``wandb.log()``
- If it's a checkpoint object, it will be saved using ``wandb.log_artifact()``.
"""
def __init__(
self,
logdir: str,
queue: Queue,
exclude: List[str],
to_config: List[str],
*args,
**kwargs,
):
import wandb
self._wandb = wandb
os.chdir(logdir)
self.queue = queue
self._exclude = set(exclude)
self._to_config = set(to_config)
self.args = args
self.kwargs = kwargs
self._trial_name = self.kwargs.get("name", "unknown")
self._logdir = logdir
def run(self):
# Since we're running in a separate process already, use threads.
os.environ["WANDB_START_METHOD"] = "thread"
run = self._wandb.init(*self.args, **self.kwargs)
run.config.trial_log_path = self._logdir
_run_wandb_process_run_info_hook(run)
while True:
item_type, item_content = self.queue.get()
if item_type == _QueueItem.END:
break
if item_type == _QueueItem.CHECKPOINT:
self._handle_checkpoint(item_content)
continue
assert item_type == _QueueItem.RESULT
log, config_update = self._handle_result(item_content)
try:
self._wandb.config.update(config_update, allow_val_change=True)
self._wandb.log(log, step=log.get(TRAINING_ITERATION))
except urllib.error.HTTPError as e:
# Ignore HTTPError. Missing a few data points is not a
# big issue, as long as things eventually recover.
logger.warning("Failed to log result to w&b: {}".format(str(e)))
except FileNotFoundError as e:
logger.error(
"FileNotFoundError: Did not log result to Weights & Biases. "
"Possible cause: relative file path used instead of absolute path. "
"Error: %s",
e,
)
self._wandb.finish()
def _handle_checkpoint(self, checkpoint_path: str):
artifact = self._wandb.Artifact(
name=f"checkpoint_{self._trial_name}", type="model"
)
artifact.add_dir(checkpoint_path)
self._wandb.log_artifact(artifact)
def _handle_result(self, result: Dict) -> Tuple[Dict, Dict]:
config_update = result.get("config", {}).copy()
log = {}
flat_result = flatten_dict(result, delimiter="/")
for k, v in flat_result.items():
if any(k.startswith(item + "/") or k == item for item in self._exclude):
continue
elif any(k.startswith(item + "/") or k == item for item in self._to_config):
config_update[k] = v
elif not _is_allowed_type(v):
continue
else:
log[k] = v
config_update.pop("callbacks", None) # Remove callbacks
return log, config_update
@PublicAPI(stability="alpha")
class WandbLoggerCallback(LoggerCallback):
"""WandbLoggerCallback
Weights and biases (https://www.wandb.ai/) is a tool for experiment
tracking, model optimization, and dataset versioning. This Ray Tune
``LoggerCallback`` sends metrics to Wandb for automatic tracking and
visualization.
Example:
.. testcode::
import random
from ray import tune
from ray.air.integrations.wandb import WandbLoggerCallback
def train_func(config):
offset = random.random() / 5
for epoch in range(2, config["epochs"]):
acc = 1 - (2 + config["lr"]) ** -epoch - random.random() / epoch - offset
loss = (2 + config["lr"]) ** -epoch + random.random() / epoch + offset
train.report({"acc": acc, "loss": loss})
tuner = tune.Tuner(
train_func,
param_space={
"lr": tune.grid_search([0.001, 0.01, 0.1, 1.0]),
"epochs": 10,
},
run_config=tune.RunConfig(
callbacks=[WandbLoggerCallback(project="Optimization_Project")]
),
)
results = tuner.fit()
.. testoutput::
:hide:
...
Args:
project: Name of the Wandb project. Mandatory.
group: Name of the Wandb group. Defaults to the trainable
name.
api_key_file: Path to file containing the Wandb API KEY. This
file only needs to be present on the node running the Tune script
if using the WandbLogger.
api_key: Wandb API Key. Alternative to setting ``api_key_file``.
excludes: List of metrics and config that should be excluded from
the log.
log_config: Boolean indicating if the ``config`` parameter of
the ``results`` dict should be logged. This makes sense if
parameters will change during training, e.g. with
PopulationBasedTraining. Defaults to False.
upload_checkpoints: If ``True``, model checkpoints will be uploaded to
Wandb as artifacts. Defaults to ``False``.
save_checkpoints: Deprecated alias of ``upload_checkpoints``. Defaults to
``False``.
upload_timeout: Maximum time in seconds to wait for pending uploads to
wandb when the experiment ends. Defaults to the Ray Train default
sync timeout.
video_kwargs: Dictionary of keyword arguments passed to wandb.Video()
when logging videos. Videos have to be logged as 5D numpy arrays
to be affected by this parameter. For valid keyword arguments, see
https://docs.wandb.ai/ref/python/data-types/video/. Defaults to ``None``.
image_kwargs: Dictionary of keyword arguments passed to wandb.Image()
when logging images. Images have to be logged as 3D or 4D numpy arrays
to be affected by this parameter. For valid keyword arguments, see
https://docs.wandb.ai/ref/python/data-types/image/. Defaults to ``None``.
**kwargs: The keyword arguments will be passed to ``wandb.init()``.
Wandb's ``group``, ``run_id`` and ``run_name`` are automatically selected
by Tune, but can be overwritten by filling out the respective configuration
values.
Please see here for all other valid configuration settings:
https://docs.wandb.ai/ref/python/init/
""" # noqa: E501
# Do not log these result keys
_exclude_results = ["done", "should_checkpoint"]
AUTO_CONFIG_KEYS = [
"trial_id",
"experiment_tag",
"node_ip",
"experiment_id",
"hostname",
"pid",
"date",
]
"""Results that are saved with `wandb.config` instead of `wandb.log`."""
_logger_actor_cls = _WandbLoggingActor
def __init__(
self,
project: Optional[str] = None,
group: Optional[str] = None,
api_key_file: Optional[str] = None,
api_key: Optional[str] = None,
excludes: Optional[List[str]] = None,
log_config: bool = False,
upload_checkpoints: bool = False,
save_checkpoints: bool = False,
upload_timeout: int = DEFAULT_SYNC_TIMEOUT,
video_kwargs: Optional[dict] = None,
image_kwargs: Optional[dict] = None,
**kwargs,
):
if not wandb:
raise RuntimeError(
"Wandb was not found - please install with `pip install wandb`"
)
if save_checkpoints:
warnings.warn(
"`save_checkpoints` is deprecated. Use `upload_checkpoints` instead.",
DeprecationWarning,
)
upload_checkpoints = save_checkpoints
self.project = project
self.group = group
self.api_key_path = api_key_file
self.api_key = api_key
self.excludes = excludes or []
self.log_config = log_config
self.upload_checkpoints = upload_checkpoints
self._upload_timeout = upload_timeout
self.video_kwargs = video_kwargs or {}
self.image_kwargs = image_kwargs or {}
self.kwargs = kwargs
self._remote_logger_class = None
self._trial_logging_actors: Dict[
"Trial", ray.actor.ActorHandle[_WandbLoggingActor]
] = {}
self._trial_logging_futures: Dict["Trial", ray.ObjectRef] = {}
self._logging_future_to_trial: Dict[ray.ObjectRef, "Trial"] = {}
self._trial_queues: Dict["Trial", Queue] = {}
def setup(self, *args, **kwargs):
self.api_key_file = (
os.path.expanduser(self.api_key_path) if self.api_key_path else None
)
_set_api_key(self.api_key_file, self.api_key)
self.project = _get_wandb_project(self.project)
if not self.project:
raise ValueError(
"Please pass the project name as argument or through "
f"the {WANDB_PROJECT_ENV_VAR} environment variable."
)
if not self.group and os.environ.get(WANDB_GROUP_ENV_VAR):
self.group = os.environ.get(WANDB_GROUP_ENV_VAR)
def log_trial_start(self, trial: "Trial"):
config = trial.config.copy()
config.pop("callbacks", None) # Remove callbacks
exclude_results = self._exclude_results.copy()
# Additional excludes
exclude_results += self.excludes
# Log config keys on each result?
if not self.log_config:
exclude_results += ["config"]
# Fill trial ID and name
trial_id = trial.trial_id if trial else None
trial_name = str(trial) if trial else None
# Project name for Wandb
wandb_project = self.project
# Grouping
wandb_group = self.group or trial.experiment_dir_name if trial else None
# remove unpickleable items!
config = _clean_log(config)
config = {
key: value for key, value in config.items() if key not in self.excludes
}
wandb_init_kwargs = dict(
id=trial_id,
name=trial_name,
resume=False,
reinit=True,
allow_val_change=True,
group=wandb_group,
project=wandb_project,
config=config,
)
wandb_init_kwargs.update(self.kwargs)
self._start_logging_actor(trial, exclude_results, **wandb_init_kwargs)
def _start_logging_actor(
self, trial: "Trial", exclude_results: List[str], **wandb_init_kwargs
):
# Reuse actor if one already exists.
# This can happen if the trial is restarted.
if trial in self._trial_logging_futures:
return
if not self._remote_logger_class:
env_vars = {}
# API key env variable is not set if authenticating through `wandb login`
if WANDB_ENV_VAR in os.environ:
env_vars[WANDB_ENV_VAR] = os.environ[WANDB_ENV_VAR]
self._remote_logger_class = ray.remote(
num_cpus=0,
**_force_on_current_node(),
runtime_env={"env_vars": env_vars},
max_restarts=-1,
max_task_retries=-1,
)(self._logger_actor_cls)
self._trial_queues[trial] = Queue(
actor_options={
"num_cpus": 0,
**_force_on_current_node(),
"max_restarts": -1,
"max_task_retries": -1,
}
)
self._trial_logging_actors[trial] = self._remote_logger_class.remote(
logdir=trial.local_path,
queue=self._trial_queues[trial],
exclude=exclude_results,
to_config=self.AUTO_CONFIG_KEYS,
**wandb_init_kwargs,
)
logging_future = self._trial_logging_actors[trial].run.remote()
self._trial_logging_futures[trial] = logging_future
self._logging_future_to_trial[logging_future] = trial
def _signal_logging_actor_stop(self, trial: "Trial"):
self._trial_queues[trial].put((_QueueItem.END, None))
def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
if trial not in self._trial_logging_actors:
self.log_trial_start(trial)
result = _clean_log(
result, video_kwargs=self.video_kwargs, image_kwargs=self.image_kwargs
)
self._trial_queues[trial].put((_QueueItem.RESULT, result))
def log_trial_save(self, trial: "Trial"):
if self.upload_checkpoints and trial.checkpoint:
checkpoint_root = None
if isinstance(trial.checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
checkpoint_root = trial.checkpoint.path
if checkpoint_root:
self._trial_queues[trial].put((_QueueItem.CHECKPOINT, checkpoint_root))
def log_trial_end(self, trial: "Trial", failed: bool = False):
self._signal_logging_actor_stop(trial=trial)
self._cleanup_logging_actors()
def _cleanup_logging_actor(self, trial: "Trial"):
del self._trial_queues[trial]
del self._trial_logging_futures[trial]
ray.kill(self._trial_logging_actors[trial])
del self._trial_logging_actors[trial]
def _cleanup_logging_actors(self, timeout: int = 0, kill_on_timeout: bool = False):
"""Clean up logging actors that have finished uploading to wandb.
Waits for `timeout` seconds to collect finished logging actors.
Args:
timeout: The number of seconds to wait. Defaults to 0 to clean up
any immediate logging actors during the run.
This is set to a timeout threshold to wait for pending uploads
on experiment end.
kill_on_timeout: Whether or not to kill and cleanup the logging actor if
it hasn't finished within the timeout.
"""
futures = list(self._trial_logging_futures.values())
done, remaining = ray.wait(futures, num_returns=len(futures), timeout=timeout)
for ready_future in done:
finished_trial = self._logging_future_to_trial.pop(ready_future)
self._cleanup_logging_actor(finished_trial)
if kill_on_timeout:
for remaining_future in remaining:
trial = self._logging_future_to_trial.pop(remaining_future)
self._cleanup_logging_actor(trial)
def on_experiment_end(self, trials: List["Trial"], **info):
"""Wait for the actors to finish their call to `wandb.finish`.
This includes uploading all logs + artifacts to wandb."""
self._cleanup_logging_actors(timeout=self._upload_timeout, kill_on_timeout=True)
def __del__(self):
if ray.is_initialized():
for trial in list(self._trial_logging_actors):
self._signal_logging_actor_stop(trial=trial)
self._cleanup_logging_actors(timeout=2, kill_on_timeout=True)
self._trial_logging_actors = {}
self._trial_logging_futures = {}
self._logging_future_to_trial = {}
self._trial_queues = {}