diff --git a/.github/workflows/training-e2e.yaml b/.github/workflows/training-e2e.yaml
index 9033f119..cc505ed3 100644
--- a/.github/workflows/training-e2e.yaml
+++ b/.github/workflows/training-e2e.yaml
@@ -96,9 +96,6 @@ jobs:
         env:
           ANSIBLE_CONFIG: ./main/training/tests/ansible.cfg
 
-      - name: Wait for 4 minutes
-        run: sleep 240
-
       # - name: Setup tmate session
       #   uses: mxschmitt/action-tmate@v3.18
       #   timeout-minutes: 20
diff --git a/training/tests/provision/playbook.yml b/training/tests/provision/playbook.yml
index 1181e34c..90d0936f 100644
--- a/training/tests/provision/playbook.yml
+++ b/training/tests/provision/playbook.yml
@@ -6,12 +6,12 @@
   gather_facts: false
 
   tasks:
-  
+
   - name: Wait until the instance is ready
     ansible.builtin.wait_for_connection:
     delay: 15
     timeout: 180
-  
+
   - name: Gather facts for first time
     ansible.builtin.setup:
 
@@ -42,7 +42,7 @@
     ansible.builtin.shell: |
       podman pull "quay.io/ai-lab/{{ image_name }}:latest" \
         --authfile=/etc/containers/auth.json \
-        --arch amd64 
+        --arch amd64
 
         # --retry=3 \
         # --retry-delay=15 \
@@ -112,6 +112,6 @@
     delegate_to: localhost
 
   - name: Reboot
-    ansible.builtin.shell: systemctl reboot
-    ignore_errors: true
-    ignore_unreachable: true
+    ansible.builtin.reboot:
+      reboot_timeout: 300
+      reboot_msg: "Reboot initiated by Ansible"