SageMaker Job hangs after python script finishes

We have a step functions pipeline that executes a bunch of SageMaker training and processing jobs in series. However, in the last couple weeks it happened four times that one of these jobs failed (timed out), even though the underlying python script finished successfully.

So what happens is that the python script reaches the last print statement of the main function (the print statement is the last line of code that needs to be executed), but then the SageMaker job doesn't report "SUCCESS". Instead, the job continues to run seemingly indefinitely (CPU and GPU utilization drop to zero) until it times out. It seems to happen randomly, and when re-starting the pipeline it runs just fine. Has anyone experienced the same issue and found a solution? Can I pass an explicit "Script finished successfully" message to the SageMaker job?

Note: when this issue happened the first time the underlying script or the docker image that I use hadn't been changed in months.