diff --git a/python/ray/tune/examples/mlflow_example.py b/python/ray/tune/examples/mlflow_example.py index 68be24139..368726c0b 100644 --- a/python/ray/tune/examples/mlflow_example.py +++ b/python/ray/tune/examples/mlflow_example.py @@ -41,7 +41,9 @@ if __name__ == "__main__": num_samples=5, loggers=DEFAULT_LOGGERS + (MLFLowLogger, ), config={ - "mlflow_experiment_id": experiment_id, + "logger_config": { + "mlflow_experiment_id": experiment_id, + }, "width": tune.sample_from( lambda spec: 10 + int(90 * random.random())), "height": tune.sample_from(lambda spec: int(100 * random.random())) diff --git a/python/ray/tune/logger.py b/python/ray/tune/logger.py index d2fae3723..a1edf9cbd 100644 --- a/python/ray/tune/logger.py +++ b/python/ray/tune/logger.py @@ -78,9 +78,12 @@ class MLFLowLogger(Logger): """ def _init(self): + logger_config = self.config.get("logger_config", {}) from mlflow.tracking import MlflowClient - client = MlflowClient() - run = client.create_run(self.config.get("mlflow_experiment_id")) + client = MlflowClient( + tracking_uri=logger_config.get("mlflow_tracking_uri"), + registry_uri=logger_config.get("mlflow_registry_uri")) + run = client.create_run(logger_config.get("mlflow_experiment_id")) self._run_id = run.info.run_id for key, value in self.config.items(): client.log_param(self._run_id, key, value) diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index b2a3c1972..b21c2fff7 100644 --- a/rllib/agents/trainer.py +++ b/rllib/agents/trainer.py @@ -370,6 +370,10 @@ COMMON_CONFIG: TrainerConfigDict = { "replay_mode": "independent", }, + # === Logger === + # Define logger-specific configuration to be used inside Logger + "logger_config": {}, + # === Replay Settings === # The number of contiguous environment steps to replay at once. This may # be set to greater than 1 to support recurrent models.