From 0aeb6c0d341f79b007426c4d0f113d3808fc0628 Mon Sep 17 00:00:00 2001
From: Ce Gao <cegao@tensorchord.ai>
Date: Wed, 2 Aug 2023 11:07:10 +0800
Subject: [PATCH] feat(cli): Support command flag in deploy (#95)

Signed-off-by: Ce Gao <cegao@tensorchord.ai>
---
 mdz/README.md         | 4 +++-
 mdz/pkg/cmd/deploy.go | 6 ++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/mdz/README.md b/mdz/README.md
index a220607..6996112 100644
--- a/mdz/README.md
+++ b/mdz/README.md
@@ -13,7 +13,9 @@ CLI for OpenModelZ.
 
 ## Installation
 
-TODO
+```
+pip install openmodelz
+```
 
 ## CLI Reference
 
diff --git a/mdz/pkg/cmd/deploy.go b/mdz/pkg/cmd/deploy.go
index e050e7d..6ce08bf 100644
--- a/mdz/pkg/cmd/deploy.go
+++ b/mdz/pkg/cmd/deploy.go
@@ -20,6 +20,7 @@ var (
 	deployName        string
 	deployGPU         int
 	deployNodeLabel   []string
+	deployCommand     string
 )
 
 // deployCmd represents the deploy command
@@ -53,6 +54,7 @@ func init() {
 	deployCmd.Flags().IntVar(&deployGPU, "gpu", 0, "Number of GPUs")
 	deployCmd.Flags().StringVar(&deployName, "name", "", "Name of inference")
 	deployCmd.Flags().StringSliceVarP(&deployNodeLabel, "node-labels", "l", []string{}, "Node labels")
+	deployCmd.Flags().StringVar(&deployCommand, "command", "", "Command to run")
 }
 
 func commandDeploy(cmd *cobra.Command, args []string) error {
@@ -87,6 +89,10 @@ func commandDeploy(cmd *cobra.Command, args []string) error {
 		},
 	}
 
+	if deployCommand != "" {
+		inf.Spec.Command = &deployCommand
+	}
+
 	if len(deployNodeLabel) > 0 {
 		inf.Spec.Constraints = []string{}
 		for _, label := range deployNodeLabel {