Skip to content

Commit

Permalink
Begin implementation of straw cli
Browse files Browse the repository at this point in the history
  • Loading branch information
HoKim98 committed Nov 1, 2023
1 parent 4f0e5db commit 9fdedcd
Show file tree
Hide file tree
Showing 13 changed files with 581 additions and 49 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ members = [
"kiss/manager",
"kiss/monitor",
"straw/api",
"straw/cli",
"straw/provider",
"vine/api",
"vine/bastion",
Expand Down
2 changes: 1 addition & 1 deletion dash/controller/src/validator/pipe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ impl<'namespace, 'kube> PipeValidator<'namespace, 'kube> {
}

async fn validate_exec_straw(&self, exec: StrawPipe) -> Result<StrawPipe> {
let session = StrawSession::new(self.kube.clone());
let session = StrawSession::new(self.kube.clone(), Some(self.namespace.into()));
session.create(&exec).await.map(|()| exec)
}
}
4 changes: 3 additions & 1 deletion dash/pipe/functions/ai/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@ straw = ["kube", "straw-api"]
[dependencies]
ark-core-k8s = { path = "../../../../ark/core/k8s", features = ["data"] }
dash-pipe-provider = { path = "../../provider", features = ["pyo3"] }
straw-api = { path = "../../../../straw/api", optional = true }
straw-api = { path = "../../../../straw/api", optional = true, features = [
"plugin",
] }

anyhow = { workspace = true }
async-trait = { workspace = true }
Expand Down
2 changes: 1 addition & 1 deletion dash/pipe/functions/ai/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ impl ::dash_pipe_provider::FunctionBuilder for Function {
ai_model_kind: kind,
} = args;

let code = self::plugin::Plugin::new().load_code(model)?;
let code = self::plugin::PluginBuilder::new().load_code(model)?;

Ok(Self {
tick: Python::with_gil(|py| {
Expand Down
38 changes: 13 additions & 25 deletions dash/pipe/functions/ai/src/plugin/mod.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
use anyhow::{anyhow, Result};
use ark_core_k8s::data::Url;

pub struct Plugin<'a> {
pub struct PluginBuilder<'a> {
loaders: &'a [ModelLoader<'a>],
}

impl<'a> Plugin<'a> {
impl<'a> PluginBuilder<'a> {
pub const fn new() -> Self {
Self {
loaders: &[ModelLoader {
Expand All @@ -16,25 +16,22 @@ impl<'a> Plugin<'a> {
}

pub fn load_code(&self, model: &Url) -> Result<&'a str> {
self.loaders
.iter()
.find(|&loader| loader.scheme == model.scheme())
self.try_load(model)
.map(|loader| loader.code)
.ok_or_else(|| anyhow!("unsupported model URL scheme: {model}"))
}

fn try_load(&self, model: &Url) -> Option<&ModelLoader<'a>> {
let scheme = model.scheme();
self.loaders.iter().find(|&loader| loader.scheme == scheme)
}
}

#[cfg(feature = "straw")]
impl ::straw_api::plugin::PluginBuilder for Plugin<'static> {
fn try_build(
&self,
scheme: &str,
) -> Result<Option<Box<dyn ::straw_api::plugin::Plugin>>> {
Ok(self
.loaders
.iter()
.find(|&loader| loader.scheme == scheme)
.map(|loader| Box::new(*loader) as Box<dyn ::straw_api::plugin::Plugin>))
impl ::straw_api::plugin::PluginBuilder for PluginBuilder<'static> {
fn try_build(&self, url: &Url) -> Option<::straw_api::plugin::DynPlugin> {
self.try_load(url)
.map(|loader| Box::new(*loader) as ::straw_api::plugin::DynPlugin)
}
}

Expand All @@ -45,13 +42,4 @@ pub struct ModelLoader<'a> {
}

#[cfg(feature = "straw")]
#[::async_trait::async_trait]
impl<'a> ::straw_api::plugin::Plugin for ModelLoader<'a> {
async fn create(&self, client: &::kube::Client) -> Result<()> {
todo!()
}

async fn delete(&self, client: &::kube::Client) -> Result<()> {
todo!()
}
}
impl<'a> ::straw_api::plugin::PluginDaemon for ModelLoader<'a> {}
14 changes: 11 additions & 3 deletions straw/api/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,20 @@ version = { workspace = true }

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[features]
plugin = ["anyhow", "async-trait", "kube"]

[dependencies]
ark-core-k8s = { path = "../../ark/core/k8s", features = ["data"] }

anyhow = { workspace = true }
async-trait = { workspace = true }
anyhow = { workspace = true, optional = true }
async-trait = { workspace = true, optional = true }
k8s-openapi = { workspace = true }
kube = { workspace = true, features = ["client", "derive"] }
kube = { workspace = true, optional = true, features = [
"client",
"runtime",
"rustls-tls",
"ws",
] }
schemars = { workspace = true }
serde = { workspace = true }
5 changes: 5 additions & 0 deletions straw/api/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,2 +1,7 @@
pub mod pipe;
#[cfg(feature = "plugin")]
pub mod plugin;

pub mod name {
pub const NAME: &str = "straw";
}
Loading

0 comments on commit 9fdedcd

Please sign in to comment.