Merge pull request #441 from Modalities/wandb_entity

added entity parameter for w&b logging

Author: le1nux

src/modalities/config/config.py

@@ -492,6 +492,7 @@ class EvaluationResultToDiscSubscriberConfig(BaseModel):
 
 class WandBEvaluationResultSubscriberConfig(BaseModel):
     global_rank: int
+    entity: Optional[str] = None
     project: str
     experiment_id: str
     mode: WandbMode

src/modalities/logging_broker/subscriber_impl/results_subscriber.py

@@ -65,14 +65,16 @@ def __init__(
         project: str,
         experiment_id: str,
         mode: WandbMode,
-        logging_directory: Path,
+        logging_directory: Path | None,
         config_file_path: Path,
+        entity: str | None = None,
     ) -> None:
         super().__init__()
 
         with open(config_file_path, "r", encoding="utf-8") as file:
             config = yaml.safe_load(file)
         self.run = wandb.init(
+            entity=entity,
             project=project,
             name=experiment_id,
             mode=mode.value.lower(),
@@ -81,7 +83,7 @@ def __init__(
             settings=wandb.Settings(init_timeout=120),
         )
 
-        self.run.log_artifact(config_file_path, name=f"config_{wandb.run.id}", type="config")
+        self.run.log_artifact(config_file_path, name=f"config_{self.run.id}", type="config")
 
     def consume_dict(self, message_dict: dict[str, Any]):
         for k, v in message_dict.items():

src/modalities/logging_broker/subscriber_impl/subscriber_factory.py

@@ -68,6 +68,7 @@ def get_wandb_result_subscriber(
         mode: WandbMode,
         config_file_path: Path,
         directory: Optional[Path] = None,
+        entity: Optional[str] = None,
     ) -> WandBEvaluationResultSubscriber:
         if global_rank == 0 and (mode != WandbMode.DISABLED):
             if directory is not None:
@@ -88,7 +89,12 @@ def get_wandb_result_subscriber(
                 absolute_dir = None
 
             result_subscriber = WandBEvaluationResultSubscriber(
-                project, experiment_id, mode, absolute_dir, config_file_path
+                project=project,
+                experiment_id=experiment_id,
+                mode=mode,
+                logging_directory=absolute_dir,
+                config_file_path=config_file_path,
+                entity=entity,
             )
         else:
             result_subscriber = ResultsSubscriberFactory.get_dummy_result_subscriber()