Skip to content
This repository was archived by the owner on Jan 27, 2026. It is now read-only.

Commit 5653659

Browse files
JannikStJannikSt
authored
refactor variable namings in storage and scheduling, add storage var validation (#361)
1 parent 2b9fa15 commit 5653659

8 files changed

Lines changed: 129 additions & 38 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/orchestrator/src/api/routes/heartbeat.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,11 @@ mod tests {
141141
name: "test".to_string(),
142142
..Default::default()
143143
};
144-
app_state.store_context.task_store.add_task(task.into());
144+
let task = match task.try_into() {
145+
Ok(task) => task,
146+
Err(e) => panic!("Failed to convert TaskRequest to Task: {}", e),
147+
};
148+
app_state.store_context.task_store.add_task(task);
145149

146150
let req = test::TestRequest::post()
147151
.uri("/heartbeat")

‎crates/orchestrator/src/api/routes/storage.rs‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -113,11 +113,11 @@ async fn request_upload(
113113
match plugin.get_node_group(address) {
114114
Ok(Some(group)) => {
115115
group_id = Some(group.id.clone());
116-
file_name = file_name.replace("${node_group_id}", &group.id);
116+
file_name = file_name.replace("${NODE_GROUP_ID}", &group.id);
117117
file_name =
118-
file_name.replace("${node_group_size}", &group.nodes.len().to_string());
118+
file_name.replace("${NODE_GROUP_SIZE}", &group.nodes.len().to_string());
119119
let idx = plugin.get_idx_in_group(&group, address).unwrap();
120-
file_name = file_name.replace("${node_group_index}", &idx.to_string());
120+
file_name = file_name.replace("${NODE_GROUP_INDEX}", &idx.to_string());
121121
}
122122
Ok(None) => {
123123
log::warn!(
@@ -174,11 +174,11 @@ async fn request_upload(
174174
};
175175
let file_number = upload_count.saturating_sub(1);
176176

177-
if file_name.contains("${upload_count}") {
178-
file_name = file_name.replace("${upload_count}", &upload_count.to_string());
177+
if file_name.contains("${TOTAL_UPLOAD_COUNT_AFTER}") {
178+
file_name = file_name.replace("${TOTAL_UPLOAD_COUNT_AFTER}", &upload_count.to_string());
179179
}
180-
if file_name.contains("${file_number}") {
181-
file_name = file_name.replace("${file_number}", &file_number.to_string());
180+
if file_name.contains("${CURRENT_FILE_INDEX}") {
181+
file_name = file_name.replace("${CURRENT_FILE_INDEX}", &file_number.to_string());
182182
}
183183

184184
let file_size = &request_upload.file_size;
@@ -287,7 +287,7 @@ pub fn storage_routes() -> Scope {
287287

288288
fn generate_file_name(template: &str, original_name: &str) -> String {
289289
let mut file_name = template.to_string();
290-
file_name = file_name.replace("${original_name}", original_name);
290+
file_name = file_name.replace("${ORIGINAL_NAME}", original_name);
291291
file_name
292292
}
293293

@@ -308,7 +308,7 @@ mod tests {
308308

309309
#[tokio::test]
310310
async fn test_generate_file_name() {
311-
let template = "test/${original_name}";
311+
let template = "test/${ORIGINAL_NAME}";
312312
let original_name = "test";
313313
let file_name = generate_file_name(template, original_name);
314314
assert_eq!(file_name, "test/test");
@@ -323,7 +323,7 @@ mod tests {
323323
image: "test-image".to_string(),
324324
name: "test-task".to_string(),
325325
storage_config: Some(StorageConfig {
326-
file_name_template: Some("model_123/user_uploads/${original_name}".to_string()),
326+
file_name_template: Some("model_123/user_uploads/${ORIGINAL_NAME}".to_string()),
327327
}),
328328
..Default::default()
329329
};
@@ -439,7 +439,7 @@ mod tests {
439439
name: "test-task".to_string(),
440440
storage_config: Some(StorageConfig {
441441
file_name_template: Some(
442-
"model_xyz/dataset_1/${node_group_id}-${node_group_size}-${node_group_index}-${upload_count}-${file_number}.parquet".to_string(),
442+
"model_xyz/dataset_1/${NODE_GROUP_ID}-${NODE_GROUP_SIZE}-${NODE_GROUP_INDEX}-${TOTAL_UPLOAD_COUNT_AFTER}-${CURRENT_FILE_INDEX}.parquet".to_string(),
443443
),
444444
}),
445445
..Default::default()
@@ -572,7 +572,8 @@ mod tests {
572572
name: "test-task".to_string(),
573573
storage_config: Some(StorageConfig {
574574
file_name_template: Some(
575-
"model_xyz/dataset_1/${upload_count}-${file_number}.parquet".to_string(),
575+
"model_xyz/dataset_1/${TOTAL_UPLOAD_COUNT_AFTER}-${CURRENT_FILE_INDEX}.parquet"
576+
.to_string(),
576577
),
577578
}),
578579
..Default::default()

‎crates/orchestrator/src/api/routes/task.rs‎

Lines changed: 53 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,16 @@ async fn get_all_tasks(app_state: Data<AppState>) -> HttpResponse {
1212
let tasks = task_store.get_all_tasks();
1313
HttpResponse::Ok().json(json!({"success": true, "tasks": tasks}))
1414
}
15-
1615
async fn create_task(task: web::Json<TaskRequest>, app_state: Data<AppState>) -> HttpResponse {
17-
let task = Task::from(task.into_inner());
16+
let task = match Task::try_from(task.into_inner()) {
17+
Ok(task) => task,
18+
Err(e) => {
19+
return HttpResponse::BadRequest().json(json!({
20+
"success": false,
21+
"error": e
22+
}))
23+
}
24+
};
1825
let task_store = app_state.store_context.task_store.clone();
1926

2027
task_store.add_task(task.clone());
@@ -74,6 +81,38 @@ mod tests {
7481
assert_eq!(json["task"]["image"], "test");
7582
}
7683

84+
#[actix_web::test]
85+
async fn test_create_task_with_invalid_storage_config() {
86+
let app_state = create_test_app_state().await;
87+
let app = test::init_service(
88+
App::new()
89+
.app_data(app_state.clone())
90+
.route("/tasks", post().to(create_task)),
91+
)
92+
.await;
93+
94+
let payload = TaskRequest {
95+
image: "test".to_string(),
96+
name: "test".to_string(),
97+
storage_config: Some(shared::models::task::StorageConfig {
98+
file_name_template: Some("${INVALID_VAR}".to_string()),
99+
}),
100+
..Default::default()
101+
};
102+
103+
let req = test::TestRequest::post()
104+
.uri("/tasks")
105+
.set_json(payload)
106+
.to_request();
107+
let resp = test::call_service(&app, req).await;
108+
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
109+
110+
let body = test::read_body(resp).await;
111+
let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
112+
assert_eq!(json["success"], serde_json::Value::Bool(false));
113+
assert!(json["error"].is_string());
114+
}
115+
77116
#[actix_web::test]
78117
async fn get_task_when_no_task_exists() {
79118
let app_state = create_test_app_state().await;
@@ -102,7 +141,8 @@ mod tests {
102141
name: "test".to_string(),
103142
..Default::default()
104143
}
105-
.into();
144+
.try_into()
145+
.unwrap();
106146

107147
let task_store = app_state.store_context.task_store.clone();
108148
task_store.add_task(task);
@@ -135,7 +175,8 @@ mod tests {
135175
name: "test".to_string(),
136176
..Default::default()
137177
}
138-
.into();
178+
.try_into()
179+
.unwrap();
139180

140181
let task_store = app_state.store_context.task_store.clone();
141182
task_store.add_task(task.clone());
@@ -175,7 +216,9 @@ mod tests {
175216
name: "test1".to_string(),
176217
..Default::default()
177218
}
178-
.into();
219+
.try_into()
220+
.unwrap();
221+
179222
task_store.add_task(task1.clone());
180223

181224
// Wait briefly to ensure different timestamps
@@ -187,7 +230,9 @@ mod tests {
187230
name: "test2".to_string(),
188231
..Default::default()
189232
}
190-
.into();
233+
.try_into()
234+
.unwrap();
235+
191236
task_store.add_task(task2.clone());
192237

193238
// Test get all tasks - should be sorted by created_at descending
@@ -216,7 +261,8 @@ mod tests {
216261
name: format!("test{}", i),
217262
..Default::default()
218263
}
219-
.into();
264+
.try_into()
265+
.unwrap();
220266
task_store.add_task(task);
221267
thread::sleep(Duration::from_millis(10));
222268
}

‎crates/orchestrator/src/plugins/node_groups/scheduler_impl.rs‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ impl SchedulerPlugin for NodeGroupsPlugin {
9999

100100
// Temporary hack to get the upload count
101101
let pattern = format!("upload:{}:{}:*", node_address, group.id);
102-
let upload_count = match self.store.client.get_connection() {
102+
let total_upload_count = match self.store.client.get_connection() {
103103
Ok(mut conn) => {
104104
let mut keys: Vec<String> = Vec::new();
105105
match conn.scan_match(&pattern) {
@@ -121,7 +121,10 @@ impl SchedulerPlugin for NodeGroupsPlugin {
121121
}
122122
};
123123
// File number starts with 0 for the first file, while filecount is at 1
124-
let file_number = upload_count.parse::<u32>().unwrap_or(0).saturating_sub(1);
124+
let last_file_idx = total_upload_count
125+
.parse::<u32>()
126+
.unwrap_or(0)
127+
.saturating_sub(1);
125128

126129
let mut env_vars = task_clone.env_vars.unwrap_or_default();
127130
env_vars.insert("GROUP_INDEX".to_string(), idx.to_string());
@@ -131,8 +134,8 @@ impl SchedulerPlugin for NodeGroupsPlugin {
131134
.replace("${GROUP_SIZE}", &group.nodes.len().to_string())
132135
.replace("${NEXT_P2P_ADDRESS}", &next_p2p_id)
133136
.replace("${GROUP_ID}", &group.id)
134-
.replace("${UPLOAD_COUNT}", &upload_count.to_string())
135-
.replace("${FILE_NUMBER}", &file_number.to_string());
137+
.replace("${TOTAL_UPLOAD_COUNT}", &total_upload_count.to_string())
138+
.replace("${LAST_FILE_IDX}", &last_file_idx.to_string());
136139

137140
*value = new_value;
138141
}
@@ -144,8 +147,8 @@ impl SchedulerPlugin for NodeGroupsPlugin {
144147
.replace("${GROUP_SIZE}", &group.nodes.len().to_string())
145148
.replace("${NEXT_P2P_ADDRESS}", &next_p2p_id)
146149
.replace("${GROUP_ID}", &group.id)
147-
.replace("${UPLOAD_COUNT}", &upload_count.to_string())
148-
.replace("${FILE_NUMBER}", &file_number.to_string())
150+
.replace("${TOTAL_UPLOAD_COUNT}", &total_upload_count.to_string())
151+
.replace("${LAST_FILE_IDX}", &last_file_idx.to_string())
149152
})
150153
.collect::<Vec<String>>()
151154
});

‎crates/orchestrator/src/plugins/node_groups/tests.rs‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -400,8 +400,11 @@ async fn test_group_scheduling() {
400400
env_vars.insert("RANK".to_string(), "${GROUP_INDEX}".to_string());
401401
env_vars.insert("WORLD_SIZE".to_string(), "${GROUP_SIZE}".to_string());
402402
env_vars.insert("GROUP_ID".to_string(), "${GROUP_ID}".to_string());
403-
env_vars.insert("UPLOAD_COUNT".to_string(), "${UPLOAD_COUNT}".to_string());
404-
env_vars.insert("FILE_NUMBER".to_string(), "${FILE_NUMBER}".to_string());
403+
env_vars.insert(
404+
"TOTAL_UPLOAD_COUNT".to_string(),
405+
"${TOTAL_UPLOAD_COUNT}".to_string(),
406+
);
407+
env_vars.insert("LAST_FILE_IDX".to_string(), "${LAST_FILE_IDX}".to_string());
405408

406409
let task1 = Task {
407410
id: Uuid::new_v4(),
@@ -419,9 +422,9 @@ async fn test_group_scheduling() {
419422
"--group-id".to_string(),
420423
"${GROUP_ID}".to_string(),
421424
"--upload-count".to_string(),
422-
"${UPLOAD_COUNT}".to_string(),
425+
"${TOTAL_UPLOAD_COUNT}".to_string(),
423426
"--file-number".to_string(),
424-
"${FILE_NUMBER}".to_string(),
427+
"${LAST_FILE_IDX}".to_string(),
425428
]),
426429
state: TaskState::PENDING,
427430
created_at: 0,
@@ -471,8 +474,8 @@ async fn test_group_scheduling() {
471474
assert_eq!(env_vars_1.get("WORLD_SIZE").unwrap(), "2");
472475
assert_eq!(task_node_1.args.as_ref().unwrap()[3], "model/Qwen3-14B-0.2");
473476
assert_ne!(env_vars_1.get("GROUP_ID").unwrap(), "${GROUP_ID}");
474-
assert_eq!(env_vars_1.get("UPLOAD_COUNT").unwrap(), "1");
475-
assert_eq!(env_vars_1.get("FILE_NUMBER").unwrap(), "0");
477+
assert_eq!(env_vars_1.get("TOTAL_UPLOAD_COUNT").unwrap(), "1");
478+
assert_eq!(env_vars_1.get("LAST_FILE_IDX").unwrap(), "0");
476479
assert_eq!(task_node_1.args.as_ref().unwrap()[9], "1"); // Check upload count in args
477480

478481
assert_eq!(filtered_tasks_2.len(), 1);
@@ -483,8 +486,8 @@ async fn test_group_scheduling() {
483486
assert_eq!(env_vars_2.get("WORLD_SIZE").unwrap(), "2");
484487
assert_eq!(task_node_2.args.as_ref().unwrap()[3], "model/Qwen3-14B-1.2");
485488
assert_ne!(env_vars_2.get("GROUP_ID").unwrap(), "${GROUP_ID}");
486-
assert_eq!(env_vars_2.get("UPLOAD_COUNT").unwrap(), "0");
487-
assert_eq!(env_vars_2.get("FILE_NUMBER").unwrap(), "0");
489+
assert_eq!(env_vars_2.get("TOTAL_UPLOAD_COUNT").unwrap(), "0");
490+
assert_eq!(env_vars_2.get("LAST_FILE_IDX").unwrap(), "0");
488491
assert_eq!(task_node_2.args.as_ref().unwrap()[9], "0"); // Check upload count in args
489492

490493
assert_eq!(task_node_1.id, task_node_2.id);

‎crates/shared/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,3 +32,4 @@ google-cloud-storage = "0.24.0"
3232
base64 = "0.22.1"
3333
chrono = { workspace = true, features = ["serde"] }
3434
async-trait = "0.1.88"
35+
regex = "1.11.1"

‎crates/shared/src/models/task.rs‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -93,9 +93,41 @@ pub struct StorageConfig {
9393
pub file_name_template: Option<String>,
9494
}
9595

96-
impl From<TaskRequest> for Task {
97-
fn from(request: TaskRequest) -> Self {
98-
Task {
96+
impl StorageConfig {
97+
pub fn validate(&self) -> Result<(), String> {
98+
if let Some(template) = &self.file_name_template {
99+
let valid_vars = [
100+
"${ORIGINAL_NAME}",
101+
"${NODE_GROUP_ID}",
102+
"${NODE_GROUP_SIZE}",
103+
"${NODE_GROUP_INDEX}",
104+
"${TOTAL_UPLOAD_COUNT_AFTER}",
105+
"${CURRENT_FILE_INDEX}",
106+
];
107+
108+
let re = regex::Regex::new(r"\$\{[^}]+\}").unwrap();
109+
for cap in re.find_iter(template) {
110+
let var = cap.as_str();
111+
if !valid_vars.contains(&var) {
112+
return Err(format!(
113+
"Storage config template contains invalid variable: {}",
114+
var
115+
));
116+
}
117+
}
118+
}
119+
Ok(())
120+
}
121+
}
122+
impl TryFrom<TaskRequest> for Task {
123+
type Error = String;
124+
125+
fn try_from(request: TaskRequest) -> Result<Self, Self::Error> {
126+
if let Some(storage_config) = &request.storage_config {
127+
storage_config.validate()?;
128+
}
129+
130+
Ok(Task {
99131
id: Uuid::new_v4(),
100132
image: request.image,
101133
name: request.name,
@@ -107,7 +139,7 @@ impl From<TaskRequest> for Task {
107139
updated_at: None,
108140
scheduling_config: request.scheduling_config,
109141
storage_config: request.storage_config,
110-
}
142+
})
111143
}
112144
}
113145

0 commit comments

Comments
 (0)