wehub-resource-sync
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from .utils import check_wandb_availability
|
|
from .monitor import Monitor
|
|
|
|
import deepspeed.comm as dist
|
|
|
|
|
|
class WandbMonitor(Monitor):
|
|
|
|
def __init__(self, wandb_config):
|
|
super().__init__(wandb_config)
|
|
check_wandb_availability()
|
|
import wandb
|
|
|
|
self.enabled = wandb_config.enabled
|
|
self.group = wandb_config.group
|
|
self.team = wandb_config.team
|
|
self.project = wandb_config.project
|
|
|
|
if self.enabled and dist.get_rank() == 0:
|
|
wandb.init(project=self.project, group=self.group, entity=self.team)
|
|
|
|
def log(self, data, step=None, commit=None):
|
|
if self.enabled and dist.get_rank() == 0:
|
|
import wandb
|
|
return wandb.log(data, step=step, commit=commit)
|
|
|
|
def write_events(self, event_list):
|
|
if self.enabled and dist.get_rank() == 0:
|
|
for event in event_list:
|
|
label = event[0]
|
|
value = event[1]
|
|
step = event[2]
|
|
self.log({label: value}, step=step)
|