[Bug]: Lambda instance becomes ureachable after `dstack server` restart

[jvstme]

Steps to reproduce

1. Create an instance using the lambda backend, wait until it becomes idle or busy.
2. Restart dstack server

Actual behaviour

The instance becomes unreachable and never recovers. If it was running a job, the job is terminated. dstack-shim no longer runs on the instance.

The first shim health check attempt fails with this error:

           DEBUG    dstack._internal.server.background.tasks.process_instances:747 Check instance cloud-0 status. shim health: Can't request shim:          
                    ('Connection aborted.', RemoteDisconnected('Remote end closed connection without response'))

All other health checks fail with this error:

           DEBUG    dstack._internal.server.background.tasks.process_instances:747 Check instance cloud-0 status. shim health: Can't request shim:          
                    ('Connection aborted.', ConnectionResetError(104, 'Connection reset by peer'))

Expected behaviour

The instance remains idle or busy.

dstack version

master

Server logs

Additional information

No response

[Bihan]

The reason for this issue is currently SSH session does not return and waits for the shim process to complete. To test this hypothesis test is conducted as below:

Testing Steps:
Modified _run_ssh_command() to print testing logs as below:

def _run_ssh_command(hostname: str, ssh_private_key: str, command: str):
    with tempfile.NamedTemporaryFile("w+", 0o600) as f:
        f.write(ssh_private_key)
        f.flush()
        subprocess.run(
            [
                "ssh",
                "-F",
                "none",
                "-o",
                "StrictHostKeyChecking=no",
                "-i",
                f.name,
                f"hotaisle@{hostname}",
                command,
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        print(f"SSH command completed successfully to {hostname}: {command[:50]}...")

Result: Summary Logs

INFO     The dstack server (no version) is running at http://127.0.0.1:3000/                                                                                          
...
...                                                       
[12:19:06] DEBUG    dstack._internal.server.app:247 Processed request POST http://127.0.0.1:3000/api/project/bihan/runs/get in 0.013133s. Status: 200                           
[12:19:08] DEBUG    dstack._internal.server.background.tasks.process_runs:153 run(b275ff)amd-test: processing run                                                               
SSH command completed successfully to 23.183.40.84: sudo apt-get update...
...
...                                                              
                                                                    
           INFO     dstack._internal.server.background.tasks.process_running_jobs:620 job(00b7a9)amd-test-0-0: now is PULLING                                                   
...
...                                                              
           INFO     dstack._internal.server.background.tasks.process_runs:375 run(b275ff)amd-test: run status has changed PROVISIONING -> RUNNING                               
...
...                                                                   
           DEBUG    dstack._internal.server.background.tasks.process_instances:771 Check instance amd-test-0 status. shim health: Service is OK                                 
[12:20:19] DEBUG    dstack._internal.server.background.tasks.process_running_jobs:337 job(00b7a9)amd-test-0-0: process running job, age=0:01:49.052079                          
[12:20:25] DEBUG    dstack._internal.server.background.tasks.process_runs:153 run(b275ff)amd-test: processing run                                                               

...
...                         
           INFO     dstack._internal.server.services.jobs:360 job(00b7a9)amd-test-0-0: job status is TERMINATED, reason: TERMINATED_BY_USER                                     
[12:22:01] DEBUG    dstack._internal.server.background.tasks.process_runs:153 run(b275ff)amd-test: processing run                                                               
           INFO     dstack._internal.server.services.runs:1127 run(b275ff)amd-test: run status has changed TERMINATING -> TERMINATED, reason: STOPPED_BY_USER                   

...
...                               
^CSSH command completed successfully to 23.183.40.84: sudo sh -c '/bin/sh -c '"'"' grep -q nvidia /etc/d...

Findings:

1. _run_ssh_command() prints SSH command completed successfully to 23.183.40.84: sudo apt-get update...
   Here SSH session returns after compeleting "apt-get update".

2. _run_ssh_command() prints SSH command completed successfully to 23.183.40.84: sudo sh -c '/bin/sh -c '"'"' grep -q nvidia /etc/d...
   Here SSH session returns only after user intervenes with (CTRL+C).

Conclusion:
Indefinitely running shim command causes SSH session to hang. The SSH session hangs waiting for the shim process to complete, but the shim never completes because it's designed to run continuously.

The shim process is linked to the SSH session, so if the SSH session dies, the shim would die too. With daemonization the shim process can be intentionally unlinked from the SSH session, so the shim continues running even after the SSH session ends.