diff --git a/cmd/envd/main.go b/cmd/envd/main.go index 1cc910e51..5537c7ba7 100644 --- a/cmd/envd/main.go +++ b/cmd/envd/main.go @@ -19,17 +19,11 @@ import ( "os" "github.com/cockroachdb/errors" - _ "github.com/moby/buildkit/client/connhelper/dockercontainer" - _ "github.com/moby/buildkit/client/connhelper/kubepod" - _ "github.com/moby/buildkit/client/connhelper/podmancontainer" - "github.com/sirupsen/logrus" - "github.com/spf13/viper" cli "github.com/urfave/cli/v2" "go.starlark.net/starlark" "go.starlark.net/syntax" - "github.com/tensorchord/envd/pkg/flag" - "github.com/tensorchord/envd/pkg/home" + "github.com/tensorchord/envd/pkg/app" "github.com/tensorchord/envd/pkg/version" ) @@ -38,61 +32,8 @@ func run(args []string) (bool, error) { fmt.Println(c.App.Name, version.Package, c.App.Version, version.Revision) } - // TODO(gaocegege): Enclose the app, maybe create the struct envdApp. - app := cli.NewApp() - app.EnableBashCompletion = true - app.Name = "envd" - app.Usage = "Build tools for data scientists" - app.Version = version.GetVersion().String() - app.Flags = []cli.Flag{ - &cli.BoolFlag{ - Name: "debug", - Usage: "enable debug output in logs", - }, - &cli.StringFlag{ - Name: flag.FlagBuildkitdImage, - Usage: "docker image to use for buildkitd", - Value: "docker.io/moby/buildkit:v0.10.3", - }, - &cli.StringFlag{ - Name: flag.FlagBuildkitdContainer, - Usage: "buildkitd container to use for buildkitd", - Value: "envd_buildkitd", - }, - } - - app.Commands = []*cli.Command{ - CommandBootstrap, - CommandBuild, - CommandDestroy, - CommandGet, - CommandPause, - CommandResume, - CommandUp, - CommandVersion, - } - - // Deal with debug flag. - var debugEnabled bool - - app.Before = func(context *cli.Context) error { - debugEnabled = context.Bool("debug") - - logrus.SetFormatter(&logrus.TextFormatter{FullTimestamp: true}) - if debugEnabled { - logrus.SetLevel(logrus.DebugLevel) - } - - if err := home.Initialize(); err != nil { - return errors.Wrap(err, "failed to initialize home manager") - } - - // TODO(gaocegege): Add a config struct to keep them. - viper.Set(flag.FlagBuildkitdContainer, context.String(flag.FlagBuildkitdContainer)) - viper.Set(flag.FlagBuildkitdImage, context.String(flag.FlagBuildkitdImage)) - return nil - } - return debugEnabled, app.Run(args) + app := app.New() + return app.Debug, app.Run(args) } func handleErr(debug bool, err error) { diff --git a/pkg/app/app.go b/pkg/app/app.go new file mode 100644 index 000000000..b2a524e71 --- /dev/null +++ b/pkg/app/app.go @@ -0,0 +1,95 @@ +// Copyright 2022 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 app + +import ( + "github.com/cockroachdb/errors" + _ "github.com/moby/buildkit/client/connhelper/dockercontainer" + _ "github.com/moby/buildkit/client/connhelper/kubepod" + _ "github.com/moby/buildkit/client/connhelper/podmancontainer" + "github.com/sirupsen/logrus" + "github.com/spf13/viper" + cli "github.com/urfave/cli/v2" + + "github.com/tensorchord/envd/pkg/flag" + "github.com/tensorchord/envd/pkg/home" + "github.com/tensorchord/envd/pkg/version" +) + +type EnvdApp struct { + cli.App + Debug bool +} + +func New() EnvdApp { + internalApp := cli.NewApp() + internalApp.EnableBashCompletion = true + internalApp.Name = "envd" + internalApp.Usage = "Build tools for data scientists" + internalApp.Version = version.GetVersion().String() + internalApp.Flags = []cli.Flag{ + &cli.BoolFlag{ + Name: "debug", + Usage: "enable debug output in logs", + }, + &cli.StringFlag{ + Name: flag.FlagBuildkitdImage, + Usage: "docker image to use for buildkitd", + Value: "docker.io/moby/buildkit:v0.10.3", + }, + &cli.StringFlag{ + Name: flag.FlagBuildkitdContainer, + Usage: "buildkitd container to use for buildkitd", + Value: "envd_buildkitd", + }, + } + + internalApp.Commands = []*cli.Command{ + CommandBootstrap, + CommandBuild, + CommandDestroy, + CommandGet, + CommandPause, + CommandResume, + CommandUp, + CommandVersion, + } + + // Deal with debug flag. + var debugEnabled bool + + internalApp.Before = func(context *cli.Context) error { + debugEnabled = context.Bool("debug") + + logrus.SetFormatter(&logrus.TextFormatter{FullTimestamp: true}) + if debugEnabled { + logrus.SetLevel(logrus.DebugLevel) + } + + if err := home.Initialize(); err != nil { + return errors.Wrap(err, "failed to initialize home manager") + } + + // TODO(gaocegege): Add a config struct to keep them. + viper.Set(flag.FlagBuildkitdContainer, context.String(flag.FlagBuildkitdContainer)) + viper.Set(flag.FlagBuildkitdImage, context.String(flag.FlagBuildkitdImage)) + return nil + } + + return EnvdApp{ + App: *internalApp, + Debug: debugEnabled, + } +} diff --git a/cmd/envd/bootstrap.go b/pkg/app/bootstrap.go similarity index 99% rename from cmd/envd/bootstrap.go rename to pkg/app/bootstrap.go index 9232ee7b8..dd2077511 100644 --- a/cmd/envd/bootstrap.go +++ b/pkg/app/bootstrap.go @@ -11,7 +11,7 @@ // 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 main +package app import ( "github.com/cockroachdb/errors" diff --git a/cmd/envd/build.go b/pkg/app/build.go similarity index 99% rename from cmd/envd/build.go rename to pkg/app/build.go index 4e9f8fe2c..5dab0f863 100644 --- a/cmd/envd/build.go +++ b/pkg/app/build.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "path/filepath" diff --git a/cmd/envd/build_test.go b/pkg/app/build_test.go similarity index 88% rename from cmd/envd/build_test.go rename to pkg/app/build_test.go index cc5b4f051..e195740f0 100644 --- a/cmd/envd/build_test.go +++ b/pkg/app/build_test.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "context" @@ -31,6 +31,9 @@ var _ = Describe("build command", func() { } BeforeEach(func() { Expect(home.Initialize()).NotTo(HaveOccurred()) + app := New() + err := app.Run([]string{"envd.test", "--debug", "bootstrap"}) + Expect(err).NotTo(HaveOccurred()) cli, err := docker.NewClient(context.TODO()) Expect(err).NotTo(HaveOccurred()) _, err = cli.Destroy(context.TODO(), buildContext) @@ -38,7 +41,8 @@ var _ = Describe("build command", func() { }) When("given the right arguments", func() { It("should build successfully", func() { - _, err := run(args) + app := New() + err := app.Run(args) Expect(err).NotTo(HaveOccurred()) }) }) diff --git a/cmd/envd/destroy.go b/pkg/app/destroy.go similarity index 99% rename from cmd/envd/destroy.go rename to pkg/app/destroy.go index 34e8d30a3..9e9bd341e 100644 --- a/cmd/envd/destroy.go +++ b/pkg/app/destroy.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "path/filepath" diff --git a/cmd/envd/get.go b/pkg/app/get.go similarity index 98% rename from cmd/envd/get.go rename to pkg/app/get.go index da20a32dc..5654f100d 100644 --- a/cmd/envd/get.go +++ b/pkg/app/get.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( cli "github.com/urfave/cli/v2" diff --git a/cmd/envd/get_env.go b/pkg/app/get_env.go similarity index 99% rename from cmd/envd/get_env.go rename to pkg/app/get_env.go index 3eb6a836e..6abf2b2ba 100644 --- a/cmd/envd/get_env.go +++ b/pkg/app/get_env.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "fmt" diff --git a/cmd/envd/get_env_dep.go b/pkg/app/get_env_dep.go similarity index 99% rename from cmd/envd/get_env_dep.go rename to pkg/app/get_env_dep.go index 39ff0cb6d..3158685ef 100644 --- a/cmd/envd/get_env_dep.go +++ b/pkg/app/get_env_dep.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "io" diff --git a/cmd/envd/get_image.go b/pkg/app/get_image.go similarity index 99% rename from cmd/envd/get_image.go rename to pkg/app/get_image.go index cfb12c7e8..ed0e68c83 100644 --- a/cmd/envd/get_image.go +++ b/pkg/app/get_image.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "io" diff --git a/cmd/envd/get_image_dep.go b/pkg/app/get_image_dep.go similarity index 99% rename from cmd/envd/get_image_dep.go rename to pkg/app/get_image_dep.go index e1fd29881..50fc6e3d0 100644 --- a/cmd/envd/get_image_dep.go +++ b/pkg/app/get_image_dep.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "os" diff --git a/cmd/envd/pause.go b/pkg/app/pause.go similarity index 99% rename from cmd/envd/pause.go rename to pkg/app/pause.go index 6d15bcf38..16fe81a29 100644 --- a/cmd/envd/pause.go +++ b/pkg/app/pause.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "github.com/cockroachdb/errors" diff --git a/cmd/envd/resume.go b/pkg/app/resume.go similarity index 99% rename from cmd/envd/resume.go rename to pkg/app/resume.go index ddd0abe03..c7c129b1e 100644 --- a/cmd/envd/resume.go +++ b/pkg/app/resume.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "github.com/cockroachdb/errors" diff --git a/cmd/envd/suite_test.go b/pkg/app/suite_test.go similarity index 98% rename from cmd/envd/suite_test.go rename to pkg/app/suite_test.go index 380eb9dc3..c87a03779 100644 --- a/cmd/envd/suite_test.go +++ b/pkg/app/suite_test.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "testing" diff --git a/cmd/envd/testdata/build.envd b/pkg/app/testdata/build.envd similarity index 100% rename from cmd/envd/testdata/build.envd rename to pkg/app/testdata/build.envd diff --git a/cmd/envd/up.go b/pkg/app/up.go similarity index 99% rename from cmd/envd/up.go rename to pkg/app/up.go index aebcf4886..38b6d89ff 100644 --- a/cmd/envd/up.go +++ b/pkg/app/up.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "fmt" diff --git a/cmd/envd/up_test.go b/pkg/app/up_test.go similarity index 88% rename from cmd/envd/up_test.go rename to pkg/app/up_test.go index 2385e4c67..022c420db 100644 --- a/cmd/envd/up_test.go +++ b/pkg/app/up_test.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "context" @@ -31,6 +31,9 @@ var _ = Describe("up command", func() { } BeforeEach(func() { Expect(home.Initialize()).NotTo(HaveOccurred()) + app := New() + err := app.Run([]string{"envd.test", "--debug", "bootstrap"}) + Expect(err).NotTo(HaveOccurred()) cli, err := docker.NewClient(context.TODO()) Expect(err).NotTo(HaveOccurred()) _, err = cli.Destroy(context.TODO(), buildContext) @@ -38,12 +41,13 @@ var _ = Describe("up command", func() { }) When("given the right arguments", func() { It("should up and destroy successfully", func() { - _, err := run(args) + app := New() + err := app.Run(args) Expect(err).NotTo(HaveOccurred()) destroyArgs := []string{ "envd.test", "--debug", "destroy", "--path", buildContext, } - _, err = run(destroyArgs) + err = app.Run(destroyArgs) Expect(err).NotTo(HaveOccurred()) }) }) diff --git a/cmd/envd/version.go b/pkg/app/version.go similarity index 99% rename from cmd/envd/version.go rename to pkg/app/version.go index 966eab1b4..3dc21838e 100644 --- a/cmd/envd/version.go +++ b/pkg/app/version.go @@ -12,10 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package app import ( "fmt" + "github.com/tensorchord/envd/pkg/version" "github.com/urfave/cli/v2" ) diff --git a/pkg/buildkitd/buildkitd.go b/pkg/buildkitd/buildkitd.go index 71e9b4b26..3b990d6c3 100644 --- a/pkg/buildkitd/buildkitd.go +++ b/pkg/buildkitd/buildkitd.go @@ -112,7 +112,7 @@ func (c *generalClient) maybeStart(ctx context.Context, return "", err } } - c.logger.Debug("container is running, check if it's ready...") + c.logger.Debugf("container is running, check if it's ready at %s...", c.BuildkitdAddr()) if err := c.waitUntilConnected(ctx, connectingTimeout); err != nil { return "", errors.Wrap(err, "failed to connect to buildkitd") @@ -133,7 +133,7 @@ func (c generalClient) waitUntilConnected( case <-time.After(interval): connected, err := c.connected(ctxTimeout) if err != nil { - logrus.Debug("failed to connect to buildkitd") + logrus.Debugf("failed to connect to buildkitd: %s", err.Error()) continue } if !connected { diff --git a/pkg/shell/zsh_test.go b/pkg/shell/zsh_test.go index 2367e31b8..c1ce32a07 100644 --- a/pkg/shell/zsh_test.go +++ b/pkg/shell/zsh_test.go @@ -15,8 +15,10 @@ package shell import ( + "os" "path/filepath" + "github.com/adrg/xdg" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/tensorchord/envd/pkg/home" @@ -31,8 +33,7 @@ var _ = Describe("zsh manager", func() { AfterEach(func() { // Cleanup the home cache. Expect(home.Initialize()).NotTo(HaveOccurred()) - m := home.GetManager() - Expect(fileutil.RemoveAll(m.CacheDir())).NotTo(HaveOccurred()) + Expect(os.RemoveAll(filepath.Join(xdg.CacheHome, "envd/cache.status"))).NotTo(HaveOccurred()) }) When("cached", func() { It("should skip", func() {