diff --git a/cmd/cmd.go b/cmd/cmd.go index b8c9c640..b43f7578 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -400,19 +400,43 @@ func loadOrUnloadModel(cmd *cobra.Command, opts *runOptions) error { return client.Generate(cmd.Context(), req, func(api.GenerateResponse) error { return nil }) } -func StopHandler(cmd *cobra.Command, args []string) error { +func stopModel(cmd *cobra.Command, model string) error { opts := &runOptions{ - Model: args[0], + Model: model, KeepAlive: &api.Duration{Duration: 0}, } + if err := loadOrUnloadModel(cmd, opts); err != nil { if strings.Contains(err.Error(), "not found") { - return fmt.Errorf("couldn't find model \"%s\" to stop", args[0]) + return fmt.Errorf("couldn't find model \"%s\" to stop", model) } } return nil } +func StopHandler(cmd *cobra.Command, args []string) error { + if args[0] == "all" { + client, err := api.ClientFromEnvironment() + if err != nil { + return err + } + + models, err := client.List(cmd.Context()) + if err != nil { + return err + } + + for _, m := range models.Models { + err = stopModel(cmd, m.Name) + if err != nil { + return err + } + } + return nil + } + return stopModel(cmd, args[0]) +} + func RunHandler(cmd *cobra.Command, args []string) error { interactive := true @@ -1382,7 +1406,7 @@ func NewCLI() *cobra.Command { runCmd.Flags().String("format", "", "Response format (e.g. json)") stopCmd := &cobra.Command{ - Use: "stop MODEL", + Use: "stop MODEL|all", Short: "Stop a running model", Args: cobra.ExactArgs(1), PreRunE: checkServerHeartbeat,