Skip to content

Auto activate uv venv with pip installed #2666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
May 22, 2025
Merged

Auto activate uv venv with pip installed #2666

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2c202b8
Auto activate uv venv with pip installed
r4victor May 20, 2025
c475f0b
Fix tests
r4victor May 20, 2025
8119da1
Merge branch 'master' into issue_2625_uv_pip
r4victor May 21, 2025
5e712a5
Bump base_image = "0.8"
r4victor May 22, 2025
e280ec7
Fix tests
r4victor May 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/dstack/_internal/core/backends/base/compute.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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}"
),
]

Expand Down
3 changes: 2 additions & 1 deletion src/dstack/_internal/core/models/configurations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
RUN_PRIOTIRY_MIN = 0
RUN_PRIOTIRY_MAX = 100
RUN_PRIORITY_DEFAULT = 0
DEFAULT_REPO_DIR = "/workflow"


class RunConfigurationType(str, Enum):
Expand Down Expand Up @@ -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 `"."` '
)
),
Expand Down
3 changes: 2 additions & 1 deletion src/dstack/_internal/core/models/runs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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 `"."`.'
)
),
Expand Down
16 changes: 15 additions & 1 deletion src/dstack/_internal/server/services/jobs/configurators/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 []
Expand All @@ -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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(nit) This means pip won't work if our base image is explicitly specified in image. I've seen some users do it, possibly because we advertise it in the docs (1, 2, and more).

Maybe we could set up the venv in the Dockerfile, so that our images already come with a working pip?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The current idea is to have .venv in the repo dir so that uv does not issue warnings about the wrong venv. And I don't want to hardcode the repo path in the image since we had plans to make it configurable.

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",
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I assume this .venv can turn up in git status, as it is not necessarily included in .gitignore. Any chance we could create it outside of the repo?

Copy link
Collaborator Author

@r4victor r4victor May 21, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I considered creating a global uv venv outside the repo. The problem with that is that having an active venv outside the repo would trigger uv warnings about the active repo not being used cause it defaults to .venv in the project directory. Not critical, but quite confusing. And since we want to support uv-first flow, I think we should optimize for this case.

As regards to .venv in git, it will never appear in git since it's always excluded by .venv/.gitignore.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Users may still be confused by a new directory in their repository, especially if they don't use uv or Python. But I guess we can live with that if there are no good alternatives

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe, let's see how it works

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())):
Expand Down
4 changes: 3 additions & 1 deletion src/dstack/_internal/server/services/jobs/configurators/extensions/cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from typing import List

from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR


class CursorDesktop:
def __init__(
Expand Down Expand Up @@ -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 ''",
]
4 changes: 3 additions & 1 deletion src/dstack/_internal/server/services/jobs/configurators/extensions/vscode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from typing import List

from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR


class VSCodeDesktop:
def __init__(
Expand Down Expand Up @@ -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 ''",
]
2 changes: 1 addition & 1 deletion src/dstack/version.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
__version__ = "0.0.0"
__is_release__ = False
base_image = "0.7"
base_image = "0.8"
2 changes: 1 addition & 1 deletion src/tests/_internal/server/background/tasks/test_process_running_jobs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
14 changes: 10 additions & 4 deletions src/tests/_internal/server/routers/test_runs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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" '
Expand All @@ -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",
Expand Down Expand Up @@ -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" '
Expand All @@ -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",
Expand Down
10 changes: 9 additions & 1 deletion src/tests/_internal/server/services/jobs/configurators/test_task.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
Loading