From bab6bc40bb62c153c7a7f2deb5222a1f8bffbce9 Mon Sep 17 00:00:00 2001
From: Arun Isaac
Date: Wed, 2 Oct 2024 01:35:37 +0100
Subject: slurm-api: Allow requesting a certain number of CPUs.

* ravanan/slurm-api.scm (submit-job): Accept cpus argument.
* ravanan/command-line-tool.scm (run-command-line-tool): Pass cpus
argument.
(build-command-line-tool-script)[cores]: New function.
Set runtime.cores based on CPU allocation.
---
 ravanan/command-line-tool.scm | 12 +++++++++++-
 ravanan/slurm-api.scm         |  9 ++++++---
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/ravanan/command-line-tool.scm b/ravanan/command-line-tool.scm
index 6a48a20..205ab48 100644
--- a/ravanan/command-line-tool.scm
+++ b/ravanan/command-line-tool.scm
@@ -451,6 +451,7 @@ path."
                                           ,store-data-file))
                                        stdout-file
                                        stderr-file
+                                       1
                                        name
                                        script
                                        #:api-endpoint slurm-api-endpoint
@@ -607,6 +608,15 @@ named @var{name} with @var{inputs} using tools from Guix manifest
                       (assoc-ref initial-work-dir-requirement
                                  "listing")))
 
+  (define (cores batch-system)
+    (case batch-system
+      ((slurm-api)
+       #~(getenv "SLURM_CPUS_ON_NODE"))
+      ((single-machine)
+       #~(total-processor-count))
+      (else
+       (assertion-violation batch-system "Unknown batch system"))))
+
   (define stdout-filename
     (cond
      ;; stdout filename or expression is specified.
@@ -927,7 +937,7 @@ directory of the workflow."
               (call-with-temporary-directory
                (lambda (inputs-directory)
                  (let ((inputs #$(copy-input-files-gexp (canonicalize-json inputs)))
-                       (runtime `(("cores" . ,(total-processor-count)))))
+                       (runtime `(("cores" . ,#$(cores batch-system)))))
 
                    ;; Set environment defined by workflow.
                    (for-each (match-lambda
diff --git a/ravanan/slurm-api.scm b/ravanan/slurm-api.scm
index fdc89fc..15dc92a 100644
--- a/ravanan/slurm-api.scm
+++ b/ravanan/slurm-api.scm
@@ -62,12 +62,14 @@ document and pass in as the body of the HTTP request."
                       #:body (call-with-output-bytevector
                               (cut scm->json body-scm <>))))
 
-(define* (submit-job environment stdout-file stderr-file name script
+(define* (submit-job environment stdout-file stderr-file cpus name script
                      #:key api-endpoint jwt)
   "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."
+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))
@@ -78,7 +80,8 @@ are files in which to write the stdout and stderr of the job respectively."
                               environment)))
       ("current_working_directory" . "/")
       ("standard_output" . ,stdout-file)
-      ("standard_error" . ,stderr-file)))
+      ("standard_error" . ,stderr-file)
+      ("minimum_cpus" . ,cpus)))
   
   (let ((response (slurm-http-post api-endpoint
                                    jwt
-- 
cgit v1.2.3