How do you schedule reboots for Ubuntu lab machines running training jobs?

Set a regular maintenance window rather than relying on the machines being idle. Get it agreed with whoever owns the lab and your supervisor, then publish the dates well ahead of time so users can plan their jobs around them.

You don't have to reboot every time the window comes around, but when maintenance is needed, make the reboot part of the policy. Otherwise a persistent login or a long-running job could keep a machine from being restarted indefinitely.

Communicate in more than one place. Post the planned reboot and expected maintenance window where lab users normally get announcements, then use wall to warn anyone actively connected to the machines shortly before the reboot.

Don't count on a terminal warning alone. Someone may have a training script running in tmux without watching the session.

Stagger the reboots across the maintenance window instead of restarting the whole lab at once. That leaves some capacity available while other machines are coming back up, assuming users can work around the individual restart times.

A scheduled task can automate the warnings and reboots, but make the timetable predictable rather than surprising people with an arbitrary restart.

An idle check sounds convenient, but define what idle means first. A remote session can be unattended while a Python process is still doing useful work, and checking only for logged-in users won't tell you whether a training job is safe to interrupt.

I'd use idle detection as a way to find a convenient earlier opportunity, not as the only condition for an overdue reboot. Tell users when the maintenance window is and that running jobs may be interrupted.