Support custom run names

cli/dstack/_internal/api/repos.py

@@ -46,15 +46,15 @@ def get_local_repo_credentials(
 
     if identity_file is not None:  # must fail if key is invalid
         try:  # user provided ssh key
-            return test_remote_repo_credentials(
+            return check_remote_repo_credentials(
                 repo_data, RepoProtocol.SSH, identity_file=identity_file
             )
         except GitCommandError:
             pass
 
     if oauth_token is not None:
         try:  # user provided oauth token
-            return test_remote_repo_credentials(
+            return check_remote_repo_credentials(
                 repo_data, RepoProtocol.HTTPS, oauth_token=oauth_token
             )
         except GitCommandError:
@@ -63,7 +63,7 @@ def get_local_repo_credentials(
     identities = get_host_config(original_hostname or repo_data.repo_host_name).get("identityfile")
     if identities:  # must fail if key is invalid
         try:  # key from ssh config
-            return test_remote_repo_credentials(
+            return check_remote_repo_credentials(
                 repo_data, RepoProtocol.SSH, identity_file=identities[0]
             )
         except GitCommandError:
@@ -75,22 +75,22 @@ def get_local_repo_credentials(
         oauth_token = gh_hosts.get(repo_data.repo_host_name, {}).get("oauth_token")
         if oauth_token is not None:
             try:  # token from gh config
-                return test_remote_repo_credentials(
+                return check_remote_repo_credentials(
                     repo_data, RepoProtocol.HTTPS, oauth_token=oauth_token
                 )
             except GitCommandError:
                 pass
 
     if os.path.exists(default_ssh_key):
         try:  # default user key
-            return test_remote_repo_credentials(
+            return check_remote_repo_credentials(
                 repo_data, RepoProtocol.SSH, identity_file=default_ssh_key
             )
         except GitCommandError:
             pass
 
 
-def test_remote_repo_credentials(
+def check_remote_repo_credentials(
     repo_data: RemoteRepoData,
     protocol: RepoProtocol,
     *,

cli/dstack/_internal/backend/base/__init__.py

@@ -88,6 +88,10 @@ def list_job_heads(self, repo_id: str, run_name: Optional[str] = None) -> List[J
     def delete_job_head(self, repo_id: str, job_id: str):
         pass
 
+    @abstractmethod
+    def delete_run_jobs(self, repo_id: str, run_name: str):
+        pass
+
     @abstractmethod
     def list_run_heads(
         self,
@@ -262,8 +266,8 @@ def logging(self) -> Logging:
     def predict_instance_type(self, job: Job) -> Optional[InstanceType]:
         return base_jobs.predict_job_instance(self.compute(), job)
 
-    def create_run(self, repo_id: str) -> str:
-        return base_runs.create_run(self.storage())
+    def create_run(self, repo_id: str, run_name: Optional[str]) -> str:
+        return base_runs.create_run(self.storage(), run_name)
 
     def create_job(self, job: Job):
         base_jobs.create_job(self.storage(), job)
@@ -293,6 +297,9 @@ def list_job_heads(self, repo_id: str, run_name: Optional[str] = None) -> List[J
     def delete_job_head(self, repo_id: str, job_id: str):
         base_jobs.delete_job_head(self.storage(), repo_id, job_id)
 
+    def delete_run_jobs(self, repo_id: str, run_name: str):
+        base_jobs.delete_jobs(self.storage(), repo_id, run_name)
+
     def list_run_heads(
         self,
         repo_id: str,

cli/dstack/_internal/backend/base/jobs.py

@@ -96,6 +96,13 @@ def delete_job_head(storage: Storage, repo_id: str, job_id: str):
         storage.delete_object(job_head_key)
 
 
+def delete_jobs(storage: Storage, repo_id: str, run_name: str):
+    job_key_run_prefix = _get_jobs_filenames_prefix(repo_id, run_name)
+    jobs_keys = storage.list_objects(job_key_run_prefix)
+    for job_key in jobs_keys:
+        storage.delete_object(job_key)
+
+
 def predict_job_instance(
     compute: Compute,
     job: Job,
@@ -176,9 +183,9 @@ def stop_job(
     if new_status is not None and new_status != job.status:
         job.status = new_status
         update_job(storage, job)
-        if new_status.is_finished():
+        if new_status in [JobStatus.TERMINATED, JobStatus.ABORTED]:
             if runner is not None:
-                runners.stop_runner(compute, runner)
+                runners.terminate_runner(compute, runner)
 
 
 def update_job_submission(job: Job):

cli/dstack/_internal/backend/base/runners.py

@@ -33,7 +33,7 @@ def delete_runner(storage: Storage, runner: Runner):
     storage.delete_object(_get_runner_filename(runner.runner_id))
 
 
-def stop_runner(compute: Compute, runner: Runner):
+def terminate_runner(compute: Compute, runner: Runner):
     if runner.request_id:
         if runner.resources.spot:
             compute.cancel_spot_request(runner)

cli/dstack/_internal/backend/base/runs.py

@@ -1,4 +1,5 @@
-from typing import List
+import re
+from typing import List, Optional
 
 import yaml
 
@@ -7,29 +8,20 @@
 from dstack._internal.backend.base.storage import Storage
 from dstack._internal.core.app import AppHead
 from dstack._internal.core.artifact import ArtifactHead
+from dstack._internal.core.error import BackendValueError
 from dstack._internal.core.job import JobErrorCode, JobHead, JobStatus
 from dstack._internal.core.run import RequestStatus, RunHead, generate_remote_run_name_prefix
 
 
-def create_run(
-    storage: Storage,
-) -> str:
-    name = generate_remote_run_name_prefix()
-    run_name_index = _next_run_name_index(storage, name)
-    run_name = f"{name}-{run_name_index}"
-    return run_name
-
-
-def _next_run_name_index(storage: Storage, run_name: str) -> int:
-    count = 0
-    key = f"run_names/{run_name}.yaml"
+def create_run(storage: Storage, run_name: Optional[str]) -> str:
+    if run_name is None:
+        return _generate_random_run_name(storage)
+    _validate_run_name(run_name)
+    key = _get_run_name_filename(run_name)
     obj = storage.get_object(key)
     if obj is None:
         storage.put_object(key=key, content=yaml.dump({"count": 1}))
-        return 1
-    count = yaml.load(obj, yaml.FullLoader)["count"]
-    storage.put_object(key=key, content=yaml.dump({"count": count + 1}))
-    return count + 1
+    return run_name
 
 
 def get_run_heads(
@@ -55,6 +47,33 @@ def get_run_heads(
     return run_heads
 
 
+def _generate_random_run_name(storage: Storage):
+    name = generate_remote_run_name_prefix()
+    run_name_index = _next_run_name_index(storage, name)
+    return f"{name}-{run_name_index}"
+
+
+def _next_run_name_index(storage: Storage, run_name: str) -> int:
+    count = 0
+    key = _get_run_name_filename(run_name)
+    obj = storage.get_object(key)
+    if obj is None:
+        storage.put_object(key=key, content=yaml.dump({"count": 1}))
+        return 1
+    count = yaml.load(obj, yaml.FullLoader)["count"]
+    storage.put_object(key=key, content=yaml.dump({"count": count + 1}))
+    return count + 1
+
+
+def _validate_run_name(run_name: str):
+    if re.match(r"^[a-zA-Z0-9_-]{5,100}$", run_name) is None:
+        raise BackendValueError(
+            "Invalid run name. "
+            "Run name may include alphanumeric characters, dashes, and underscores, "
+            "and its length should be between 5 and 100 characters."
+        )
+
+
 def _create_run(
     storage: Storage,
     compute: Compute,
@@ -183,3 +202,7 @@ def _update_run(
                     jobs.update_job(storage, job)
         run.status = job_head.status
     run.job_heads.append(job_head)
+
+
+def _get_run_name_filename(run_name: str) -> str:
+    return f"run_names/{run_name}.yaml"

cli/dstack/_internal/backend/local/compute.py

@@ -3,6 +3,7 @@
 from dstack._internal.backend.base.compute import Compute, choose_instance_type
 from dstack._internal.backend.local import runners
 from dstack._internal.backend.local.config import LocalConfig
+from dstack._internal.core.error import BackendValueError
 from dstack._internal.core.instance import InstanceType, LaunchedInstanceInfo
 from dstack._internal.core.job import Job
 from dstack._internal.core.request import RequestHead
@@ -29,11 +30,14 @@ def run_instance(self, job: Job, instance_type: InstanceType) -> LaunchedInstanc
         return LaunchedInstanceInfo(request_id=pid, location=None)
 
     def restart_instance(self, job: Job):
+        if runners.get_container(job.run_name) is None:
+            raise BackendValueError("Container not found")
         pid = runners.start_runner_process(self.backend_config, job.runner_id)
         return LaunchedInstanceInfo(request_id=pid, location=None)
 
     def terminate_instance(self, runner: Runner):
         runners.stop_process(runner.request_id)
+        runners.remove_container(runner.job.run_name)
 
     def cancel_spot_request(self, runner: Runner):
-        runners.stop_process(runner.request_id)
+        pass

cli/dstack/_internal/backend/local/runners.py

@@ -7,12 +7,16 @@
 from typing import Optional
 
 import cpuinfo
+import docker.errors
 import psutil
 import requests
 import yaml
+from docker.models.containers import Container
 from psutil import NoSuchProcess
 from tqdm import tqdm
 
+import docker
+
 from dstack import version
 from dstack._internal.backend.base.config import BACKEND_CONFIG_FILENAME, RUNNER_CONFIG_FILENAME
 from dstack._internal.backend.local.config import LocalConfig
@@ -212,3 +216,17 @@ def is_running(request_id: str) -> bool:
         return True
     except NoSuchProcess:
         return False
+
+
+def get_container(container_name: str) -> Optional[Container]:
+    client = docker.from_env()
+    try:
+        return client.containers.get(container_name)
+    except docker.errors.NotFound:
+        return None
+
+
+def remove_container(container_name: str):
+    container = get_container(container_name)
+    if container is not None:
+        container.remove()

cli/dstack/_internal/cli/commands/build/__init__.py

@@ -39,13 +39,13 @@ def register(self):
             type=str,
             dest="file_name",
         )
-        add_project_argument(self._parser)
         self._parser.add_argument(
             "-y",
             "--yes",
             help="Do not ask for plan confirmation",
             action="store_true",
         )
+        add_project_argument(self._parser)
         self._parser.add_argument(
             "--profile",
             metavar="PROFILE",

cli/dstack/_internal/cli/commands/ls/__init__.py

@@ -32,7 +32,6 @@ def register(self):
             nargs="?",
             default="",
         )
-
         self._parser.add_argument(
             "-r", "--recursive", help="Show all files recursively", action="store_true"
         )

cli/dstack/_internal/cli/commands/rm/__init__.py

@@ -37,22 +37,12 @@ def _command(self, args: Namespace):
             and (args.yes or Confirm.ask(f"[red]Delete the run '{args.run_name}'?[/]"))
         ) or (args.all and (args.yes or Confirm.ask("[red]Delete all runs?[/]"))):
             hub_client = get_hub_client(project_name=args.project)
-            deleted_run = False
-            job_heads = hub_client.list_job_heads(args.run_name)
-            if job_heads:
-                finished_job_heads = []
-                for job_head in job_heads:
-                    if job_head.status.is_finished():
-                        finished_job_heads.append(job_head)
-                    elif args.run_name:
-                        console.print("The run is not finished yet. Stop the run first.")
-                        exit(1)
-                for job_head in finished_job_heads:
-                    hub_client.delete_job_head(job_head.job_id)
-                    deleted_run = True
-            if args.run_name and not deleted_run:
+            run_heads = hub_client.list_run_heads(args.run_name)
+            if len(run_heads) == 0 and args.run_name:
                 console.print(f"Cannot find the run '{args.run_name}'")
                 exit(1)
+            for run_head in run_heads:
+                hub_client.delete_run(run_head.run_name)
             console.print(f"[grey58]OK[/]")
         else:
             if not args.run_name and not args.all: