From b0fb87ea044a05a3ee34bbbcb0972456094c3f12 Mon Sep 17 00:00:00 2001
From: Radostin Stoyanov <rstoyanov@fedoraproject.org>
Date: Mon, 10 Nov 2025 12:55:17 +0000
Subject: [PATCH] criu: add support for --gpu-device-map

This option allows to specify a GPU device map that will be used during
restore. The provided map must contain all checkpointed GPU devices
in the format "oldUuid=newUuid,oldUuid=newUuid,..." that will be
used to remap old devices to new devices.

Example of migrating all process from GPU_0 to GPU_1:

i=0
for uuid in $(nvidia-smi --list-gpus | grep -oP 'UUID: \K[^)]+'); do
    export GPU_$i=$uuid
    i=$((i+1))
done

criu restore --gpu-device-map=$GPU_0=$GPU_1,$GPU_1=$GPU_0,$GPU_2=$GPU_2,$GPU_3=$GPU_3

Signed-off-by: Radostin Stoyanov <rstoyanov@fedoraproject.org>
---
 Documentation/criu.txt     |  5 +++++
 criu/config.c              |  4 ++++
 criu/crtools.c             |  3 +++
 criu/include/cr_options.h  |  5 +++++
 images/rpc.proto           |  1 +
 plugins/cuda/cuda_plugin.c | 25 +++++++++++++++++++++----
 6 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/Documentation/criu.txt b/Documentation/criu.txt
index 0c9a9e527..c28ec4698 100644
--- a/Documentation/criu.txt
+++ b/Documentation/criu.txt
@@ -720,6 +720,11 @@ The 'mode' may be one of the following:
     Required when dumped with this option. Refer to this option in the section
     on dumping for more details.
 
+*--gpu-device-map*::
+    Optional GPU device map used during restore to remap old devices onto new ones.
+    For more information, please refer to the documentation for the *--device-map*
+    option of the *cuda-checkpoint* tool.
+
 *check*
 ~~~~~~~
 Checks whether the kernel supports the features needed by *criu* to
diff --git a/criu/config.c b/criu/config.c
index d7ef3f8e8..7bc87b709 100644
--- a/criu/config.c
+++ b/criu/config.c
@@ -705,6 +705,7 @@ int parse_options(int argc, char **argv, bool *usage_error, bool *has_exec_cmd,
 		BOOL_OPT("unprivileged", &opts.unprivileged),
 		BOOL_OPT("ghost-fiemap", &opts.ghost_fiemap),
 		BOOL_OPT(OPT_ALLOW_UPROBES, &opts.allow_uprobes),
+		{ "gpu-device-map", required_argument, 0, 1101 },
 		{},
 	};
 
@@ -1045,6 +1046,9 @@ int parse_options(int argc, char **argv, bool *usage_error, bool *has_exec_cmd,
 				return 1;
 			}
 			break;
+		case 1101:
+			SET_CHAR_OPTS(gpu_device_map, optarg);
+			break;
 		case 'V':
 			pr_msg("Version: %s\n", CRIU_VERSION);
 			if (strcmp(CRIU_GITID, "0"))
diff --git a/criu/crtools.c b/criu/crtools.c
index 4dc55a065..aad9a2184 100644
--- a/criu/crtools.c
+++ b/criu/crtools.c
@@ -521,6 +521,9 @@ usage:
 	       "  --skip-file-rwx-check\n"
 	       "			Skip checking file permissions\n"
 	       "			(r/w/x for u/g/o) on restore.\n"
+	       "  --gpu-device-map\n"
+	       "			Optional GPU device map used during restore\n"
+	       "			to remap old devices onto new ones.\n"
 	       "\n"
 	       "Check options:\n"
 	       "  Without options, \"criu check\" checks availability of absolutely required\n"
diff --git a/criu/include/cr_options.h b/criu/include/cr_options.h
index 8c5707b41..d461fccfa 100644
--- a/criu/include/cr_options.h
+++ b/criu/include/cr_options.h
@@ -247,6 +247,11 @@ struct cr_options {
 	 * explicitly request it as it comes with many limitations.
 	 */
 	int unprivileged;
+	/*
+	 * Optional GPU device map used during restore to manually remap old devices
+	 * onto new ones. See `cuda-checkpoint --help` for more information.
+	 */
+	char *gpu_device_map;
 };
 
 extern struct cr_options opts;
diff --git a/images/rpc.proto b/images/rpc.proto
index 1a4722a9c..8db173d3a 100644
--- a/images/rpc.proto
+++ b/images/rpc.proto
@@ -145,6 +145,7 @@ message criu_opts {
 	optional bool			leave_stopped		= 69;
 	optional bool			display_stats		= 70;
 	optional bool			log_to_stderr		= 71;
+	optional string			gpu_device_map		= 72;
 /*	optional bool			check_mounts		= 128;	*/
 }
 
diff --git a/plugins/cuda/cuda_plugin.c b/plugins/cuda/cuda_plugin.c
index 9ccb04224..d5e69935b 100644
--- a/plugins/cuda/cuda_plugin.c
+++ b/plugins/cuda/cuda_plugin.c
@@ -261,15 +261,32 @@ static int cuda_process_checkpoint_action(int pid, const char *action, unsigned
 {
 	char pid_buf[16];
 	char timeout_buf[16];
+	int args_idx = 5;
 
 	snprintf(pid_buf, sizeof(pid_buf), "%d", pid);
 
-	const char *args[] = { CUDA_CHECKPOINT, "--action", action, "--pid", pid_buf, NULL /* --timeout */,
-			       NULL /* timeout_val */, NULL };
+	const char *args[] = {
+		CUDA_CHECKPOINT,
+		"--action", action,
+		"--pid", pid_buf,
+		NULL /* --timeout */,
+		NULL /* timeout_val */,
+		NULL /* --device-map */,
+		NULL /* device_map_val */,
+		NULL
+	};
+
 	if (timeout > 0) {
 		snprintf(timeout_buf, sizeof(timeout_buf), "%d", timeout);
-		args[5] = "--timeout";
-		args[6] = timeout_buf;
+		args[args_idx] = "--timeout";
+		args[args_idx+1] = timeout_buf;
+		args_idx += 2;
+	}
+
+	if (opts.gpu_device_map && strncmp(action, ACTION_RESTORE, strlen(ACTION_RESTORE)) == 0) {
+		pr_debug("opts.gpu_device_map: %s\n", opts.gpu_device_map);
+		args[args_idx] = "--device-map";
+		args[args_idx+1] = opts.gpu_device_map;
 	}
 
 	return launch_cuda_checkpoint(args, msg_buf, buf_size);