Factory: Improved Shutdown of Local Workers (#4196)

dthain · btovar · commit 256e17b21f8b · 2025-09-25T12:37:16.000-04:00
* Add batch_queue_remove_mode_t parameter to allow for friendly and unfriendly kill modes.

Conflicts:
	batch_job/src/batch_queue_internal.h

* format

* Modify factory to shutdown workers in parallel.
And, for local workers, send SIGKILL after 30s until gone.

* Adjust makeflow to use batch_queue_remove_mode_t.

* Extend worker kill from vine_factory to work_queue_factory.

* Remove mistaken prune

* Remove unused int count

* Remove unused int count

* After waiting 30s, send the kill signal every 5s until worker exited.
diff --git a/batch_job/src/batch_queue.c b/batch_job/src/batch_queue.c
@@ -244,9 +244,9 @@ batch_queue_id_t batch_queue_wait_timeout(struct batch_queue *q, struct batch_jo
 	return q->module->wait(q, info, stoptime);
 }
 
-int batch_queue_remove(struct batch_queue *q, batch_queue_id_t jobid)
+int batch_queue_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
-	return q->module->remove(q, jobid);
+	return q->module->remove(q, jobid, mode);
 }
 
 /* vim: set noexpandtab tabstop=8: */
diff --git a/batch_job/src/batch_queue.h b/batch_job/src/batch_queue.h
@@ -84,6 +84,12 @@ typedef enum {
 	BATCH_QUEUE_TYPE_UNKNOWN = -1         /**< An invalid batch queue type. */
 } batch_queue_type_t;
 
+/** Indicates how aggressively to remove a batch job. */
+typedef enum {
+	BATCH_QUEUE_REMOVE_MODE_FRIENDLY,     /**< Remove a batch job with a friendly signal that allows it to clean up. */
+	BATCH_QUEUE_REMOVE_MODE_KILL,	      /**< Remove a batch job by sending a kill signal that terminates quickly. */
+} batch_queue_remove_mode_t;
+
 /** Create a new batch queue.
 @param type The type of the queue.
 @param ssl_key_file The location of the queue manager's ssl key file, if it has one.
@@ -133,7 +139,7 @@ You must still call @ref batch_queue_wait to wait for the removal to complete.
 @param jobid The job to be removed.
 @return Greater than zero if the job exists and was removed, zero otherwise.
 */
-int batch_queue_remove(struct batch_queue *q, batch_queue_id_t jobid);
+int batch_queue_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode );
 
 /** Converts a string into a batch queue type.
 @param str A string listing all of the known batch queue types (which changes over time.)
diff --git a/batch_job/src/batch_queue_amazon.c b/batch_job/src/batch_queue_amazon.c
@@ -789,7 +789,7 @@ kill the local ssh process forcibly, and then we save
 the Amazon instance and delete other expired instances.
 */
 
-static int batch_queue_amazon_remove(struct batch_queue *q, batch_queue_id_t jobid)
+static int batch_queue_amazon_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
 	struct batch_queue_amazon_info *info;
 
diff --git a/batch_job/src/batch_queue_cluster.c b/batch_job/src/batch_queue_cluster.c
@@ -393,7 +393,7 @@ static batch_queue_id_t batch_queue_cluster_wait(struct batch_queue *q, struct b
 	return -1;
 }
 
-static int batch_queue_cluster_remove(struct batch_queue *q, batch_queue_id_t jobid)
+static int batch_queue_cluster_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
 	struct batch_job_info *info;
 
diff --git a/batch_job/src/batch_queue_condor.c b/batch_job/src/batch_queue_condor.c
@@ -358,7 +358,7 @@ static batch_queue_id_t batch_queue_condor_wait(struct batch_queue *q, struct ba
 	return -1;
 }
 
-static int batch_queue_condor_remove(struct batch_queue *q, batch_queue_id_t jobid)
+static int batch_queue_condor_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
 	char *command = string_format("condor_rm %" PRIbjid, jobid);
 
diff --git a/batch_job/src/batch_queue_dryrun.c b/batch_job/src/batch_queue_dryrun.c
@@ -15,8 +15,8 @@ See the file COPYING for details.
 #include "debug.h"
 #include "process.h"
 #include "macros.h"
-#include "stringtools.h"
 #include "path.h"
+#include "stringtools.h"
 #include "xxmalloc.h"
 
 static batch_queue_id_t batch_queue_dryrun_submit(struct batch_queue *q, struct batch_job *bt)
@@ -84,7 +84,7 @@ static batch_queue_id_t batch_queue_dryrun_wait(struct batch_queue *q, struct ba
 	}
 }
 
-static int batch_queue_dryrun_remove(struct batch_queue *q, batch_queue_id_t jobid)
+static int batch_queue_dryrun_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
 	return 0;
 }
diff --git a/batch_job/src/batch_queue_flux.c b/batch_job/src/batch_queue_flux.c
@@ -258,7 +258,7 @@ static batch_queue_id_t batch_queue_flux_wait(struct batch_queue *q, struct batc
 	return batch_queue_flux_wait_jobid(q, info_out, stoptime, 0);
 }
 
-static int batch_queue_flux_remove(struct batch_queue *q, batch_queue_id_t jobid)
+static int batch_queue_flux_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
 	struct flux_job_info *info = itable_lookup(batch_queue_jobid_info_table, jobid);
 	if (!info) {
diff --git a/batch_job/src/batch_queue_internal.h b/batch_job/src/batch_queue_internal.h
@@ -32,7 +32,8 @@ struct batch_queue_module {
 
 	batch_queue_id_t (*submit) (struct batch_queue *Q, struct batch_job *bt );
 	batch_queue_id_t (*wait) (struct batch_queue *Q, struct batch_job_info *info, time_t stoptime);
-	int (*remove) (struct batch_queue *Q, batch_queue_id_t id);
+
+	int (*remove) (struct batch_queue *Q, batch_queue_id_t id, batch_queue_remove_mode_t mode );
 };
 
 struct batch_queue {
diff --git a/batch_job/src/batch_queue_k8s.c b/batch_job/src/batch_queue_k8s.c
@@ -349,7 +349,7 @@ static batch_queue_id_t batch_queue_k8s_submit(struct batch_queue *q, struct bat
 	return jobid;
 }
 
-static int batch_queue_k8s_remove(struct batch_queue *q, batch_queue_id_t jobid)
+static int batch_queue_k8s_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
 
 	pid_t pid = fork();
@@ -408,7 +408,7 @@ static void batch_queue_k8s_handle_complete_task(char *pod_id, int job_id,
 	free(info);
 
 	list_remove(running_pod_lst, pod_id);
-	batch_queue_k8s_remove(q, job_id);
+	batch_queue_k8s_remove(q, job_id, BATCH_QUEUE_REMOVE_MODE_FRIENDLY);
 
 	if (curr_k8s_job_info->is_running == 0) {
 		process_wait(timeout);
diff --git a/batch_job/src/batch_queue_local.c b/batch_job/src/batch_queue_local.c
@@ -107,10 +107,16 @@ static batch_queue_id_t batch_queue_local_wait(struct batch_queue *q, struct bat
 	}
 }
 
-static int batch_queue_local_remove(struct batch_queue *q, batch_queue_id_t jobid)
+/* Note that batch_queue_remove triggers a job to exit, but it must still be returned by batch_queue_wait */
+
+static int batch_queue_local_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
-	int max_wait = 5; // maximum seconds we wish to wait for a given process
-	process_kill_waitpid(jobid, max_wait);
+	if (mode == BATCH_QUEUE_REMOVE_MODE_FRIENDLY) {
+		kill(jobid, SIGQUIT);
+	} else {
+		kill(jobid, SIGKILL);
+	}
+
 	return 0;
 }
 
diff --git a/batch_job/src/batch_queue_vine.c b/batch_job/src/batch_queue_vine.c
@@ -163,7 +163,7 @@ static batch_queue_id_t batch_queue_vine_wait(struct batch_queue *q, struct batc
 	}
 }
 
-static int batch_queue_vine_remove(struct batch_queue *q, batch_queue_id_t jobid)
+static int batch_queue_vine_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
 	return 0;
 }
diff --git a/batch_job/src/batch_queue_work_queue.c b/batch_job/src/batch_queue_work_queue.c
@@ -146,7 +146,7 @@ static batch_queue_id_t batch_queue_wq_wait(struct batch_queue *q, struct batch_
 	}
 }
 
-static int batch_queue_wq_remove(struct batch_queue *q, batch_queue_id_t jobid)
+static int batch_queue_wq_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)
 {
 	return 0;
 }
diff --git a/batch_job/src/vine_factory.c b/batch_job/src/vine_factory.c
@@ -151,6 +151,10 @@ static int single_shot = 0;
 // A pretend port number to unqiuely identify the factory in catalog updates.
 static int pretend_port = 0;
 
+// When killing a worker, how long to wait between sending SIGQUIT and SIGKILL
+static int worker_kill_timeout = 30;
+
+
 /*
 In a signal handler, only a limited number of functions are safe to
 invoke, so we construct a string and emit it with a low-level write.
@@ -620,20 +624,54 @@ static int submit_workers( struct batch_queue *queue, struct itable *job_table,
 	return i;
 }
 
-void remove_all_workers( struct batch_queue *queue, struct itable *job_table )
+void wait_all_workers( struct batch_queue *queue, struct itable *job_table, time_t stoptime)
+{
+	debug(D_VINE,"waiting for workers to exit...");
+	while(itable_size(job_table)) {
+		struct batch_job_info info;
+		uint64_t jobid = batch_queue_wait_timeout(queue,&info,stoptime);
+		if(jobid>0) {
+			debug(D_VINE,"worker job %"PRId64" completed",jobid);
+			itable_remove(job_table,jobid);
+		} else if(jobid<=0) {
+			/* indicates no more jobs or time expired. */
+			break;
+		}
+	}
+}
+
+void remove_all_workers( struct batch_queue *queue, struct itable *job_table, batch_queue_remove_mode_t mode )
 {
 	uint64_t jobid;
 	void *value;
 
-	debug(D_VINE,"removing all remaining worker jobs...");
-	int count = itable_size(job_table);
+	debug(D_VINE,"removing all workers...");
+
 	itable_firstkey(job_table);
 	while(itable_nextkey(job_table,&jobid,&value)) {
-		debug(D_VINE,"removing job %"PRId64,jobid);
-		batch_queue_remove(queue,jobid);
+		debug(D_VINE,"removing worker job %"PRId64,jobid);
+		batch_queue_remove(queue,jobid,mode);
 	}
-	debug(D_VINE,"%d workers removed.",count);
+}
 
+void remove_and_wait_all_workers( struct batch_queue *queue, struct itable *job_table )
+{
+	remove_all_workers(queue,job_table,BATCH_QUEUE_REMOVE_MODE_FRIENDLY);
+
+	/*
+	For non-local batch queues, we skip this be because our ability to
+	wait on removed jobs is unreliable.  For local jobs, we follow up
+	the friendly remove after 30s with an unfriendly kill every 5s,
+	and we wait for children to exit to ensure cleanup is complete.
+	*/
+
+	if(batch_queue_get_type(queue)==BATCH_QUEUE_TYPE_LOCAL) {
+		wait_all_workers(queue,job_table,time(0)+worker_kill_timeout);
+		while(itable_size(job_table)) {
+			remove_all_workers(queue,job_table,BATCH_QUEUE_REMOVE_MODE_KILL);
+			wait_all_workers(queue,job_table,time(0)+5);
+		}
+	}
 }
 
 void print_stats(struct jx *j) {
@@ -1176,7 +1214,7 @@ static void mainloop( struct batch_queue *queue )
 	}
 
 	printf("removing %d workers...\n",itable_size(job_table));
-	remove_all_workers(queue,job_table);
+	remove_and_wait_all_workers(queue,job_table);
 	printf("all workers removed.\n");
 	itable_delete(job_table);
 }
diff --git a/batch_job/src/work_queue_factory.c b/batch_job/src/work_queue_factory.c
@@ -139,6 +139,10 @@ static int disable_afs_check = 0;
 // A pretend port number to unqiuely identify the factory in catalog updates.
 static int pretend_port = 0;
 
+// When killing a worker, how long to wait between sending SIGQUIT and SIGKILL
+static int worker_kill_timeout = 30;
+
+
 /*
 In a signal handler, only a limited number of functions are safe to
 invoke, so we construct a string and emit it with a low-level write.
@@ -587,20 +591,53 @@ static int submit_workers( struct batch_queue *queue, struct itable *job_table,
 	return i;
 }
 
-void remove_all_workers( struct batch_queue *queue, struct itable *job_table )
+void wait_all_workers( struct batch_queue *queue, struct itable *job_table, time_t stoptime)
+{
+	debug(D_WQ,"waiting for workers to exit...");
+	while(itable_size(job_table)) {
+		struct batch_job_info info;
+		uint64_t jobid = batch_queue_wait_timeout(queue,&info,stoptime);
+		if(jobid>0) {
+			debug(D_WQ,"worker job %"PRId64" completed",jobid);
+			itable_remove(job_table,jobid);
+		} else if(jobid<=0) {
+			/* indicates no more jobs or time expired. */
+			break;
+		}
+	}
+}
+
+void remove_all_workers( struct batch_queue *queue, struct itable *job_table, batch_queue_remove_mode_t mode )
 {
 	uint64_t jobid;
 	void *value;
 
-	debug(D_WQ,"removing all remaining worker jobs...");
-	int count = itable_size(job_table);
+	debug(D_WQ,"removing all workers...");
+
 	itable_firstkey(job_table);
 	while(itable_nextkey(job_table,&jobid,&value)) {
-		debug(D_WQ,"removing job %"PRId64,jobid);
-		batch_queue_remove(queue,jobid);
+		debug(D_WQ,"removing worker job %"PRId64,jobid);
+		batch_queue_remove(queue,jobid,mode);
 	}
-	debug(D_WQ,"%d workers removed.",count);
+}
 
+void remove_and_wait_all_workers( struct batch_queue *queue, struct itable *job_table )
+{
+	remove_all_workers(queue,job_table,BATCH_QUEUE_REMOVE_MODE_FRIENDLY);
+
+	/*
+	For non-local batch queues, we skip this be because our ability to
+	wait on removed jobs is unreliable.  For local jobs, we follow up
+	the friendly remove after 30s with an unfriendly kill, and we wait
+	for children to exit to ensure cleanup is complete.
+	*/
+
+	if(batch_queue_get_type(queue)==BATCH_QUEUE_TYPE_LOCAL) {
+		while(itable_size(job_table)) {
+			wait_all_workers(queue,job_table,time(0)+worker_kill_timeout);
+			remove_all_workers(queue,job_table,BATCH_QUEUE_REMOVE_MODE_KILL);
+		}
+	}
 }
 
 void print_stats(struct jx *j) {
@@ -1127,7 +1164,7 @@ static void mainloop( struct batch_queue *queue )
 	}
 
 	printf("removing %d workers...\n",itable_size(job_table));
-	remove_all_workers(queue,job_table);
+	remove_and_wait_all_workers(queue,job_table);
 	printf("all workers removed.\n");
 	itable_delete(job_table);
 }
diff --git a/makeflow/src/makeflow.c b/makeflow/src/makeflow.c
@@ -391,7 +391,7 @@ static void makeflow_abort_job( struct dag *d, struct dag_node *n, struct batch_
 {
 	printf("aborting %s job %" PRIu64 "\n", name, jobid);
 
-	batch_queue_remove(q, jobid);
+	batch_queue_remove(q, jobid, BATCH_QUEUE_REMOVE_MODE_FRIENDLY);
 	makeflow_hook_node_abort(n);
 	makeflow_log_state_change(d, n, DAG_NODE_STATE_ABORTED);
 	makeflow_clean_node(d,q,n);
@@ -438,10 +438,10 @@ void makeflow_node_reset( struct dag *d, struct dag_node *n )
 
 	if(n->state == DAG_NODE_STATE_RUNNING) {
 		if(n->local_job && local_queue) {
-			batch_queue_remove(local_queue, n->jobid);
+			batch_queue_remove(local_queue, n->jobid, BATCH_QUEUE_REMOVE_MODE_FRIENDLY);
 			itable_remove(d->local_job_table, n->jobid);
 		} else {
-			batch_queue_remove(remote_queue, n->jobid);
+			batch_queue_remove(remote_queue, n->jobid, BATCH_QUEUE_REMOVE_MODE_FRIENDLY);
 			itable_remove(d->remote_job_table, n->jobid);
 		}
 	}

Original file line number	Diff line number	Diff line change
`@@ -244,9 +244,9 @@ batch_queue_id_t batch_queue_wait_timeout(struct batch_queue *q, struct batch_jo`
`244`	`244`	`return q->module->wait(q, info, stoptime);`
`245`	`245`	`}`
`246`	`246`
`247`		`-int batch_queue_remove(struct batch_queue *q, batch_queue_id_t jobid)`
	`247`	`+int batch_queue_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)`
`248`	`248`	`{`
`249`		`- return q->module->remove(q, jobid);`
	`249`	`+ return q->module->remove(q, jobid, mode);`
`250`	`250`	`}`
`251`	`251`
`252`	`252`	`/* vim: set noexpandtab tabstop=8: */`
Original file line number	Diff line number	Diff line change
`@@ -789,7 +789,7 @@ kill the local ssh process forcibly, and then we save`
`789`	`789`	`the Amazon instance and delete other expired instances.`
`790`	`790`	`*/`
`791`	`791`
`792`		`-static int batch_queue_amazon_remove(struct batch_queue *q, batch_queue_id_t jobid)`
	`792`	`+static int batch_queue_amazon_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)`
`793`	`793`	`{`
`794`	`794`	`struct batch_queue_amazon_info *info;`
`795`	`795`
Original file line number	Diff line number	Diff line change
`@@ -393,7 +393,7 @@ static batch_queue_id_t batch_queue_cluster_wait(struct batch_queue *q, struct b`
`393`	`393`	`return -1;`
`394`	`394`	`}`
`395`	`395`
`396`		`-static int batch_queue_cluster_remove(struct batch_queue *q, batch_queue_id_t jobid)`
	`396`	`+static int batch_queue_cluster_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)`
`397`	`397`	`{`
`398`	`398`	`struct batch_job_info *info;`
`399`	`399`
Original file line number	Diff line number	Diff line change
`@@ -358,7 +358,7 @@ static batch_queue_id_t batch_queue_condor_wait(struct batch_queue *q, struct ba`
`358`	`358`	`return -1;`
`359`	`359`	`}`
`360`	`360`
`361`		`-static int batch_queue_condor_remove(struct batch_queue *q, batch_queue_id_t jobid)`
	`361`	`+static int batch_queue_condor_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)`
`362`	`362`	`{`
`363`	`363`	`char *command = string_format("condor_rm %" PRIbjid, jobid);`
`364`	`364`
Original file line number	Diff line number	Diff line change
`@@ -15,8 +15,8 @@ See the file COPYING for details.`
`15`	`15`	`#include "debug.h"`
`16`	`16`	`#include "process.h"`
`17`	`17`	`#include "macros.h"`
`18`		`-#include "stringtools.h"`
`19`	`18`	`#include "path.h"`
	`19`	`+#include "stringtools.h"`
`20`	`20`	`#include "xxmalloc.h"`
`21`	`21`
`22`	`22`	`static batch_queue_id_t batch_queue_dryrun_submit(struct batch_queue q, struct batch_job bt)`
`@@ -84,7 +84,7 @@ static batch_queue_id_t batch_queue_dryrun_wait(struct batch_queue *q, struct ba`
`84`	`84`	`}`
`85`	`85`	`}`
`86`	`86`
`87`		`-static int batch_queue_dryrun_remove(struct batch_queue *q, batch_queue_id_t jobid)`
	`87`	`+static int batch_queue_dryrun_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)`
`88`	`88`	`{`
`89`	`89`	`return 0;`
`90`	`90`	`}`
Original file line number	Diff line number	Diff line change
`@@ -258,7 +258,7 @@ static batch_queue_id_t batch_queue_flux_wait(struct batch_queue *q, struct batc`
`258`	`258`	`return batch_queue_flux_wait_jobid(q, info_out, stoptime, 0);`
`259`	`259`	`}`
`260`	`260`
`261`		`-static int batch_queue_flux_remove(struct batch_queue *q, batch_queue_id_t jobid)`
	`261`	`+static int batch_queue_flux_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)`
`262`	`262`	`{`
`263`	`263`	`struct flux_job_info *info = itable_lookup(batch_queue_jobid_info_table, jobid);`
`264`	`264`	`if (!info) {`
Original file line number	Diff line number	Diff line change
`@@ -349,7 +349,7 @@ static batch_queue_id_t batch_queue_k8s_submit(struct batch_queue *q, struct bat`
`349`	`349`	`return jobid;`
`350`	`350`	`}`
`351`	`351`
`352`		`-static int batch_queue_k8s_remove(struct batch_queue *q, batch_queue_id_t jobid)`
	`352`	`+static int batch_queue_k8s_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)`
`353`	`353`	`{`
`354`	`354`
`355`	`355`	`pid_t pid = fork();`
`@@ -408,7 +408,7 @@ static void batch_queue_k8s_handle_complete_task(char *pod_id, int job_id,`
`408`	`408`	`free(info);`
`409`	`409`
`410`	`410`	`list_remove(running_pod_lst, pod_id);`
`411`		`- batch_queue_k8s_remove(q, job_id);`
	`411`	`+ batch_queue_k8s_remove(q, job_id, BATCH_QUEUE_REMOVE_MODE_FRIENDLY);`
`412`	`412`
`413`	`413`	`if (curr_k8s_job_info->is_running == 0) {`
`414`	`414`	`process_wait(timeout);`
Original file line number	Diff line number	Diff line change
`@@ -107,10 +107,16 @@ static batch_queue_id_t batch_queue_local_wait(struct batch_queue *q, struct bat`
`107`	`107`	`}`
`108`	`108`	`}`
`109`	`109`
`110`		`-static int batch_queue_local_remove(struct batch_queue *q, batch_queue_id_t jobid)`
	`110`	`+/* Note that batch_queue_remove triggers a job to exit, but it must still be returned by batch_queue_wait */`
	`111`	`+`
	`112`	`+static int batch_queue_local_remove(struct batch_queue *q, batch_queue_id_t jobid, batch_queue_remove_mode_t mode)`
`111`	`113`	`{`
`112`		`- int max_wait = 5; // maximum seconds we wish to wait for a given process`
`113`		`- process_kill_waitpid(jobid, max_wait);`
	`114`	`+ if (mode == BATCH_QUEUE_REMOVE_MODE_FRIENDLY) {`
	`115`	`+ kill(jobid, SIGQUIT);`
	`116`	`+ } else {`
	`117`	`+ kill(jobid, SIGKILL);`
	`118`	`+ }`
	`119`	`+`
`114`	`120`	`return 0;`
`115`	`121`	`}`
`116`	`122`