Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit f64f835

Browse files
spiralispirali
committed
cpu management
1 parent 7144980 commit f64f835

24 files changed

Lines changed: 474 additions & 92 deletions

‎Cargo.lock‎

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

‎src/client/commands/submit.rs‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,26 @@ use tako::messages::common::ProgramDefinition;
66
use crate::client::globalsettings::GlobalSettings;
77
use crate::client::job::print_job_detail;
88

9+
use crate::client::resources::parse_cpu_request;
910
use crate::common::arraydef::ArrayDef;
1011
use crate::rpc_call;
1112
use crate::transfer::connection::ClientConnection;
1213
use crate::transfer::messages::{FromClientMessage, JobType, SubmitRequest, ToClientMessage};
1314
use anyhow::anyhow;
15+
use std::str::FromStr;
16+
use tako::common::resources::{CpuRequest, ResourceRequest};
17+
18+
// I am wrapping CpuRequest + implementing FromStr as I am not able to provide
19+
// own parser function into clap??
20+
struct ArgCpuRequest(CpuRequest);
21+
22+
impl FromStr for ArgCpuRequest {
23+
type Err = anyhow::Error;
24+
25+
fn from_str(s: &str) -> Result<Self, Self::Err> {
26+
parse_cpu_request(s).map(ArgCpuRequest)
27+
}
28+
}
1429

1530
#[derive(Clap)]
1631
#[clap(setting = clap::AppSettings::ColoredHelp)]
@@ -21,15 +36,31 @@ pub struct SubmitOpts {
2136
#[clap(long)]
2237
array: Option<ArrayDef>,
2338

39+
#[clap(long, default_value = "1")]
40+
cpus: ArgCpuRequest,
41+
2442
#[clap(long)]
2543
name: Option<String>,
44+
45+
#[clap(long)]
46+
pin: bool,
47+
}
48+
49+
impl SubmitOpts {
50+
fn resource_request(&self) -> ResourceRequest {
51+
ResourceRequest::new(self.cpus.0.clone())
52+
}
2653
}
2754

2855
pub async fn submit_computation(
2956
gsettings: &GlobalSettings,
3057
connection: &mut ClientConnection,
3158
opts: SubmitOpts,
3259
) -> crate::Result<()> {
60+
let resources = opts.resource_request();
61+
62+
resources.validate()?;
63+
3364
let name = if let Some(name) = opts.name {
3465
validate_name(name)?
3566
} else {
@@ -47,7 +78,6 @@ pub async fn submit_computation(
4778
let submit_cwd = std::env::current_dir().unwrap();
4879
let stdout = submit_cwd.join("stdout.%{JOB_ID}.%{TASK_ID}");
4980
let stderr = submit_cwd.join("stderr.%{JOB_ID}.%{TASK_ID}");
50-
5181
let message = FromClientMessage::Submit(SubmitRequest {
5282
job_type,
5383
name,
@@ -59,6 +89,8 @@ pub async fn submit_computation(
5989
stderr: Some(stderr),
6090
cwd: None,
6191
},
92+
resources,
93+
pin: opts.pin,
6294
});
6395
let response = rpc_call!(connection, message, ToClientMessage::SubmitResponse(r) => r).await?;
6496
print_job_detail(gsettings, response.job, true, false);

‎src/client/job.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use cli_table::format::Justify;
22
use cli_table::{print_stdout, Cell, CellStruct, Color, Style, Table};
33

44
use crate::client::globalsettings::GlobalSettings;
5+
use crate::client::resources::cpu_request_to_string;
56
use crate::server::job::{JobTaskCounters, JobTaskInfo, JobTaskState};
67
use crate::transfer::messages::{JobDetail, JobInfo, JobType};
78
use crate::JobTaskCount;
@@ -161,6 +162,18 @@ pub fn print_job_detail(
161162

162163
rows.push(vec!["Tasks".cell().bold(true), n_tasks.cell()]);
163164

165+
let resources = cpu_request_to_string(job.resources.cpus());
166+
167+
rows.push(vec![
168+
"Resources".cell().bold(true),
169+
if job.pin {
170+
format!("{} [pin]", resources)
171+
} else {
172+
resources
173+
}
174+
.cell(),
175+
]);
176+
164177
// TODO: Each argument on own line, after the bug in cli-table is fixed
165178
let program_def = job.program_def;
166179
rows.push(vec![

‎src/client/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
pub mod commands;
22
pub mod globalsettings;
33
mod job;
4+
pub mod resources;
45
pub mod utils;
56
pub mod worker;

‎src/client/resources.rs‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
use crate::common::parser::{p_uint, NomResult};
2+
use anyhow::anyhow;
3+
use nom::branch::alt;
4+
use nom::bytes::complete::tag;
5+
use nom::character::complete::multispace1;
6+
use nom::combinator::{all_consuming, map, map_res, opt};
7+
use nom::sequence::{preceded, tuple};
8+
use tako::common::resources::CpuRequest;
9+
10+
fn p_cpu_request(input: &str) -> NomResult<CpuRequest> {
11+
alt((
12+
map(tag("all"), |_| CpuRequest::All),
13+
map_res(
14+
tuple((
15+
p_uint,
16+
opt(preceded(
17+
multispace1,
18+
alt((tag("compact!"), tag("compact"), tag("scatter"))),
19+
)),
20+
)),
21+
|(count, policy)| {
22+
if count == 0 {
23+
return Err(nom::Err::Error("xxx"));
24+
}
25+
Ok(match policy {
26+
None | Some("compact") => CpuRequest::Compact(count),
27+
Some("compact!") => CpuRequest::ForceCompact(count),
28+
Some("scatter") => CpuRequest::Scatter(count),
29+
_ => unreachable!(),
30+
})
31+
},
32+
),
33+
))(input)
34+
}
35+
36+
pub fn parse_cpu_request(input: &str) -> anyhow::Result<CpuRequest> {
37+
all_consuming(p_cpu_request)(input)
38+
.map(|r| r.1)
39+
.map_err(|e| anyhow!(e.to_string()))
40+
}
41+
42+
pub fn cpu_request_to_string(cr: &CpuRequest) -> String {
43+
match cr {
44+
CpuRequest::Compact(n_cpus) => {
45+
format!("{} compact", *n_cpus)
46+
}
47+
CpuRequest::ForceCompact(n_cpus) => {
48+
format!("{} compact!", *n_cpus)
49+
}
50+
CpuRequest::Scatter(n_cpus) => {
51+
format!("{} scatter", *n_cpus)
52+
}
53+
CpuRequest::All => "all".to_string(),
54+
}
55+
}
56+
57+
#[cfg(test)]
58+
mod test {
59+
use super::*;
60+
61+
#[test]
62+
fn test_parse_cpu_request() {
63+
assert_eq!(parse_cpu_request("all").unwrap(), CpuRequest::All);
64+
assert_eq!(parse_cpu_request("10").unwrap(), CpuRequest::Compact(10));
65+
assert_eq!(
66+
parse_cpu_request("5 compact").unwrap(),
67+
CpuRequest::Compact(5)
68+
);
69+
assert_eq!(
70+
parse_cpu_request("351 compact!").unwrap(),
71+
CpuRequest::ForceCompact(351)
72+
);
73+
assert_eq!(
74+
parse_cpu_request("10 scatter").unwrap(),
75+
CpuRequest::Scatter(10)
76+
);
77+
}
78+
}

‎src/client/worker.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ pub fn print_worker_info(workers: Vec<WorkerInfo>, gsettings: &GlobalSettings) {
1515
Some(_) => "OFFLINE".cell().foreground_color(Some(Color::Red)),
1616
},
1717
w.configuration.hostname.cell(),
18-
w.configuration.n_cpus.cell(),
18+
w.configuration.resources.summary().cell(),
1919
w.configuration
2020
.extra
2121
.get("MANAGER")
@@ -39,7 +39,7 @@ pub fn print_worker_info(workers: Vec<WorkerInfo>, gsettings: &GlobalSettings) {
3939
"Id".cell(),
4040
"State".cell().bold(true),
4141
"Hostname".cell().bold(true),
42-
"# cpus".cell().bold(true),
42+
"Resources".cell().bold(true),
4343
"Manager".cell().bold(true),
4444
"Manager Job Id".cell().bold(true),
4545
]);

‎src/common/arrayparser.rs‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,11 @@
11
use anyhow::anyhow;
22
use nom::bytes::complete::tag;
3-
use nom::character::complete::digit1;
43
use nom::combinator::all_consuming;
54
use nom::combinator::{map, map_res, opt};
6-
use nom::error::VerboseError;
75
use nom::sequence::{preceded, tuple};
8-
use nom::IResult;
96

107
use crate::common::arraydef::{ArrayDef, TaskIdRange};
11-
12-
pub type NomResult<'a, Ret> = IResult<&'a str, Ret, VerboseError<&'a str>>;
13-
14-
fn p_uint(input: &str) -> NomResult<u32> {
15-
map_res(digit1, |digit_str: &str| digit_str.parse::<u32>())(input)
16-
}
8+
use crate::common::parser::{p_uint, NomResult};
179

1810
fn p_task_id_range(input: &str) -> NomResult<TaskIdRange> {
1911
map_res(

‎src/common/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ pub mod arraydef;
44
pub mod arrayparser;
55
pub mod error;
66
pub mod fsutils;
7+
pub mod parser;
78
pub mod serverdir;
89
pub mod setup;
910
pub mod wrapped;

‎src/common/parser.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
use nom::character::complete::digit1;
2+
use nom::combinator::map_res;
3+
use nom::error::VerboseError;
4+
use nom::IResult;
5+
6+
pub type NomResult<'a, Ret> = IResult<&'a str, Ret, VerboseError<&'a str>>;
7+
8+
pub fn p_uint(input: &str) -> NomResult<u32> {
9+
map_res(digit1, |digit_str: &str| digit_str.parse::<u32>())(input)
10+
}

‎src/server/client.rs‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use std::sync::Arc;
44

55
use futures::{Sink, SinkExt, Stream, StreamExt};
66
use orion::kdf::SecretKey;
7-
use tako::messages::common::ProgramDefinition;
7+
use tako::messages::common::{LauncherDefinition, ProgramDefinition};
88
use tako::messages::gateway::{
99
CancelTasks, FromGatewayMessage, NewTasksMessage, StopWorkerRequest, TaskDef, ToGatewayMessage,
1010
};
@@ -241,16 +241,25 @@ async fn handle_submit(
241241
tako_ref: &TakoServer,
242242
message: SubmitRequest,
243243
) -> ToClientMessage {
244+
if message.resources.validate().is_err() {
245+
return ToClientMessage::Error("Invalid resource request".to_string());
246+
}
247+
244248
let make_task = |job_id, task_id, tako_id| {
245-
let program_def = make_program_def_for_task(&message.spec, job_id, task_id);
246-
let body = rmp_serde::to_vec_named(&program_def).unwrap();
249+
let program = make_program_def_for_task(&message.spec, job_id, task_id);
250+
let launcher_def = LauncherDefinition {
251+
program,
252+
pin: message.pin,
253+
};
254+
let body = rmp_serde::to_vec_named(&launcher_def).unwrap();
247255
TaskDef {
248256
id: tako_id,
249257
type_id: 0,
250258
body,
251259
keep: false,
252260
observe: true,
253261
n_outputs: 0,
262+
resources: message.resources.clone(),
254263
}
255264
};
256265
let (task_defs, job_detail) = {
@@ -275,6 +284,8 @@ async fn handle_submit(
275284
tako_base_id,
276285
message.name.clone(),
277286
message.spec,
287+
message.resources,
288+
message.pin,
278289
);
279290
let job_detail = job.make_job_detail(false);
280291
state.add_job(job);

0 commit comments

Comments
 (0)