Skip to content

Commit

Permalink
Override W&B data on a resumed training (#595)
Browse files Browse the repository at this point in the history
* Override W&B data on a resumed training

* Suggestions
  • Loading branch information
vrigal authored May 15, 2024
1 parent adb890e commit ea95bc0
Showing 1 changed file with 33 additions and 10 deletions.
43 changes: 33 additions & 10 deletions tracking/translations_parser/publishers.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import csv
import logging
import os
from abc import ABC
from collections import defaultdict
from pathlib import Path
Expand Down Expand Up @@ -99,27 +100,49 @@ def open(self, parser) -> None:
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(
self.project,
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.")

Expand Down

0 comments on commit ea95bc0

Please sign in to comment.