Set mlflow run name from itwinai

Author: eldoo

Fires/_utilities/logger_itwinai.py

@@ -205,7 +205,8 @@ def __init__(
 			savedir: str, 
 			name: str = "itwinai", 
 			version:Optional[int | None] = None,
-			experiment_name:str = "experiment_name"):
+			experiment_name:str = "experiment_name",
+			run_name:Optional[str | None] = None):
 		super().__init__()
 		self.logger = MLFlowLogger(
 			savedir=savedir,

Fires/_utilities/utils_trainer.py

@@ -120,7 +120,7 @@ def get_itwinai_loggers() -> LoggersCollection:
 	_loggers = []
 
 	# define Itwinai MLFlow logger
-	_itwinai_mlflow_logger = Itwinai_MLFLogger(experiment_name=os.getenv('MLFLOW_EXPERIMENT_NAME'), tracking_uri=os.getenv('MLFLOW_TRACKING_URI'), log_freq='epoch')
+	_itwinai_mlflow_logger = Itwinai_MLFLogger(experiment_name=os.getenv('MLFLOW_EXPERIMENT_NAME'), run_name=os.getenv('MLFLOW_RUN_NAME'), tracking_uri=os.getenv('MLFLOW_TRACKING_URI'), log_freq='epoch')
 	_loggers.append(_itwinai_mlflow_logger)
 
 	# define Itwinai Provenance logger

phase_training_100.py

@@ -374,10 +374,6 @@ def main():
 		global_rank = trainer.global_rank
 		print(f" | Global rank {global_rank}")
 
-		if not global_rank: 
-			# change MLFLow run name
-			mlflow.set_tag('mlflow.runName', os.getenv('MLFLOW_RUN_NAME'))
-
 		# fit the model
 		trainer.fit(
 			model=model,