diff --git a/src/dstack/_internal/cli/commands/preset.py b/src/dstack/_internal/cli/commands/preset.py index c006a4121..d4ed0128f 100644 --- a/src/dstack/_internal/cli/commands/preset.py +++ b/src/dstack/_internal/cli/commands/preset.py @@ -33,16 +33,19 @@ from dstack._internal.cli.services.presets.export import export_preset from dstack._internal.cli.services.presets.output import get_presets_table, print_presets from dstack._internal.cli.services.presets.session import ( + get_presets_dir, list_preset_sessions, load_preset_session, load_resumable_session, resolve_session_ref, + session_process_alive, ) from dstack._internal.cli.services.presets.store import ( PresetStore, load_preset_configuration, parse_preset_configuration, ) +from dstack._internal.cli.services.presets.workspace import remove_agent_workspace from dstack._internal.cli.services.profile import ( register_profile_args, ) @@ -367,10 +370,12 @@ def _delete(self, args: argparse.Namespace) -> None: preset_ids = [args.preset] description = f"preset [code]{args.preset}[/]" else: - if preset is None: - raise CLIError(f"Preset {args.preset!r} does not exist") - preset_ids = [preset.id] - description = f"preset [code]{preset.id}[/] for [code]{preset.base}[/]" + if preset is not None: + preset_ids = [preset.id] + description = f"preset [code]{preset.id}[/] for [code]{preset.base}[/]" + else: + preset_ids = [_creation_id(args.preset)] + description = f"preset [code]{preset_ids[0]}[/]" else: target = args.base or args.repo presets = _filter_presets(store.list(), base=args.base, repo=args.repo) @@ -380,14 +385,52 @@ def _delete(self, args: argparse.Namespace) -> None: preset_ids = [preset.id for preset in presets] count = f"{len(presets)} preset{'s' if len(presets) != 1 else ''}" description = f"{count} for [code]{target}[/]" + # Checked before the prompt, so a creation that cannot be deleted says + # so instead of asking first. + for preset_id in preset_ids: + _check_creation_not_in_use(preset_id) if not args.yes and not confirm_ask(f"Delete {description}?"): console.print("\nExiting...") return for preset_id in preset_ids: + # The workspace alias is a symlink outside the preset directory, so + # it has to go before the directory does. + with suppress(CLIError): + remove_agent_workspace(load_preset_session(preset_id)) store.delete(preset_id) console.print(f"Deleted {description}") +def _creation_id(ref: str) -> str: + """The creation a reference names. A directory whose state cannot be read + resolves too, since that is a state the user needs to delete.""" + preset_id = resolve_session_ref(ref) + if not (get_presets_dir() / preset_id).is_dir(): + raise CLIError(f"Preset {ref!r} does not exist") + return preset_id + + +def _check_creation_not_in_use(preset_id: str) -> None: + """Raises while the creation's process is alive, and while its runs may still + be up, because deleting the session drops `runs.jsonl`, the only record of + which runs it started.""" + try: + session = load_preset_session(preset_id) + except CLIError: + return + state = session.read_state() + if state is None or state.status != "running": + return + stop = f"Stop it with `dstack preset stop {preset_id}`" + if session_process_alive(state): + raise CLIError(f"Failed to delete preset {preset_id}. Preset creation is in use. {stop}") + if session.runs_path.is_file() and session.runs_path.stat().st_size > 0: + raise CLIError( + f"Failed to delete preset {preset_id}." + f" Preset creation was interrupted and its runs may still be active. {stop}" + ) + + def _get_unfinished_preset(ref: str) -> Optional[Preset]: """The preset as its creation session knows it, for any state but verified.""" try: diff --git a/src/dstack/_internal/cli/services/presets/store.py b/src/dstack/_internal/cli/services/presets/store.py index f069f16a6..25e01e004 100644 --- a/src/dstack/_internal/cli/services/presets/store.py +++ b/src/dstack/_internal/cli/services/presets/store.py @@ -120,12 +120,13 @@ def release_name(self, name: str) -> VerifiedPreset | None: return detached def delete(self, preset_id: str) -> bool: - # Deletes by path so a corrupt preset file is still removable. + # Deletes by path so a corrupt preset file, or a creation that never + # saved one, is still removable. _validate_preset_id(preset_id) if not self.root.exists(): return False directory = self.root / preset_id - if not (directory / "preset.yml").is_file(): + if directory.parent != self.root or directory.is_symlink() or not directory.is_dir(): return False shutil.rmtree(directory) return True @@ -241,7 +242,9 @@ def _relative_to_preset_dir(local_path: str, directory: Path) -> str: def _validate_preset_id(preset_id: str) -> None: - if not preset_id or preset_id.startswith(".") or any(char in preset_id for char in "/\\"): + # `:` is rejected so a Windows drive-relative reference (`D:x`) cannot name a + # directory outside the store, or another one inside it. + if not preset_id or preset_id.startswith(".") or any(char in preset_id for char in "/\\:"): raise CLIError(f"Invalid preset ID: {preset_id!r}") diff --git a/src/dstack/_internal/cli/services/presets/workspace.py b/src/dstack/_internal/cli/services/presets/workspace.py index acb336731..c10613175 100644 --- a/src/dstack/_internal/cli/services/presets/workspace.py +++ b/src/dstack/_internal/cli/services/presets/workspace.py @@ -126,7 +126,9 @@ def remove_agent_workspace(session: PresetSession) -> None: if alias and alias != workspace and Path(alias).is_symlink(): with suppress(OSError): os.unlink(alias) - if workspace: + # The recorded path is trusted only when it is the one this session created, + # so a hand-edited or copied `session.json` cannot point the delete elsewhere. + if workspace and Path(workspace) == session.path / "workspace": shutil.rmtree(workspace, ignore_errors=True) diff --git a/src/tests/_internal/cli/commands/test_preset.py b/src/tests/_internal/cli/commands/test_preset.py index 130b6f64b..26cda776b 100644 --- a/src/tests/_internal/cli/commands/test_preset.py +++ b/src/tests/_internal/cli/commands/test_preset.py @@ -8,8 +8,10 @@ import pytest from dstack._internal.cli.commands.preset import _check_stdin_configuration_confirmable +from dstack._internal.cli.models.preset_agent import PresetSessionWorkspace from dstack._internal.cli.services.presets import output as presets_utils from dstack._internal.cli.services.presets.store import PresetStore +from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.errors import CLIError from dstack._internal.utils.common import render_datetime_as_api from tests._internal.cli.common import ( @@ -230,6 +232,102 @@ def test_deletes_preset_without_api_client(self, tmp_path): assert PresetStore(tmp_path / ".dstack" / "presets").list() == [] + def test_deletes_an_interrupted_creation_that_never_saved_a_preset(self, tmp_path, capsys): + session_dir = self._session(tmp_path, status="interrupted") + + assert run_dstack_cli(["preset", "delete", "smoke", "-y"], home_dir=tmp_path) == 0 + + assert not session_dir.exists() + assert "Deleted preset ab12cd34" in capsys.readouterr().out + + def test_refuses_to_delete_a_running_creation(self, tmp_path, capsys): + session_dir = self._session(tmp_path, status="running") + + with patch( + "dstack._internal.cli.commands.preset.session_process_alive", return_value=True + ): + exit_code = run_dstack_cli(["preset", "delete", "ab12cd34", "-y"], home_dir=tmp_path) + + assert exit_code != 0 + assert "Preset creation is in use" in capsys.readouterr().out + assert session_dir.exists() + + def test_refuses_to_delete_an_orphaned_creation_that_recorded_runs(self, tmp_path, capsys): + # The agent died without stopping its runs: `runs.jsonl` is the only + # record of them, so deleting it would strand them. + session_dir = self._session(tmp_path, status="running") + session_dir.joinpath("runs.jsonl").write_text('{"name": "qwen-1"}\n') + + exit_code = run_dstack_cli(["preset", "delete", "ab12cd34", "-y"], home_dir=tmp_path) + + assert exit_code != 0 + assert "Preset creation was interrupted" in capsys.readouterr().out + assert session_dir.exists() + + def test_deletes_an_orphaned_creation_without_runs(self, tmp_path): + session_dir = self._session(tmp_path, status="running") + + assert run_dstack_cli(["preset", "delete", "ab12cd34", "-y"], home_dir=tmp_path) == 0 + + assert not session_dir.exists() + + def test_deletes_a_creation_whose_state_cannot_be_read(self, tmp_path): + session_dir = tmp_path / ".dstack" / "presets" / "ab12cd34" + session_dir.mkdir(parents=True) + session_dir.joinpath("session.json").write_text("{") + + assert run_dstack_cli(["preset", "delete", "ab12cd34", "-y"], home_dir=tmp_path) == 0 + + assert not session_dir.exists() + + @pytest.mark.skipif(IS_WINDOWS, reason="workspace alias symlinks are POSIX-only") + def test_delete_removes_the_workspace_alias_but_nothing_outside_the_session(self, tmp_path): + session_dir = self._session(tmp_path, status="interrupted") + workspace = session_dir / "workspace" + workspace.mkdir() + alias = tmp_path / "dpe-alias" + alias.symlink_to(workspace, target_is_directory=True) + outside = tmp_path / "outside" + outside.mkdir() + (outside / "keep.txt").write_text("keep") + session_dir.joinpath("session.json").write_text( + json.dumps( + get_session_state( + status="interrupted", + run=get_session_run( + workspace=PresetSessionWorkspace(path=str(outside), alias=str(alias)) + ), + ).model_dump(mode="json") + ) + ) + + assert run_dstack_cli(["preset", "delete", "ab12cd34", "-y"], home_dir=tmp_path) == 0 + + assert not session_dir.exists() + # A dangling symlink is gone from exists() but still on disk. + assert not alias.is_symlink() + # The recorded path pointed outside the session, so it must survive. + assert (outside / "keep.txt").read_text() == "keep" + + def _session(self, tmp_path, *, status: str): + session_dir = tmp_path / ".dstack" / "presets" / "ab12cd34" + session_dir.mkdir(parents=True) + session_dir.joinpath("session.json").write_text( + json.dumps( + get_session_state( + status=status, + name="smoke", + run=get_session_run( + workspace=PresetSessionWorkspace( + path=str(session_dir / "workspace"), + alias=str(session_dir / "workspace"), + ) + ), + ).model_dump(mode="json") + ) + ) + return session_dir + def test_gets_complete_preset_as_json_without_api_client(self, tmp_path, capsys): preset = get_preset() PresetStore(tmp_path / ".dstack" / "presets").save(preset) diff --git a/src/tests/_internal/cli/services/presets/test_store.py b/src/tests/_internal/cli/services/presets/test_store.py index 0fababf9b..0ed4b2cdf 100644 --- a/src/tests/_internal/cli/services/presets/test_store.py +++ b/src/tests/_internal/cli/services/presets/test_store.py @@ -7,7 +7,8 @@ from dstack._internal.cli.services.presets import store as store_module from dstack._internal.cli.services.presets.store import PresetStore -from dstack._internal.core.errors import ConfigurationError +from dstack._internal.compat import IS_WINDOWS +from dstack._internal.core.errors import CLIError, ConfigurationError from dstack._internal.core.models.envs import EnvSentinel from dstack._internal.core.models.files import FilePathMapping from dstack._internal.core.models.presets import PresetConfiguration @@ -45,16 +46,40 @@ def test_saving_same_id_overwrites_existing_preset(self, tmp_path: Path): assert store.get(preset.id) == updated - def test_ignores_directories_without_a_preset_file(self, tmp_path: Path): + def test_ignores_directories_without_a_preset_file_but_keeps_them_deletable( + self, tmp_path: Path + ): root = tmp_path / "presets" store = PresetStore(root) - # A creation-session directory, or anything else that is not a preset. + # A creation session that never saved a preset. (root / "ab12cd34").mkdir(parents=True) (root / "ab12cd34" / "session.json").write_text("{}") assert store.list() == [] + assert store.delete("ab12cd34") is True + assert not (root / "ab12cd34").exists() assert store.delete("ab12cd34") is False + def test_refuses_an_id_that_could_name_a_directory_outside_the_store(self, tmp_path: Path): + store = PresetStore(tmp_path / "presets") + + # `D:x` is a drive-relative path on Windows, and `..` escapes anywhere. + for preset_id in ["D:x", "..", "../presets", "a/b"]: + with pytest.raises(CLIError, match="Invalid preset ID"): + store.delete(preset_id) + + @pytest.mark.skipif(IS_WINDOWS, reason="symlink creation requires privileges on Windows") + def test_refuses_to_delete_through_a_symlinked_directory(self, tmp_path: Path): + root = tmp_path / "presets" + root.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + (outside / "keep.txt").write_text("keep") + (root / "ab12cd34").symlink_to(outside, target_is_directory=True) + + assert PresetStore(root).delete("ab12cd34") is False + assert (outside / "keep.txt").read_text() == "keep" + def test_skips_invalid_preset_on_list_but_keeps_it_deletable(self, tmp_path: Path, capsys): store = PresetStore(tmp_path / "presets") valid = get_preset().model_copy(update={"id": "01234567"})