diff --git a/containers/spindle-slurm-ubuntu/testing-plugin/conf/slurm.conf b/containers/spindle-slurm-ubuntu/testing-plugin/conf/slurm.conf index abf060d5..e4298157 100644 --- a/containers/spindle-slurm-ubuntu/testing-plugin/conf/slurm.conf +++ b/containers/spindle-slurm-ubuntu/testing-plugin/conf/slurm.conf @@ -34,6 +34,7 @@ JobAcctGatherFrequency=30 AccountingStorageType=accounting_storage/slurmdbd AccountingStorageHost=slurm-db AccountingStoragePort=6819 +PrologFlags=Contain NodeName=slurm-node-1 NodeAddr=slurm-node-1 CPUs=3 RealMemory=1000 State=UNKNOWN NodeName=slurm-node-2 NodeAddr=slurm-node-2 CPUs=3 RealMemory=1000 State=UNKNOWN NodeName=slurm-node-3 NodeAddr=slurm-node-3 CPUs=3 RealMemory=1000 State=UNKNOWN diff --git a/doc/slurm_plugin.md b/doc/slurm_plugin.md new file mode 100644 index 00000000..ab8765c1 --- /dev/null +++ b/doc/slurm_plugin.md @@ -0,0 +1,115 @@ +Spindle Slurm plugin +==================== + +The Spindle Slurm plugin integrates Spindle into Slurm through the +SPANK interface as an alternative launch mechanism to the srun wrapper. +It adds the ability to launch job steps using `srun --spindle`. + +## Building and configuring the plugin + +Configure Spindle with `--enable-slurm-plugin`: + +```bash +./configure --with-rm=slurm-plugin --enable-slurm-plugin [--with-slurm-dir=/path/to/slurm] ... +make +make install +``` + +Refer to `INSTALL` for more details on configuring Spindle. + +After installation of Spindle, the plugin is installed at +`$PREFIX/lib/libspindleslurm.so`. It is registered with Slurm by adding the +following line to `/etc/slurm/plugstack.conf`: + +``` +required /path/to/spindle/lib/libspindleslurm.so +``` + +## Session launch modes + +The manner in which Spindle sessions are started varies depending on +the configuration of Spindle and of Slurm. + +When starting a session, the plugin must arrange for Spindle to start +on each compute node before any step runs within the allocation. +The most straightforward way to do this is to configure the cluster +to run job prologs at allocation time. If your `slurm.conf` includes +`PrologFlags=Alloc` (or another flag that implies it: `Contain`, +`RunInJob`, `X11`, `ForceRequeueOnFail`, or `NoHold`), then sessions +will be started on each node of the allocation at the time the allocation +is made. + +If `PrologFlags=Alloc` or a related setting is *not* used, one of two +mechanisms is used to start the job on every node: + +**RSH launch**: Spindle can use RSH/SSH to launch daemons from the +frontend (FE) process. To use the RSH launch mode, the cluster must be configured +such that passwordless ssh can be used to run commands on every compute +node within the allocation without any interactive user input. +This mode is enabled by configuring Spindle with: + +```bash +./configure --with-rm=slurm-plugin --enable-slurm-plugin --with-rsh-launch [--with-rsh-cmd=/usr/bin/ssh] ... +``` + +**Dummy srun fallback**: If neither `PrologFlags=Alloc` nor RSH launch is available, +Spindle will fall back on using a dummy `srun` invocation to force the prolog +to run on every compute node of the allocation. Note that this has the side-effect +of consuming step 0, so that the user's first step will instead be numbered 1. + +## Using Spindle through the Slurm plugin + +### Per-step mode: `--spindle` + +Add `--spindle` to any `srun` command to use Spindle for that step. +Spindle daemons start before the application runs and shut down when +the step finishes. + +```bash +srun --spindle ./my_application +``` + +Additional arguments can be passed to Spindle as an optional value of the argument `--spindle`: + +```bash +srun --spindle="--level=low" ./my_application +``` + +### Session mode: `--spindle-session` + +Session mode shares a Spindle session across multiple steps. +The use of sessions in the Slurm plugin differs from its use with +the other launchers. Unlike the other launchers, sessions are *not* +started with `spindle --start-session`. Rather, an additional argument +`--spindle-session` is added to `salloc` and `sbatch`. + +To use a session, include `--spindle-session` when creating the allocation: + +```bash +salloc --spindle-session ... +``` + +Then run steps with `--spindle`: + +```bash +srun --spindle ./app1 +srun --spindle ./app2 +srun --spindle ./app3 +``` + +All steps within the allocation will run in the same Spindle session. +When the allocation exits, the session will terminate automatically. + +Sessions can be used with an `sbatch` script as shown below: + +```bash +#!/bin/bash +#SBATCH --spindle-session +#SBATCH -N 4 +#SBATCH -n 4 + +srun --spindle ./app1 +srun --spindle ./app2 +srun --spindle ./app3 +``` + diff --git a/src/client/client_comlib/client_api.c b/src/client/client_comlib/client_api.c index cec31051..07735226 100644 --- a/src/client/client_comlib/client_api.c +++ b/src/client/client_comlib/client_api.c @@ -40,42 +40,31 @@ static struct lock_t comm_lock; int send_cachepath_query( int fd, char **chosen_realized_cachepath, char **chosen_parsed_cachepath){ - int retries = 0, max_retries = 1000, rc = 0; - struct timespec delay_between_retries = { .tv_sec = 0, .tv_nsec = 1000000 }; + int rc = 0; ldcs_message_t message; - char buffer[MAX_PATH_LEN+1]; + char buffer[2*(MAX_PATH_LEN+1)]; buffer[MAX_PATH_LEN] = '\0'; - do{ - message.header.type = LDCS_MSG_CHOSEN_CACHEPATH_REQUEST; - message.header.len = 0; - message.data = buffer; - - COMM_LOCK; - - debug_printf3("sending message of type: CHOSEN_CACHEPATH_REQUEST.\n" ); - rc = client_send_msg(fd, &message); - if( rc != 0 ){ - return rc; - } - rc = client_recv_msg_static(fd, &message, LDCS_READ_BLOCK); - if( rc != 0 ){ - return rc; - } - - COMM_UNLOCK; - - if( message.header.type == LDCS_MSG_NO_CACHEPATH_CONSENSUS_YET ){ - if( retries++ >= max_retries ){ - break; - } - nanosleep( &delay_between_retries, NULL ); - continue; - } - break; + message.header.type = LDCS_MSG_CHOSEN_CACHEPATH_REQUEST; + message.header.len = 0; + message.data = buffer; + + debug_printf3("sending message of type: CHOSEN_CACHEPATH_REQUEST.\n" ); + COMM_LOCK; - }while( 1 ); + rc = client_send_msg(fd, &message); + if( rc != 0 ){ + COMM_UNLOCK; + return rc; + } + rc = client_recv_msg_static(fd, &message, LDCS_READ_BLOCK); + if( rc != 0 ){ + COMM_UNLOCK; + return rc; + } + + COMM_UNLOCK; if (message.header.type != LDCS_MSG_CHOSEN_CACHEPATH || message.header.len > MAX_PATH_LEN) { err_printf("Got unexpected message of type %d\n", (int) message.header.type); diff --git a/src/fe/startup/spindle_fe.cc b/src/fe/startup/spindle_fe.cc index a038e201..0d5c0817 100644 --- a/src/fe/startup/spindle_fe.cc +++ b/src/fe/startup/spindle_fe.cc @@ -41,7 +41,6 @@ static const char *logging_file = NULL; #endif static const char spindle_bootstrap[] = LIBEXECDIR "/spindle_bootstrap"; static bool sendAndWaitForAlive(); -static void determineCachepathConsensus(); #define STARTUP_TIMEOUT 60 @@ -433,7 +432,6 @@ int spindleInitFE(const char **hosts, spindle_args_t *params) /* Wait for servers to indicate startup */ sendAndWaitForAlive(); - determineCachepathConsensus(); return 0; } @@ -490,17 +488,6 @@ void markRSHPidReapedFE() clear_fe_rsh_pid(); } -static void determineCachepathConsensus( void ){ - ldcs_message_t consensus_req_msg; - consensus_req_msg.header.type = LDCS_MSG_REQUEST_CACHEPATH_CONSENSUS; - consensus_req_msg.header.len = 0; - consensus_req_msg.data = NULL; - int result = ldcs_audit_server_fe_broadcast(&consensus_req_msg, NULL); - if (result == -1) { - debug_printf("Failure sending cachepath consensus message\n"); - } -} - static bool sendAndWaitForAlive() { int result; diff --git a/src/include/ldcs_api.h b/src/include/ldcs_api.h index 0bcd8f40..2b261a56 100644 --- a/src/include/ldcs_api.h +++ b/src/include/ldcs_api.h @@ -85,10 +85,8 @@ typedef enum { LDCS_MSG_PICKONE_RESP, LDCS_MSG_ALIVE_REQ, LDCS_MSG_ALIVE_RESP, - LDCS_MSG_REQUEST_CACHEPATH_CONSENSUS, LDCS_MSG_CHOSEN_CACHEPATH_REQUEST, LDCS_MSG_CHOSEN_CACHEPATH, - LDCS_MSG_NO_CACHEPATH_CONSENSUS_YET, LDCS_MSG_UNKNOWN } ldcs_message_ids_t; diff --git a/src/server/auditserver/ldcs_audit_server_handlers.c b/src/server/auditserver/ldcs_audit_server_handlers.c index 1afdf2e8..6b3a6b1c 100644 --- a/src/server/auditserver/ldcs_audit_server_handlers.c +++ b/src/server/auditserver/ldcs_audit_server_handlers.c @@ -181,7 +181,6 @@ static int handle_setup_alias(ldcs_process_data_t *procdata, char *pathname, cha static int handle_client_dirlists_req(ldcs_process_data_t *procdata, int nc); static int handle_close_client_query(ldcs_process_data_t *procdata, int nc); static int handle_alive_msg(ldcs_process_data_t *procdata, ldcs_message_t *msg); -static int handle_cachepath_consensus(ldcs_process_data_t *procdata, ldcs_message_t *msg); static int handle_chosen_cachepath_request(ldcs_process_data_t *procdata, int nc); extern void getValidCachePathByIndex( uint64_t validBitIdx, char **realizedCachePath, char **parsedCachePath, char **symbolicCachePath ); @@ -1997,8 +1996,6 @@ int handle_server_message(ldcs_process_data_t *procdata, node_peer_t peer, ldcs_ case LDCS_MSG_ALIVE_REQ: case LDCS_MSG_ALIVE_RESP: return handle_alive_msg(procdata, msg); - case LDCS_MSG_REQUEST_CACHEPATH_CONSENSUS: - return handle_cachepath_consensus(procdata, msg); default: err_printf("Received unexpected message from node: %d\n", (int) msg->header.type); assert(0); @@ -2961,35 +2958,26 @@ static int handle_client_pickone_msg(ldcs_process_data_t *procdata, int nc, ldcs } /** - * Handle LDCS_MSG_REQUEST_CACHEPATH_CONSENSUS to determine which cachepaths are - * available across all of the servers. + * Determine which cachepaths are available across all of the servers. */ -static int cachepath_consensus_reached; -static int handle_cachepath_consensus(ldcs_process_data_t *procdata, ldcs_message_t *msg){ - +int handle_cachepath_consensus(ldcs_process_data_t *procdata) +{ int num_children = ldcs_audit_server_md_get_num_children(procdata); - debug_printf( "Processing REQUEST_CACHEPATH_CONSENSUS.\n" ); + debug_printf( "Calculating cachepath consensus.\n" ); debug_printf3( " procdata->cachepath_bitidx = %#"PRIx64"\n", procdata->cachepath_bitidx ); debug_printf3( " procdata->cachepaths = %s\n", procdata->cachepaths ); debug_printf3( " procdata->cachepath = %s [should be null]\n", procdata->cachepath ); debug_printf3( " procdata->commpath = %s\n", procdata->commpath ); debug_printf3( " num_children = %d\n", num_children ); - if (num_children) { - spindle_broadcast(procdata, msg); - debug_printf3( "Successfully broadcast REQUEST_CACHEPATH_CONSENSUS\n" ); - msgbundle_force_flush(procdata); - debug_printf3( "Successfully flushed the broadcast of REQUEST_CACHEPATH_CONSENSUS\n" ); - } - ldcs_audit_server_md_allreduce_AND( &procdata->cachepath_bitidx ); debug_printf3( "The consensus value for procdata->cachepath_bitidx is: %#"PRIx64"\n", procdata->cachepath_bitidx ); - if( procdata->cachepath_bitidx == 0 ){ - err_printf("No valid cachepath path available. Falling back to \"commpath\" path (%s).\n", procdata->commpath); + if( procdata->cachepath_bitidx == 0 ) { + debug_printf("No valid cachepath path available. Falling back to \"commpath\" path (%s).\n", procdata->commpath); procdata->cachepath = procdata->commpath; - }else{ + } else { getValidCachePathByIndex( procdata->cachepath_bitidx, &procdata->cachepath, &procdata->parsed_cachepath, @@ -3003,14 +2991,10 @@ static int handle_cachepath_consensus(ldcs_process_data_t *procdata, ldcs_messag ldcs_audit_server_filemngt_init(procdata->cachepath, procdata->commpath); test_printf(" cachepath=%s\n", procdata->cachepath); - cachepath_consensus_reached = 1; return 0; } -/** - * Handle LDCS_MSG_CHOSEN_CACHEPATH_REQUEST - */ -static int handle_chosen_cachepath_request(ldcs_process_data_t *procdata, int nc){ +static int handle_chosen_cachepath_request(ldcs_process_data_t *procdata, int nc) { ldcs_message_t msg; int connid; ldcs_client_t *client; @@ -3022,18 +3006,12 @@ static int handle_chosen_cachepath_request(ldcs_process_data_t *procdata, int nc return 0; - if( cachepath_consensus_reached ){ - msg.header.type = LDCS_MSG_CHOSEN_CACHEPATH; - msg.header.len = strlen(procdata->cachepath) + 1 + strlen(procdata->parsed_cachepath) + 1; - msg.data = calloc( 1, msg.header.len ); - strcpy( msg.data, procdata->cachepath ); - strcpy( &msg.data[ strlen(procdata->cachepath)+1 ], procdata->parsed_cachepath ); - }else{ - msg.header.type = LDCS_MSG_NO_CACHEPATH_CONSENSUS_YET; - msg.header.len = 0; - msg.data = NULL; - } - + msg.header.type = LDCS_MSG_CHOSEN_CACHEPATH; + msg.header.len = strlen(procdata->cachepath) + 1 + strlen(procdata->parsed_cachepath) + 1; + msg.data = calloc( 1, msg.header.len ); + strcpy( msg.data, procdata->cachepath ); + strcpy( &msg.data[ strlen(procdata->cachepath)+1 ], procdata->parsed_cachepath ); + ldcs_send_msg(connid, &msg); free( msg.data ); procdata->server_stat.clientmsg.cnt++; diff --git a/src/server/auditserver/ldcs_audit_server_handlers.h b/src/server/auditserver/ldcs_audit_server_handlers.h index 8702316c..790e52b3 100644 --- a/src/server/auditserver/ldcs_audit_server_handlers.h +++ b/src/server/auditserver/ldcs_audit_server_handlers.h @@ -25,6 +25,7 @@ int handle_server_error(ldcs_process_data_t *procdata, node_peer_t peer); int handle_client_message(ldcs_process_data_t *procdata, int nc, ldcs_message_t *msg); int handle_client_start(ldcs_process_data_t *procdata, int nc); int handle_client_end(ldcs_process_data_t *procdata, int nc); +int handle_cachepath_consensus(ldcs_process_data_t *procdata); int exit_note_cb(int infd, int serverid, void *data); diff --git a/src/server/auditserver/ldcs_audit_server_process.c b/src/server/auditserver/ldcs_audit_server_process.c index b6ecdbff..7e32337d 100644 --- a/src/server/auditserver/ldcs_audit_server_process.c +++ b/src/server/auditserver/ldcs_audit_server_process.c @@ -137,7 +137,7 @@ void stopprofile() int ldcs_audit_server_process(spindle_args_t *args) { - int serverid, fd; + int serverid, fd, result; startprofile(args); @@ -198,7 +198,20 @@ int ldcs_audit_server_process(spindle_args_t *args) if (ldcs_process_data.opts & OPT_PROCCLEAN) init_cleanup_proc(ldcs_process_data.cachepath, ldcs_process_data.commpath); - debug_printf3("Initializing connections for clients at %s and %lu\n", + /* Calculate location of cache */ + debug_printf2("Calculating cache path location\n"); + determineValidCachePaths( + &ldcs_process_data.cachepath_bitidx, + ldcs_process_data.cachepaths, + ldcs_process_data.number ); + result = handle_cachepath_consensus(&ldcs_process_data); + if (result == -1) { + err_printf("Could not determine cachepath consensus\n"); + return -1; + } + + /* Setup connections for clients to start connecting */ + debug_printf2("Initializing connections for clients at %s and %lu\n", ldcs_process_data.commpath, (unsigned long) ldcs_process_data.number); serverid = ldcs_create_server(ldcs_process_data.commpath, ldcs_process_data.number); if (serverid == -1) { @@ -231,10 +244,7 @@ int ldcs_audit_server_process(spindle_args_t *args) if (fd != -1) { ldcs_listen_register_fd(fd, serverid, forceExitCB, (void *) &ldcs_process_data); } - determineValidCachePaths( - &ldcs_process_data.cachepath_bitidx, - ldcs_process_data.cachepaths, - ldcs_process_data.number ); + return 0; } diff --git a/src/server/comlib/ldcs_api_util.c b/src/server/comlib/ldcs_api_util.c index af101815..28a3d90c 100644 --- a/src/server/comlib/ldcs_api_util.c +++ b/src/server/comlib/ldcs_api_util.c @@ -91,10 +91,8 @@ char* _message_type_to_str (ldcs_message_ids_t type) { STR_CASE(LDCS_MSG_PICKONE_RESP); STR_CASE(LDCS_MSG_ALIVE_REQ); STR_CASE(LDCS_MSG_ALIVE_RESP); - STR_CASE(LDCS_MSG_REQUEST_CACHEPATH_CONSENSUS); STR_CASE(LDCS_MSG_CHOSEN_CACHEPATH_REQUEST); STR_CASE(LDCS_MSG_CHOSEN_CACHEPATH); - STR_CASE(LDCS_MSG_NO_CACHEPATH_CONSENSUS_YET); STR_CASE(LDCS_MSG_UNKNOWN); } return "unknown"; diff --git a/src/slurm_plugin/plugin_utils.c b/src/slurm_plugin/plugin_utils.c index f77531e4..87c2a37d 100644 --- a/src/slurm_plugin/plugin_utils.c +++ b/src/slurm_plugin/plugin_utils.c @@ -254,11 +254,11 @@ int isFEHost(char **hostlist, unsigned int num_hosts) int feresult = -1; for (i = 0; i < num_hosts; i++) { - if (!last_host || strcmp(hostlist[i], last_host) == 1) { + if (!last_host || strcmp(hostlist[i], last_host) > 0) { last_host = hostlist[i]; } } - sdprintf(2, "last_host = %s\n", last_host ? last_host : NULL); + sdprintf(2, "last_host = %s\n", last_host ? last_host : "(null)"); if (!last_host) { error = errno; sdprintf(1, "ERROR: Could not get current system's hostname: %s\n", strerror(error)); @@ -515,6 +515,15 @@ int signalSpankSessionEnd(spindle_args_t *params) char *unique_file = NULL; +void cleanup_unique_file() +{ + if (unique_file) { + unlink(unique_file); + free(unique_file); + unique_file = NULL; + } +} + #define UNIQUE_FILE_NAME "spindle_unique" int isBEProc(spindle_args_t *params, unsigned int exit_phase) @@ -687,8 +696,10 @@ void push_env(spank_t spank, saved_env_t **env) e->new_spindledebug = readSpankEnv(spank, "SPINDLE_DEBUG"); e->old_spindledebug = getenv("SPINDLE_DEBUG"); - if (e->new_pwd) - chdir(e->new_pwd); + if (e->new_pwd) { + if (chdir(e->new_pwd) == -1) + sdprintf(1, "WARNING: Could not chdir to %s: %s\n", e->new_pwd, strerror(errno)); + } if (e->new_home) setenv("HOME", e->new_home, 1); @@ -727,8 +738,10 @@ void pop_env(saved_env_t *env) else unsetenv("SPINDLE_DEBUG"); - if (env->old_pwd) - chdir(env->old_pwd); + if (env->old_pwd) { + if (chdir(env->old_pwd) == -1) + sdprintf(1, "WARNING: Could not chdir to %s: %s\n", env->old_pwd, strerror(errno)); + } if (env->new_home) free(env->new_home); @@ -815,7 +828,7 @@ int dropPrivilegeAndRun(dpr_function_t func, uid_t uid, void *input, char **outp exit(-1); } if (output_len) { - result = safe_write(pipe_fds[1], output_str, output_len+1); + result = safe_write(pipe_fds[1], child_output_str, output_len+1); if (result != output_len+1) { error = errno; fprintf(stderr, "Spindle error. Could not write result string to pipe: %s\n", strerror(error)); @@ -989,7 +1002,11 @@ pid_t grandchild_fork() int result, fork_result = -1; pipe_fds[0] = pipe_fds[1] = -1; - pipe(pipe_fds); + result = pipe(pipe_fds); + if (result == -1) { + sdprintf(1, "ERROR: pipe() failed in grandchild_fork. Aborting spindle\n"); + return -1; + } child_pid = fork(); if (child_pid == -1) { diff --git a/src/slurm_plugin/plugin_utils.h b/src/slurm_plugin/plugin_utils.h index 0ac2c8c8..a29b9805 100644 --- a/src/slurm_plugin/plugin_utils.h +++ b/src/slurm_plugin/plugin_utils.h @@ -45,6 +45,7 @@ char **getHostsParse(unsigned int num_hosts, const char *shortlist); int isFEHost(char **hostlist, unsigned int num_hosts); extern char *unique_file; +void cleanup_unique_file(); int isBEProc(spindle_args_t *params, unsigned int exit_phase); int doesFEExitSocketExist(spindle_args_t *params); diff --git a/src/slurm_plugin/slurm_plugin.c b/src/slurm_plugin/slurm_plugin.c index 0989299c..aa27f6bb 100644 --- a/src/slurm_plugin/slurm_plugin.c +++ b/src/slurm_plugin/slurm_plugin.c @@ -27,6 +27,7 @@ Place, Suite 330, Boston, MA 02111-1307 USA #include #include +#include #include #include "spindle_launch.h" @@ -102,6 +103,7 @@ static __thread spank_t current_spank; static const char *user_options = NULL; static int enable_spindle = 0; static int start_session = 0; +static int prolog_alloc_mode = 0; extern char *parse_location(char *loc, number_t number); extern char *realize(char *path); @@ -156,7 +158,14 @@ static int should_use_session(spank_t spank) { if (session_env) return 1; err = spank_option_getopt(spank, &session_option, NULL); - return (err == ESPANK_SUCCESS); + return (err == ESPANK_SUCCESS); + } + + /* In the allocator, we are running in the same process that + * handled the command line arguments and can check the + * flag directly. */ + if (context == S_CTX_ALLOCATOR) { + return start_session; } return 0; @@ -168,6 +177,23 @@ int slurm_spank_init(spank_t spank, int ac, char *argv[]) { if (context == S_CTX_ALLOCATOR) { spank_option_register(spank, &session_option); } + +#if defined(PROLOG_FLAG_ALLOC) + /* Check whether Slurm is configured to run the prolog at + * allocation time (PROLOG_FLAG_ALLOC), or the default mode + * where the prolog runs on the first step. */ + if (!prolog_alloc_mode) { + slurm_conf_t *conf = NULL; + if (slurm_load_ctl_conf(0, &conf) == SLURM_SUCCESS) { + if (conf->prolog_flags & PROLOG_FLAG_ALLOC) + prolog_alloc_mode = 1; + slurm_free_ctl_conf(conf); + } else { + sdprintf(1, "Could not read Slurm config, falling back to non-prolog launch.\n"); + } + } +#endif + return 0; } @@ -186,6 +212,17 @@ int slurm_spank_init_post_opt(spank_t spank, int ac, char *argv[]) { if (start_session) { setenv(SPANK_SPINDLE_USE_SESSION, "1", 1); } + + /* With PrologFlags=Alloc, forward env vars to job control here; + them. Without PrologFlags=Alloc, this forwarding happens later + in local context (srun). */ + if (prolog_alloc_mode && start_session) { + int result = forward_environment_to_job_control(spank); + if (result == -1) { + slurm_error("ERROR: Spindle plugin error. Unable to forward environment variables to job control.\n"); + return result; + } + } } return 0; } @@ -314,10 +351,13 @@ int slurm_spank_local_user_init(spank_t spank, int ac, char *argv[]) goto done; } - use_session = should_use_session(spank); + use_session = should_use_session(spank); if (!use_session) goto done; + if (prolog_alloc_mode) + goto done; + result = process_spindle_args(spank, ac, argv, ¶ms, NULL, NULL, use_session); if (result == -1) { slurm_error("ERROR: Spindle plugin error. Could not process spindle args in local user init.\n"); @@ -384,15 +424,18 @@ int slurm_spank_job_prolog(spank_t spank, int ac, char *argv[]) { return 0; - envVal = getenv("SPANK_SPINDLE_RSHLAUNCH"); - if (envVal && strcmp(envVal, "1") == 0) - return 0; + if (!prolog_alloc_mode) { + envVal = getenv("SPANK_SPINDLE_RSHLAUNCH"); + if (envVal && strcmp(envVal, "1") == 0) + return 0; + } // The prolog starts in the user's home directory. // Change to $SLURM_JOB_WORK_DIR so logs go to right place. work_dir = getenv("SLURM_JOB_WORK_DIR"); if (work_dir) { - chdir(work_dir); + if (chdir(work_dir) == -1) + sdprintf(1, "WARNING: Could not chdir to %s: %s\n", work_dir, strerror(errno)); } err = spank_get_item(spank, S_JOB_UID, &userid); @@ -505,11 +548,13 @@ int slurm_spank_task_init(spank_t spank, int site_argc, char *site_argv[]) } if (params.opts & OPT_OFF) { + pop_env(env); return 0; } - /* When using a session without RSHLAUNCH, handle start in job prolog, not here. */ - if ((!use_session) || (params.opts & OPT_RSHLAUNCH)) { + /* When using a session without RSHLAUNCH, handle start in job prolog, not here. + With PrologFlags=Alloc, session+RSHLAUNCH is also handled in the prolog. */ + if ((!use_session) || ((params.opts & OPT_RSHLAUNCH) && !prolog_alloc_mode)) { start_params.spank = spank; start_params.site_argc = site_argc; start_params.site_argv = site_argv; @@ -560,8 +605,9 @@ static int handleStart(void *params, char **output_str) return 0; } - // Only initialize a session once - if (use_session && (args.opts & OPT_RSHLAUNCH)) { + /* Only initialize a session once. In prolog context (S_CTX_JOB_SCRIPT), + * there is no step yet, so skip the step ID check. */ + if (use_session && (args.opts & OPT_RSHLAUNCH) && spank_context() != S_CTX_JOB_SCRIPT) { err = get_stepid(spank, &stepid); if (err != ESPANK_SUCCESS) { slurm_error("ERROR: Spindle plugin error. Could not get step id."); @@ -1175,9 +1221,7 @@ static int launchBE(spank_t spank, spindle_args_t *params) else sdprintf(1, "spindleRunBE completed. Session finishing.\n"); - if (unique_file) unlink(unique_file); - free(unique_file); - unique_file = NULL; + cleanup_unique_file(); exit(result);