diff --git a/src/plugins/mpi/pmix/pmixp_spawn.c b/src/plugins/mpi/pmix/pmixp_spawn.c index ef3be1e6b62af992422b7a9ff1ce1a5f5d03f16c..9c608cf9b89d99b28ecda99684d32a828549beeb 100644 --- a/src/plugins/mpi/pmix/pmixp_spawn.c +++ b/src/plugins/mpi/pmix/pmixp_spawn.c @@ -867,7 +867,8 @@ void _populate_srun_argv(const pmix_app_t *app, int app_argc, char **argv, } int _validate_launch(const pmix_proc_t *proc, const pmix_app_t apps[], - size_t napps, size_t **_out_ids, size_t *_out_nids) { + size_t napps, size_t **_out_ids, size_t *_out_nids, + char **_out_vrm_id) { // Make sure we do not launch "nothing" if (!napps) return 0; @@ -929,6 +930,11 @@ int _validate_launch(const pmix_proc_t *proc, const pmix_app_t apps[], return 0; } + // Fetch vrm job id + char *vrm_id = strtok(NULL, ","); + *_out_vrm_id = (char *)malloc(strlen(vrm_id) + 1); + strcpy(*_out_vrm_id, vrm_id); + // Ignore nonce strtok(NULL, ","); @@ -947,7 +953,7 @@ int _validate_launch(const pmix_proc_t *proc, const pmix_app_t apps[], } void _launch_app(const pmix_proc_t *parent, const pmix_app_t *app, - size_t dynamicId) { + size_t dynamicId, char *vrm_job_id) { int app_argc; char **argv = _create_srun_argv(app, &app_argc); _populate_srun_argv(app, app_argc, argv, dynamicId); @@ -972,7 +978,7 @@ void _launch_app(const pmix_proc_t *parent, const pmix_app_t *app, char **env = env_array_copy(app->env); env_array_append_fmt(&env, "HOME", "/root"); env_array_append_fmt(&env, "DPM_AGENT_PORT", "25000"); - env_array_append_fmt(&env, "SLURM_VRM_JOBID", "%u", pmixp_info_jobid()); + env_array_append_fmt(&env, "SLURM_VRM_JOBID", "%s", vrm_job_id); env_array_append_fmt(&env, "DPM_PMIX_DYNAMIC_ID", "%zu", dynamicId); execve(SLURM_PREFIX "/bin/srun", argv, env); abort(); @@ -982,8 +988,10 @@ int pmixp_spawn(const pmix_proc_t *proc, const pmix_app_t apps[], size_t napps, pmix_spawn_cbfunc_t cbfunc, void *cbdata) { // First ask for permission size_t *dynamicIds; + char *vrm_job_id = NULL; size_t nDynamicIds; - if (!_validate_launch(proc, apps, napps, &dynamicIds, &nDynamicIds)) + if (!_validate_launch(proc, apps, napps, &dynamicIds, &nDynamicIds, + &vrm_job_id)) return 0; size_t currentApp = 0; @@ -1008,11 +1016,12 @@ int pmixp_spawn(const pmix_proc_t *proc, const pmix_app_t apps[], size_t napps, __pid_t pid = fork(); if (pid == 0) { - _launch_app(proc, app, dynamicIds[currentApp]); + _launch_app(proc, app, dynamicIds[currentApp], vrm_job_id); } currentApp++; } free(dynamicIds); + free(vrm_job_id); fclose(ptr); return 1; } \ No newline at end of file