update mlflow_data_extraction.py

Aarhus-Psychiatry-Research · May 15, 2024 · ba1932d · ba1932d
1 parent 1a7b580
commit ba1932d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/psycop/common/global_utils/mlflow/mlflow_data_extraction.py b/psycop/common/global_utils/mlflow/mlflow_data_extraction.py
@@ -82,6 +82,10 @@ def get_all_metrics_for_experiment(self, experiment_name: str) -> MlflowAllMetri
 
         return MlflowAllMetricsFrame(frame=metrics_df, allow_extra_columns=False)
 
+    def get_run(self, experiment_name: str, run_name: str) -> PsycopMlflowRun:
+        runs = self._get_mlflow_runs_by_experiment(experiment_name=experiment_name)
+        return next(run for run in runs if run.info.run_name == run_name)
+
     def get_best_run_from_experiment(
         self, experiment_name: str, metric: str, larger_is_better: bool = True
     ) -> PsycopMlflowRun: