Skip to content

Commit

Permalink
Remove host_info.json on instance deploy if exists (#1479)
Browse files Browse the repository at this point in the history
Ensures that shim will create a fresh file at the first start

Fixes: #1465
  • Loading branch information
un-def authored Aug 1, 2024
1 parent 1845241 commit 5aebcee
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 2 deletions.
19 changes: 17 additions & 2 deletions src/dstack/_internal/core/backends/remote/provisioning.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@

DSTACK_SHIM_ENV_FILE = "dstack-shim.env"

HOST_INFO_FILE = "host_info.json"


def sftp_upload(client: paramiko.SSHClient, path: str, body: str) -> None:
try:
Expand Down Expand Up @@ -139,22 +141,35 @@ def check_dstack_shim_service(client: paramiko.SSHClient):
raise ProvisioningError(f"The dstack-shim service doesn't start: {line.strip()}")


def remove_host_info_if_exists(client: paramiko.SSHClient, working_dir: str) -> None:
file_path = f"{working_dir}/{HOST_INFO_FILE}"
try:
_, _, stderr = client.exec_command(
f"sudo test -e {file_path} && sudo rm {file_path}", timeout=10
)
err = stderr.read().decode().strip()
if err:
logger.debug(f"{HOST_INFO_FILE} hasn't been removed: %s", err)
except (paramiko.SSHException, OSError) as e:
raise ProvisioningError(f"remove_host_info_if_exists failed: {e}")


def get_host_info(client: paramiko.SSHClient, working_dir: str) -> Dict[str, Any]:
# wait host_info
retries = 60
iter_delay = 3
for _ in range(retries):
try:
_, stdout, stderr = client.exec_command(
f"sudo cat {working_dir}/host_info.json", timeout=10
f"sudo cat {working_dir}/{HOST_INFO_FILE}", timeout=10
)
err = stderr.read().decode().strip()
if err:
logger.debug("Retry after error: %s", err)
time.sleep(iter_delay)
continue
except (paramiko.SSHException, OSError) as e:
logger.debug("Cannot run `cat host_info.json` in the remote instance: %s", e)
logger.debug(f"Cannot run `cat {HOST_INFO_FILE}` in the remote instance: %s", e)
else:
try:
host_info_json = stdout.read()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
get_paramiko_connection,
get_shim_healthcheck,
host_info_to_instance_type,
remove_host_info_if_exists,
run_pre_start_commands,
run_shim_as_systemd_service,
upload_envs,
Expand Down Expand Up @@ -169,6 +170,9 @@ def deploy_instance(
upload_envs(client, DSTACK_WORKING_DIR, shim_envs)
logger.debug("The dstack-shim environment variables have been installed")

# Ensure host info file does not exist
remove_host_info_if_exists(client, DSTACK_WORKING_DIR)

# Run dstack-shim as a systemd service
run_shim_as_systemd_service(
client=client,
Expand Down

0 comments on commit 5aebcee

Please sign in to comment.