pytorch support

I just tried this out on PyTorch and it seems to work for the cuda state but I'm hitting issues with criu when saving the parent process. It seems like the issue is with saving the nvidia driver in criu.

Are there any plans to expand support for this with criu for common ML frameworks?

~/D/torch-criu (main)> third-party/cuda-checkpoint/bin/x86_64_Linux/cuda-checkpoint --toggle --pid 125704
~/D/torch-criu (main)> sudo criu dump --shell-job --images-dir demo --tree 125704
Error (criu/files-ext.c:94): Can't dump file 19 of that type [20666] (chr 195:255)
Error (criu/cr-dump.c:1669): Dump files (pid: 125704) failed with -1
Error (criu/cr-dump.c:2093): Dumping FAILED.

There's no longer an active cuda process after toggling but still seems to have access to a nvidia device.

~/D/torch-criu (main) [1]> nvidia-smi
Mon Apr 29 17:11:21 2024       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.67                 Driver Version: 550.67         CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3090        On  |   00000000:06:00.0 Off |                  N/A |
|  0%   32C    P8             16W /  350W |       5MiB /  24576MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA GeForce RTX 3090        On  |   00000000:07:00.0 Off |                  N/A |
|  0%   34C    P8             21W /  350W |       5MiB /  24576MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

The file that failed to save seems to be nvidia.

~/D/torch-criu (main)> grep 195 /proc/devices 
195 nvidia
195 nvidia-modeset
195 nvidiactl

Test script

import time
import os
import torch

device = torch.device("cuda")

a = torch.tensor(10, device=device)

print(os.getpid())
time.sleep(1000)

~/D/torch-criu (main)> criu -V
Version: 3.18
GitID: v3.18

Reproduce the failure with the provided Python script, cuda-checkpoint toggle command, CRIU 3.18, and demo images directory; start with CRIU's files-ext.c error for the NVIDIA character device. Done would mean the PyTorch process can be dumped without the reported NVIDIA-device failure.