Override W&B data on a resumed training (#595)

* Override W&B data on a resumed training

* Suggestions
This commit is contained in:
Valentin Rigal 2024-05-15 20:35:54 +02:00 коммит произвёл GitHub
Родитель adb890e14c
Коммит ea95bc0cb1
Не найден ключ, соответствующий данной подписи
Идентификатор ключа GPG: B5690EEEBB952194
1 изменённых файлов: 33 добавлений и 10 удалений

Просмотреть файл

@ -1,5 +1,6 @@
import csv
import logging
import os
from abc import ABC
from collections import defaultdict
from pathlib import Path
@ -99,8 +100,9 @@ class WandB(Publisher):
config.update(self.extra_kwargs.pop("config", {}))
try:
# Check if a W&B run already exists with this name
project = next(filter(lambda p: p.name == self.project, wandb.Api().projects()), None)
# Check if a W&B run already exists with this name
existing_runs = []
if project and (name := self.extra_kwargs.get("name")):
existing_runs = list(
wandb.Api().runs(
@ -108,18 +110,39 @@ class WandB(Publisher):
filters={"display_name": name, "group": self.extra_kwargs.get("group")},
)
)
if len(existing_runs) > 0:
if len(existing_runs) == 0:
# Start a new W&B run
self.wandb = wandb.init(
project=self.project,
config=config,
**self.extra_kwargs,
)
return
elif len(existing_runs) == 1:
run = existing_runs[0]
# Avoid overriding an existing run on a first training, this should not happen
if int(os.environ.get("RUN_ID", 0)) < 1:
logger.warning(
f"This run already exists on W&B: {existing_runs}. No data will be published."
f"A W&B run already exists with name '{name}': {run}. No data will be published."
)
return
# Start a W&B run
self.wandb = wandb.init(
project=self.project,
config=config,
**self.extra_kwargs,
)
# Resume an existing run
logger.info(
f"Training has been resumed from an earlier run wit name '{name}', "
f"continue W&B publication with run {run}."
)
self.wandb = wandb.init(
project=self.project,
config=config,
id=run.id,
resume="must",
**self.extra_kwargs,
)
else:
logger.warning(
f"Multiple W&B runs already exist with name '{name}': {existing_runs}. No data will be published."
)
return
except Exception as e:
logger.error(f"WandB client could not be initialized: {e}. No data will be published.")