How can I checkpoint and restart a TFC program?

[leakec]

I have programs that take a while to compile and/or optimize. After I compile and/or optimize, I'd like to save their state so I can restart it at a different time. How can I do that?

[leakec]

TFC version 1.0.4 contains a new feature in TFCUtils called criuCheckpoint. It can be used like this

from tfc.utils.TFCUtils import criuCheckpoint
criuCheckpoint()

There are some arguments to criuCheckpoint the user can use to modify the behavior. See the criuCheckpoint API for more details.

This will create a checkpoint in the criu_checkpoint with CRIU directory that contains the data needed to restart the program from the same point using CRIU. A message will be printed when running this function that will instruct the user how to restart at a later time.

Warnings/known limitations

• The user must ensure no external sockets are used. This happens most commonly with matplotlib. Either disable matplotlib plotting or plot with a different program such as plotly.
• The user must ensure no GPU memory is used. There is some support for AMD GPUs with CRIU, but no NVIDIA support at the moment. Unless the user has built JAX for their GPU and is purposely running computations on the GPU, this is not an issue.
• The user_mode is only available in CRIU 3.18, and is not full featured yet. Hence, use this mode with caution.