Auto activate uv venv with pip installed (#2666)

* Auto activate uv venv with pip installed

* Fix tests

* Bump base_image = "0.8"

* Fix tests

Author: r4victor

src/dstack/_internal/core/backends/base/compute.py

@@ -19,6 +19,7 @@
     DSTACK_RUNNER_SSH_PORT,
     DSTACK_SHIM_HTTP_PORT,
 )
+from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR
 from dstack._internal.core.models.gateways import (
     GatewayComputeConfiguration,
     GatewayProvisioningData,
@@ -754,7 +755,7 @@ def get_docker_commands(
             f" --ssh-port {DSTACK_RUNNER_SSH_PORT}"
             " --temp-dir /tmp/runner"
             " --home-dir /root"
-            " --working-dir /workflow"
+            f" --working-dir {DEFAULT_REPO_DIR}"
         ),
     ]
 

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

@@ -26,6 +26,7 @@
 RUN_PRIOTIRY_MIN = 0
 RUN_PRIOTIRY_MAX = 100
 RUN_PRIORITY_DEFAULT = 0
+DEFAULT_REPO_DIR = "/workflow"
 
 
 class RunConfigurationType(str, Enum):
@@ -181,7 +182,7 @@ class BaseRunConfiguration(CoreModel):
         Field(
             description=(
                 "The path to the working directory inside the container."
-                " It's specified relative to the repository directory (`/workflow`) and should be inside it."
+                f" It's specified relative to the repository directory (`{DEFAULT_REPO_DIR}`) and should be inside it."
                 ' Defaults to `"."` '
             )
         ),

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

@@ -8,6 +8,7 @@
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import ApplyAction, CoreModel, NetworkMode, RegistryAuth
 from dstack._internal.core.models.configurations import (
+    DEFAULT_REPO_DIR,
     AnyRunConfiguration,
     RunConfiguration,
 )
@@ -338,7 +339,7 @@ class RunSpec(CoreModel):
         Field(
             description=(
                 "The path to the working directory inside the container."
-                " It's specified relative to the repository directory (`/workflow`) and should be inside it."
+                f" It's specified relative to the repository directory (`{DEFAULT_REPO_DIR}`) and should be inside it."
                 ' Defaults to `"."`.'
             )
         ),

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

@@ -10,6 +10,7 @@
 from dstack._internal.core.errors import DockerRegistryError, ServerClientError
 from dstack._internal.core.models.common import RegistryAuth
 from dstack._internal.core.models.configurations import (
+    DEFAULT_REPO_DIR,
     PortMapping,
     PythonVersion,
     RunConfigurationType,
@@ -149,7 +150,8 @@ async def _commands(self) -> List[str]:
             commands = self.run_spec.configuration.commands
         elif shell_commands := self._shell_commands():
             entrypoint = [self._shell(), "-i", "-c"]
-            commands = [_join_shell_commands(shell_commands)]
+            dstack_image_commands = self._dstack_image_commands()
+            commands = [_join_shell_commands(dstack_image_commands + shell_commands)]
         else:  # custom docker image without commands
             image_config = await self._get_image_config()
             entrypoint = image_config.entrypoint or []
@@ -164,6 +166,18 @@ async def _commands(self) -> List[str]:
 
         return result
 
+    def _dstack_image_commands(self) -> List[str]:
+        if (
+            self.run_spec.configuration.image is not None
+            or self.run_spec.configuration.entrypoint is not None
+        ):
+            return []
+        return [
+            f"uv venv --prompt workflow --seed {DEFAULT_REPO_DIR}/.venv > /dev/null 2>&1",
+            f"echo 'source {DEFAULT_REPO_DIR}/.venv/bin/activate' >> ~/.bashrc",
+            f"source {DEFAULT_REPO_DIR}/.venv/bin/activate",
+        ]
+
     def _app_specs(self) -> List[AppSpec]:
         specs = []
         for i, pm in enumerate(filter_reserved_ports(self._ports())):

src/dstack/_internal/server/services/jobs/configurators/extensions/cursor.py

@@ -1,5 +1,7 @@
 from typing import List
 
+from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR
+
 
 class CursorDesktop:
     def __init__(
@@ -37,6 +39,6 @@ def get_print_readme_commands(self) -> List[str]:
         return [
             "echo To open in Cursor, use link below:",
             "echo ''",
-            f"echo '  cursor://vscode-remote/ssh-remote+{self.run_name}/workflow'",  # TODO use $REPO_DIR
+            f"echo '  cursor://vscode-remote/ssh-remote+{self.run_name}{DEFAULT_REPO_DIR}'",  # TODO use $REPO_DIR
             "echo ''",
         ]

src/dstack/_internal/server/services/jobs/configurators/extensions/vscode.py

@@ -1,5 +1,7 @@
 from typing import List
 
+from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR
+
 
 class VSCodeDesktop:
     def __init__(
@@ -37,6 +39,6 @@ def get_print_readme_commands(self) -> List[str]:
         return [
             "echo To open in VS Code Desktop, use link below:",
             "echo ''",
-            f"echo '  vscode://vscode-remote/ssh-remote+{self.run_name}/workflow'",  # TODO use $REPO_DIR
+            f"echo '  vscode://vscode-remote/ssh-remote+{self.run_name}{DEFAULT_REPO_DIR}'",  # TODO use $REPO_DIR
             "echo ''",
         ]

src/dstack/version.py

@@ -1,3 +1,3 @@
 __version__ = "0.0.0"
 __is_release__ = False
-base_image = "0.7"
+base_image = "0.8"

src/tests/_internal/server/background/tasks/test_process_running_jobs.py

@@ -330,7 +330,7 @@ async def test_provisioning_shim_with_volumes(
             name="test-run-0-0",
             registry_username="",
             registry_password="",
-            image_name="dstackai/base:py3.13-0.7-cuda-12.1",
+            image_name="dstackai/base:py3.13-0.8-cuda-12.1",
             container_user="root",
             privileged=privileged,
             gpu=None,

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

@@ -166,7 +166,10 @@ def get_dev_env_run_plan_dict(
                         "/bin/bash",
                         "-i",
                         "-c",
-                        "(echo pip install ipykernel... && "
+                        "uv venv --prompt workflow --seed /workflow/.venv > /dev/null 2>&1"
+                        " && echo 'source /workflow/.venv/bin/activate' >> ~/.bashrc"
+                        " && source /workflow/.venv/bin/activate"
+                        " && (echo pip install ipykernel... && "
                         "pip install -q --no-cache-dir "
                         'ipykernel 2> /dev/null) || echo "no '
                         'pip, ipykernel was not installed" '
@@ -181,7 +184,7 @@ def get_dev_env_run_plan_dict(
                     ],
                     "env": {},
                     "home_dir": "/root",
-                    "image_name": "dstackai/base:py3.13-0.7-cuda-12.1",
+                    "image_name": "dstackai/base:py3.13-0.8-cuda-12.1",
                     "user": None,
                     "privileged": privileged,
                     "job_name": f"{run_name}-0-0",
@@ -322,7 +325,10 @@ def get_dev_env_run_dict(
                         "/bin/bash",
                         "-i",
                         "-c",
-                        "(echo pip install ipykernel... && "
+                        "uv venv --prompt workflow --seed /workflow/.venv > /dev/null 2>&1"
+                        " && echo 'source /workflow/.venv/bin/activate' >> ~/.bashrc"
+                        " && source /workflow/.venv/bin/activate"
+                        " && (echo pip install ipykernel... && "
                         "pip install -q --no-cache-dir "
                         'ipykernel 2> /dev/null) || echo "no '
                         'pip, ipykernel was not installed" '
@@ -337,7 +343,7 @@ def get_dev_env_run_dict(
                     ],
                     "env": {},
                     "home_dir": "/root",
-                    "image_name": "dstackai/base:py3.13-0.7-cuda-12.1",
+                    "image_name": "dstackai/base:py3.13-0.8-cuda-12.1",
                     "user": None,
                     "privileged": privileged,
                     "job_name": f"{run_name}-0-0",

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

@@ -94,7 +94,15 @@ async def test_with_commands_no_image(self, shell: Optional[str], expected_shell
 
         job_specs = await configurator.get_job_specs(replica_num=0)
 
-        assert job_specs[0].commands == [expected_shell, "-i", "-c", "sleep inf"]
+        assert job_specs[0].commands == [
+            expected_shell,
+            "-i",
+            "-c",
+            "uv venv --prompt workflow --seed /workflow/.venv > /dev/null 2>&1"
+            " && echo 'source /workflow/.venv/bin/activate' >> ~/.bashrc"
+            " && source /workflow/.venv/bin/activate"
+            " && sleep inf",
+        ]
 
     async def test_no_commands(self, image_config_mock: ImageConfig):
         image_config_mock.entrypoint = ["/entrypoint.sh"]