Skip to content
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.

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

feat: support uv #1990

Merged
merged 2 commits into from
Mar 7, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
20 changes: 15 additions & 5 deletions e2e/language/python_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,30 +16,29 @@ package language

import (
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"

"github.com/tensorchord/envd/e2e"
)

var _ = Describe("python", Ordered, func() {
testcase := "e2e"
It("Should build packages successfully", func() {
exampleName := "python/packages"
testcase := "e2e"
e := e2e.NewExample(e2e.BuildContextDirWithName(exampleName), testcase)
e.BuildImage(true)()
e.RunContainer()()
e.DestroyContainer()()
})
It("Should build requirements successfully", func() {
exampleName := "python/requirements"
testcase := "e2e"
e := e2e.NewExample(e2e.BuildContextDirWithName(exampleName), testcase)
e.BuildImage(true)()
e.RunContainer()()
e.DestroyContainer()()
})
It("Should build hybrid successfully", func() {
exampleName := "python/hybrid"
testcase := "e2e"
e := e2e.NewExample(e2e.BuildContextDirWithName(exampleName), testcase)
e.BuildImage(true)()
e.RunContainer()()
Expand All @@ -48,7 +47,6 @@ var _ = Describe("python", Ordered, func() {

It("Should build conda with channel successfully", func() {
exampleName := "python/conda"
testcase := "e2e"
e := e2e.NewExample(e2e.BuildContextDirWithName(exampleName), testcase)
e.BuildImage(true)()
e.RunContainer()()
Expand All @@ -57,10 +55,22 @@ var _ = Describe("python", Ordered, func() {

It("Should build conda with separate channel setting successfully", func() {
exampleName := "python/conda_channel"
testcase := "e2e"
e := e2e.NewExample(e2e.BuildContextDirWithName(exampleName), testcase)
e.BuildImage(true)()
e.RunContainer()()
e.DestroyContainer()()
})

Describe("Should build uv with Python successfully", func() {
exampleName := "python/uv"
e := e2e.NewExample(e2e.BuildContextDirWithName(exampleName), testcase)
BeforeAll(e.BuildImage(true))
BeforeEach(e.RunContainer())
It("Should have Python installed", func() {
res, err := e.ExecRuntimeCommand("uv-python")
Expect(err).To(BeNil())
Expect(res).To(ContainSubstring("python"))
})
AfterEach(e.DestroyContainer())
})
})
12 changes: 12 additions & 0 deletions e2e/language/testdata/python/uv/build.envd
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
# syntax=v1


def build():
base(dev=True)
install.uv()
shell("fish")
runtime.command(
commands={
"uv-python": "uv python list --only-installed",
}
)
15 changes: 15 additions & 0 deletions envd/api/v1/install.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,21 @@ def conda(use_mamba: bool = False):
"""


def uv(python_version: str = "3.11"):
"""Install UV (an extremely fast Python package and project manager).

`uv` is much faster than `conda`. Choose this one instead of `conda` if you don't
need any machine learning packages.

This doesn't support installing Python packages through `install.python_packages`
because that part should be managed by `uv`. You can run `uv sync` in the `envd`
environment to install all the dependencies.

Args:
python_version (str): install this Python version through UV
"""


def r_lang():
"""Install R Lang."""

Expand Down
1 change: 1 addition & 0 deletions pkg/lang/frontend/starlark/v1/install/const.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ const (
// language
rulePython = "install.python"
ruleConda = "install.conda"
ruleUV = "install.uv"
ruleRLang = "install.r_lang"
ruleJulia = "install.julia"

Expand Down
15 changes: 15 additions & 0 deletions pkg/lang/frontend/starlark/v1/install/install.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ var Module = &starlarkstruct.Module{
// language
"python": starlark.NewBuiltin(rulePython, ruleFuncPython),
"conda": starlark.NewBuiltin(ruleConda, ruleFuncConda),
"uv": starlark.NewBuiltin(ruleUV, ruleFuncUV),
"r_lang": starlark.NewBuiltin(ruleRLang, ruleFuncRLang),
"julia": starlark.NewBuiltin(ruleJulia, ruleFuncJulia),
// packages
Expand Down Expand Up @@ -77,6 +78,20 @@ func ruleFuncConda(thread *starlark.Thread, _ *starlark.Builtin,
return starlark.None, nil
}

func ruleFuncUV(thread *starlark.Thread, _ *starlark.Builtin,
args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
logger.Debugf("rule `%s` is invoked", ruleUV)

pythonVersion := ir.PythonVersionDefault

if err := starlark.UnpackArgs(ruleUV, args, kwargs, "python_version?", &pythonVersion); err != nil {
return nil, err
}

ir.UV(pythonVersion)
return starlark.None, nil
}

func ruleFuncRLang(thread *starlark.Thread, _ *starlark.Builtin,
args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
logger.Debugf("rule `%s` is invoked", ruleRLang)
Expand Down
4 changes: 4 additions & 0 deletions pkg/lang/ir/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,10 @@ type CondaConfig struct {
UseMicroMamba bool
}

type UVConfig struct {
PythonVersion string
}

type GitConfig struct {
Name string
Email string
Expand Down
4 changes: 4 additions & 0 deletions pkg/lang/ir/v1/compile.go
Original file line number Diff line number Diff line change
Expand Up @@ -388,6 +388,10 @@ func (g *generalGraph) CompileLLB(uid, gid int) (llb.State, error) {
return llb.State{}, errors.Wrap(err, "failed to compile shell")
}
prompt := g.compilePrompt(shell)
if g.UVConfig != nil {
// re-install uv Python for dev user
prompt = g.compileUVPython(prompt)
}
entrypoint, err := g.compileEntrypoint(prompt)
if err != nil {
return llb.State{}, errors.Wrap(err, "failed to compile entrypoint")
Expand Down
8 changes: 8 additions & 0 deletions pkg/lang/ir/v1/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,14 @@ func Conda(mamba bool) {
}
}

func UV(pythonVersion string) {
g := DefaultGraph.(*generalGraph)

g.UVConfig = &ir.UVConfig{
PythonVersion: pythonVersion,
}
}

func RLang() {
g := DefaultGraph.(*generalGraph)

Expand Down
11 changes: 4 additions & 7 deletions pkg/lang/ir/v1/system.go
Original file line number Diff line number Diff line change
Expand Up @@ -61,10 +61,7 @@ func (g generalGraph) compileRun(root llb.State) llb.State {

cmdStr := fmt.Sprintf("bash -c '%s'", sb.String())
logrus.WithField("command", cmdStr).Debug("compile run command")
// Mount the build context into the build process.
// TODO(gaocegege): Maybe we should make it readonly,
// but these cases then cannot be supported:
// run(commands=["git clone xx.git"])
// mount host here is read-only
run := root.Dir(workingDir).Run(llb.Shlex(cmdStr))
if execGroup.MountHost {
run.AddMount(workingDir, llb.Local(flag.FlagBuildContext))
Expand Down Expand Up @@ -207,6 +204,9 @@ func (g *generalGraph) compileLanguagePackages(root llb.State) llb.State {
channel := g.compileCondaChannel(pack)
pack = g.compileCondaPackages(channel)
}
if g.UVConfig != nil {
pack = g.compileUV(pack)
}

for _, language := range g.Languages {
switch language.Name {
Expand Down Expand Up @@ -271,9 +271,6 @@ func (g *generalGraph) compileBaseImage() (llb.State, error) {
if err != nil {
return llb.State{}, errors.Wrapf(err, "failed to get the image config, check if the image(%s) exists", g.Image)
}
if err != nil {
return llb.State{}, errors.Wrap(err, "failed to get the image metadata")
}

// Set the environment variables to RuntimeEnviron to keep it in the resulting image.
logger.Logger.Debugf("inherit envs from base image: %s", config.Env)
Expand Down
1 change: 1 addition & 0 deletions pkg/lang/ir/v1/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ type generalGraph struct {
*ir.JupyterConfig
*ir.GitConfig
*ir.CondaConfig
*ir.UVConfig
*ir.RStudioServerConfig

Writer compileui.Writer `json:"-"`
Expand Down
1 change: 0 additions & 1 deletion pkg/lang/ir/v1/util.go
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,6 @@ func GetDefaultGraphHash() string {
// GetCUDAImage finds the correct CUDA base image
// refer to https://hub.docker.com/r/nvidia/cuda/tags
func GetCUDAImage(image string, cuda *string, cudnn string, dev bool) string {
// TODO: support CUDA 10
target := "runtime"
if dev {
target = "devel"
Expand Down
53 changes: 53 additions & 0 deletions pkg/lang/ir/v1/uv.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
// Copyright 2025 The envd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package v1

import "github.com/moby/buildkit/client/llb"

const (
uvVersion = "0.6.5"
)

func (g generalGraph) compileUV(root llb.State) llb.State {
if g.UVConfig == nil {
return root
}
// uv configurations
g.RuntimeEnviron["UV_LINK_MODE"] = "copy"
g.RuntimeEnviron["UV_PYTHON_PREFERENCE"] = "only-managed"

base := llb.Image(builderImage)
builder := base.Run(
llb.Shlexf(`sh -c "wget -qO- https://github.com/astral-sh/uv/releases/download/%s/uv-$(uname -m)-unknown-linux-gnu.tar.gz | tar -xz --strip-components=1 -C /tmp || exit 1"`, uvVersion),
llb.WithCustomNamef("[internal] download uv %s", uvVersion),
).Root()

root = root.File(
llb.Copy(builder, "/tmp/uv", "/usr/bin/uv"), llb.WithCustomName("[internal] install uv")).
File(llb.Copy(builder, "/tmp/uvx", "/usr/bin/uvx"), llb.WithCustomName("[internal] install uvx"))
return g.compileUVPython(root)
}

func (g generalGraph) compileUVPython(root llb.State) llb.State {
if g.UVConfig == nil {
return root
}

root = root.Run(
llb.Shlexf(`uv python install %s`, g.UVConfig.PythonVersion),
llb.WithCustomNamef("[internal] install uv Python=%s", g.UVConfig.PythonVersion),
).Root()
return root
}
2 changes: 1 addition & 1 deletion pkg/metrics/docker_collector.go
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ func (c *dockerCollector) Watch(ctx context.Context, cid string) chan Metrics {
err := c.client.Stats(ctx, cid, stats, c.done)
if err != nil {
logrus.WithField("cid", cid).
WithError(err).Error("error occurred in dockerCollecter.Watch")
WithError(err).Error("error occurred in dockerCollector.Watch")
}
c.running = false
}()
Expand Down
1 change: 0 additions & 1 deletion pkg/types/envd.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,6 @@ var BaseEnvironment = []struct {
{"PATH", DefaultSystemPath},
{"LANG", "en_US.UTF-8"},
{"LC_ALL", "en_US.UTF-8"},
{"UV_LINK_MODE", "copy"}, // uv link-mode for installing Python packages
}
var BaseAptPackage = []string{
"bash-static",
Expand Down
Loading