chore: import upstream snapshot with attribution

python/ray/tune/logger/__init__.py

@@ -0,0 +1,15 @@
+from ray.tune.logger.csv import CSVLoggerCallback
+from ray.tune.logger.json import JsonLoggerCallback
+from ray.tune.logger.logger import (
+    LoggerCallback,
+    pretty_print,
+)
+from ray.tune.logger.tensorboardx import TBXLoggerCallback
+
+__all__ = [
+    "LoggerCallback",
+    "pretty_print",
+    "CSVLoggerCallback",
+    "JsonLoggerCallback",
+    "TBXLoggerCallback",
+]

python/ray/tune/logger/aim.py

@@ -0,0 +1,185 @@
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+import numpy as np
+
+from ray.air.constants import TRAINING_ITERATION
+from ray.tune.logger.logger import LoggerCallback
+from ray.tune.result import TIME_TOTAL_S, TIMESTEPS_TOTAL
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial
+
+try:
+    from aim.sdk import Repo, Run
+except ImportError:
+    Repo, Run = None, None
+
+logger = logging.getLogger(__name__)
+
+VALID_SUMMARY_TYPES = [int, float, np.float32, np.float64, np.int32, np.int64]
+
+
+@PublicAPI
+class AimLoggerCallback(LoggerCallback):
+    """Aim Logger: logs metrics in Aim format.
+
+    Aim is an open-source, self-hosted ML experiment tracking tool.
+    It's good at tracking lots (thousands) of training runs, and it allows you to
+    compare them with a performant and well-designed UI.
+
+    Source: https://github.com/aimhubio/aim
+    """
+
+    VALID_HPARAMS = (str, bool, int, float, list, type(None))
+    VALID_NP_HPARAMS = (np.bool_, np.float32, np.float64, np.int32, np.int64)
+
+    def __init__(
+        self,
+        repo: Optional[Union[str, "Repo"]] = None,
+        experiment_name: Optional[str] = None,
+        metrics: Optional[List[str]] = None,
+        **aim_run_kwargs,
+    ):
+        """Initialize the Aim logger callback.
+
+        Args:
+            repo: Aim repository directory or a `Repo` object that the Run object
+                will log results to. If not provided, a default repo will be set
+                up in the experiment directory (one level above trial directories).
+            experiment_name: Sets the `experiment` property of each Run object,
+                which is the experiment name associated with it. Can be used
+                later to query runs/sequences. If not provided, the default
+                will be the Tune experiment name set by `RunConfig(name=...)`.
+            metrics: List of metric names (out of the metrics reported by Tune)
+                to track in Aim. If no metric are specified, log everything
+                that is reported.
+            **aim_run_kwargs: Additional arguments that will be passed when
+                creating the individual `Run` objects for each trial. For the
+                full list of arguments, please see the Aim documentation:
+                https://aimstack.readthedocs.io/en/latest/refs/sdk.html
+        """
+        assert Run is not None, (
+            "aim must be installed!. You can install aim with"
+            " the command: `pip install aim`."
+        )
+        self._repo_path = repo
+        self._experiment_name = experiment_name
+        if not (bool(metrics) or metrics is None):
+            raise ValueError(
+                "`metrics` must either contain at least one metric name, or be None, "
+                "in which case all reported metrics will be logged to the aim repo."
+            )
+        self._metrics = metrics
+        self._aim_run_kwargs = aim_run_kwargs
+        self._trial_to_run: Dict["Trial", Run] = {}
+
+    def _create_run(self, trial: "Trial") -> Run:
+        """Initializes an Aim Run object for a given trial.
+
+        Args:
+            trial: The Tune trial that aim will track as a Run.
+
+        Returns:
+            Run: The created aim run for a specific trial.
+        """
+        experiment_dir = trial.local_experiment_path
+        run = Run(
+            repo=self._repo_path or experiment_dir,
+            experiment=self._experiment_name or trial.experiment_dir_name,
+            **self._aim_run_kwargs,
+        )
+        # Attach a few useful trial properties
+        run["trial_id"] = trial.trial_id
+        run["trial_log_dir"] = trial.path
+        trial_ip = trial.get_ray_actor_ip()
+        if trial_ip:
+            run["trial_ip"] = trial_ip
+        return run
+
+    def log_trial_start(self, trial: "Trial"):
+        if trial in self._trial_to_run:
+            # Cleanup an existing run if the trial has been restarted
+            self._trial_to_run[trial].close()
+
+        trial.init_local_path()
+        self._trial_to_run[trial] = self._create_run(trial)
+
+        if trial.evaluated_params:
+            self._log_trial_hparams(trial)
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        tmp_result = result.copy()
+
+        step = result.get(TIMESTEPS_TOTAL, None) or result[TRAINING_ITERATION]
+
+        for k in ["config", "pid", "timestamp", TIME_TOTAL_S, TRAINING_ITERATION]:
+            tmp_result.pop(k, None)  # not useful to log these
+
+        # `context` and `epoch` are special keys that users can report,
+        # which are treated as special aim metrics/configurations.
+        context = tmp_result.pop("context", None)
+        epoch = tmp_result.pop("epoch", None)
+
+        trial_run = self._trial_to_run[trial]
+        path = ["ray", "tune"]
+
+        flat_result = flatten_dict(tmp_result, delimiter="/")
+        valid_result = {}
+
+        for attr, value in flat_result.items():
+            if self._metrics and attr not in self._metrics:
+                continue
+
+            full_attr = "/".join(path + [attr])
+            if isinstance(value, tuple(VALID_SUMMARY_TYPES)) and not (
+                np.isnan(value) or np.isinf(value)
+            ):
+                valid_result[attr] = value
+                trial_run.track(
+                    value=value,
+                    name=full_attr,
+                    epoch=epoch,
+                    step=step,
+                    context=context,
+                )
+            elif (isinstance(value, (list, tuple, set)) and len(value) > 0) or (
+                isinstance(value, np.ndarray) and value.size > 0
+            ):
+                valid_result[attr] = value
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        trial_run = self._trial_to_run.pop(trial)
+        trial_run.close()
+
+    def _log_trial_hparams(self, trial: "Trial"):
+        params = flatten_dict(trial.evaluated_params, delimiter="/")
+        flat_params = flatten_dict(params)
+
+        scrubbed_params = {
+            k: v for k, v in flat_params.items() if isinstance(v, self.VALID_HPARAMS)
+        }
+
+        np_params = {
+            k: v.tolist()
+            for k, v in flat_params.items()
+            if isinstance(v, self.VALID_NP_HPARAMS)
+        }
+
+        scrubbed_params.update(np_params)
+        removed = {
+            k: v
+            for k, v in flat_params.items()
+            if not isinstance(v, self.VALID_HPARAMS + self.VALID_NP_HPARAMS)
+        }
+        if removed:
+            logger.info(
+                "Removed the following hyperparameter values when "
+                "logging to aim: %s",
+                str(removed),
+            )
+
+        run = self._trial_to_run[trial]
+        run["hparams"] = scrubbed_params

python/ray/tune/logger/comet.py

@@ -0,0 +1,3 @@
+from ray.air.integrations.comet import CometLoggerCallback
+
+CometLoggerCallback.__module__ = "ray.tune.logger.comet"

python/ray/tune/logger/csv.py

@@ -0,0 +1,79 @@
+import csv
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, TextIO
+
+from ray.air.constants import EXPR_PROGRESS_FILE
+from ray.tune.logger.logger import LoggerCallback
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI
+class CSVLoggerCallback(LoggerCallback):
+    """Logs results to progress.csv under the trial directory.
+
+    Automatically flattens nested dicts in the result dict before writing
+    to csv:
+
+        {"a": {"b": 1, "c": 2}} -> {"a/b": 1, "a/c": 2}
+
+    """
+
+    _SAVED_FILE_TEMPLATES = [EXPR_PROGRESS_FILE]
+
+    def __init__(self):
+        self._trial_continue: Dict["Trial", bool] = {}
+        self._trial_files: Dict["Trial", TextIO] = {}
+        self._trial_csv: Dict["Trial", csv.DictWriter] = {}
+
+    def _setup_trial(self, trial: "Trial"):
+        if trial in self._trial_files:
+            self._trial_files[trial].close()
+
+        # Make sure logdir exists
+        trial.init_local_path()
+        local_file_path = Path(trial.local_path, EXPR_PROGRESS_FILE)
+
+        # Resume the file from remote storage.
+        self._restore_from_remote(EXPR_PROGRESS_FILE, trial)
+
+        self._trial_continue[trial] = (
+            local_file_path.exists() and local_file_path.stat().st_size > 0
+        )
+
+        self._trial_files[trial] = local_file_path.open("at")
+        self._trial_csv[trial] = None
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        if trial not in self._trial_files:
+            self._setup_trial(trial)
+
+        tmp = result.copy()
+        tmp.pop("config", None)
+        result = flatten_dict(tmp, delimiter="/")
+
+        if not self._trial_csv[trial]:
+            self._trial_csv[trial] = csv.DictWriter(
+                self._trial_files[trial], result.keys()
+            )
+            if not self._trial_continue[trial]:
+                self._trial_csv[trial].writeheader()
+
+        self._trial_csv[trial].writerow(
+            {k: v for k, v in result.items() if k in self._trial_csv[trial].fieldnames}
+        )
+        self._trial_files[trial].flush()
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        if trial not in self._trial_files:
+            return
+
+        del self._trial_csv[trial]
+        self._trial_files[trial].close()
+        del self._trial_files[trial]

python/ray/tune/logger/json.py

@@ -0,0 +1,78 @@
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, TextIO
+
+import ray.cloudpickle as cloudpickle
+from ray.air.constants import EXPR_PARAM_FILE, EXPR_PARAM_PICKLE_FILE, EXPR_RESULT_FILE
+from ray.tune.logger.logger import LoggerCallback
+from ray.tune.utils.util import SafeFallbackEncoder
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI
+class JsonLoggerCallback(LoggerCallback):
+    """Logs trial results in json format.
+
+    Also writes to a results file and param.json file when results or
+    configurations are updated. Experiments must be executed with the
+    JsonLoggerCallback to be compatible with the ExperimentAnalysis tool.
+    """
+
+    _SAVED_FILE_TEMPLATES = [EXPR_RESULT_FILE, EXPR_PARAM_FILE, EXPR_PARAM_PICKLE_FILE]
+
+    def __init__(self):
+        self._trial_configs: Dict["Trial", Dict] = {}
+        self._trial_files: Dict["Trial", TextIO] = {}
+
+    def log_trial_start(self, trial: "Trial"):
+        if trial in self._trial_files:
+            self._trial_files[trial].close()
+
+        # Update config
+        self.update_config(trial, trial.config)
+
+        # Make sure logdir exists
+        trial.init_local_path()
+        local_file = Path(trial.local_path, EXPR_RESULT_FILE)
+
+        # Resume the file from remote storage.
+        self._restore_from_remote(EXPR_RESULT_FILE, trial)
+
+        self._trial_files[trial] = local_file.open("at")
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        if trial not in self._trial_files:
+            self.log_trial_start(trial)
+        json.dump(result, self._trial_files[trial], cls=SafeFallbackEncoder)
+        self._trial_files[trial].write("\n")
+        self._trial_files[trial].flush()
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        if trial not in self._trial_files:
+            return
+
+        self._trial_files[trial].close()
+        del self._trial_files[trial]
+
+    def update_config(self, trial: "Trial", config: Dict):
+        self._trial_configs[trial] = config
+
+        config_out = Path(trial.local_path, EXPR_PARAM_FILE)
+        with config_out.open("w") as f:
+            json.dump(
+                self._trial_configs[trial],
+                f,
+                indent=2,
+                sort_keys=True,
+                cls=SafeFallbackEncoder,
+            )
+
+        config_pkl = Path(trial.local_path, EXPR_PARAM_PICKLE_FILE)
+        with config_pkl.open("wb") as f:
+            cloudpickle.dump(self._trial_configs[trial], f)

python/ray/tune/logger/logger.py

@@ -0,0 +1,155 @@
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional, Set
+
+import pyarrow
+import yaml
+
+from ray.air._internal.json import SafeFallbackEncoder
+from ray.tune.callback import Callback
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+# Apply flow style for sequences of this length
+_SEQUENCE_LEN_FLOW_STYLE = 3
+
+
+@PublicAPI
+class LoggerCallback(Callback):
+    """Base class for experiment-level logger callbacks
+
+    This base class defines a general interface for logging events,
+    like trial starts, restores, ends, checkpoint saves, and receiving
+    trial results.
+
+    Callbacks implementing this interface should make sure that logging
+    utilities are cleaned up properly on trial termination, i.e. when
+    ``log_trial_end`` is received. This includes e.g. closing files.
+    """
+
+    def log_trial_start(self, trial: "Trial"):
+        """Handle logging when a trial starts.
+
+        Args:
+            trial: Trial object.
+        """
+        pass
+
+    def log_trial_restore(self, trial: "Trial"):
+        """Handle logging when a trial restores.
+
+        Args:
+            trial: Trial object.
+        """
+        pass
+
+    def log_trial_save(self, trial: "Trial"):
+        """Handle logging when a trial saves a checkpoint.
+
+        Args:
+            trial: Trial object.
+        """
+        pass
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        """Handle logging when a trial reports a result.
+
+        Args:
+            iteration: Iteration of the experiment that this result belongs to.
+            trial: Trial object.
+            result: Result dictionary.
+        """
+        pass
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        """Handle logging when a trial ends.
+
+        Args:
+            trial: Trial object.
+            failed: True if the Trial finished gracefully, False if
+                it failed (e.g. when it raised an exception).
+        """
+        pass
+
+    def on_trial_result(
+        self,
+        iteration: int,
+        trials: List["Trial"],
+        trial: "Trial",
+        result: Dict,
+        **info,
+    ):
+        self.log_trial_result(iteration, trial, result)
+
+    def on_trial_start(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_start(trial)
+
+    def on_trial_restore(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_restore(trial)
+
+    def on_trial_save(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_save(trial)
+
+    def on_trial_complete(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_end(trial, failed=False)
+
+    def on_trial_error(
+        self, iteration: int, trials: List["Trial"], trial: "Trial", **info
+    ):
+        self.log_trial_end(trial, failed=True)
+
+    def _restore_from_remote(self, file_name: str, trial: "Trial") -> None:
+        if not trial.checkpoint:
+            # If there's no checkpoint, there's no logging artifacts to restore
+            # since we're starting from scratch.
+            return
+
+        local_file = Path(trial.local_path, file_name).as_posix()
+        remote_file = Path(trial.storage.trial_fs_path, file_name).as_posix()
+
+        try:
+            pyarrow.fs.copy_files(
+                remote_file,
+                local_file,
+                source_filesystem=trial.storage.storage_filesystem,
+            )
+            logger.debug(f"Copied {remote_file} to {local_file}")
+        except FileNotFoundError:
+            logger.warning(f"Remote file not found: {remote_file}")
+        except Exception:
+            logger.exception(f"Error downloading {remote_file}")
+
+
+class _RayDumper(yaml.SafeDumper):
+    def represent_sequence(self, tag, sequence, flow_style=None):
+        if len(sequence) > _SEQUENCE_LEN_FLOW_STYLE:
+            return super().represent_sequence(tag, sequence, flow_style=True)
+        return super().represent_sequence(tag, sequence, flow_style=flow_style)
+
+
+@DeveloperAPI
+def pretty_print(result, exclude: Optional[Set[str]] = None):
+    result = result.copy()
+    result.update(config=None)  # drop config from pretty print
+    result.update(hist_stats=None)  # drop hist_stats from pretty print
+    out = {}
+    for k, v in result.items():
+        if v is not None and (exclude is None or k not in exclude):
+            out[k] = v
+
+    cleaned = json.dumps(out, cls=SafeFallbackEncoder)
+    return yaml.dump(json.loads(cleaned), Dumper=_RayDumper, default_flow_style=False)

python/ray/tune/logger/mlflow.py

@@ -0,0 +1,3 @@
+from ray.air.integrations.mlflow import MLflowLoggerCallback
+
+MLflowLoggerCallback.__module__ = "ray.tune.logger.mlflow"

python/ray/tune/logger/tensorboardx.py

@@ -0,0 +1,178 @@
+import logging
+from typing import TYPE_CHECKING, Dict
+
+import numpy as np
+
+from ray.air.constants import TRAINING_ITERATION
+from ray.tune.logger.logger import LoggerCallback
+from ray.tune.result import TIME_TOTAL_S, TIMESTEPS_TOTAL
+from ray.tune.utils import flatten_dict
+from ray.util.annotations import PublicAPI
+from ray.util.debug import log_once
+
+if TYPE_CHECKING:
+    from ray.tune.experiment.trial import Trial  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+VALID_SUMMARY_TYPES = [int, float, np.float32, np.float64, np.int32, np.int64]
+
+
+@PublicAPI
+class TBXLoggerCallback(LoggerCallback):
+    """TensorBoardX Logger.
+
+    Note that hparams will be written only after a trial has terminated.
+    This logger automatically flattens nested dicts to show on TensorBoard:
+
+        {"a": {"b": 1, "c": 2}} -> {"a/b": 1, "a/c": 2}
+    """
+
+    _SAVED_FILE_TEMPLATES = ["events.out.tfevents.*"]
+
+    VALID_HPARAMS = (str, bool, int, float, list, type(None))
+    VALID_NP_HPARAMS = (np.bool_, np.float32, np.float64, np.int32, np.int64)
+
+    def __init__(self):
+        try:
+            from tensorboardX import SummaryWriter
+
+            self._summary_writer_cls = SummaryWriter
+        except ImportError:
+            if log_once("tbx-install"):
+                logger.info('pip install "ray[tune]" to see TensorBoard files.')
+            raise
+        self._trial_writer: Dict["Trial", SummaryWriter] = {}
+        self._trial_result: Dict["Trial", Dict] = {}
+
+    def log_trial_start(self, trial: "Trial"):
+        if trial in self._trial_writer:
+            self._trial_writer[trial].close()
+        trial.init_local_path()
+        self._trial_writer[trial] = self._summary_writer_cls(
+            trial.local_path, flush_secs=30
+        )
+        self._trial_result[trial] = {}
+
+    def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
+        if trial not in self._trial_writer:
+            self.log_trial_start(trial)
+
+        step = result.get(TIMESTEPS_TOTAL) or result[TRAINING_ITERATION]
+
+        tmp = result.copy()
+        for k in ["config", "pid", "timestamp", TIME_TOTAL_S, TRAINING_ITERATION]:
+            if k in tmp:
+                del tmp[k]  # not useful to log these
+
+        flat_result = flatten_dict(tmp, delimiter="/")
+        path = ["ray", "tune"]
+        valid_result = {}
+
+        for attr, value in flat_result.items():
+            full_attr = "/".join(path + [attr])
+            if isinstance(value, tuple(VALID_SUMMARY_TYPES)) and not np.isnan(value):
+                valid_result[full_attr] = value
+                self._trial_writer[trial].add_scalar(full_attr, value, global_step=step)
+            elif (isinstance(value, list) and len(value) > 0) or (
+                isinstance(value, np.ndarray) and value.size > 0
+            ):
+                valid_result[full_attr] = value
+
+                # Must be a single image.
+                if isinstance(value, np.ndarray) and value.ndim == 3:
+                    self._trial_writer[trial].add_image(
+                        full_attr,
+                        value,
+                        global_step=step,
+                    )
+                    continue
+
+                # Must be a batch of images.
+                if isinstance(value, np.ndarray) and value.ndim == 4:
+                    self._trial_writer[trial].add_images(
+                        full_attr,
+                        value,
+                        global_step=step,
+                    )
+                    continue
+
+                # Must be video
+                if isinstance(value, np.ndarray) and value.ndim == 5:
+                    self._trial_writer[trial].add_video(
+                        full_attr, value, global_step=step, fps=20
+                    )
+                    continue
+
+                try:
+                    self._trial_writer[trial].add_histogram(
+                        full_attr, value, global_step=step
+                    )
+                # In case TensorboardX still doesn't think it's a valid value
+                # (e.g. `[[]]`), warn and move on.
+                except (ValueError, TypeError):
+                    if log_once("invalid_tbx_value"):
+                        logger.warning(
+                            "You are trying to log an invalid value ({}={}) "
+                            "via {}!".format(full_attr, value, type(self).__name__)
+                        )
+
+        self._trial_result[trial] = valid_result
+        self._trial_writer[trial].flush()
+
+    def log_trial_end(self, trial: "Trial", failed: bool = False):
+        if trial in self._trial_writer:
+            if trial and trial.evaluated_params and self._trial_result[trial]:
+                flat_result = flatten_dict(self._trial_result[trial], delimiter="/")
+                scrubbed_result = {
+                    k: value
+                    for k, value in flat_result.items()
+                    if isinstance(value, tuple(VALID_SUMMARY_TYPES))
+                }
+                self._try_log_hparams(trial, scrubbed_result)
+            self._trial_writer[trial].close()
+            del self._trial_writer[trial]
+            del self._trial_result[trial]
+
+    def _try_log_hparams(self, trial: "Trial", result: Dict):
+        # TBX currently errors if the hparams value is None.
+        flat_params = flatten_dict(trial.evaluated_params)
+        scrubbed_params = {
+            k: v for k, v in flat_params.items() if isinstance(v, self.VALID_HPARAMS)
+        }
+
+        np_params = {
+            k: v.tolist()
+            for k, v in flat_params.items()
+            if isinstance(v, self.VALID_NP_HPARAMS)
+        }
+
+        scrubbed_params.update(np_params)
+
+        removed = {
+            k: v
+            for k, v in flat_params.items()
+            if not isinstance(v, self.VALID_HPARAMS + self.VALID_NP_HPARAMS)
+        }
+        if removed:
+            logger.info(
+                "Removed the following hyperparameter values when "
+                "logging to tensorboard: %s",
+                str(removed),
+            )
+
+        from tensorboardX.summary import hparams
+
+        try:
+            experiment_tag, session_start_tag, session_end_tag = hparams(
+                hparam_dict=scrubbed_params, metric_dict=result
+            )
+            self._trial_writer[trial].file_writer.add_summary(experiment_tag)
+            self._trial_writer[trial].file_writer.add_summary(session_start_tag)
+            self._trial_writer[trial].file_writer.add_summary(session_end_tag)
+        except Exception:
+            logger.exception(
+                "TensorboardX failed to log hparams. "
+                "This may be due to an unsupported type "
+                "in the hyperparameter values."
+            )

python/ray/tune/logger/wandb.py

@@ -0,0 +1,3 @@
+from ray.air.integrations.wandb import WandbLoggerCallback
+
+WandbLoggerCallback.__module__ = "ray.tune.logger.wandb"