From be8611f5274c01832e1d43886eef434b239cd69e Mon Sep 17 00:00:00 2001 From: Arun Isaac Date: Wed, 9 Oct 2024 23:17:57 +0100 Subject: slurm-api: Allow requesting for a specific partition. * ravanan/batch-system.scm ()[partition]: Add field. * bin/ravanan (%options): Add slurm-partition. (print-usage): Document it. (main): Initialize partition field of object. * ravanan/command-line-tool.scm (run-command-line-tool): Pass #:partition argument to submit-job. * ravanan/slurm-api.scm (submit-job): Accept #:partition argument. --- bin/ravanan | 7 ++++++- ravanan/batch-system.scm | 8 +++++--- ravanan/command-line-tool.scm | 3 ++- ravanan/slurm-api.scm | 39 ++++++++++++++++++++++----------------- 4 files changed, 35 insertions(+), 22 deletions(-) diff --git a/bin/ravanan b/bin/ravanan index 5807a8e..2bcf0db 100755 --- a/bin/ravanan +++ b/bin/ravanan @@ -57,6 +57,9 @@ exec guile --no-auto-compile -e main -s "$0" "$@" (option (list "slurm-jwt") #t #f (lambda (opt name arg result) (acons 'slurm-jwt arg result))) + (option (list "slurm-partition") #t #f + (lambda (opt name arg result) + (acons 'slurm-partition arg result))) (option (list "help") #f #t (lambda (opt name arg result) (acons 'help #t result))))) @@ -85,6 +88,7 @@ Slurm API batch system options: --slurm-api-endpoint=SLURM-API-ENDPOINT slurm API endpoint to connect to --slurm-jwt=SLURM-JWT slurm JWT to authenticate with + --slurm-partition=SLURM-PARTITION slurm partition to request " program)) @@ -152,7 +156,8 @@ files that have the token in the @verbatim{SLURM_JWT=token} format." (slurm-api-batch-system (assq-ref args 'slurm-api-endpoint) (and (assq-ref args 'slurm-jwt) - (read-jwt (assq-ref args 'slurm-jwt)))))) + (read-jwt (assq-ref args 'slurm-jwt))) + (assq-ref args 'slurm-partition)))) #:guix-daemon-socket (assq-ref args 'guix-daemon-socket)) (current-output-port) #:pretty #t) diff --git a/ravanan/batch-system.scm b/ravanan/batch-system.scm index 726bff8..830b3a0 100644 --- a/ravanan/batch-system.scm +++ b/ravanan/batch-system.scm @@ -21,10 +21,12 @@ #:export (slurm-api-batch-system slurm-api-batch-system? slurm-api-batch-system-endpoint - slurm-api-batch-system-jwt)) + slurm-api-batch-system-jwt + slurm-api-batch-system-partition)) (define-immutable-record-type - (slurm-api-batch-system endpoint jwt) + (slurm-api-batch-system endpoint jwt partition) slurm-api-batch-system? (endpoint slurm-api-batch-system-endpoint) - (jwt slurm-api-batch-system-jwt)) + (jwt slurm-api-batch-system-jwt) + (partition slurm-api-batch-system-partition)) diff --git a/ravanan/command-line-tool.scm b/ravanan/command-line-tool.scm index 4254df9..eaaf90b 100644 --- a/ravanan/command-line-tool.scm +++ b/ravanan/command-line-tool.scm @@ -447,7 +447,8 @@ the same as in @code{run-workflow} from @code{(ravanan workflow)}." name script #:api-endpoint (slurm-api-batch-system-endpoint batch-system) - #:jwt (slurm-api-batch-system-jwt batch-system)))) + #:jwt (slurm-api-batch-system-jwt batch-system) + #:partition (slurm-api-batch-system-partition batch-system)))) (format (current-error-port) "~a submitted as job ID ~a~%" script diff --git a/ravanan/slurm-api.scm b/ravanan/slurm-api.scm index 15dc92a..5d9b660 100644 --- a/ravanan/slurm-api.scm +++ b/ravanan/slurm-api.scm @@ -63,25 +63,30 @@ document and pass in as the body of the HTTP request." (cut scm->json body-scm <>)))) (define* (submit-job environment stdout-file stderr-file cpus name script - #:key api-endpoint jwt) + #:key api-endpoint jwt partition) "Submit job named @var{name} running @var{script} to slurm via @var{api-endpoint} -and authenticating using @var{jwt}. @var{environment} is an association list of -environment variables to set in the job. @var{stdout-file} and @var{stderr-file} -are files in which to write the stdout and stderr of the job respectively. -@var{cpus} is the number of CPUs (in slurm terminology, a CPU is a hyperthread; see @url{https://slurm.schedmd.com/faq.html#cpu_count, the Slurm FAQ}) to request for -the job." +and authenticating using @var{jwt}. Request slurm @var{partition} if it is not +@code{#f}. @var{environment} is an association list of environment variables to +set in the job. @var{stdout-file} and @var{stderr-file} are files in which to +write the stdout and stderr of the job respectively. @var{cpus} is the number of +CPUs (in slurm terminology, a CPU is a hyperthread; see +@url{https://slurm.schedmd.com/faq.html#cpu_count, the Slurm FAQ}) to request +for the job." (define job-spec - `(("name" . ,name) - ("script" . ,(string-append "#!/bin/bash\n" script)) - ("environment" . ,(list->vector - (map (match-lambda - ((name . value) - (string-append name "=" value))) - environment))) - ("current_working_directory" . "/") - ("standard_output" . ,stdout-file) - ("standard_error" . ,stderr-file) - ("minimum_cpus" . ,cpus))) + (append `(("name" . ,name) + ("script" . ,(string-append "#!/bin/bash\n" script)) + ("environment" . ,(list->vector + (map (match-lambda + ((name . value) + (string-append name "=" value))) + environment))) + ("current_working_directory" . "/") + ("standard_output" . ,stdout-file) + ("standard_error" . ,stderr-file) + ("minimum_cpus" . ,cpus)) + (if partition + `(("partition" . ,partition)) + '()))) (let ((response (slurm-http-post api-endpoint jwt -- cgit v1.2.3