Merge pull request #6728 from grondo/issue#6703

export `FLUX_JOB_RANKS` to housekeeping scripts

Author: mergify[bot]

doc/man7/flux-environment.rst

@@ -323,6 +323,13 @@ contains the following:
   :term:`instance owner`
 - :envvar:`FLUX_JOB_ID` - the job id that triggered the script
 - :envvar:`FLUX_JOB_USERID` - the uid of the of the job owner
+- :envvar:`FLUX_JOB_RANKS` - the broker ranks which were assigned to the
+  job for which the current script is running.  This can be used in conjunction
+  with the ``hostlist`` attribute to construct the job hostlist. For example:
+
+  .. code-block:: shell
+
+   FLUX_JOB_HOSTLIST=$(flux hostlist --nth=${FLUX_JOB_RANKS} instance)
 
 If the IMP is configured to allow other ``FLUX_`` prefixed environment
 variables to be set as described in :man5:`flux-config-security-imp`,
@@ -333,18 +340,6 @@ then the following are set to allow Flux commands to work from the script:
 - :envvar:`FLUX_EXEC_PATH`
 - :envvar:`FLUX_CONNECTOR_PATH`
 
-In addition, the following extra environment variables are set in the
-prolog and epilog for convenience:
-
-- :envvar:`FLUX_JOB_RANKS` - the broker ranks which were assigned to the
-  job for which the current prolog or epilog is running. This can be
-  used in conjunction with the ``hostlist`` attribute to construct the
-  job hostlist. For example:
-
-  .. code-block:: shell
-
-   FLUX_JOB_HOSTLIST=$(flux hostlist --nth=${FLUX_JOB_RANKS} instance)
-
 TESTING
 =======
 

src/modules/job-manager/housekeeping.c

@@ -41,6 +41,7 @@
  * Script environment:
  *   FLUX_JOB_ID - the job whose resources are running housekeeping
  *   FLUX_JOB_USERID - the UID of the job's owner
+ *   FLUX_JOB_RANKS - idset of broker ranks which were assigned to FLUX_JOB_ID
  *   FLUX_URI - the URI of the local flux broker
  *   The IMP must be configured to explicitly allow FLUX_* to pass through.
  *
@@ -161,11 +162,28 @@ static void allocation_destructor (void **item)
 
 }
 
-static int update_cmd_env (flux_cmd_t *cmd, flux_jobid_t id, uint32_t userid)
+static char *get_rlist_ranks (struct rlist *rl)
 {
-    if (flux_cmd_setenvf (cmd, 1, "FLUX_JOB_ID", "%ju", (uintmax_t)id) < 0
-        || flux_cmd_setenvf (cmd, 1, "FLUX_JOB_USERID", "%u", userid) < 0)
+    struct idset *ids;
+    if (!(ids = rlist_ranks (rl)))
+        return NULL;
+    return idset_encode (ids, IDSET_FLAG_RANGE);
+}
+
+static int update_cmd_env (flux_cmd_t *cmd,
+                           flux_jobid_t id,
+                           uint32_t userid,
+                           struct rlist *rl)
+{
+    char *ranks;
+
+    if (!(ranks = get_rlist_ranks (rl))
+        || flux_cmd_setenvf (cmd, 1, "FLUX_JOB_ID", "%ju", (uintmax_t)id) < 0
+        || flux_cmd_setenvf (cmd, 1, "FLUX_JOB_USERID", "%u", userid) < 0
+        || flux_cmd_setenvf (cmd, 1, "FLUX_JOB_RANKS", "%s", ranks) < 0)
         return -1;
+
+    free (ranks);
     return 0;
 }
 
@@ -195,7 +213,7 @@ static struct allocation *allocation_create (struct housekeeping *hk,
                                               id,
                                               "housekeeping",
                                                a))
-        || update_cmd_env (hk->cmd, id, userid) < 0
+        || update_cmd_env (hk->cmd, id, userid, a->rl) < 0
         || bulk_exec_push_cmd (a->bulk_exec, a->pending, hk->cmd, 0) < 0) {
         allocation_destroy (a);
         return NULL;

t/t2226-housekeeping.t

@@ -277,7 +277,8 @@ test_expect_success 'run a job on rank 3, wait for hk, and check environment' '
 	wait_for_running 0 &&
 	grep "^FLUX_JOB_ID=$(flux job last | flux job id --to=dec)$" env.out &&
 	grep "^FLUX_JOB_USERID=$(id -u)$" env.out &&
-	grep "^FLUX_URI=$(flux exec -r 3 flux getattr local-uri)$" env.out
+	grep "^FLUX_URI=$(flux exec -r 3 flux getattr local-uri)$" env.out &&
+	grep "^FLUX_JOB_RANKS=$(flux jobs -no {ranks} $(flux job last))" env.out
 '
 test_expect_success 'configure housekeeping to sleep forever' '
 	flux config load <<-EOT