use crate::api::ApiClient;
use crate::types::job::ExperimentDict;
use crate::Job;
pub struct Experiment {
uuid: String,
name: String,
api_client: ApiClient,
}
impl Experiment {
pub fn create_or_get(uri_or_name: &str) -> crate::Result<Self> {
let config = crate::Config::load();
let mut api_client = ApiClient::new(&config)?;
let response = api_client.create_or_get_experiment(uri_or_name)?;
Ok(Self {
uuid: response.uuid,
name: uri_or_name.to_string(),
api_client,
})
}
pub fn get_by_uri(uri: &str) -> crate::Result<Self> {
let config = crate::Config::load();
let mut api_client = ApiClient::new(&config)?;
let resource = api_client.get_resource(uri)?;
Ok(Self {
uuid: resource.uuid,
name: resource.name,
api_client,
})
}
pub fn uuid(&self) -> &str {
&self.uuid
}
pub fn name(&self) -> &str {
&self.name
}
pub fn add_job(&mut self, job: &Job) -> crate::Result<()> {
self.api_client
.add_job_to_experiment(&self.uuid, job.id())?;
Ok(())
}
pub fn get_jobs(&mut self, count: usize, status: Option<&str>) -> crate::Result<Vec<Job>> {
let page_size = std::cmp::min(count, 1_000);
let response = self
.api_client
.get_experiment_jobs(&self.uuid, page_size, status, None)?;
let mut jobs: Vec<Job> = Vec::new();
for jd in response.results {
if let Ok(job) = Job::from_uuid(&self.api_client, &jd.uuid, Some(&jd.auth_token)) {
jobs.push(job);
}
}
for page_number in 2..=response.page_count {
if jobs.len() >= count {
break;
}
let page_response = self.api_client.get_experiment_jobs(
&self.uuid,
page_size,
status,
Some(page_number),
)?;
for jd in page_response.results {
if let Ok(job) = Job::from_uuid(&self.api_client, &jd.uuid, Some(&jd.auth_token)) {
jobs.push(job);
}
}
}
jobs.truncate(count);
Ok(jobs)
}
pub fn wait(&mut self) -> crate::Result<()> {
loop {
let mut all_done = true;
let mut jobs = self.get_jobs(1000, None)?;
for job in &mut jobs {
if !job.is_finished()? {
all_done = false;
break;
}
}
if all_done {
break;
}
std::thread::sleep(std::time::Duration::from_secs(5));
}
Ok(())
}
pub fn show_experiments(count: usize) -> crate::Result<Vec<ExperimentDict>> {
let config = crate::Config::load();
let mut api_client = ApiClient::new(&config)?;
let response = api_client.get_experiments(count)?;
Ok(response.results)
}
}