From ba5449868b69058dafcda10aaf6374dad71e972f Mon Sep 17 00:00:00 2001 From: Evan Thomas Date: Tue, 29 Oct 2024 16:54:13 +0100 Subject: [PATCH] Add run status table --- migration/src/lib.rs | 2 + ...m20241029_154332_create_runstatus_table.rs | 131 ++++++++++++++++++ src/config.rs | 8 +- src/external/k8s/crd.rs | 1 - src/submissions/mod.rs | 1 + src/submissions/run_status/db.rs | 38 +++++ src/submissions/run_status/mod.rs | 1 + src/submissions/views.rs | 46 ++---- 8 files changed, 193 insertions(+), 35 deletions(-) create mode 100644 migration/src/m20241029_154332_create_runstatus_table.rs create mode 100644 src/submissions/run_status/db.rs create mode 100644 src/submissions/run_status/mod.rs diff --git a/migration/src/lib.rs b/migration/src/lib.rs index a44d074..a2b72f5 100644 --- a/migration/src/lib.rs +++ b/migration/src/lib.rs @@ -3,6 +3,7 @@ pub use sea_orm_migration::prelude::*; mod m20240926_143036_create_submission_table; mod m20241009_142236_create_system_status_table; mod m20241010_073350_create_input_objects; +mod m20241029_154332_create_runstatus_table; pub struct Migrator; @@ -13,6 +14,7 @@ impl MigratorTrait for Migrator { Box::new(m20240926_143036_create_submission_table::Migration), Box::new(m20241009_142236_create_system_status_table::Migration), Box::new(m20241010_073350_create_input_objects::Migration), + Box::new(m20241029_154332_create_runstatus_table::Migration), ] } } diff --git a/migration/src/m20241029_154332_create_runstatus_table.rs b/migration/src/m20241029_154332_create_runstatus_table.rs new file mode 100644 index 0000000..7f6b86f --- /dev/null +++ b/migration/src/m20241029_154332_create_runstatus_table.rs @@ -0,0 +1,131 @@ +use sea_orm_migration::prelude::*; +use sea_orm_migration::sea_orm::prelude::Json; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + // Create the RunStatus table + manager + .create_table( + Table::create() + .table(RunStatus::Table) + .if_not_exists() + .col(ColumnDef::new(RunStatus::Id).uuid().primary_key()) + .col(ColumnDef::new(RunStatus::SubmissionId).uuid().not_null()) + .col(ColumnDef::new(RunStatus::KubernetesPodName).string().null()) + .col(ColumnDef::new(RunStatus::Status).string().null()) + .col( + ColumnDef::new(RunStatus::IsRunning) + .boolean() + .not_null() + .default(false), + ) + .col( + ColumnDef::new(RunStatus::IsSuccessful) + .boolean() + .not_null() + .default(false), + ) + .col( + ColumnDef::new(RunStatus::IsStillKubernetesResource) + .boolean() + .not_null() + .default(false), + ) + .col(ColumnDef::new(RunStatus::TimeStarted).string().null()) + .col( + ColumnDef::new(RunStatus::Logs) + .json() + .not_null() + .default(Json::Array(vec![])), + ) + .col( + ColumnDef::new(RunStatus::TimeAddedUtc) + .date_time() + .not_null(), + ) + .col( + ColumnDef::new(RunStatus::LastUpdated) + .date_time() + .not_null(), + ) + .foreign_key( + ForeignKeyCreateStatement::new() + .name("fk_run_status_submission_id") + .from_tbl(RunStatus::Table) + .from_col(RunStatus::SubmissionId) + .to_tbl(Submissions::Table) + .to_col(Submissions::Id), + ) + .to_owned(), + ) + .await?; + + // Create indexes for the RunStatus table + manager + .create_index( + Index::create() + .name("idx_run_status_submission_id") + .table(RunStatus::Table) + .col(RunStatus::SubmissionId) + .to_owned(), + ) + .await?; + + manager + .create_index( + Index::create() + .name("idx_run_status_kubernetes_pod_name") + .table(RunStatus::Table) + .col(RunStatus::KubernetesPodName) + .to_owned(), + ) + .await?; + + manager + .create_index( + Index::create() + .name("idx_run_status_status") + .table(RunStatus::Table) + .col(RunStatus::Status) + .to_owned(), + ) + .await?; + + Ok(()) + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + // Drop the RunStatus table in the down migration + manager + .drop_table(Table::drop().table(RunStatus::Table).to_owned()) + .await?; + + Ok(()) + } +} + +#[derive(DeriveIden)] +enum RunStatus { + Table, + Id, + SubmissionId, + KubernetesPodName, + Status, + IsRunning, + IsSuccessful, + IsStillKubernetesResource, + TimeStarted, + Logs, + TimeAddedUtc, + LastUpdated, +} + +#[derive(DeriveIden)] +enum Submissions { + Table, + Id, +} diff --git a/src/config.rs b/src/config.rs index 6b7f962..bb66210 100644 --- a/src/config.rs +++ b/src/config.rs @@ -11,8 +11,8 @@ pub struct Config { pub db_name: String, pub db_prefix: String, pub db_url: Option, - pub s3_url: String, pub app_name: String, + pub s3_url: String, pub s3_bucket: String, pub s3_access_key: String, pub s3_secret_key: String, @@ -23,6 +23,8 @@ pub struct Config { pub _kube_config: PathBuf, pub kube_namespace: String, pub interval_external_services: u64, + pub submission_base_image: String, + pub submission_base_image_tag: String, pub s3_prefix: String, // Prefix within the bucket, ie. labcaller-dev pub pod_prefix: String, // What is prefixed to the pod name, ie. labcaller-dev} @@ -74,6 +76,10 @@ impl Config { .unwrap_or_else(|_| "60".to_string()) .parse() .unwrap(), + submission_base_image: env::var("SUBMISSION_BASE_IMAGE") + .expect("SUBMISSION_BASE_IMAGE must be set"), + submission_base_image_tag: env::var("SUBMISSION_BASE_IMAGE_TAG") + .expect("SUBMISSION_BASE_IMAGE_TAG must be set"), db_prefix, db_url, s3_prefix, diff --git a/src/external/k8s/crd.rs b/src/external/k8s/crd.rs index 492c099..2a36e7e 100644 --- a/src/external/k8s/crd.rs +++ b/src/external/k8s/crd.rs @@ -1,7 +1,6 @@ use kube::CustomResource; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use serde_json::Value; #[derive(CustomResource, Debug, Serialize, Deserialize, Clone, JsonSchema)] #[kube( diff --git a/src/submissions/mod.rs b/src/submissions/mod.rs index eca28b0..35f3a52 100644 --- a/src/submissions/mod.rs +++ b/src/submissions/mod.rs @@ -1,3 +1,4 @@ pub mod db; pub mod models; +pub mod run_status; pub mod views; diff --git a/src/submissions/run_status/db.rs b/src/submissions/run_status/db.rs new file mode 100644 index 0000000..277287b --- /dev/null +++ b/src/submissions/run_status/db.rs @@ -0,0 +1,38 @@ +use chrono::NaiveDateTime; +use sea_orm::entity::prelude::*; +use uuid::Uuid; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel)] +#[sea_orm(table_name = "run_status")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: Uuid, + pub submission_id: Uuid, + pub kubernetes_pod_name: Option, + pub status: Option, + pub is_running: bool, + pub is_successful: bool, + pub is_still_kubernetes_resource: bool, + pub time_started: Option, + pub logs: Json, + pub time_added_utc: NaiveDateTime, + pub last_updated: NaiveDateTime, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm( + belongs_to = "crate::submissions::db::Entity", + from = "Column::SubmissionId", + to = "crate::submissions::db::Column::Id" + )] + Submissions, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Submissions.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/src/submissions/run_status/mod.rs b/src/submissions/run_status/mod.rs new file mode 100644 index 0000000..dec1023 --- /dev/null +++ b/src/submissions/run_status/mod.rs @@ -0,0 +1 @@ +pub mod db; diff --git a/src/submissions/views.rs b/src/submissions/views.rs index 4a2cc16..00cbc58 100644 --- a/src/submissions/views.rs +++ b/src/submissions/views.rs @@ -3,10 +3,12 @@ use crate::common::filter::{apply_filters, parse_range}; use crate::common::models::FilterOptions; use crate::common::pagination::calculate_content_range; use crate::common::sort::generic_sort; -// use crate::external::k8s::crd::DictionaryField; +use crate::external::k8s::crd::{ + Environment, EnvironmentItems, TrainingWorkload, TrainingWorkloadSpec, ValueField, +}; use anyhow::Result; -use axum::debug_handler; use axum::{ + debug_handler, extract::{Path, Query, State}, http::StatusCode, response::IntoResponse, @@ -15,14 +17,12 @@ use axum::{ use axum_keycloak_auth::{ instance::KeycloakAuthInstance, layer::KeycloakAuthLayer, PassthroughMode, }; -use kube::api::PostParams; -use kube::Api; +use kube::{api::PostParams, Api}; use rand::Rng; use sea_orm::{ query::*, ActiveModelTrait, DatabaseConnection, DeleteResult, EntityTrait, IntoActiveModel, ModelTrait, SqlErr, }; -use serde_json::json; use std::sync::Arc; use uuid::Uuid; @@ -234,24 +234,6 @@ pub async fn delete_one(State(db): State, Path(id): Path, @@ -283,8 +265,11 @@ pub async fn execute_workflow( Err(_) => return StatusCode::INTERNAL_SERVER_ERROR, }; - // Create the `TrainingWorkload` custom resource instance - + let base_image = format!( + "{}:{}", + config.submission_base_image, config.submission_base_image_tag, + ); + // Create a new TrainingWorkload custom resource let training_workload = TrainingWorkload::new( &job_name, TrainingWorkloadSpec { @@ -313,16 +298,14 @@ pub async fn execute_workflow( value: id.to_string(), }, base_image: ValueField { - value: "registry.rcp.epfl.ch/rcp-test-ejthomas/dorado:0.2".to_string(), + value: base_image.clone(), }, }, }, gpu: ValueField { value: "1".to_string(), }, - image: ValueField { - value: "registry.rcp.epfl.ch/rcp-test-ejthomas/dorado:0.2".to_string(), - }, + image: ValueField { value: base_image }, image_pull_policy: ValueField { value: "Always".to_string(), }, @@ -332,10 +315,7 @@ pub async fn execute_workflow( run_as_gid: None, run_as_uid: None, run_as_user: None, - // run_as_gid: Some(ValueField { value: 1000 }), - // run_as_uid: Some(ValueField { value: 1000 }), - // run_as_user: Some(ValueField { value: true }), - service_type: None, // or Some(ValueField { value: "service_type_value".to_string() }) + service_type: None, usage: Some("Submit".to_string()), }, );