[shim] Don't revoke host SSH access of co-located tasks - #4226
Open
un-def wants to merge 1 commit into
Open
Conversation
The shim appends `cfg.HostSshKeys` to the host user's `authorized_keys` when a task starts and removes every entry matching those fingerprints when it ends. That key is the user's own SSH key, shared by all of their runs unless overridden with `dstack apply --ssh-identity`, so on an instance that holds more than one of their jobs at a time -- a fleet with `blocks` -- the first job to finish revokes the SSH access of the jobs still running. The same unconditional removal also drops a key the user added by hand, if a run happened to reuse it. The per-task add/remove deltas are replaced with a reconcile of the entries the shim owns, that is, the ones carrying the `# added by dstack-shim` marker shipped in #4206. For every host user, those entries become exactly the keys of the stored tasks that are not cleaned up yet; everything else in the file -- keys added by the user, keys added by the server when provisioning an SSH fleet, comments -- is kept verbatim. Refcounting is implicit in set membership, so there is no new state to keep or persist: `HostSshKeys` and `HostSshUser` are already a part of the task config persisted since #4220. * Honoring the marker is safe as of 0.21.3, the first release carrying both the marker and the task state file, so no shim version writes one without the other. A task restored without a state file has no config, and therefore no keys in the reconciled set, but it was started by an earlier shim, whose entries are unmarked and never touched. * `cleanupLocked()` commits `cleanedUp` before releasing the resources, so that the finishing task is already out of the reconciled set by the time its keys are reconciled. If the shim stops running in between, the state file, written last, still reports the task as not cleaned up, and the cleanup runs again, idempotently, after the restart. * The file is written to a temporary file in the same dir and renamed over the old one, preserving its mode and setting its owner, so that a failed write cannot leave the user without their keys. `~/.ssh` and the file itself are created if missing, as the server does when provisioning an SSH fleet, instead of failing on a missing file. The `.bak` copy is dropped: the rename is atomic and the entries the shim does not own are copied through untouched, so it protected nothing, while leaving a root-owned copy of the user's keys in `~/.ssh`. * A single reconcile on start drops the entries of the tasks that are gone and collapses the duplicates left by the earlier versions, which appended a key once per task using it. Fixes: #4174 Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The shim appends
cfg.HostSshKeysto the host user'sauthorized_keyswhen a task starts and removes every entry matching those fingerprints when it ends. That key is the user's own SSH key, shared by all of their runs unless overridden withdstack apply --ssh-identity, so on an instance that holds more than one of their jobs at a time -- a fleet withblocks-- the first job to finish revokes the SSH access of the jobs still running. The same unconditional removal also drops a key the user added by hand, if a run happened to reuse it.The per-task add/remove deltas are replaced with a reconcile of the entries the shim owns, that is, the ones carrying the
# added by dstack-shimmarker shipped in #4206. For every host user, those entries become exactly the keys of the stored tasks that are not cleaned up yet; everything else in the file -- keys added by the user, keys added by the server when provisioning an SSH fleet, comments -- is kept verbatim. Refcounting is implicit in set membership, so there is no new state to keep or persist:HostSshKeysandHostSshUserare already a part of the task config persisted since #4220.cleanupLocked()commitscleanedUpbefore releasing the resources, so that the finishing task is already out of the reconciled set by the time its keys are reconciled. If the shim stops running in between, the state file, written last, still reports the task as not cleaned up, and the cleanup runs again, idempotently, after the restart.~/.sshand the file itself are created if missing, as the server does when provisioning an SSH fleet, instead of failing on a missing file. The.bakcopy is dropped: the rename is atomic and the entries the shim does not own are copied through untouched, so it protected nothing, while leaving a root-owned copy of the user's keys in~/.ssh.Fixes: #4174