Skip to content

Commit a7b6e8a

Browse files
committed
add support for downloading coman squash if no squash specified
1 parent 37974f2 commit a7b6e8a

1 file changed

Lines changed: 38 additions & 2 deletions

File tree

coman/src/cscs/handlers.rs

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,11 @@ use std::{
1111
use base64::prelude::*;
1212
use color_eyre::{Result, eyre::eyre};
1313
use eyre::Context;
14+
use futures::StreamExt;
1415
use iroh::SecretKey;
1516
use itertools::Itertools;
1617
use reqwest::Url;
18+
use tokio::{fs::File, io::AsyncWriteExt};
1719

1820
use super::api_client::client::{EdfSpec, ScriptSpec};
1921
use crate::{
@@ -219,7 +221,41 @@ async fn inject_coman_squash(
219221
let config = Config::new().unwrap();
220222
let local_squash_path = match config.values.coman_squash_path.clone() {
221223
Some(path) => path,
222-
None => todo!(), //download from github for architecture
224+
None => {
225+
//download from github for architecture
226+
let system = config
227+
.values
228+
.cscs
229+
.systems
230+
.get(current_system)
231+
.ok_or(eyre!("couldn't find architecture for system {}", current_system))?;
232+
let architecture = system
233+
.architecture
234+
.first()
235+
.ok_or(eyre!("no architecture set for {}", current_system))?;
236+
let target_path = get_data_dir().join(format!("coman_{}.sqsh", architecture));
237+
if !target_path.exists() {
238+
let url = match architecture.as_str() {
239+
"arm64" => {
240+
"https://github.com/SwissDataScienceCenter/coman/releases/latest/download/coman_Linux-aarch64.sqsh"
241+
}
242+
"amd64" => {
243+
"https://github.com/SwissDataScienceCenter/coman/releases/latest/download/coman_Linux-x86_64.sqsh"
244+
}
245+
_ => {
246+
return Err(eyre!("unsupported architecture {}", architecture));
247+
}
248+
};
249+
let mut out = File::create(target_path.clone()).await?;
250+
let mut resp = reqwest::get(url).await?.bytes_stream();
251+
while let Some(chunk_result) = resp.next().await {
252+
let chunk = chunk_result?;
253+
out.write_all(&chunk).await?;
254+
}
255+
out.flush().await?;
256+
}
257+
target_path
258+
}
223259
};
224260
let target = base_path.join("coman.sqsh");
225261
let file_meta = std::fs::metadata(local_squash_path.clone())?;
@@ -409,7 +445,7 @@ pub async fn cscs_job_start(
409445
match get_access_token().await {
410446
Ok(access_token) => {
411447
let api_client = CscsApi::new(access_token.0, platform).unwrap();
412-
let config = Config::new().unwrap();
448+
let config = Config::new()?;
413449
let current_system = &system.unwrap_or(config.values.cscs.current_system);
414450
let account = account.or(config.values.cscs.account);
415451
let user_info = api_client.get_userinfo(current_system).await?;

0 commit comments

Comments
 (0)