feat: support uv

Signed-off-by: Keming <[email protected]>

e2e/language/python_test.go

@@ -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()()
@@ -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()()
@@ -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()()
 	})
+
+	It("Should build uv with Python successfully", func() {
+		exampleName := "python/uv"
+		e := e2e.NewExample(e2e.BuildContextDirWithName(exampleName), testcase)
+		e.BuildImage(true)()
+		e.RunContainer()()
+		It("Should have Python installed", func() {
+			res, err := e.ExecRuntimeCommand("uv-python")
+			Expect(err).To(BeNil())
+			Expect(res).To(ContainSubstring("python"))
+		})
+		e.DestroyContainer()()
+	})
 })

e2e/language/testdata/python/uv/build.envd

@@ -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",
+        }
+    )

envd/api/v1/install.py

@@ -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."""
 

pkg/lang/frontend/starlark/v1/install/const.go

@@ -18,6 +18,7 @@ const (
 	// language
 	rulePython = "install.python"
 	ruleConda  = "install.conda"
+	ruleUV     = "install.uv"
 	ruleRLang  = "install.r_lang"
 	ruleJulia  = "install.julia"
 

pkg/lang/frontend/starlark/v1/install/install.go

@@ -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
@@ -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)

pkg/lang/ir/types.go

@@ -61,6 +61,10 @@ type CondaConfig struct {
 	UseMicroMamba      bool
 }
 
+type UVConfig struct {
+	PythonVersion string
+}
+
 type GitConfig struct {
 	Name   string
 	Email  string

pkg/lang/ir/v1/compile.go

@@ -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")

pkg/lang/ir/v1/interface.go

@@ -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)
 

pkg/lang/ir/v1/system.go

@@ -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))
@@ -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 {
@@ -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)

pkg/lang/ir/v1/types.go

@@ -72,6 +72,7 @@ type generalGraph struct {
 	*ir.JupyterConfig
 	*ir.GitConfig
 	*ir.CondaConfig
+	*ir.UVConfig
 	*ir.RStudioServerConfig
 
 	Writer compileui.Writer `json:"-"`

pkg/lang/ir/v1/util.go

@@ -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"

pkg/lang/ir/v1/uv.go

@@ -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
+}

pkg/types/envd.go

@@ -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",