main.go 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. // Copyright 2020 gorse Project Authors
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package main
  15. import (
  16. "fmt"
  17. _ "net/http/pprof"
  18. "github.com/spf13/cobra"
  19. "github.com/zhenghaoz/gorse/base/log"
  20. "github.com/zhenghaoz/gorse/cmd/version"
  21. "github.com/zhenghaoz/gorse/worker"
  22. "go.uber.org/zap"
  23. )
  24. var workerCommand = &cobra.Command{
  25. Use: "gorse-worker",
  26. Short: "The worker node of gorse recommender system.",
  27. Run: func(cmd *cobra.Command, args []string) {
  28. // show version
  29. showVersion, _ := cmd.PersistentFlags().GetBool("version")
  30. if showVersion {
  31. fmt.Println(version.BuildInfo())
  32. return
  33. }
  34. masterHost, _ := cmd.PersistentFlags().GetString("master-host")
  35. masterPort, _ := cmd.PersistentFlags().GetInt("master-port")
  36. httpHost, _ := cmd.PersistentFlags().GetString("http-host")
  37. httpPort, _ := cmd.PersistentFlags().GetInt("http-port")
  38. workingJobs, _ := cmd.PersistentFlags().GetInt("jobs")
  39. managedModel, _ := cmd.PersistentFlags().GetBool("managed")
  40. // setup logger
  41. debug, _ := cmd.PersistentFlags().GetBool("debug")
  42. log.SetLogger(cmd.PersistentFlags(), debug)
  43. // create worker
  44. cachePath, _ := cmd.PersistentFlags().GetString("cache-path")
  45. w := worker.NewWorker(masterHost, masterPort, httpHost, httpPort, workingJobs, cachePath, managedModel)
  46. w.Serve()
  47. },
  48. }
  49. func init() {
  50. log.AddFlags(workerCommand.PersistentFlags())
  51. workerCommand.PersistentFlags().BoolP("version", "v", false, "gorse version")
  52. workerCommand.PersistentFlags().String("master-host", "127.0.0.1", "host of master node")
  53. workerCommand.PersistentFlags().Int("master-port", 8086, "port of master node")
  54. workerCommand.PersistentFlags().String("http-host", "127.0.0.1", "host for Prometheus metrics export")
  55. workerCommand.PersistentFlags().Int("http-port", 8089, "port for Prometheus metrics export")
  56. workerCommand.PersistentFlags().Bool("debug", false, "use debug log mode")
  57. workerCommand.PersistentFlags().Bool("managed", false, "enable managed mode")
  58. workerCommand.PersistentFlags().IntP("jobs", "j", 1, "number of working jobs.")
  59. workerCommand.PersistentFlags().String("cache-path", "worker_cache.data", "path of cache file")
  60. }
  61. func main() {
  62. if err := workerCommand.Execute(); err != nil {
  63. log.Logger().Fatal("failed to execute", zap.Error(err))
  64. }
  65. }