Skip to content

Fix logic for when to stream output when running models #65

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mattt merged 2 commits into from
Feb 10, 2024
Merged

Fix logic for when to stream output when running models #65

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
775ef8d
Fix logic for when to stream output when running models
mattt Feb 10, 2024
441d2df
Reorder conditionals
mattt Feb 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 20 additions & 8 deletions internal/cmd/prediction/create.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,8 +78,15 @@ var CreateCmd = &cobra.Command{
return fmt.Errorf("failed to coerce inputs: %w", err)
}

shouldWait := cmd.Flags().Changed("wait") || !cmd.Flags().Changed("no-wait")
shouldStream := !cmd.Flags().Changed("wait") && cmd.Flags().Changed("stream") || (outputSchema != nil && outputSchema.Type == "array" && outputSchema.Items.Value.Type == "string" && outputSchema.Items.Value.Format != "uri")
shouldWait := (cmd.Flags().Changed("wait") || !cmd.Flags().Changed("no-wait"))

canStream := (outputSchema != nil &&
outputSchema.Type == "array" &&
outputSchema.Items.Value.Type == "string" &&
outputSchema.Extensions["x-cog-array-type"] == "iterator" &&
outputSchema.Extensions["x-cog-array-display"] == "concatenate")
shouldStream := canStream && !cmd.Flags().Changed("wait") &&
(cmd.Flags().Changed("stream") || !cmd.Flags().Changed("no-stream"))

s.Start()
var prediction *replicate.Prediction
Expand Down Expand Up @@ -304,15 +311,20 @@ func init() {

func AddCreateFlags(cmd *cobra.Command) {
cmd.Flags().Bool("json", false, "Emit JSON")
cmd.Flags().Bool("no-wait", false, "Don't wait for prediction to complete")
cmd.Flags().Bool("web", false, "View on web")
cmd.MarkFlagsMutuallyExclusive("json", "web")

cmd.Flags().BoolP("wait", "w", true, "Wait for prediction to complete")
cmd.Flags().Bool("no-wait", false, "Don't wait for prediction to complete")
cmd.MarkFlagsMutuallyExclusive("wait", "no-wait")

cmd.Flags().Bool("stream", false, "Stream prediction output")
cmd.Flags().Bool("web", false, "View on web")
cmd.Flags().Bool("no-stream", false, "Don't stream prediction output")
cmd.MarkFlagsMutuallyExclusive("stream", "no-stream")
cmd.MarkFlagsMutuallyExclusive("stream", "wait")

cmd.Flags().String("separator", "=", "Separator between input key and value")

cmd.Flags().Bool("save", false, "Save prediction outputs to directory")
cmd.Flags().String("output-directory", "", "Output directory, defaults to ./{prediction-id}")

cmd.MarkFlagsMutuallyExclusive("json", "web")
cmd.MarkFlagsMutuallyExclusive("stream", "wait")
cmd.MarkFlagsMutuallyExclusive("wait", "no-wait")
}