jobmap: added getter for job rank counts

codes/codes-jobmap.h

@@ -12,8 +12,8 @@
  * Example:
  *
  * job  0      1    2        (<-- jobmap defined "job" IDs)
- * rank 0 1 2  0 1  0 1 2 3  (<-- jobmap defined "job local" IDs)
- * ID   0 1 2  3 4  5 6 7 8  (<-- LP relative "global" IDs)
+ * rank 0 1 2  0 1  0 1 2 3  (<-- "job local" IDs)
+ * ID   0 1 2  3 4  5 6 7 8  (<-- jobmap-defined LP relative "global" IDs)
  * LP   A B C  D E  F G H I  (<-- provided by codes-mapping)
  * */
 
@@ -78,6 +78,8 @@ int codes_jobmap_to_global_id(
 
 int codes_jobmap_get_num_jobs(struct codes_jobmap_ctx const * c);
 
+int codes_jobmap_get_num_ranks(int job_id, struct codes_jobmap_ctx const * c);
+
 #endif
 
 /*

src/util/codes-jobmap-method-impl.h

@@ -18,6 +18,7 @@ struct codes_jobmap_impl {
     struct codes_jobmap_id (*to_local) (int id, void const * ctx);
     int                    (*to_global) (struct codes_jobmap_id id, void const * ctx);
     int                    (*get_num_jobs)(void const * ctx);
+    int                    (*get_num_ranks)(int job_id, void const * ctx);
 };
 
 struct codes_jobmap_ctx {

src/util/codes-jobmap.c

@@ -72,6 +72,11 @@ int codes_jobmap_get_num_jobs(struct codes_jobmap_ctx const * c)
     return c->impl->get_num_jobs(c->ctx);
 }
 
+int codes_jobmap_get_num_ranks(int job_id, struct codes_jobmap_ctx const * c)
+{
+    return c->impl->get_num_ranks(job_id, c->ctx);
+}
+
 /*
  * Local variables:
  *  c-indent-level: 4

src/util/jobmap-impl/jobmap-dummy.c

@@ -55,17 +55,27 @@ static int jobmap_dummy_to_global(struct codes_jobmap_id id, void const * ctx)
         return -1;
 }
 
-int jobmap_dummy_get_num_jobs(void const * ctx)
+static int jobmap_dummy_get_num_jobs(void const * ctx)
 {
     return *(int const *) ctx;
 }
 
+int jobmap_dummy_get_num_ranks(int job_id, void const * ctx)
+{
+    int num_jobs = *(int const *) ctx;
+    if (job_id < 0 || job_id >= num_jobs)
+        return -1;
+    else
+        return 1;
+}
+
 struct codes_jobmap_impl jobmap_dummy_impl = {
     jobmap_dummy_configure,
     jobmap_dummy_destroy,
     jobmap_dummy_to_local,
     jobmap_dummy_to_global,
-    jobmap_dummy_get_num_jobs
+    jobmap_dummy_get_num_jobs,
+    jobmap_dummy_get_num_ranks
 };
 
 /*

src/util/jobmap-impl/jobmap-identity.c

@@ -59,12 +59,18 @@ int jobmap_identity_get_num_jobs(void const * ctx)
     return 1;
 }
 
+int jobmap_identity_get_num_ranks(int job_id, void const * ctx)
+{
+    return (job_id == 0) ? *(int const *) ctx : -1;
+}
+
 struct codes_jobmap_impl jobmap_identity_impl = {
     jobmap_identity_configure,
     jobmap_identity_destroy,
     jobmap_identity_to_local,
     jobmap_identity_to_global,
-    jobmap_identity_get_num_jobs
+    jobmap_identity_get_num_jobs,
+    jobmap_identity_get_num_ranks
 };
 
 /*

src/util/jobmap-impl/jobmap-list.c

@@ -172,7 +172,7 @@ static struct codes_jobmap_id jobmap_list_to_local(int id, void const * ctx)
     rtn.job = -1;
     rtn.rank = -1;
 
-    struct jobmap_list *lst = (struct jobmap_list*)ctx;
+    struct jobmap_list const *lst = (struct jobmap_list const *)ctx;
 
     for(int i=0; i<lst->num_jobs; i++) {
         for(int j=0; j < lst->rank_counts[i]; j++) {
@@ -189,7 +189,7 @@ static struct codes_jobmap_id jobmap_list_to_local(int id, void const * ctx)
 
 static int jobmap_list_to_global(struct codes_jobmap_id id, void const * ctx)
 {
-    struct jobmap_list *lst = (struct jobmap_list*)ctx;
+    struct jobmap_list const *lst = (struct jobmap_list*)ctx;
 
     if (id.job < lst->num_jobs)
         return lst->global_ids[id.job][id.rank];
@@ -201,7 +201,15 @@ static int jobmap_list_get_num_jobs(void const * ctx)
 {
     struct jobmap_list *lst = (struct jobmap_list*)ctx;
     return lst->num_jobs;
+}
 
+static int jobmap_list_get_num_ranks(int job_id, void const * ctx)
+{
+    struct jobmap_list const *lst = (struct jobmap_list const *) ctx;
+    if (job_id < 0 || job_id >= lst->num_jobs)
+        return -1;
+    else
+        return lst->rank_counts[job_id];
 }
 
 static void jobmap_list_destroy(void * ctx)
@@ -222,7 +230,8 @@ struct codes_jobmap_impl jobmap_list_impl = {
     jobmap_list_destroy,
     jobmap_list_to_local,
     jobmap_list_to_global,
-    jobmap_list_get_num_jobs
+    jobmap_list_get_num_jobs,
+    jobmap_list_get_num_ranks
 };
 
 /*

tests/jobmap-test.c

@@ -26,6 +26,11 @@ static int test_jobmap_identity(int num_ranks)
     if (1 != num_jobs)
         ERR("jobmap-identity: expected exactly 1 job, got %d\n", num_jobs);
 
+    int num_ranks_from_map = codes_jobmap_get_num_ranks(0, c);
+    if (num_ranks_from_map != num_ranks)
+        ERR("jobmap-identity: expected %d ranks for job 0, got %d\n",
+                num_ranks, num_ranks_from_map);
+
     struct codes_jobmap_id lid;
     int gid;
     for (int i = 0; i < num_ranks; i++) {
@@ -60,6 +65,11 @@ static int test_jobmap_list(char * fname)
     int gid, gid_expected = 0;
     struct codes_jobmap_id lid_expected, lid;
     for (int i = 0; i < num_jobs; i++) {
+        int num_ranks = codes_jobmap_get_num_ranks(i, c);
+        if (num_ranks != rank_count_per_job[i])
+            ERR("jobmap-list: expected %d ranks for job %d, got %d\n",
+                    rank_count_per_job[i], i, num_ranks);
+
         for (int j = 0; j < rank_count_per_job[i]; j++) {
             lid_expected.job = i;
             lid_expected.rank = j;