From 52ee27f3309e8ab51d2e5a95bee05b7935ef8aed Mon Sep 17 00:00:00 2001 From: "Bradford N. Roarr" Date: Thu, 27 Aug 2020 09:41:08 -0400 Subject: [PATCH] Adding parallelization via sbatch --- README.md | 2 +- src/main.rs | 142 +++++++++++++++++++++++++++++++++++++++++++--------- 2 files changed, 119 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 9d24fdc..65580d6 100644 --- a/README.md +++ b/README.md @@ -59,7 +59,7 @@ cargo clippy ## Future work -* Parallelize builds using [Slurm](https://slurm.schedmd.com/) +* ~~Parallelize builds using [Slurm](https://slurm.schedmd.com/)~~ * Support single Docker images without a manifest * Support Docker registries other than [Docker Hub](https://hub.docker.com) * Support native Singularity builds diff --git a/src/main.rs b/src/main.rs index 5cea675..2afc32b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,16 +3,17 @@ use chrono::{DateTime, Utc}; use clap::{App, Arg}; use serde::Deserialize; use std::fs::{self, File}; -use std::io::{self, Read}; +use std::io::{self, Read, Write}; use std::path::Path; -use std::process::Command; +use std::process::{Command, Stdio}; use std::time::SystemTime; #[derive(Debug)] struct Options { dry_run: bool, - force: bool, first_sync: usize, + force: bool, + parallelize: bool, } #[derive(Deserialize, Debug)] @@ -165,7 +166,110 @@ fn tags_after_timestamp( Ok(tags) } -fn sync_docker_image(image: &str, directory: &str, options: &Options) -> Result<()> { +fn slurm_sync_command( + directory: &str, + repository: &str, + image: &str, + tag: &str, + options: &Options, +) -> Result<()> { + let sif_path = format!("{}/{}/{}-{}.sif", directory, repository, image, tag); + let docker_uri = format!("docker://{}/{}:{}", repository, image, tag); + let job_name = format!("SingularitySync-{}-{}", repository, image); + let output = format!( + "/gpfs/scratch/%u/SingularitySync-{}-{}-%j.out", + repository, image + ); + let force = if options.force { "-F" } else { "" }; + + let memory = option_env!("SBATCH_MEM_PER_NODE"); + let memory = match memory { + Some(m) => m, + None => "8G", + }; + + let time = option_env!("SBATCH_TIMELIMIT"); + let time = match time { + Some(t) => t, + None => "8:00:00", + }; + + let singularity_cachedir = option_env!("SINGULARITY_CACHEDIR"); + let singularity_cachedir = match singularity_cachedir { + Some(x) => x, + None => "${HOME}/scratch/singularity", + }; + + let singularity_tmpdir = option_env!("SINGULARITY_TMPDIR"); + let singularity_tmpdir = match singularity_tmpdir { + Some(x) => x, + None => "${HOME}/scratch/tmp", + }; + + let script = [ + String::from("#!/usr/bin/env bash\n"), + format!("#SBATCH --time={}\n", time), + format!("#SBATCH --mem={}\n", memory), + format!("#SBATCH -J {}\n", job_name), + format!("#SBATCH -o {}\n", output), + format!("export SINGULARITY_CACHEDIR={}\n", singularity_cachedir), + format!("export SINGULARITY_TMPDIR={}\n", singularity_tmpdir), + format!("singularity build {} {} {}", force, sif_path, docker_uri), + ] + .concat(); + + if options.dry_run { + println!("sbatch < Result<()> { + if options.parallelize { + slurm_sync_command(directory, repository, image, tag, options)?; + return Ok(()); + } + + let sif_path = format!("{}/{}/{}-{}.sif", directory, repository, image, tag); + let docker_uri = format!("docker://{}/{}:{}", repository, image, tag); + + if options.dry_run { + let force = if options.force { "-F" } else { "" }; + let sbatch_cmd = format!("singularity build {} {} {}", force, sif_path, docker_uri); + println!("{}", sbatch_cmd); + return Ok(()); + } + + let mut command = Command::new("singularity"); + + command.arg("build"); + + if options.force { + command.arg("-F"); + } + + command.arg(sif_path).arg(docker_uri).status()?; + + Ok(()) +} + +fn sync_docker_image(directory: &str, image: &str, options: &Options) -> Result<()> { let image_split: Vec<&str> = image.rsplit('/').collect(); let image = String::from(image_split[0]); let repository = String::from(image_split[1]); @@ -190,24 +294,7 @@ fn sync_docker_image(image: &str, directory: &str, options: &Options) -> Result< }; for tag in tags_to_sync { - let sif_path = format!("{}/{}/{}-{}.sif", directory, repository, image, tag); - let docker_uri = format!("docker://{}/{}:{}", repository, image, tag); - - if options.dry_run { - let force = if options.force { "-F" } else { "" }; - let sbatch_cmd = format!("singularity build {} {} {}", force, sif_path, docker_uri); - println!("{}", sbatch_cmd); - } else { - let mut command = Command::new("singularity"); - - command.arg("build"); - - if options.force { - command.arg("-F"); - } - - command.arg(sif_path).arg(docker_uri).status()?; - } + run_sync_command(directory, &repository, &image, &tag, options)?; } Ok(()) @@ -215,7 +302,7 @@ fn sync_docker_image(image: &str, directory: &str, options: &Options) -> Result< fn sync_manifest(directory: &str, manifest: &Manifest, options: &Options) -> Result<()> { for image in &manifest.docker { - sync_docker_image(image, directory, options)?; + sync_docker_image(directory, image, options)?; } Ok(()) } @@ -257,6 +344,12 @@ fn main() -> Result<()> { .default_value("5") .help("The number of tags to pull on first sync"), ) + .arg( + Arg::with_name("parallelize") + .short("p") + .long("parallelize") + .help("Parallelize using Slurm scheduler"), + ) .version("v0.2.0") .get_matches(); @@ -266,8 +359,9 @@ fn main() -> Result<()> { let directory = String::from(matches.value_of("DIR").unwrap()); let options = Options { dry_run: matches.is_present("dry_run"), - force: matches.is_present("force"), first_sync: matches.value_of("first_sync").unwrap().parse()?, + force: matches.is_present("force"), + parallelize: matches.is_present("parallelize"), }; sync_manifest(&directory, &manifest, &options)?;