This source file includes following definitions.
- pmi_init
- pmi_get_local_ranks
- pmi_get_shmem_size
- pmi_put_key_loc
- pmi_put_key_rem
- pmi_put_double
- pmi_commit
- pmi_fence
- pmi_fini
- pmi_get_key_loc
- pmi_get_key_rem
- pmi_get_double
1
2
3
4
5
6
7
8
9
10
11 #include <stdio.h>
12 #include <slurm/pmi2.h>
13 #include "pmi2_pmap_parser.h"
14 #include "pmi2_utils.h"
15
16 static int my_rank, my_node;
17 static char * kvs_name;
18
19 void pmi_init(int *rank, int *size)
20 {
21 int spawned, appnum;
22 int rc;
23
24 *size = -1;
25 *rank = -1;
26 appnum = -1;
27
28 if (PMI2_SUCCESS != (rc = PMI2_Init(&spawned, size, rank, &appnum))) {
29 fprintf(stderr, "pmi2: PMI2_Init: error rc = %d\n", rc);
30 abort();
31 }
32 if( *size < 0 || *rank < 0 ){
33 fprintf(stderr, "pmi2: PMI2_Init: bad size=%d or rank=%d values\n",
34 *size, *rank);
35 abort();
36 }
37 my_rank = *rank;
38
39 kvs_name = (char*)malloc(PMI2_MAX_VALLEN);
40 if( kvs_name == NULL ){
41 fprintf(stderr, "pmi_init (pmi2): cannot malloc kvs name buffer\n");
42 abort();
43 }
44 rc = PMI2_Job_GetId(kvs_name, PMI2_MAX_VALLEN);
45 if( PMI2_SUCCESS != rc ){
46 fprintf(stderr, "pmi_init (pmi2): cannot get kvs name: rc = %d\n", rc);
47 abort();
48 }
49 }
50
51 void pmi_get_local_ranks(int **local_ranks, int *local_cnt)
52 {
53 int rc, found;
54
55 char *pmapping = (char*)malloc(PMI2_MAX_VALLEN);
56 if( pmapping == NULL ){
57 fprintf(stderr,"pmi_get_local_ranks (pmi2): cannot mallic for pmappig\n");
58 abort();
59 }
60
61
62 rc = PMI2_Info_GetJobAttr("PMI_process_mapping", pmapping, PMI2_MAX_VALLEN, &found);
63 if( !found || PMI2_SUCCESS != rc ) {
64 fprintf(stderr,"pmi_get_local_ranks (pmi2): PMI2_Info_GetJobAttr rc = %d\n", rc);
65 abort();
66 }
67
68 *local_ranks = NULL;
69 *local_ranks = mca_common_pmi2_parse_pmap(pmapping, my_rank, &my_node, local_cnt);
70 if (NULL == *local_ranks) {
71 fprintf(stderr,"mca_common_pmi2_parse_pmap: error\n");
72 abort();
73 }
74
75 free(pmapping);
76 }
77
78 void pmi_get_shmem_size(char *is_avail, size_t *size)
79 {
80 *is_avail = 0;
81 }
82
83 void pmi_put_key_loc(char *key, int *key_val, int key_size)
84 {
85 char *encoded = pmi_encode(key_val, key_size * sizeof(int));
86 if( NULL == encoded ){
87 fprintf(stderr, "pmi_encode: error on key: %s\n", key);
88 abort();
89 }
90 PMI2_Info_PutNodeAttr(key, encoded);
91 }
92
93 void pmi_put_key_rem(char *key, int *key_val, int key_size)
94 {
95 char *encoded = pmi_encode(key_val, key_size * sizeof(int));
96 if( NULL == encoded ){
97 fprintf(stderr, "pmi_encode: error on key: %s\n", key);
98 abort();
99 }
100
101 PMI2_KVS_Put(key, encoded);
102 }
103
104 void pmi_put_double(char *key, double val)
105 {
106 char buf[128];
107 sprintf(buf, "%lf", val);
108 PMI2_KVS_Put(key, buf);
109 }
110
111
112 void pmi_commit()
113 {
114
115 }
116
117 void pmi_fence(int collect)
118 {
119 int rc;
120
121 if (PMI2_SUCCESS != (rc = PMI2_KVS_Fence()) ) {
122 fprintf(stderr, "PMI2_KVS_Fence: error %d\n", rc);
123 abort();
124 }
125 }
126
127 void pmi_fini()
128 {
129 PMI2_Finalize();
130 }
131
132 void pmi_get_key_loc(int rank, char *key_name, int **key_val, int *key_size)
133 {
134 int found, rc;
135 size_t tmp_size;
136 char *tmp = calloc(PMI2_MAX_VALLEN, sizeof(char));
137 if( (rc = PMI2_Info_GetNodeAttr(key_name, tmp, PMI2_MAX_VALLEN, &found, 1) ) ){
138 fprintf(stderr,"PMI2_Info_GetNodeAttr: error rc = %d\n", rc);
139 abort();
140 }
141 if( !found ){
142 fprintf(stderr,"pmi_get_key_loc: key %s not found\n", key_name);
143 abort();
144 }
145
146 *key_val = (int*)pmi_decode(tmp, &tmp_size);
147 *key_size = tmp_size / sizeof(int);
148
149 if( NULL == *key_val ){
150 fprintf(stderr,"pmi_decode: cannot decode key %s\n", key_name);
151 abort();
152 }
153 free(tmp);
154 }
155
156 void pmi_get_key_rem(int rank, char *key_name, int **key_val, int *key_size)
157 {
158 int len, rc;
159 size_t tmp_size;
160
161 char *tmp = calloc(PMI2_MAX_VALLEN, sizeof(char));
162 if( (rc = PMI2_KVS_Get(kvs_name, PMI2_ID_NULL, key_name, tmp, PMI2_MAX_VALLEN, &len) ) ){
163 fprintf(stderr,"PMI2_Info_GetNodeAttr: error rc = %d\n", rc);
164 abort();
165 }
166
167 *key_val = (int*)pmi_decode(tmp, &tmp_size);
168 *key_size = tmp_size / sizeof(int);
169
170 if( NULL == *key_val ){
171 fprintf(stderr,"pmi_decode: cannot decode key %s\n", key_name);
172 abort();
173 }
174 free(tmp);
175 }
176
177 double pmi_get_double(int rank, char *key)
178 {
179 int len, rc;
180 size_t tmp_size;
181 char *tmp = calloc(PMI2_MAX_VALLEN, sizeof(char));
182 double v;
183
184 if( (rc = PMI2_KVS_Get(kvs_name, PMI2_ID_NULL, key, tmp, PMI2_MAX_VALLEN, &len) ) ){
185 fprintf(stderr,"PMI2_Info_GetNodeAttr: error rc = %d\n", rc);
186 abort();
187 }
188 sscanf(tmp, "%lf", &v);
189 free(tmp);
190 return v;
191 }