chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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 = {}
|
||||
Reference in New Issue
Block a user