Skip to content

Commit

Permalink
feat(kollama): validate args
Browse files Browse the repository at this point in the history
Signed-off-by: Neko Ayaka <[email protected]>
  • Loading branch information
nekomeowww committed Apr 17, 2024
1 parent e99fb39 commit 2c7fb80
Show file tree
Hide file tree
Showing 3 changed files with 73 additions and 44 deletions.
9 changes: 3 additions & 6 deletions internal/cli/kollama/cmd.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,12 +34,9 @@ var (
// NewCmd provides a cobra command wrapping NamespaceOptions
func NewCmd(streams genericiooptions.IOStreams) *cobra.Command {
cmd := &cobra.Command{
Use: "kollama [cmd] [args] [flags]",
Short: "CLI for Ollama Operator",
SilenceUsage: true,
RunE: func(c *cobra.Command, args []string) error {
return nil
},
Use: "kollama [cmd] [args] [flags]",
Short: "CLI for Ollama Operator",
Args: cobra.NoArgs,
}

cmd.AddCommand(NewCmdDeploy(streams))
Expand Down
70 changes: 48 additions & 22 deletions internal/cli/kollama/cmd_deploy.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,21 +21,24 @@ import (
)

const (
deployExample = `# Deploy a phi model
deployExample = `
# Deploy a phi model
$ kollama deploy phi
kollama deploy phi
or if using as kubectl plugin
$ kubectl ollama deploy phi
or
# Deploy a model with a specific image
$ kollama deploy phi --image=phi-image
kubectl ollama deploy phi
or if using as kubectl plugin
$ kubectl ollama deploy phi --image=phi-image
# Deploy a phi model in a specific namespace
# Deploy a phi model in a specific namespace
$ kollama deploy phi -n phi-namespace
kollama deploy phi -n phi-namespace
or
kubectl ollama deploy phi -n phi-namespace`
or if using as kubectl plugin
$ kubectl ollama deploy phi -n phi-namespace`
)

// CmdDeployOptions provides information required to deploy a model
Expand All @@ -47,6 +50,7 @@ type CmdDeployOptions struct {
discoveryClient discovery.DiscoveryInterface

userSpecifiedNamespace string
modelImage string

genericiooptions.IOStreams
}
Expand All @@ -64,15 +68,29 @@ func NewCmdDeploy(streams genericiooptions.IOStreams) *cobra.Command {
o := NewCmdDeployOptions(streams)

cmd := &cobra.Command{
Use: "deploy [model name] [flags]",
Short: "Deploy a model with the given name by using Ollama Operator",
Example: deployExample,
SilenceUsage: true,
Use: "deploy [model name] [flags]",
Short: "Deploy a model with the given name by using Ollama Operator",
Example: deployExample,
Args: func(cmd *cobra.Command, args []string) error {
if len(args) < 1 {
return fmt.Errorf("model name is required")
}
if args[0] == "" {
return fmt.Errorf("model name cannot be empty")
}

return nil
},
RunE: func(c *cobra.Command, args []string) error {
return o.runE(c, args)
},
}

cmd.Flags().StringVar(&o.modelImage, "image", "", ""+
"Model image to deploy. If not specified, the model name will be used as the "+
"image name (will be pulled from registry.ollama.ai/library/<model name> by "+
"default if no registry is specified), the tag will be latest.")

o.configFlags.AddFlags(cmd.Flags())
o.clientConfig = o.configFlags.ToRawKubeConfigLoader()
o.kubeConfig = lo.Must(o.clientConfig.ClientConfig())
Expand Down Expand Up @@ -109,22 +127,30 @@ func (o *CmdDeployOptions) runE(cmd *cobra.Command, args []string) error {
return ErrOllamaModelNotSupported
}

modelImage := args[0]
modelName := args[0]

modelImage, err := cmd.Flags().GetString("image")
if err != nil {
return err
}
if modelImage == "" {
modelImage = modelName
}

ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()

model, err := getOllama(ctx, o.dynamicClient, o.userSpecifiedNamespace, modelImage)
model, err := getOllama(ctx, o.dynamicClient, o.userSpecifiedNamespace, modelName)
if err != nil {
return err
}
if model != nil {
if model.Spec.Image == modelImage {
fmt.Println(modelImage, "deploy")
if model.Spec.Image == modelName {
fmt.Println(modelName, "deployed")
return nil
}

model.Spec.Image = modelImage
model.Spec.Image = modelName

unstructuredObj, err := Unstructured(model)
if err != nil {
Expand All @@ -139,7 +165,7 @@ func (o *CmdDeployOptions) runE(cmd *cobra.Command, args []string) error {
return err
}

fmt.Println(modelImage, "updated")
fmt.Println(modelName, "updated")

return nil
}
Expand All @@ -150,7 +176,7 @@ func (o *CmdDeployOptions) runE(cmd *cobra.Command, args []string) error {
Kind: "Model",
},
ObjectMeta: metav1.ObjectMeta{
Name: modelImage,
Name: modelName,
},
Spec: ollamav1.ModelSpec{
Image: modelImage,
Expand All @@ -170,7 +196,7 @@ func (o *CmdDeployOptions) runE(cmd *cobra.Command, args []string) error {
return err
}

fmt.Println(modelImage, "deployed")
fmt.Println(modelName, "deployed")

return nil
}
38 changes: 22 additions & 16 deletions internal/cli/kollama/cmd_undeploy.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,21 +19,18 @@ import (
)

const (
unDeployExample = `# Undeploy a phi model
unDeployExample = `
# Undeploy a phi model
$ kollama undeploy phi
kollama undeploy phi
# or if using as kubectl plugin
$ kubectl ollama undeploy phi
or
# Undeploy a phi model in a specific namespace
$ kollama undeploy phi -n phi-namespace
kubectl ollama undeploy phi
# Undeploy a phi model in a specific namespace
kollama undeploy phi -n phi-namespace
or
kubectl ollama undeploy phi -n phi-namespace`
# or if using as kubectl plugin
$ kubectl ollama undeploy phi -n phi-namespace`
)

type CmdUndeployOptions struct {
Expand All @@ -59,10 +56,19 @@ func NewCmdUndeploy(streams genericiooptions.IOStreams) *cobra.Command {
o := NewCmdUndeployOptions(streams)

cmd := &cobra.Command{
Use: "undeploy [model name] [flags]",
Short: "Undeploy a model with the given name by using Ollama Operator",
Example: unDeployExample,
SilenceUsage: true,
Use: "undeploy [model name] [flags]",
Short: "Undeploy a model with the given name by using Ollama Operator",
Example: unDeployExample,
Args: func(cmd *cobra.Command, args []string) error {
if len(args) < 1 {
return fmt.Errorf("model name is required")
}
if args[0] == "" {
return fmt.Errorf("model name cannot be empty")
}

return nil
},
RunE: func(c *cobra.Command, args []string) error {
return o.runE(c, args)
},
Expand Down

0 comments on commit 2c7fb80

Please sign in to comment.