Skip to main content

mlflow/
client.rs

1use crate::{
2    storage::{
3        errors::{CreateExperimentError, GetExperimentError, StorageError},
4        Storage,
5    },
6    Experiment,
7};
8use std::sync::Arc;
9
10/// An MLflow client.
11///
12/// This is the heart of this library.
13/// It allows creating and accessing [`Experiment`]s.
14pub struct Client {
15    pub(crate) storage: Arc<dyn Storage>,
16}
17
18/// Possible backends for a `Client`.
19impl Client {
20    /// Create a `Client` for a MLflow Tracking Server.
21    ///
22    /// The `url` should be something like `http://127.0.0.1:5000/api`.
23    pub fn for_server(url: &str) -> Self {
24        Client {
25            storage: Arc::new(crate::storage::Server::new(url)),
26        }
27    }
28}
29
30/// Client methods without error handling.
31impl Client {
32    pub fn create_experiment(&self, name: &str) -> Option<Experiment> {
33        match self.try_create_experiment(name) {
34            Ok(experiment) => Some(experiment),
35            Err(CreateExperimentError::AlreadyExists(_)) => None,
36            Err(err @ CreateExperimentError::Storage(_)) => {
37                panic!("{}", err);
38            }
39        }
40    }
41
42    pub fn get_experiment(&self, name: &str) -> Option<Experiment> {
43        match self.try_get_experiment(name) {
44            Ok(experiment) => Some(experiment),
45            Err(GetExperimentError::DoesNotExist(_)) => None,
46            Err(err @ GetExperimentError::Storage(_)) => {
47                panic!("{}", err);
48            }
49        }
50    }
51
52    pub fn list_experiments(&self) -> Vec<Experiment> {
53        self.try_list_experiments().unwrap()
54    }
55}
56
57/// Client methods with error handling.
58impl Client {
59    pub fn try_create_experiment(&self, name: &str) -> Result<Experiment, CreateExperimentError> {
60        let primitive = self.storage.create_experiment(name)?;
61        Ok(Experiment::new(self, primitive))
62    }
63
64    pub fn try_get_experiment(&self, name: &str) -> Result<Experiment, GetExperimentError> {
65        let primitive = self.storage.get_experiment(name)?;
66        Ok(Experiment::new(self, primitive))
67    }
68
69    pub fn try_list_experiments(&self) -> Result<Vec<Experiment>, StorageError> {
70        let primitives = self.storage.list_experiments()?;
71        Ok(primitives
72            .into_iter()
73            .map(|e| Experiment::new(self, e))
74            .collect())
75    }
76}