diff --git a/e2e/language/python_test.go b/e2e/language/python_test.go index bc07cd7d0..d9cb61456 100644 --- a/e2e/language/python_test.go +++ b/e2e/language/python_test.go @@ -16,14 +16,15 @@ 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()() @@ -31,7 +32,6 @@ var _ = Describe("python", Ordered, func() { }) It("Should build requirements successfully", func() { exampleName := "python/requirements" - testcase := "e2e" e := e2e.NewExample(e2e.BuildContextDirWithName(exampleName), testcase) e.BuildImage(true)() e.RunContainer()() @@ -39,7 +39,6 @@ var _ = Describe("python", Ordered, func() { }) 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()() }) + + 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()) + }) }) diff --git a/e2e/language/testdata/python/uv/build.envd b/e2e/language/testdata/python/uv/build.envd new file mode 100644 index 000000000..03d500668 --- /dev/null +++ b/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", + } + ) diff --git a/envd/api/v1/install.py b/envd/api/v1/install.py index 9cd3dbc69..6fcee7bd5 100644 --- a/envd/api/v1/install.py +++ b/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.""" diff --git a/pkg/lang/frontend/starlark/v1/install/const.go b/pkg/lang/frontend/starlark/v1/install/const.go index 937d8fe49..37fde4941 100644 --- a/pkg/lang/frontend/starlark/v1/install/const.go +++ b/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" diff --git a/pkg/lang/frontend/starlark/v1/install/install.go b/pkg/lang/frontend/starlark/v1/install/install.go index 3f0b46b25..2bd9c2b3e 100644 --- a/pkg/lang/frontend/starlark/v1/install/install.go +++ b/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) diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go index f2a8f21a8..a0da2f1c2 100644 --- a/pkg/lang/ir/types.go +++ b/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 diff --git a/pkg/lang/ir/v1/compile.go b/pkg/lang/ir/v1/compile.go index a36a59781..e5e88bdf3 100644 --- a/pkg/lang/ir/v1/compile.go +++ b/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") diff --git a/pkg/lang/ir/v1/interface.go b/pkg/lang/ir/v1/interface.go index 09354f378..95ecc6e6c 100644 --- a/pkg/lang/ir/v1/interface.go +++ b/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) diff --git a/pkg/lang/ir/v1/system.go b/pkg/lang/ir/v1/system.go index c1ae9df0f..94288f5c4 100644 --- a/pkg/lang/ir/v1/system.go +++ b/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) diff --git a/pkg/lang/ir/v1/types.go b/pkg/lang/ir/v1/types.go index 2a6f7073e..394cb0ed0 100644 --- a/pkg/lang/ir/v1/types.go +++ b/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:"-"` diff --git a/pkg/lang/ir/v1/util.go b/pkg/lang/ir/v1/util.go index 660d7198a..d32547ee4 100644 --- a/pkg/lang/ir/v1/util.go +++ b/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" diff --git a/pkg/lang/ir/v1/uv.go b/pkg/lang/ir/v1/uv.go new file mode 100644 index 000000000..768b98eac --- /dev/null +++ b/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 +} diff --git a/pkg/metrics/docker_collector.go b/pkg/metrics/docker_collector.go index 8fd523980..eb3a2e036 100644 --- a/pkg/metrics/docker_collector.go +++ b/pkg/metrics/docker_collector.go @@ -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 }() diff --git a/pkg/types/envd.go b/pkg/types/envd.go index 1b850c57b..cc9d3fee1 100644 --- a/pkg/types/envd.go +++ b/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",