diff --git a/src/main.rs b/src/main.rs index 4b320da..891d805 100644 --- a/src/main.rs +++ b/src/main.rs @@ -235,14 +235,12 @@ async fn run( args, ), Command::OfferTemplate => { - let template = offer_template::template(&runtime_config)?; - io::stdout().write_all(template.as_ref())?; - return Ok(()); - } - Command::Test => { - offer_template::gpu_detection(&runtime_config)?; + let offer_template = RUNTIME::offer_template(&runtime_config)?; + let offer_template = serde_json::to_string_pretty(&offer_template)?; + io::stdout().write_all(offer_template.as_bytes())?; return Ok(()); } + Command::Test => return RUNTIME::test(&runtime_config), }; let runtime_config = Box::pin(runtime_config); diff --git a/src/offer_template.rs b/src/offer_template.rs index d57f965..f7b1d0f 100644 --- a/src/offer_template.rs +++ b/src/offer_template.rs @@ -1,16 +1,9 @@ use crate::process::RuntimeConfig; -use gpu_detection::model::Gpu; +use gpu_detection::model::Gpu; use gpu_detection::GpuDetectionError; -use serde::{Deserialize, Serialize}; -use std::borrow::Cow; -use std::collections::BTreeMap; +use ya_agreement_utils::OfferTemplate; -#[derive(Deserialize, Serialize)] -struct OfferTemplate { - properties: BTreeMap, - constraints: String, -} pub(crate) fn gpu_detection(config: &CONFIG) -> anyhow::Result> { match gpu_detection::GpuDetection::init() { Ok(gpu_detection) => { @@ -22,18 +15,8 @@ pub(crate) fn gpu_detection(config: &CONFIG) -> anyhow::R } } -pub(crate) fn template( - config: &CONFIG, -) -> anyhow::Result> { +pub(crate) fn template(_config: &CONFIG) -> anyhow::Result { let offer_template = include_bytes!("offer-template.json"); - let mut template: OfferTemplate = serde_json::from_slice(offer_template.as_ref())?; - - if let Some(gpu) = gpu_detection(config)? { - let gpu = serde_json::value::to_value(gpu)?; - template - .properties - .insert("golem.!exp.gap-35.v1.inf.gpu".into(), gpu); - } - - Ok(Cow::Owned(serde_json::to_vec_pretty(&template)?)) + let template: OfferTemplate = serde_json::from_slice(offer_template.as_ref())?; + Ok(template) } diff --git a/src/process.rs b/src/process.rs index 7bd6e92..c85c5af 100644 --- a/src/process.rs +++ b/src/process.rs @@ -1,10 +1,9 @@ use anyhow::Context; use async_trait::async_trait; - use futures::TryFutureExt; use serde::de::DeserializeOwned; - use serde_json::Value; + use std::cell::RefCell; use std::env::current_exe; use std::fmt::Debug; @@ -15,6 +14,10 @@ use std::process::ExitStatus; use std::rc::Rc; use std::task::Poll; +use ya_agreement_utils::OfferTemplate; + +use crate::offer_template::{self, gpu_detection}; + pub mod automatic; pub mod dummy; pub mod win; @@ -40,6 +43,22 @@ pub(crate) trait Runtime: Sized { async fn stop(&mut self) -> anyhow::Result<()>; async fn wait(&mut self) -> std::io::Result; + + fn test(config: &Self::CONFIG) -> anyhow::Result<()> { + gpu_detection(config).context("Testing runtime failed. Unable to detect GPU.")?; + Ok(()) + } + + fn offer_template(config: &Self::CONFIG) -> anyhow::Result { + let mut template = offer_template::template(config)?; + if let Some(gpu) = gpu_detection(config) + .context("Generating offer template failed. Unable to detect GPU.")? + { + let gpu = serde_json::value::to_value(gpu)?; + template.set_property("golem.!exp.gap-35.v1.inf.gpu", gpu); + } + Ok(template) + } } pub(crate) trait RuntimeConfig: DeserializeOwned + Default + Debug + Clone { diff --git a/src/process/dummy.rs b/src/process/dummy.rs index 45787d4..b1c012a 100644 --- a/src/process/dummy.rs +++ b/src/process/dummy.rs @@ -8,6 +8,10 @@ use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::process::{Child, Command}; use tokio::sync::Mutex; +use ya_agreement_utils::OfferTemplate; + +use crate::offer_template; + use super::{Runtime, RuntimeConfig}; #[derive(Clone)] @@ -81,4 +85,12 @@ impl Runtime for Dummy { let mut child = self.child.lock().await; child.wait().await } + + fn test(_config: &Self::CONFIG) -> anyhow::Result<()> { + Ok(()) + } + + fn offer_template(config: &Self::CONFIG) -> anyhow::Result { + offer_template::template(config) + } }