use mlmd::metadata::{EventType, ExecutionState, PropertyType};
use mlmd::MetadataStore;
use structopt::StructOpt;
#[derive(StructOpt)]
struct Opt {
database_uri: String,
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
env_logger::init();
let opt = Opt::from_args();
let mut store = MetadataStore::connect(&opt.database_uri).await?;
println!("Create ArtifactTypes, e.g., Data and Model");
let data_type_id = store
.put_artifact_type("DataSet")
.property("day", PropertyType::Int)
.property("split", PropertyType::String)
.execute()
.await?;
let model_type_id = store
.put_artifact_type("SavedModel")
.property("version", PropertyType::Int)
.property("name", PropertyType::String)
.execute()
.await?;
println!("Create an ExecutionType, e.g., Trainer");
let trainer_type_id = store.put_execution_type("Trainer").execute().await?;
println!("Create an input artifact of type DataSet");
let data_artifact_id = store
.post_artifact(data_type_id)
.uri("path/to/data")
.property("day", 1i32)
.property("split", "train")
.execute()
.await?;
println!("Register the Execution of a Trainer run");
let trainer_run_id = store
.post_execution(trainer_type_id)
.state(ExecutionState::Running)
.execute()
.await?;
println!("Define the input event and record it in the metadata store");
store
.put_event(trainer_run_id, data_artifact_id)
.ty(EventType::DeclaredInput)
.execute()
.await?;
println!("Declare the output artifact of type SavedModel");
let model_artifact_id = store
.post_artifact(model_type_id)
.uri("path/to/model")
.property("version", 1i32)
.property("name", "MNIST-v1")
.execute()
.await?;
println!("Declare the output event and submit it to the Metadata Store");
store
.put_event(trainer_run_id, model_artifact_id)
.ty(EventType::DeclaredOutput)
.execute()
.await?;
println!("Mark the execution as completed");
store
.put_execution(trainer_run_id)
.state(ExecutionState::Complete)
.execute()
.await?;
println!("Create a ContextType, e.g., Experiment with a note property");
let experiment_type_id = store
.put_context_type("Experiment")
.property("note", PropertyType::String)
.execute()
.await?;
println!("Group the model and the trainer run to an experiment");
let experiment_id = store
.post_context(experiment_type_id, &format!("exp.{}", std::process::id()))
.property("note", "My first experiment")
.execute()
.await?;
store
.put_attribution(experiment_id, model_artifact_id)
.execute()
.await?;
store
.put_association(experiment_id, trainer_run_id)
.execute()
.await?;
Ok(())
}