Add an option to load W&B API key

Author: saridormi

compute_metrics.py

@@ -1,6 +1,6 @@
 import logging
 from pprint import pprint
-
+import os
 import hydra
 import pandas as pd
 import wandb
@@ -34,6 +34,9 @@ def main(cfg: MetricsConfig):
     print(f"==== Using config ====\n{OmegaConf.to_yaml(cfg)}")
 
     if cfg.logger.use_wandb:
+        if cfg.logger.use_api_key:
+            with open(hydra.utils.to_absolute_path("wandb_api_key.txt"), "r") as f:
+                os.environ["WANDB_API_KEY"] = f.read().strip()
         wandb.Table.MAX_ROWS = 50000
         run: wandb.wandb_sdk.wandb_run.Run = wandb.init(
             project=cfg.logger.project,

conf/eval_config.py

@@ -49,11 +49,13 @@ class WandbEvalConfig:
         use_wandb: Whether W&B will be used for logging or not.
         project: Name of project this run will appear in.
         load_artifact: Whether model checkpoint should be loaded from W&B artifact or not.
+        use_api_key: True to read an API key from a local file (expected to be stored in `wandb_api_key.txt`).
     """
 
     use_wandb: bool = True
     project: str = "commit_message_completion"
     load_artifact: bool = True
+    use_api_key: bool = False
     artifact_config: ArtifactEvalConfig = field(default_factory=ArtifactEvalConfig)
 
 

conf/metrics_config.py

@@ -39,11 +39,13 @@ class WandbMetricConfig:
         use_wandb: Whether W&B will be used for logging or not.
         project: Name of project this run will appear in.
         load_artifact: Whether model predictions should be loaded from W&B artifact or not.
+        use_api_key: True to read an API key from a local file (expected to be stored in `wandb_api_key.txt`).
     """
 
     use_wandb: bool = True
     project: str = "commit_message_completion"
     load_artifact: bool = True
+    use_api_key: bool = False
     artifact_config: ArtifactMetricConfig = field(default_factory=ArtifactMetricConfig)
 
 

conf/train_config.py

@@ -50,11 +50,13 @@ class WandbTrainConfig:
         use_wandb: Whether W&B will be used for logging or not.
         project: Name of a project this run will appear in.
         save_artifact: True to load model checkpoints to W&B as artifacts, False otherwise.
+        use_api_key: True to read an API key from a local file (expected to be stored in `wandb_api_key.txt`).
     """
 
     use_wandb: bool = True
     project: str = "commit_message_completion"
     save_artifact: bool = True
+    use_api_key: bool = False
 
 
 @dataclass

eval.py

@@ -54,6 +54,9 @@ def main(cfg: EvalConfig) -> None:
     )
 
     if cfg.logger.use_wandb and cfg.logger.load_artifact:
+        if cfg.logger.use_api_key:
+            with open(hydra.utils.to_absolute_path("wandb_api_key.txt"), "r") as f:
+                os.environ["WANDB_API_KEY"] = f.read().strip()
         artifact_name = f"{cfg.logger.artifact_config.project}/{run_name}:{cfg.logger.artifact_config.version}"
         artifact = trainer_logger.experiment.use_artifact(artifact_name)
         if "tags" in artifact.metadata:

train.py

@@ -95,6 +95,9 @@ def main(cfg: TrainConfig) -> None:
 
     # logger
     if cfg.logger.use_wandb:
+        if cfg.logger.use_api_key:
+            with open(hydra.utils.to_absolute_path("wandb_api_key.txt"), "r") as f:
+                os.environ["WANDB_API_KEY"] = f.read().strip()
         trainer_logger = pl.loggers.WandbLogger(
             name=run_name,
             project=cfg.logger.project,