18149755418

Committed 01 Oct 2025 02:45AM UTC coverage: 42.475% (-0.3%) from 42.787%

Build # 18149755418

Build Type

push

github

Committed by

web-flow

Commit Message

chore(deps): bump github.com/go-viper/mapstructure/v2 from 2.3.0 to 2.4.0 (#2034)

chore(deps): bump github.com/go-viper/mapstructure/v2

Bumps [github.com/go-viper/mapstructure/v2](https://github.com/go-viper/mapstructure) from 2.3.0 to 2.4.0.
- [Release notes](https://github.com/go-viper/mapstructure/releases)
- [Changelog](https://github.com/go-viper/mapstructure/blob/main/CHANGELOG.md)
- [Commits](https://github.com/go-viper/mapstructure/compare/v2.3.0...v2.4.0)

---
updated-dependencies:
- dependency-name: github.com/go-viper/mapstructure/v2
  dependency-version: 2.4.0
  dependency-type: indirect
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

Run Details

5182 of 12200 relevant lines covered (42.48%)

165.16 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

43.51

/pkg/envd/docker.go

// Copyright 2023 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 envd

import (
        "context"
        "encoding/json"
        "fmt"
        "path/filepath"
        "runtime"
        "strconv"
        "strings"
        "time"

        "github.com/cockroachdb/errors"
        "github.com/containerd/errdefs"
        "github.com/docker/cli/opts"
        "github.com/docker/docker/api/types/container"
        "github.com/docker/docker/api/types/filters"
        dockerimage "github.com/docker/docker/api/types/image"
        "github.com/docker/docker/api/types/mount"
        "github.com/docker/docker/client"
        "github.com/docker/go-connections/nat"
        dockerutils "github.com/docker/go-units"
        "github.com/sirupsen/logrus"

        envdconfig "github.com/tensorchord/envd/pkg/config"
        "github.com/tensorchord/envd/pkg/lang/ir"
        "github.com/tensorchord/envd/pkg/lang/version"
        "github.com/tensorchord/envd/pkg/ssh"
        sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
        "github.com/tensorchord/envd/pkg/types"
        "github.com/tensorchord/envd/pkg/util/fileutil"
        "github.com/tensorchord/envd/pkg/util/netutil"
)

type dockerEngine struct {
        *client.Client
}

func dockerFilters(gpu bool) filters.Args {
        f := filters.NewArgs()
        f.Add("label", fmt.Sprintf("%s=%s", types.ImageLabelVendor, types.ImageVendorEnvd))
        if gpu {
                f.Add("label", fmt.Sprintf("%s=true", types.ImageLabelGPU))
        }
        return f
}

func dockerFiltersWithName(name string) filters.Args {
        f := filters.NewArgs()
        f.Add("reference", name)
        return f
}

func (e dockerEngine) ListImage(ctx context.Context) ([]types.EnvdImage, error) {
        images, err := e.ImageList(ctx, dockerimage.ListOptions{
                Filters: dockerFilters(false),
        })
        if err != nil {
                return nil, errors.Wrap(err, "failed to get the images")
        }

        envdImgs := make([]types.EnvdImage, 0)
        for _, img := range images {
                envdImg, err := types.NewImageFromSummary(img)
                if err != nil {
                        return nil, errors.Wrapf(err, "failed to create envd image `%s` from the docker image", img.ID)
                }
                envdImgs = append(envdImgs, *envdImg)
        }
        return envdImgs, nil
}

func (e dockerEngine) GetEnvironment(ctx context.Context, env string) (*types.EnvdEnvironment, error) {
        ctrs, err := e.ContainerList(ctx, container.ListOptions{
                Filters: dockerFiltersWithName(env),
        })
        if err != nil {
                return nil, errors.Wrapf(err, "failed to get container: %s", env)
        }

        if len(ctrs) <= 0 {
                return nil, errors.Newf("can not find the container: %s", env)
        }

        environment, err := types.NewEnvironmentFromContainer(ctrs[0])
        if err != nil {
                return nil, errors.Wrap(err, "failed to create env from the container")
        }
        return environment, nil
}

func (e dockerEngine) ListEnvironment(
        ctx context.Context) ([]types.EnvdEnvironment, error) {
        ctrs, err := e.ContainerList(ctx, container.ListOptions{
                Filters: dockerFilters(false),
        })
        if err != nil {
                return nil, errors.Wrap(err, "failed to list containers")
        }

        envs := make([]types.EnvdEnvironment, 0)
        for _, ctr := range ctrs {
                env, err := types.NewEnvironmentFromContainer(ctr)
                if err != nil {
                        return nil, errors.Wrap(err, "failed to create env from the container")
                }
                envs = append(envs, *env)
        }
        return envs, nil
}

func (e dockerEngine) PauseEnvironment(ctx context.Context, env string) (string, error) {
        logger := logrus.WithFields(logrus.Fields{
                "env": env,
        })
        logger.Debug("pausing environment")
        err := e.ContainerPause(ctx, env)
        if err != nil {
                errCause := errors.UnwrapAll(err).Error()
                switch {
                case strings.Contains(errCause, "is already paused"):
                        logger.Debug("container is already paused, there is no need to pause it again")
                        return "", nil
                case strings.Contains(errCause, "No such container"):
                        logger.Debug("container is not found, there is no need to pause it")
                        return "", errors.New("container not found")
                default:
                        return "", errors.Wrap(err, "failed to pause container")
                }
        }
        return env, nil
}

func (e dockerEngine) ResumeEnvironment(ctx context.Context, env string) (string, error) {
        logger := logrus.WithFields(logrus.Fields{
                "env": env,
        })
        logger.Debug("resuming environment")

        err := e.ContainerUnpause(ctx, env)
        if err != nil {
                errCause := errors.UnwrapAll(err).Error()
                switch {
                case strings.Contains(errCause, "is not paused"):
                        logger.Debug("container is not paused, there is no need to resume")
                        return "", nil
                case strings.Contains(errCause, "No such container"):
                        logger.Debug("container is not found, there is no need to resume it")
                        return "", errors.New("container not found")
                default:
                        return "", errors.Wrap(err, "failed to resume container")
                }
        }
        return env, nil
}

// ListImageDependency gets the dependencies of the given environment.
func (e dockerEngine) ListImageDependency(ctx context.Context, image string) (*types.Dependency, error) {
        logger := logrus.WithFields(logrus.Fields{
                "image": image,
        })
        logger.Debug("getting dependencies")
        images, err := e.ImageList(ctx, dockerimage.ListOptions{
                Filters: dockerFiltersWithName(image),
        })
        if err != nil {
                return nil, err
        }
        if len(images) == 0 {
                return nil, errors.Errorf("image %s not found", image)
        }

        img := images[0]
        dep, err := types.NewDependencyFromImageSummary(img)
        if err != nil {
                return nil, errors.Wrap(err, "failed to create dependency from image")
        }
        return dep, nil
}

func (e dockerEngine) getVerFromImageLabel(ctx context.Context, env string) (version.Getter, error) {
        ctr, err := e.GetImage(ctx, env)
        if err != nil {
                return nil, errors.Wrapf(err, "failed to inspect image: %s", env)
        }
        ver, ok := ctr.Labels[types.ImageLabelSyntaxVer]
        if !ok {
                return version.NewByVersion("v1"), nil
        }
        return version.NewByVersion(ver), nil
}

func (e dockerEngine) listEnvGeneralGraph(ctx context.Context, env string, g ir.Graph) (ir.Graph, error) {
        ctr, err := e.GetImage(ctx, env)
        if err != nil {
                return nil, errors.Wrapf(err, "failed to inspect image: %s", env)
        }
        code, ok := ctr.Labels[types.GeneralGraphCode]
        if !ok {
                return nil, errors.Newf("failed to get runtime graph label from image: %s", env)
        }
        logrus.WithField("env", env).Debugf("general graph: %s", code)
        newg, err := g.GeneralGraphFromLabel([]byte(code))
        if err != nil {
                return nil, errors.Wrapf(err, "failed to create runtime graph from the image: %s", env)
        }
        return newg, err
}

func (e dockerEngine) ListEnvRuntimeGraph(ctx context.Context, env string) (*ir.RuntimeGraph, error) {
        ctr, err := e.ContainerInspect(ctx, env)
        if err != nil {
                return nil, errors.Wrapf(err, "failed to inspect container: %s", env)
        }
        code, ok := ctr.Config.Labels[types.RuntimeGraphCode]
        if !ok {
                return nil, errors.Newf("failed to get runtime graph label from container: %s", env)
        }
        logrus.WithField("env", env).Debugf("runtime graph: %s", code)
        rg := ir.RuntimeGraph{}
        err = rg.Load([]byte(code))
        if err != nil {
                return nil, errors.Wrapf(err, "failed to create runtime graph from the container: %s", env)
        }
        return &rg, err
}

// ListEnvDependency gets the dependencies of the given environment.
func (e dockerEngine) ListEnvDependency(
        ctx context.Context, env string) (*types.Dependency, error) {
        logger := logrus.WithFields(logrus.Fields{
                "env": env,
        })
        logger.Debug("getting dependencies")
        ctr, err := e.ContainerInspect(ctx, env)
        if err != nil {
                return nil, errors.Wrap(err, "failed to get container")
        }
        dep, err := types.NewDependencyFromContainerJSON(ctr)
        if err != nil {
                return nil, errors.Wrap(err, "failed to create dependency from the container")
        }
        return dep, nil
}

func (e dockerEngine) ListEnvPortBinding(ctx context.Context, env string) ([]types.PortBinding, error) {
        logrus.WithField("env", env).Debug("getting env port bindings")
        ctr, err := e.ContainerInspect(ctx, env)
        if err != nil {
                return nil, errors.Wrap(err, "failed to get container")
        }

        ports, err := types.NewPortBindingFromContainerJSON(ctr)
        if err != nil {
                return nil, errors.Wrap(err, "failed to get port binding from container json")
        }

        return ports, nil
}

func (e dockerEngine) GetInfo(ctx context.Context) (*types.EnvdInfo, error) {
        info, err := e.Info(ctx)
        if err != nil {
                return nil, errors.Wrap(err, "failed to get docker client info")
        }
        return &types.EnvdInfo{
                Info: info,
        }, nil
}

func (e dockerEngine) CleanEnvdIfExists(ctx context.Context, name string, force bool) error {
        created, err := e.Exists(ctx, name)
        if err != nil {
                return err
        }
        if !created {
                return nil
        }
        return e.ContainerRemove(ctx, name, container.RemoveOptions{Force: force})
}

func (e dockerEngine) Exists(ctx context.Context, cname string) (bool, error) {
        _, err := e.ContainerInspect(ctx, cname)
        if err != nil {
                if errdefs.IsNotFound(err) {
                        return false, nil
                }
                return false, err
        }
        return true, nil
}

func (e dockerEngine) IsRunning(ctx context.Context, cname string) (bool, error) {
        container, err := e.ContainerInspect(ctx, cname)
        if err != nil {
                if errdefs.IsNotFound(err) {
                        return false, nil
                }
                return false, err
        }
        return container.State.Running, nil
}

func (e dockerEngine) GenerateSSHConfig(name, iface, privateKeyPath string,
        startResult *StartResult) (sshconfig.EntryOptions, error) {
        eo := sshconfig.EntryOptions{
                Name:               name,
                IFace:              iface,
                Port:               startResult.SSHPort,
                PrivateKeyPath:     privateKeyPath,
                EnableHostKeyCheck: false,
                EnableAgentForward: true,
        }
        return eo, nil
}

func (e dockerEngine) Attach(name, iface, privateKeyPath string,
        startResult *StartResult, g ir.Graph) error {
        opt := ssh.DefaultOptions()
        opt.Server = iface
        opt.PrivateKeyPath = privateKeyPath
        opt.Port = startResult.SSHPort
        sshClient, err := ssh.NewClient(opt)
        if err != nil {
                return errors.Wrap(err, "failed to create the ssh client")
        }
        opt.Server = iface

        if err := sshClient.Attach(); err != nil {
                return errors.Wrap(err, "failed to attach to the container")
        }
        return nil
}

// StartEnvd creates the container for the given tag and container name.
func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartResult, error) {
        logger := logrus.WithFields(logrus.Fields{
                "tag":           so.Image,
                "environment":   so.EnvironmentName,
                "gpu-set":       so.GPUSet,
                "shm":           so.ShmSize,
                "cpu":           so.NumCPU,
                "cpu-set":       so.CPUSet,
                "memory":        so.NumMem,
                "build-context": so.BuildContext,
        })

        bar := InitProgressBar(5)
        defer bar.Finish()
        bar.UpdateTitle("configure the environment")

        if len(so.GPUSet) > 0 {
                nvruntimeExists, err := e.GPUEnabled(ctx)
                if err != nil {
                        return nil, errors.Wrap(err, "failed to check if nvidia-runtime is installed")
                }
                if !nvruntimeExists {
                        return nil, errors.New("GPU is required but nvidia container runtime is not installed, please refer to https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker")
                }
        }

        sshPortInHost, err := netutil.GetFreePort()
        if err != nil {
                return nil, errors.Wrap(err, "failed to get a free port")
        }

        err = e.CleanEnvdIfExists(ctx, so.EnvironmentName, so.Forced)
        if err != nil {
                return nil, errors.Wrap(err, "failed to clean the envd environment")
        }
        config := &container.Config{
                Image:        so.Image,
                ExposedPorts: nat.PortSet{},
        }
        base := fileutil.EnvdHomeDir(filepath.Base(so.BuildContext))
        config.WorkingDir = base

        var g ir.Graph
        if so.Image == "" {
                g = so.DockerSource.Graph
        } else {
                // Use specified image as the running image
                getter, err := e.getVerFromImageLabel(ctx, so.Image)
                if err != nil {
                        return nil, errors.Wrap(err, "failed to get the version from the image label")
                }
                defaultGraph := getter.GetDefaultGraph()
                if err != nil {
                        return nil, errors.Wrap(err, "failed to get the graph from the image")
                }
                g, err = e.listEnvGeneralGraph(ctx, so.Image, defaultGraph)
                if err != nil {
                        return nil, errors.Wrap(err, "failed to get the graph from the image")
                }
                so.DockerSource.Graph = g
        }

        if so.DockerSource == nil || so.DockerSource.Graph == nil {
                return nil, errors.New("failed to get the docker-specific options")
        }

        mountOption := make([]mount.Mount, 0,
                len(so.DockerSource.MountOptions)+len(g.GetMount())+1)
        for _, option := range so.DockerSource.MountOptions {
                mStr := strings.Split(option, ":")
                if len(mStr) != 2 {
                        return nil, errors.Newf("Invalid mount options %s", option)
                }

                logger.WithFields(logrus.Fields{
                        "mount-path":     mStr[0],
                        "container-path": mStr[1],
                }).Debug("setting up container working directory")
                mountOption = append(mountOption, mount.Mount{
                        Type:   mount.TypeBind,
                        Source: mStr[0],
                        Target: mStr[1],
                })
        }
        for _, m := range g.GetMount() {
                logger.WithFields(logrus.Fields{
                        "mount-path":     m.Source,
                        "container-path": m.Destination,
                }).Debug("setting up declared mount directory")
                mountOption = append(mountOption, mount.Mount{
                        Type:   mount.TypeBind,
                        Source: m.Source,
                        Target: m.Destination,
                })
        }

        mountOption = append(mountOption, mount.Mount{
                Type:   mount.TypeBind,
                Source: so.BuildContext,
                Target: base,
        })

        logger.WithFields(logrus.Fields{
                "mount-path":  so.BuildContext,
                "working-dir": base,
        }).Debug("setting up container working directory")

        rp := container.RestartPolicy{
                Name: "always",
        }
        hostConfig := &container.HostConfig{
                PortBindings:  nat.PortMap{},
                Mounts:        mountOption,
                RestartPolicy: rp,
        }

        // shared memory size
        if so.ShmSize > 0 {
                hostConfig.ShmSize = int64(so.ShmSize) * 1024 * 1024
        }
        // resource
        if len(so.NumCPU) > 0 {
                cpu, err := strconv.ParseFloat(so.NumCPU, 64)
                if err != nil {
                        logger.Infof("parse `cpu` error: %v, ignore this argument", err)
                } else if runtime.GOOS == "windows" {
                        hostConfig.NanoCPUs = int64(cpu * 10e9)
                } else {
                        // refer to https://docs.docker.com/config/containers/resource_constraints/#configure-the-default-cfs-scheduler
                        // CPU quota and CPU period only work for UNIX platform
                        hostConfig.CPUQuota = int64(cpu * 10e5)
                }
        }
        if len(so.CPUSet) > 0 {
                hostConfig.CpusetCpus = so.CPUSet
        }
        if len(so.NumMem) > 0 {
                mem, err := dockerutils.RAMInBytes(so.NumMem)
                if err != nil {
                        logger.WithError(err).Info("parse `memory` error, ignore this argument")
                } else {
                        hostConfig.Memory = mem
                }
        }

        // Configure ssh port.
        natPort := nat.Port(fmt.Sprintf("%d/tcp", envdconfig.SSHPortInContainer))
        hostConfig.PortBindings[natPort] = []nat.PortBinding{
                {
                        HostIP:   so.SshdHost,
                        HostPort: strconv.Itoa(sshPortInHost),
                },
        }

        var jupyterPortInHost int
        // TODO(gaocegege): Avoid specific logic to set the port.
        // Add a func to builder to generate all the ports from the build process.
        if g.GetJupyterConfig() != nil {
                jc := g.GetJupyterConfig()
                if jc.Port != 0 {
                        jupyterPortInHost = int(jc.Port)
                } else {
                        var err error
                        jupyterPortInHost, err = netutil.GetFreePort()
                        if err != nil {
                                return nil, errors.Wrap(err, "failed to get a free port")
                        }
                }
                natPort := nat.Port(fmt.Sprintf("%d/tcp", envdconfig.JupyterPortInContainer))
                hostConfig.PortBindings[natPort] = []nat.PortBinding{
                        {
                                HostIP:   Localhost,
                                HostPort: strconv.Itoa(jupyterPortInHost),
                        },
                }
                config.ExposedPorts[natPort] = struct{}{}
        }
        var rStudioPortInHost int
        if g.GetRStudioServerConfig() != nil {
                var err error
                rStudioPortInHost, err = netutil.GetFreePort()
                if err != nil {
                        return nil, errors.Wrap(err, "failed to get a free port")
                }
                natPort := nat.Port(fmt.Sprintf("%d/tcp", envdconfig.RStudioServerPortInContainer))
                hostConfig.PortBindings[natPort] = []nat.PortBinding{
                        {
                                HostIP:   Localhost,
                                HostPort: strconv.Itoa(rStudioPortInHost),
                        },
                }
                config.ExposedPorts[natPort] = struct{}{}
        }

        if len(g.GetExposedPorts()) > 0 {

                for _, item := range g.GetExposedPorts() {
                        var err error
                        if item.HostPort == 0 {
                                item.HostPort, err = netutil.GetFreePort()
                                if err != nil {
                                        return nil, errors.Wrap(err, "failed to get a free port")
                                }
                        }
                        natPort := nat.Port(fmt.Sprintf("%d/tcp", item.EnvdPort))
                        hostConfig.PortBindings[natPort] = []nat.PortBinding{
                                {
                                        HostIP:   item.ListeningAddr,
                                        HostPort: strconv.Itoa(item.HostPort),
                                },
                        }
                        config.ExposedPorts[natPort] = struct{}{}
                }
        }

        if len(so.GPUSet) > 0 {
                logger.Debug("GPU is enabled.")
                hostConfig.DeviceRequests, err = deviceRequests(so.GPUSet)
                if err != nil {
                        return nil, errors.Wrap(err, "failed to parse gpu-set flag")

                }
        }

        config.Labels = e.labels(g, so.EnvironmentName,
                sshPortInHost, jupyterPortInHost, rStudioPortInHost)

        logger = logger.WithFields(logrus.Fields{
                "entrypoint":  config.Entrypoint,
                "working-dir": config.WorkingDir,
        })
        logger.Debugf("starting %s container", so.EnvironmentName)

        bar.UpdateTitle("create the environment")
        resp, err := e.ContainerCreate(ctx, config, hostConfig, nil, nil, so.EnvironmentName)
        if err != nil {
                return nil, errors.Wrap(err, "failed to create the container")
        }

        for _, w := range resp.Warnings {
                logger.Warnf("run with warnings: %s", w)
        }

        bar.UpdateTitle("start the environment")
        if err := e.ContainerStart(
                ctx, resp.ID, container.StartOptions{}); err != nil {
                errCause := errors.UnwrapAll(err)
                // Hack to check if the port is already allocated.
                if strings.Contains(errCause.Error(), "port is already allocated") {
                        logger.WithError(err).Debug("failed to allocate the port")
                        return nil, errors.New("port is already allocated in the host")
                }
                return nil, errors.Wrap(err, "failed to run the container")
        }

        container, err := e.ContainerInspect(ctx, resp.ID)
        if err != nil {
                return nil, errors.Wrap(err, "failed to inspect the container")
        }

        bar.UpdateTitle("wait for the environment to start")
        if err := e.WaitUntilRunning(
                ctx, container.Name, so.Timeout); err != nil {
                return nil, errors.Wrap(err, "failed to wait until the container is running")
        }

        bar.UpdateTitle("attach the environment")
        result := &StartResult{
                SSHPort: sshPortInHost,
                Address: container.NetworkSettings.IPAddress,
                Name:    container.Name,
        }
        return result, nil
}

func (e dockerEngine) Destroy(ctx context.Context, name string) (string, error) {
        logger := logrus.WithField("container", name)

        ctr, err := e.ContainerInspect(ctx, name)
        if err != nil {
                errCause := errors.UnwrapAll(err).Error()
                if strings.Contains(errCause, "No such container") {
                        // If the container is not found, it is already destroyed or the name is wrong.
                        logger.Infof("cannot find container %s, maybe it's already destroyed or the name is wrong", name)
                        return "", nil
                }
                return "", errors.Wrap(err, "failed to inspect the container")
        }

        // Refer to https://docs.docker.com/engine/reference/commandline/container_kill/
        if err := e.ContainerKill(ctx, name, "KILL"); err != nil {
                errCause := errors.UnwrapAll(err).Error()
                switch {
                case strings.Contains(errCause, "is not running"):
                        // If the container is not running, there is no need to kill it.
                        logger.Debug("container is not running, there is no need to kill it")
                case strings.Contains(errCause, "No such container"):
                        // If the container is not found, it is already destroyed or the name is wrong.
                        logger.Infof("cannot find container %s, maybe it's already destroyed or the name is wrong", name)
                        return "", nil
                default:
                        return "", errors.Wrap(err, "failed to kill the container")
                }
        }

        if err := e.ContainerRemove(ctx, name, container.RemoveOptions{}); err != nil {
                return "", errors.Wrap(err, "failed to remove the container")
        }

        if _, err := e.ImageRemove(ctx, ctr.Image, dockerimage.RemoveOptions{}); err != nil {
                return "", errors.Errorf("remove image %s failed: %w", ctr.Image, err)
        }
        logger.Infof("image(%s) is destroyed", ctr.Image)

        return name, nil
}

func (e dockerEngine) WaitUntilRunning(ctx context.Context,
        name string, timeout time.Duration) error {
        logger := logrus.WithField("container", name)
        logger.Debug("waiting to start")

        // First, wait for the container to be marked as started.
        ctxTimeout, cancel := context.WithTimeout(ctx, timeout)
        defer cancel()
        for {
                select {
                case <-time.After(waitingInterval):
                        isRunning, err := e.IsRunning(ctxTimeout, name)
                        if err != nil {
                                // Has not yet started. Keep waiting.
                                return errors.Wrap(err, "failed to check if container is running")
                        }
                        if isRunning {
                                logger.Debug("the container is running")
                                return nil
                        }

                case <-ctxTimeout.Done():
                        container, err := e.ContainerInspect(ctx, name)
                        if err != nil {
                                logger.Debugf("failed to inspect container %s", name)
                        }
                        state, err := json.Marshal(container.State)
                        if err != nil {
                                logger.Debug("failed to marshal container state")
                        }
                        logger.Debugf("container state: %s", state)
                        return errors.Errorf("timeout %s: container did not start", timeout)
                }
        }
}

func (e dockerEngine) GPUEnabled(ctx context.Context) (bool, error) {
        info, err := e.GetInfo(ctx)
        if err != nil {
                return false, errors.Wrap(err, "failed to get docker info")
        }
        logrus.WithField("info", info).Debug("docker info")
        nv := info.Runtimes["nvidia"]
        if nv.Path != "" {
                return true, nil
        } else if strings.HasSuffix(info.KernelVersion, "WSL2") {
                logrus.Warn("We couldn't detect if your runtime support GPU on WSL2, we will continue to run your environment.")
                return true, nil
        }
        return false, nil
}

func (e dockerEngine) GetImage(ctx context.Context, imageName string) (types.EnvdImage, error) {
        images, err := e.ImageList(ctx, dockerimage.ListOptions{
                Filters: dockerFiltersWithName(imageName),
        })
        if err != nil {
                return types.EnvdImage{}, err
        }
        if len(images) == 0 {
                return types.EnvdImage{},
                        errors.Errorf("image %s not found", imageName)
        }
        img, err := types.NewImageFromSummary(images[0])
        if err != nil {
                return types.EnvdImage{}, err
        }
        return *img, nil
}

func (e dockerEngine) PruneImage(ctx context.Context) (dockerimage.PruneReport, error) {
        pruneReport, err := e.ImagesPrune(ctx, filters.Args{})
        if err != nil {
                return dockerimage.PruneReport{}, errors.Wrap(err, "failed to prune images")
        }
        return pruneReport, nil
}

func deviceRequests(value string) ([]container.DeviceRequest, error) {
        if value == "" {
                return nil, nil
        }
        if !strings.Contains(value, "capabilities=") {
                value += ",\"capabilities=nvidia,compute,compat32,graphics,utility,video,display\""
        }
        if !strings.Contains(value, "driver=") {
                value += ",driver=nvidia"
        }
        opt := opts.GpuOpts{}
        if err := opt.Set(value); err != nil {
                return nil, err
        }
        return opt.Value(), nil
}

func (e dockerEngine) labels(g ir.Graph, name string,
        sshPortInHost, jupyterPortInHost, rstudioServerPortInHost int) map[string]string {
        res := make(map[string]string)
        res[types.ContainerLabelName] = name
        res[types.ContainerLabelSSHPort] = strconv.Itoa(sshPortInHost)
        if g.GetJupyterConfig() != nil {
                res[types.ContainerLabelJupyterAddr] =
                        fmt.Sprintf("http://%s:%d", Localhost, jupyterPortInHost)
        }
        if g.GetRStudioServerConfig() != nil {
                res[types.ContainerLabelRStudioServerAddr] =
                        fmt.Sprintf("http://%s:%d", Localhost, rstudioServerPortInHost)
        }

        return res
}

tensorchord / envd / 18149755418

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous