Skip to content

Commit

Permalink
adding common shared service things and starting to use them in artif…
Browse files Browse the repository at this point in the history 
…act service
  • Loading branch information
@wild-endeavor
wild-endeavor committed Oct 11, 2023
1 parent 3f88f17 commit 4f044e0
Show file tree
Hide file tree
Showing 6 changed files with 488 additions and 83 deletions.
162 changes: 79 additions & 83 deletions flyteartifacts/cmd/main.go
View file
Original file line number Diff line number Diff line change
@@ -1,96 +1,92 @@
package main

import (
"flag"
"fmt"
"github.com/flyteorg/flyte/flyteartifacts/pkg/configuration"
"github.com/flyteorg/flyte/flyteartifacts/pkg/server"
"github.com/flyteorg/flyte/flytestdlib/config"
"github.com/flyteorg/flyte/flytestdlib/config/viper"
"github.com/flyteorg/flyte/flytestdlib/logger"
"github.com/golang/glog"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"google.golang.org/grpc"
"os"

"context"
sharedCmd "github.com/flyteorg/flyte/flyteartifacts/cmd/shared"
"github.com/flyteorg/flyte/flytestdlib/logger"

_ "net/http/pprof" // Required to serve application.
)

var (
cfgFile string
configAccessor = viper.NewAccessor(config.Options{})
)

var serveCmd = &cobra.Command{
Use: "serve",
Short: "Launches the Flyte artifacts server",
RunE: func(cmd *cobra.Command, args []string) error {
ctx := context.Background()
cfg := configuration.ApplicationConfig.GetConfig().(*configuration.ApplicationConfiguration)
fmt.Printf("cfg: [%+v]\n", cfg)
opts := make([]grpc.ServerOption, 0)
return server.Serve(ctx, opts...)
},
}

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
Use: "artifacts",
Short: "Fill in later",
Long: `
To get started run the serve subcommand which will start a server on localhost:50051
`,
PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
return initConfig(cmd.Flags())
},
}

func init() {
pflag.CommandLine.AddGoFlagSet(flag.CommandLine)

// Add persistent flags - persistent flags persist through all sub-commands
RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is ./artifact_config.yaml)")

// Allow viper to read the value of the flags
configAccessor.InitializePflags(RootCmd.PersistentFlags())

// Command information
RootCmd.AddCommand(serveCmd)

err := flag.CommandLine.Parse([]string{})
if err != nil {
fmt.Println(err)
os.Exit(-1)
}

}

func initConfig(flags *pflag.FlagSet) error {
configAccessor = viper.NewAccessor(config.Options{
SearchPaths: []string{cfgFile, "./artifact_config.yaml", ".", "/etc/flyte/config", "$GOPATH/src/github.com/flyteorg/flyte/flyteartifacts"},
StrictMode: false,
})

logger.Infof(context.TODO(), "Using config file: %v", configAccessor.ConfigFilesUsed())

configAccessor.InitializePflags(flags)

err := flag.CommandLine.Parse([]string{})
if err != nil {
fmt.Println(err)
os.Exit(-1)
}

return configAccessor.UpdateConfig(context.TODO())
}
//
//var (
// cfgFile string
// configAccessor = viper.NewAccessor(config.Options{})
//)
//
//var serveCmd = &cobra.Command{
// Use: "serve",
// Short: "Launches the Flyte artifacts server",
// RunE: func(cmd *cobra.Command, args []string) error {
// ctx := context.Background()
// cfg := configuration.ApplicationConfig.GetConfig().(*configuration.ApplicationConfiguration)
// fmt.Printf("cfg: [%+v]\n", cfg)
// opts := make([]grpc.ServerOption, 0)
// return server.Serve(ctx, opts...)
// },
//}
//
//// RootCmd represents the base command when called without any subcommands
//var RootCmd = &cobra.Command{
// Use: "artifacts",
// Short: "Fill in later",
// Long: `
//To get started run the serve subcommand which will start a server on localhost:50051
//`,
// PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
// return initConfig(cmd.Flags())
// },
//}
//
//func init() {
// pflag.CommandLine.AddGoFlagSet(flag.CommandLine)
//
// // Add persistent flags - persistent flags persist through all sub-commands
// RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is ./artifact_config.yaml)")
//
// // Allow viper to read the value of the flags
// configAccessor.InitializePflags(RootCmd.PersistentFlags())
//
// // Command information
// RootCmd.AddCommand(serveCmd)
//
// err := flag.CommandLine.Parse([]string{})
// if err != nil {
// fmt.Println(err)
// os.Exit(-1)
// }
//
//}
//
//func initConfig(flags *pflag.FlagSet) error {
// configAccessor = viper.NewAccessor(config.Options{
// SearchPaths: []string{cfgFile, "./artifact_config.yaml", ".", "/etc/flyte/config", "$GOPATH/src/github.com/flyteorg/flyte/flyteartifacts"},
// StrictMode: false,
// })
//
// logger.Infof(context.TODO(), "Using config file: %v", configAccessor.ConfigFilesUsed())
//
// configAccessor.InitializePflags(flags)
//
// err := flag.CommandLine.Parse([]string{})
// if err != nil {
// fmt.Println(err)
// os.Exit(-1)
// }
//
// return configAccessor.UpdateConfig(context.TODO())
//}

func main() {
glog.V(2).Info("Beginning Flyte Artifacts Service")
if err := RootCmd.Execute(); err != nil {
fmt.Println(err)
ctx := context.Background()
logger.Infof(ctx, "Beginning Flyte Artifacts Service")
rootCmd := sharedCmd.NewRootCmd("artifacts")
//if err := RootCmd.ExecuteC(); err != nil {
// fmt.Println(err)
// panic(err)
//}
err := rootCmd.ExecuteContext(ctx)
if err != nil {
panic(err)
}
}
108 changes: 108 additions & 0 deletions flyteartifacts/cmd/shared/root.go
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
package shared

import (
"context"
"flag"
"fmt"
sharedCfg "github.com/flyteorg/flyte/flyteartifacts/pkg/configuration/shared"
"github.com/flyteorg/flyte/flytestdlib/config"
"github.com/flyteorg/flyte/flytestdlib/config/viper"
"github.com/flyteorg/flyte/flytestdlib/contextutils"
"github.com/flyteorg/flyte/flytestdlib/logger"
"github.com/flyteorg/flyte/flytestdlib/profutils"
"github.com/flyteorg/flyte/flytestdlib/promutils/labeled"
"github.com/flyteorg/flyte/flytestdlib/storage"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"os"
)

var (
cfgFile string
configAccessor = viper.NewAccessor(config.Options{})
)

//var XXRootCmd = &cobra.Command{
// Use: "artifacts",
// Short: "Fill in later",
// PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
// return initConfig(cmd.Flags())
// },
//}

// NewRootCmd represents the base command when called without any subcommands
func NewRootCmd(rootUse string, grpcHook GrpcRegistrationHook, httpHook HttpRegistrationHook) *cobra.Command {

rootCmd := &cobra.Command{
Use: rootUse,
Short: "Short description",
Long: "Long description to be filled in later",
PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
err := initConfig(cmd, args)
if err != nil {
return err
}

go func() {
ctx := context.Background()
sharedConfig := sharedCfg.SharedServerConfig.GetConfig().(*sharedCfg.ServerConfiguration)
err := profutils.StartProfilingServerWithDefaultHandlers(ctx,
sharedConfig.Metrics.Port.Port, nil)
if err != nil {
logger.Panicf(ctx, "Failed to Start profiling and metrics server. Error: %v", err)
}
}()

return nil
},
}

initSubCommands(rootCmd, grpcHook, httpHook)
return rootCmd
}

func init() {
// Set Keys
labeled.SetMetricKeys(contextutils.AppNameKey, contextutils.ProjectKey,
contextutils.DomainKey, storage.FailureTypeLabel)
}

func initConfig(cmd *cobra.Command, _ []string) error {
configAccessor = viper.NewAccessor(config.Options{
SearchPaths: []string{cfgFile, ".", "/etc/flyte/config", "$GOPATH/src/github.com/flyteorg/flyte/flyteartifacts"},
StrictMode: false,
})

fmt.Println("Using config file: ", configAccessor.ConfigFilesUsed())

// persistent flags were initially bound to the root command so we must bind to the same command to avoid
// overriding those initial ones. We need to traverse up to the root command and initialize pflags for that.
rootCmd := cmd
for rootCmd.Parent() != nil {
rootCmd = rootCmd.Parent()
}

configAccessor.InitializePflags(rootCmd.PersistentFlags())

return configAccessor.UpdateConfig(context.TODO())
}

func initSubCommands(rootCmd *cobra.Command, grpcHook GrpcRegistrationHook, httpHook HttpRegistrationHook) {
// allows ` --logtostderr` to work
pflag.CommandLine.AddGoFlagSet(flag.CommandLine)

// Add persistent flags - persistent flags persist through all sub-commands
rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "artifact_config.yaml", "config file (default is ./artifact_config.yaml)")

rootCmd.AddCommand(viper.GetConfigCommand())
rootCmd.AddCommand(NewServeCmd(rootCmd.Use, grpcHook, httpHook))

// Allow viper to read the value of the flags
configAccessor.InitializePflags(rootCmd.PersistentFlags())

err := flag.CommandLine.Parse([]string{})
if err != nil {
fmt.Println(err)
os.Exit(-1)
}
}
Loading

0 comments on commit 4f044e0

Please sign in to comment.