Merge pull request #279 from jameshcorbett/getrabbit-jobid-memo-lookups

`getrabbit`: jobid memo lookups
flux-framework · Feb 18, 2025 · 8a68a59 · 8a68a59
2 parents 67a423a + 53285b7
commit 8a68a59
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/src/cmd/flux-getrabbit.py b/src/cmd/flux-getrabbit.py
@@ -55,15 +55,19 @@ def map_to_rabbits(args, handle, mapping):
         for jobid in args.jobids:
             try:
                 nodelist = (
-                    flux.job.job_list_id(handle, JobID(jobid), ["nodelist"])
+                    flux.job.job_list_id(handle, JobID(jobid), ["annotations"])
                     .get_jobinfo()
-                    .nodelist
+                    .user.rabbits
                 )
             except FileNotFoundError:
                 sys.exit(f"Could not find job {jobid}")
             except Exception as exc:
                 sys.exit(f"Lookup of job {jobid} failed: {exc}")
-            args.computes.append(nodelist)
+            # if there was no `user.rabbits` annotation, a custom class is returned
+            # that converts to an empty string
+            if not str(nodelist):
+                sys.exit(f"No rabbits found for job {jobid}")
+            hlist.append(nodelist)
     aggregated_computes = Hostlist()
     for computes in args.computes:
         aggregated_computes.append(computes)
@@ -94,7 +98,7 @@ def main():
             pass
         if path is None:
             sys.exit(
-                "Flux is misconfigured, 'rabbit.mapping' key not set in"
+                "Flux is misconfigured, 'rabbit.mapping' key not set in "
                 "current instance or root/system instance"
             )
     try:

diff --git a/t/t2001-getrabbit.t b/t/t2001-getrabbit.t
@@ -77,15 +77,23 @@ test_expect_success 'flux rabbitmapping works in nested instances' '
 '
 
 test_expect_success 'flux rabbitmapping works on jobids' '
-    echo "{\"computes\": {\"$(hostname)\": \"rabbit101\"}}" > local_rabbitmapping &&
+    echo "{\"computes\": {\"$(hostname)\": \"rabbit219\"}}" > local_rabbitmapping &&
     echo "
 [rabbit]
 mapping = \"$(pwd)/local_rabbitmapping\"
     " | flux config load &&
-    jobid=$(flux submit -n1 hostname) &&
-    $CMD -j ${jobid} &&
+    jobid=$(flux submit -n1 sleep 1) &&
+    flux job memo ${jobid} rabbits=rabbit101
     test $($CMD -j ${jobid}) = rabbit101 &&
-    test $($CMD -j ${jobid} -c $(hostname)) = rabbit101
+    jobid=$(flux submit -n1 sleep 1) &&
+    flux job memo ${jobid} rabbits=rabbit[202-210] &&
+    test $($CMD -j ${jobid}) = rabbit[202-210] &&
+    test $($CMD -j ${jobid} -c $(hostname)) = rabbit[202-210,219]
+'
+
+test_expect_success 'flux rabbitmapping fails when job has no rabbit memo' '
+    jobid=$(flux submit -n1 hostname) &&
+    test_must_fail $CMD -j ${jobid}
 '
 
 test_done