File size: 1,592 Bytes
24c2665 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
from typing import Union, List
from verl.utils.tracking import Tracking
class ReasonRLTracking(Tracking):
def __init__(self, project_name, experiment_name, default_backend: Union[str, List[str]] = 'console', config=None, resume='never', run_id=None, tags: List[str] = None):
if isinstance(default_backend, str):
default_backend = [default_backend]
for backend in default_backend:
if backend == 'tracking':
import warnings
warnings.warn("`tracking` logger is deprecated. use `wandb` instead.", DeprecationWarning)
else:
assert backend in self.supported_backend, f'{backend} is not supported'
self.logger = {}
if 'tracking' in default_backend or 'wandb' in default_backend:
import wandb
wandb_kwargs = {}
if resume == 'must':
wandb_kwargs = {'resume': 'must', 'id': run_id}
elif resume == 'allow':
wandb_kwargs = {'resume': 'allow', 'id': run_id}
if tags is not None:
wandb_kwargs['tags'] = tags
run = wandb.init(project=project_name, settings=wandb.Settings(start_method="thread"), name=experiment_name, config=config, **wandb_kwargs)
self.run_id = run.id
self.logger['wandb'] = wandb
if 'console' in default_backend:
from verl.utils.logger.aggregate_logger import LocalLogger
self.console_logger = LocalLogger(print_to_console=True)
self.logger['console'] = self.console_logger
|