Skip to content

[shim] Don't revoke host SSH access of co-located tasks - #4226

Open
un-def wants to merge 1 commit into
masterfrom
issue_4174_shim_fix_colocated_jobs_keys_revocation
Open

[shim] Don't revoke host SSH access of co-located tasks#4226
un-def wants to merge 1 commit into
masterfrom
issue_4174_shim_fix_colocated_jobs_keys_revocation

Conversation

@un-def

@un-def un-def commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator

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

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: A finishing job revokes host SSH access for still-running co-located jobs

1 participant