Add shell run property (#2542)

Closes: #2530

Author: un-def

src/dstack/_internal/core/models/configurations.py

@@ -1,6 +1,7 @@
 import re
 from collections import Counter
 from enum import Enum
+from pathlib import PurePosixPath
 from typing import Any, Dict, List, Optional, Union
 
 from pydantic import Field, ValidationError, conint, constr, root_validator, validator
@@ -210,6 +211,16 @@ class BaseRunConfiguration(CoreModel):
         Env,
         Field(description="The mapping or the list of environment variables"),
     ] = Env()
+    shell: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The shell used to run commands."
+                " Allowed values are `sh`, `bash`, or an absolute path, e.g., `/usr/bin/zsh`."
+                " Defaults to `/bin/sh` if the `image` is specified, `/bin/bash` otherwise"
+            )
+        ),
+    ] = None
     # deprecated since 0.18.31; task, service -- no effect; dev-environment -- executed right before `init`
     setup: CommandsList = []
     resources: Annotated[
@@ -244,6 +255,17 @@ def validate_user(cls, v) -> Optional[str]:
         UnixUser.parse(v)
         return v
 
+    @validator("shell")
+    def validate_shell(cls, v) -> Optional[str]:
+        if v is None:
+            return None
+        if v in ["sh", "bash"]:
+            return v
+        path = PurePosixPath(v)
+        if path.is_absolute():
+            return v
+        raise ValueError("The value must be `sh`, `bash`, or an absolute path")
+
 
 class BaseRunConfigurationWithPorts(BaseRunConfiguration):
     ports: Annotated[
@@ -261,7 +283,7 @@ def convert_ports(cls, v) -> PortMapping:
 
 
 class BaseRunConfigurationWithCommands(BaseRunConfiguration):
-    commands: Annotated[CommandsList, Field(description="The bash commands to run")] = []
+    commands: Annotated[CommandsList, Field(description="The shell commands to run")] = []
 
     @root_validator
     def check_image_or_commands_present(cls, values):
@@ -276,7 +298,7 @@ class DevEnvironmentConfigurationParams(CoreModel):
         Field(description="The IDE to run. Supported values include `vscode` and `cursor`"),
     ]
     version: Annotated[Optional[str], Field(description="The version of the IDE")] = None
-    init: Annotated[CommandsList, Field(description="The bash commands to run on startup")] = []
+    init: Annotated[CommandsList, Field(description="The shell commands to run on startup")] = []
     inactivity_duration: Annotated[
         Optional[Union[Literal["off"], int, bool, str]],
         Field(

src/dstack/_internal/server/services/jobs/configurators/base.py

@@ -1,6 +1,7 @@
 import shlex
 import sys
 from abc import ABC, abstractmethod
+from pathlib import PurePosixPath
 from typing import Dict, List, Optional, Union
 
 from cachetools import TTLCache, cached
@@ -131,16 +132,24 @@ async def _get_job_spec(
         )
         return job_spec
 
+    def _shell(self) -> str:
+        shell = self.run_spec.configuration.shell
+        if shell is not None:
+            path = PurePosixPath(shell)
+            if path.is_absolute():
+                return shell
+            return str("/bin" / path)
+        if self.run_spec.configuration.image is None:  # dstackai/base
+            return "/bin/bash"
+        return "/bin/sh"
+
     async def _commands(self) -> List[str]:
         if self.run_spec.configuration.entrypoint is not None:  # docker-like format
             entrypoint = shlex.split(self.run_spec.configuration.entrypoint)
             commands = self.run_spec.configuration.commands
-        elif self.run_spec.configuration.image is None:  # dstackai/base
-            entrypoint = ["/bin/bash", "-i", "-c"]
-            commands = [_join_shell_commands(self._shell_commands())]
-        elif self._shell_commands():  # custom docker image with shell commands
-            entrypoint = ["/bin/sh", "-i", "-c"]
-            commands = [_join_shell_commands(self._shell_commands())]
+        elif shell_commands := self._shell_commands():
+            entrypoint = [self._shell(), "-i", "-c"]
+            commands = [_join_shell_commands(shell_commands)]
         else:  # custom docker image without commands
             image_config = await self._get_image_config()
             entrypoint = image_config.entrypoint or []

src/dstack/api/server/_runs.py

@@ -118,6 +118,8 @@ def _get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
         profile_excludes.add("tags")
     if isinstance(configuration, ServiceConfiguration) and not configuration.rate_limits:
         configuration_excludes["rate_limits"] = True
+    if configuration.shell is None:
+        configuration_excludes["shell"] = True
 
     if configuration_excludes:
         spec_excludes["configuration"] = configuration_excludes

src/tests/_internal/core/models/test_configurations.py

@@ -52,6 +52,26 @@ def test_conf(replicas: Any, scaling: Optional[Any] = None):
                 )
             )
 
+    @pytest.mark.parametrize("shell", [None, "sh", "bash", "/usr/bin/zsh"])
+    def test_shell_valid(self, shell: Optional[str]):
+        conf = {
+            "type": "task",
+            "shell": shell,
+            "commands": ["sleep inf"],
+        }
+        assert parse_run_configuration(conf).shell == shell
+
+    def test_shell_invalid(self):
+        conf = {
+            "type": "task",
+            "shell": "zsh",
+            "commands": ["sleep inf"],
+        }
+        with pytest.raises(
+            ConfigurationError, match="The value must be `sh`, `bash`, or an absolute path"
+        ):
+            parse_run_configuration(conf)
+
 
 def test_registry_auth_hashable():
     """

src/tests/_internal/server/routers/test_runs.py

@@ -93,6 +93,7 @@ def get_dev_env_run_plan_dict(
                 "version": None,
                 "image": None,
                 "user": None,
+                "shell": None,
                 "privileged": privileged,
                 "init": [],
                 "ports": [],
@@ -247,6 +248,7 @@ def get_dev_env_run_dict(
                 "version": None,
                 "image": None,
                 "user": None,
+                "shell": None,
                 "privileged": privileged,
                 "init": [],
                 "ports": [],

src/tests/_internal/server/services/jobs/configurators/test_task.py

@@ -1,17 +1,19 @@
+from typing import Optional
 from unittest.mock import patch
 
 import pytest
 
 from dstack._internal.core.models.configurations import TaskConfiguration
 from dstack._internal.core.models.runs import JobSSHKey
+from dstack._internal.server.services.docker import ImageConfig
 from dstack._internal.server.services.jobs.configurators.task import TaskJobConfigurator
 from dstack._internal.server.testing.common import get_run_spec
 
 
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("image_config_mock")
-class TestTaskJobConfigurator:
-    async def test_ssh_key_single_node(self):
+class TestSSHKey:
+    async def test_single_node(self):
         configuration = TaskConfiguration(nodes=1, image="debian")
         run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration)
         configurator = TaskJobConfigurator(run_spec)
@@ -21,7 +23,7 @@ async def test_ssh_key_single_node(self):
         assert len(job_specs) == 1
         assert job_specs[0].ssh_key is None
 
-    async def test_ssh_key_multi_node(self):
+    async def test_multi_node(self):
         configuration = TaskConfiguration(nodes=2, image="debian")
         run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration)
         configurator = TaskJobConfigurator(run_spec)
@@ -33,3 +35,74 @@ async def test_ssh_key_multi_node(self):
         assert len(job_specs) == 2
         assert job_specs[0].ssh_key == JobSSHKey(private="private1", public="public1")
         assert job_specs[1].ssh_key == JobSSHKey(private="private1", public="public1")
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("image_config_mock")
+class TestCommands:
+    @pytest.mark.parametrize(
+        ["commands", "expected_commands"],
+        [
+            pytest.param([], ["/entrypoint.sh", "-v"], id="no-commands"),
+            pytest.param(["-x", "-u"], ["/entrypoint.sh", "-v", "-x", "-u"], id="with-commands"),
+        ],
+    )
+    async def test_with_entrypoint(self, commands: list[str], expected_commands: list[str]):
+        configuration = TaskConfiguration(
+            image="debian",
+            entrypoint="/entrypoint.sh -v",
+            commands=commands,
+        )
+        run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration)
+        configurator = TaskJobConfigurator(run_spec)
+
+        job_specs = await configurator.get_job_specs(replica_num=0)
+
+        assert job_specs[0].commands == expected_commands
+
+    @pytest.mark.parametrize(
+        ["shell", "expected_shell"],
+        [
+            pytest.param(None, "/bin/sh", id="default-shell"),
+            pytest.param("sh", "/bin/sh", id="sh"),
+            pytest.param("bash", "/bin/bash", id="bash"),
+            pytest.param("/usr/bin/zsh", "/usr/bin/zsh", id="custom-shell"),
+        ],
+    )
+    async def test_with_commands_and_image(self, shell: Optional[str], expected_shell: str):
+        configuration = TaskConfiguration(image="debian", commands=["sleep inf"], shell=shell)
+        run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration)
+        configurator = TaskJobConfigurator(run_spec)
+
+        job_specs = await configurator.get_job_specs(replica_num=0)
+
+        assert job_specs[0].commands == [expected_shell, "-i", "-c", "sleep inf"]
+
+    @pytest.mark.parametrize(
+        ["shell", "expected_shell"],
+        [
+            pytest.param(None, "/bin/bash", id="default-shell"),
+            pytest.param("sh", "/bin/sh", id="sh"),
+            pytest.param("bash", "/bin/bash", id="bash"),
+            pytest.param("/usr/bin/zsh", "/usr/bin/zsh", id="custom-shell"),
+        ],
+    )
+    async def test_with_commands_no_image(self, shell: Optional[str], expected_shell: str):
+        configuration = TaskConfiguration(commands=["sleep inf"], shell=shell)
+        run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration)
+        configurator = TaskJobConfigurator(run_spec)
+
+        job_specs = await configurator.get_job_specs(replica_num=0)
+
+        assert job_specs[0].commands == [expected_shell, "-i", "-c", "sleep inf"]
+
+    async def test_no_commands(self, image_config_mock: ImageConfig):
+        image_config_mock.entrypoint = ["/entrypoint.sh"]
+        image_config_mock.cmd = ["-f", "-x"]
+        configuration = TaskConfiguration(image="debian")
+        run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration)
+        configurator = TaskJobConfigurator(run_spec)
+
+        job_specs = await configurator.get_job_specs(replica_num=0)
+
+        assert job_specs[0].commands == ["/entrypoint.sh", "-f", "-x"]