diff --git a/ansible/expand_ubuntu_disk_workflow.yml b/ansible/expand_ubuntu_disk_workflow.yml
index 56b09d7..1284be2 100644
--- a/ansible/expand_ubuntu_disk_workflow.yml
+++ b/ansible/expand_ubuntu_disk_workflow.yml
@@ -15,16 +15,21 @@
         ansible_password: "{{ ansible_password }}"
         ansible_become: true
         ansible_become_password: "{{ ansible_become_password }}"
-        ansible_ssh_extra_args: "-o StrictHostKeyChecking=no"
+        ansible_ssh_common_args: "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"
 
-- name: Stage 2 - Wait for SSH
-  hosts: ubuntu_resize_group
+- name: Stage 2 - Wait for SSH port
+  hosts: localhost
+  connection: local
   gather_facts: false
 
   tasks:
-    - name: Wait for SSH connection
-      ansible.builtin.wait_for_connection:
-        timeout: 300
+    - name: Wait for SSH port on Ubuntu VM
+      ansible.builtin.wait_for:
+        host: "{{ vm_ip }}"
+        port: 22
+        timeout: 60
+        sleep: 2
+
 
 - name: Stage 3 - Read current guest disk size
   hosts: ubuntu_resize_group