WQ+Vine: Track Specific Cores Used

dttools/src/Makefile

@@ -123,6 +123,7 @@ SOURCES = \
 	url_encode.c \
 	username.c \
 	uuid.c \
+	xpu_tracker.c \
 	xxmalloc.c \
 
 HEADERS_PUBLIC = \
@@ -174,6 +175,7 @@ HEADERS_PUBLIC = \
 	text_list.h \
 	timestamp.h \
 	unlink_recursive.h \
+	xpu_tracker.h \
 	xxmalloc.h \
 
 LIBRARIES = libdttools.a

dttools/src/xpu_tracker.c

@@ -0,0 +1,110 @@
+/*
+Copyright (C) 2026- The University of Notre Dame
+This software is distributed under the GNU General Public License.
+See the file COPYING for details.
+*/
+
+#include "xpu_tracker.h"
+#include "buffer.h"
+#include "debug.h"
+
+struct xpu_tracker {
+	const char *name;  // name of units: cpus, gpus, etc.
+	int count;	   // number of available units
+	int *unit_to_task; // array indicating task assigned to each unit.
+};
+
+/*
+Create a new tracker for units of type "name".
+*/
+
+struct xpu_tracker *xpu_tracker_create(const char *name, int count)
+{
+	struct xpu_tracker *x = malloc(sizeof(*x));
+	x->name = strdup(name);
+	x->count = count;
+	x->unit_to_task = calloc(count, sizeof(int));
+	return x;
+}
+
+/*
+Display the units associated with each task to the debug log
+*/
+
+void xpu_tracker_debug(struct xpu_tracker *x)
+{
+	buffer_t b;
+	buffer_init(&b);
+	buffer_putfstring(&b, "%s assigned to tasks: [ ", x->name);
+	int i;
+	for (i = 0; i < x->count; i++) {
+		buffer_putfstring(&b, "%d ", x->unit_to_task[i]);
+	}
+	buffer_putfstring(&b, " ]");
+	debug(D_VINE, "%s", buffer_tostring(&b));
+	buffer_free(&b);
+}
+
+/*
+Free all of the units associated with this taskid.
+*/
+
+void xpu_tracker_free(struct xpu_tracker *x, int taskid)
+{
+	int i;
+	for (i = 0; i < x->count; i++) {
+		if (x->unit_to_task[i] == taskid) {
+			x->unit_to_task[i] = 0;
+		}
+	}
+}
+
+/*
+Allocate n specific units to the given task.
+This assumes the total number of units has been
+accurately tracked: this function will fatal()
+if not enough are available.
+*/
+
+void xpu_tracker_alloc(struct xpu_tracker *x, int n, int taskid)
+{
+	int i;
+	for (i = 0; i < x->count && n > 0; i++) {
+		if (x->unit_to_task[i] == 0) {
+			x->unit_to_task[i] = taskid;
+			n--;
+		}
+	}
+
+	if (n > 0)
+		fatal("xpu_tracker_alloc: accounting error: ran out of %ss to assign!", x->name);
+
+	xpu_tracker_debug(x);
+}
+
+/*
+Return a string representing the units allocated to taskid.
+For example, if units 1 and 3 are allocated, return "1,3"
+This string must be freed after use.
+*/
+
+char *xpu_tracker_to_string(struct xpu_tracker *x, int taskid)
+{
+	int i;
+	int first = 1;
+	buffer_t b;
+	buffer_init(&b);
+	for (i = 0; i < x->count; i++) {
+		if (x->unit_to_task[i] == taskid) {
+			if (first) {
+				first = 0;
+			} else {
+				buffer_putfstring(&b, ",");
+			}
+			buffer_putfstring(&b, "%d", i);
+		}
+	}
+	char *str = strdup(buffer_tostring(&b));
+	buffer_free(&b);
+	return str;
+}

dttools/src/xpu_tracker.h

@@ -0,0 +1,22 @@
+/*
+Copyright (C) 2026- The University of Notre Dame
+This software is distributed under the GNU General Public License.
+See the file COPYING for details.
+*/
+
+/*
+This module tracks how and array of processing units (cpus, gpus, dpus)
+is assigned to specific tasks, so that we can indicate to running tasks
+exactly which processing units should be used.
+*/
+
+#ifndef XPU_TRACKER_H
+#define XPU_TRACKER_H
+
+struct xpu_tracker * xpu_tracker_create( const char *name, int count );
+void xpu_tracker_debug( struct xpu_tracker *x );
+void xpu_tracker_alloc( struct xpu_tracker *x, int n, int taskid );
+void xpu_tracker_free( struct xpu_tracker *x, int taskid );
+char * xpu_tracker_to_string( struct xpu_tracker *x, int taskid );
+
+#endif

taskvine/src/worker/Makefile

@@ -12,7 +12,6 @@ SOURCES = \
 	vine_transfer_server.c \
 	vine_process.c \
 	vine_watcher.c \
-	vine_gpus.c \
 	vine_workspace.c \
 	vine_worker_options.c \
 	vine_worker.c

taskvine/src/worker/vine_process.c

@@ -5,7 +5,6 @@ See the file COPYING for details.
 */
 
 #include "vine_process.h"
-#include "vine_gpus.h"
 #include "vine_manager.h"
 #include "vine_protocol.h"
 #include "vine_sandbox.h"
@@ -29,6 +28,7 @@ See the file COPYING for details.
 #include "stringtools.h"
 #include "timestamp.h"
 #include "trash.h"
+#include "xpu_tracker.h"
 #include "xxmalloc.h"
 
 #include "jx.h"
@@ -50,6 +50,8 @@ See the file COPYING for details.
 #include <sys/wait.h>
 
 extern struct vine_cache *cache_manager;
+extern struct xpu_tracker *core_tracker;
+extern struct xpu_tracker *gpu_tracker;
 
 /*
 Give the letter code used for the process sandbox dir.
@@ -179,10 +181,18 @@ static void set_integer_env_var(struct vine_process *p, const char *name, int64_
 	free(value_str);
 }
 
+static void set_tracker_env_var(struct vine_process *p, const char *name, struct xpu_tracker *tracker)
+{
+	char *str = xpu_tracker_to_string(tracker, p->task->task_id);
+	vine_task_set_env_var(p->task, name, str);
+	free(str);
+}
+
 static void set_resources_vars(struct vine_process *p)
 {
 	if (p->task->resources_requested->cores > 0) {
 		set_integer_env_var(p, "CORES", p->task->resources_requested->cores);
+		set_tracker_env_var(p, "CORES_LIST", core_tracker);
 		set_integer_env_var(p, "OMP_NUM_THREADS", p->task->resources_requested->cores);
 		set_integer_env_var(p, "OPENBLAS_NUM_THREADS", p->task->resources_requested->cores);
 		set_integer_env_var(p, "VECLIB_NUM_THREADS", p->task->resources_requested->cores);
@@ -200,9 +210,8 @@ static void set_resources_vars(struct vine_process *p)
 
 	if (p->task->resources_requested->gpus > 0) {
 		set_integer_env_var(p, "GPUS", p->task->resources_requested->gpus);
-		char *str = vine_gpus_to_string(p->task->task_id);
-		vine_task_set_env_var(p->task, "CUDA_VISIBLE_DEVICES", str);
-		free(str);
+		set_tracker_env_var(p, "GPUS_LIST", gpu_tracker);
+		set_tracker_env_var(p, "CUDA_VISIBLE_DEVICES", gpu_tracker);
 	}
 }