databricks-cli/cmd/workspace/serving-endpoints/serving-endpoints.go

1236 lines
34 KiB
Go
Executable File

// Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
package serving_endpoints
import (
"fmt"
"time"
"github.com/databricks/cli/cmd/root"
"github.com/databricks/cli/libs/cmdio"
"github.com/databricks/cli/libs/flags"
"github.com/databricks/databricks-sdk-go/service/serving"
"github.com/spf13/cobra"
)
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var cmdOverrides []func(*cobra.Command)
func New() *cobra.Command {
cmd := &cobra.Command{
Use: "serving-endpoints",
Short: `The Serving Endpoints API allows you to create, update, and delete model serving endpoints.`,
Long: `The Serving Endpoints API allows you to create, update, and delete model
serving endpoints.
You can use a serving endpoint to serve models from the Databricks Model
Registry or from Unity Catalog. Endpoints expose the underlying models as
scalable REST API endpoints using serverless compute. This means the endpoints
and associated compute resources are fully managed by Databricks and will not
appear in your cloud account. A serving endpoint can consist of one or more
MLflow models from the Databricks Model Registry, called served entities. A
serving endpoint can have at most ten served entities. You can configure
traffic settings to define how requests should be routed to your served
entities behind an endpoint. Additionally, you can configure the scale of
resources that should be applied to each served entity.`,
GroupID: "serving",
Annotations: map[string]string{
"package": "serving",
},
}
// Add methods
cmd.AddCommand(newBuildLogs())
cmd.AddCommand(newCreate())
cmd.AddCommand(newDelete())
cmd.AddCommand(newExportMetrics())
cmd.AddCommand(newGet())
cmd.AddCommand(newGetOpenApi())
cmd.AddCommand(newGetPermissionLevels())
cmd.AddCommand(newGetPermissions())
cmd.AddCommand(newList())
cmd.AddCommand(newLogs())
cmd.AddCommand(newPatch())
cmd.AddCommand(newPut())
cmd.AddCommand(newPutAiGateway())
cmd.AddCommand(newQuery())
cmd.AddCommand(newSetPermissions())
cmd.AddCommand(newUpdateConfig())
cmd.AddCommand(newUpdatePermissions())
// Apply optional overrides to this command.
for _, fn := range cmdOverrides {
fn(cmd)
}
return cmd
}
// start build-logs command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var buildLogsOverrides []func(
*cobra.Command,
*serving.BuildLogsRequest,
)
func newBuildLogs() *cobra.Command {
cmd := &cobra.Command{}
var buildLogsReq serving.BuildLogsRequest
// TODO: short flags
cmd.Use = "build-logs NAME SERVED_MODEL_NAME"
cmd.Short = `Get build logs for a served model.`
cmd.Long = `Get build logs for a served model.
Retrieves the build logs associated with the provided served model.
Arguments:
NAME: The name of the serving endpoint that the served model belongs to. This
field is required.
SERVED_MODEL_NAME: The name of the served model that build logs will be retrieved for. This
field is required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(2)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
buildLogsReq.Name = args[0]
buildLogsReq.ServedModelName = args[1]
response, err := w.ServingEndpoints.BuildLogs(ctx, buildLogsReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range buildLogsOverrides {
fn(cmd, &buildLogsReq)
}
return cmd
}
// start create command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var createOverrides []func(
*cobra.Command,
*serving.CreateServingEndpoint,
)
func newCreate() *cobra.Command {
cmd := &cobra.Command{}
var createReq serving.CreateServingEndpoint
var createJson flags.JsonFlag
var createSkipWait bool
var createTimeout time.Duration
cmd.Flags().BoolVar(&createSkipWait, "no-wait", createSkipWait, `do not wait to reach NOT_UPDATING state`)
cmd.Flags().DurationVar(&createTimeout, "timeout", 20*time.Minute, `maximum amount of time to reach NOT_UPDATING state`)
// TODO: short flags
cmd.Flags().Var(&createJson, "json", `either inline JSON string or @path/to/file.json with request body`)
// TODO: complex arg: ai_gateway
// TODO: array: rate_limits
cmd.Flags().BoolVar(&createReq.RouteOptimized, "route-optimized", createReq.RouteOptimized, `Enable route optimization for the serving endpoint.`)
// TODO: array: tags
cmd.Use = "create"
cmd.Short = `Create a new serving endpoint.`
cmd.Long = `Create a new serving endpoint.`
cmd.Annotations = make(map[string]string)
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
if cmd.Flags().Changed("json") {
diags := createJson.Unmarshal(&createReq)
if diags.HasError() {
return diags.Error()
}
if len(diags) > 0 {
err := cmdio.RenderDiagnosticsToErrorOut(ctx, diags)
if err != nil {
return err
}
}
} else {
return fmt.Errorf("please provide command input in JSON format by specifying the --json flag")
}
wait, err := w.ServingEndpoints.Create(ctx, createReq)
if err != nil {
return err
}
if createSkipWait {
return cmdio.Render(ctx, wait.Response)
}
spinner := cmdio.Spinner(ctx)
info, err := wait.OnProgress(func(i *serving.ServingEndpointDetailed) {
status := i.State.ConfigUpdate
statusMessage := fmt.Sprintf("current status: %s", status)
spinner <- statusMessage
}).GetWithTimeout(createTimeout)
close(spinner)
if err != nil {
return err
}
return cmdio.Render(ctx, info)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range createOverrides {
fn(cmd, &createReq)
}
return cmd
}
// start delete command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var deleteOverrides []func(
*cobra.Command,
*serving.DeleteServingEndpointRequest,
)
func newDelete() *cobra.Command {
cmd := &cobra.Command{}
var deleteReq serving.DeleteServingEndpointRequest
// TODO: short flags
cmd.Use = "delete NAME"
cmd.Short = `Delete a serving endpoint.`
cmd.Long = `Delete a serving endpoint.
Arguments:
NAME: The name of the serving endpoint. This field is required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
deleteReq.Name = args[0]
err = w.ServingEndpoints.Delete(ctx, deleteReq)
if err != nil {
return err
}
return nil
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range deleteOverrides {
fn(cmd, &deleteReq)
}
return cmd
}
// start export-metrics command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var exportMetricsOverrides []func(
*cobra.Command,
*serving.ExportMetricsRequest,
)
func newExportMetrics() *cobra.Command {
cmd := &cobra.Command{}
var exportMetricsReq serving.ExportMetricsRequest
// TODO: short flags
cmd.Use = "export-metrics NAME"
cmd.Short = `Get metrics of a serving endpoint.`
cmd.Long = `Get metrics of a serving endpoint.
Retrieves the metrics associated with the provided serving endpoint in either
Prometheus or OpenMetrics exposition format.
Arguments:
NAME: The name of the serving endpoint to retrieve metrics for. This field is
required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
exportMetricsReq.Name = args[0]
response, err := w.ServingEndpoints.ExportMetrics(ctx, exportMetricsReq)
if err != nil {
return err
}
defer response.Contents.Close()
return cmdio.Render(ctx, response.Contents)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range exportMetricsOverrides {
fn(cmd, &exportMetricsReq)
}
return cmd
}
// start get command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var getOverrides []func(
*cobra.Command,
*serving.GetServingEndpointRequest,
)
func newGet() *cobra.Command {
cmd := &cobra.Command{}
var getReq serving.GetServingEndpointRequest
// TODO: short flags
cmd.Use = "get NAME"
cmd.Short = `Get a single serving endpoint.`
cmd.Long = `Get a single serving endpoint.
Retrieves the details for a single serving endpoint.
Arguments:
NAME: The name of the serving endpoint. This field is required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
getReq.Name = args[0]
response, err := w.ServingEndpoints.Get(ctx, getReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range getOverrides {
fn(cmd, &getReq)
}
return cmd
}
// start get-open-api command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var getOpenApiOverrides []func(
*cobra.Command,
*serving.GetOpenApiRequest,
)
func newGetOpenApi() *cobra.Command {
cmd := &cobra.Command{}
var getOpenApiReq serving.GetOpenApiRequest
// TODO: short flags
cmd.Use = "get-open-api NAME"
cmd.Short = `Get the schema for a serving endpoint.`
cmd.Long = `Get the schema for a serving endpoint.
Get the query schema of the serving endpoint in OpenAPI format. The schema
contains information for the supported paths, input and output format and
datatypes.
Arguments:
NAME: The name of the serving endpoint that the served model belongs to. This
field is required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
getOpenApiReq.Name = args[0]
err = w.ServingEndpoints.GetOpenApi(ctx, getOpenApiReq)
if err != nil {
return err
}
return nil
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range getOpenApiOverrides {
fn(cmd, &getOpenApiReq)
}
return cmd
}
// start get-permission-levels command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var getPermissionLevelsOverrides []func(
*cobra.Command,
*serving.GetServingEndpointPermissionLevelsRequest,
)
func newGetPermissionLevels() *cobra.Command {
cmd := &cobra.Command{}
var getPermissionLevelsReq serving.GetServingEndpointPermissionLevelsRequest
// TODO: short flags
cmd.Use = "get-permission-levels SERVING_ENDPOINT_ID"
cmd.Short = `Get serving endpoint permission levels.`
cmd.Long = `Get serving endpoint permission levels.
Gets the permission levels that a user can have on an object.
Arguments:
SERVING_ENDPOINT_ID: The serving endpoint for which to get or manage permissions.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
getPermissionLevelsReq.ServingEndpointId = args[0]
response, err := w.ServingEndpoints.GetPermissionLevels(ctx, getPermissionLevelsReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range getPermissionLevelsOverrides {
fn(cmd, &getPermissionLevelsReq)
}
return cmd
}
// start get-permissions command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var getPermissionsOverrides []func(
*cobra.Command,
*serving.GetServingEndpointPermissionsRequest,
)
func newGetPermissions() *cobra.Command {
cmd := &cobra.Command{}
var getPermissionsReq serving.GetServingEndpointPermissionsRequest
// TODO: short flags
cmd.Use = "get-permissions SERVING_ENDPOINT_ID"
cmd.Short = `Get serving endpoint permissions.`
cmd.Long = `Get serving endpoint permissions.
Gets the permissions of a serving endpoint. Serving endpoints can inherit
permissions from their root object.
Arguments:
SERVING_ENDPOINT_ID: The serving endpoint for which to get or manage permissions.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
getPermissionsReq.ServingEndpointId = args[0]
response, err := w.ServingEndpoints.GetPermissions(ctx, getPermissionsReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range getPermissionsOverrides {
fn(cmd, &getPermissionsReq)
}
return cmd
}
// start list command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var listOverrides []func(
*cobra.Command,
)
func newList() *cobra.Command {
cmd := &cobra.Command{}
cmd.Use = "list"
cmd.Short = `Get all serving endpoints.`
cmd.Long = `Get all serving endpoints.`
cmd.Annotations = make(map[string]string)
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
response := w.ServingEndpoints.List(ctx)
return cmdio.RenderIterator(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range listOverrides {
fn(cmd)
}
return cmd
}
// start logs command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var logsOverrides []func(
*cobra.Command,
*serving.LogsRequest,
)
func newLogs() *cobra.Command {
cmd := &cobra.Command{}
var logsReq serving.LogsRequest
// TODO: short flags
cmd.Use = "logs NAME SERVED_MODEL_NAME"
cmd.Short = `Get the latest logs for a served model.`
cmd.Long = `Get the latest logs for a served model.
Retrieves the service logs associated with the provided served model.
Arguments:
NAME: The name of the serving endpoint that the served model belongs to. This
field is required.
SERVED_MODEL_NAME: The name of the served model that logs will be retrieved for. This field
is required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(2)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
logsReq.Name = args[0]
logsReq.ServedModelName = args[1]
response, err := w.ServingEndpoints.Logs(ctx, logsReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range logsOverrides {
fn(cmd, &logsReq)
}
return cmd
}
// start patch command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var patchOverrides []func(
*cobra.Command,
*serving.PatchServingEndpointTags,
)
func newPatch() *cobra.Command {
cmd := &cobra.Command{}
var patchReq serving.PatchServingEndpointTags
var patchJson flags.JsonFlag
// TODO: short flags
cmd.Flags().Var(&patchJson, "json", `either inline JSON string or @path/to/file.json with request body`)
// TODO: array: add_tags
// TODO: array: delete_tags
cmd.Use = "patch NAME"
cmd.Short = `Update tags of a serving endpoint.`
cmd.Long = `Update tags of a serving endpoint.
Used to batch add and delete tags from a serving endpoint with a single API
call.
Arguments:
NAME: The name of the serving endpoint who's tags to patch. This field is
required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
if cmd.Flags().Changed("json") {
diags := patchJson.Unmarshal(&patchReq)
if diags.HasError() {
return diags.Error()
}
if len(diags) > 0 {
err := cmdio.RenderDiagnosticsToErrorOut(ctx, diags)
if err != nil {
return err
}
}
}
patchReq.Name = args[0]
response, err := w.ServingEndpoints.Patch(ctx, patchReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range patchOverrides {
fn(cmd, &patchReq)
}
return cmd
}
// start put command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var putOverrides []func(
*cobra.Command,
*serving.PutRequest,
)
func newPut() *cobra.Command {
cmd := &cobra.Command{}
var putReq serving.PutRequest
var putJson flags.JsonFlag
// TODO: short flags
cmd.Flags().Var(&putJson, "json", `either inline JSON string or @path/to/file.json with request body`)
// TODO: array: rate_limits
cmd.Use = "put NAME"
cmd.Short = `Update rate limits of a serving endpoint.`
cmd.Long = `Update rate limits of a serving endpoint.
Used to update the rate limits of a serving endpoint. NOTE: Only foundation
model endpoints are currently supported. For external models, use AI Gateway
to manage rate limits.
Arguments:
NAME: The name of the serving endpoint whose rate limits are being updated. This
field is required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
if cmd.Flags().Changed("json") {
diags := putJson.Unmarshal(&putReq)
if diags.HasError() {
return diags.Error()
}
if len(diags) > 0 {
err := cmdio.RenderDiagnosticsToErrorOut(ctx, diags)
if err != nil {
return err
}
}
}
putReq.Name = args[0]
response, err := w.ServingEndpoints.Put(ctx, putReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range putOverrides {
fn(cmd, &putReq)
}
return cmd
}
// start put-ai-gateway command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var putAiGatewayOverrides []func(
*cobra.Command,
*serving.PutAiGatewayRequest,
)
func newPutAiGateway() *cobra.Command {
cmd := &cobra.Command{}
var putAiGatewayReq serving.PutAiGatewayRequest
var putAiGatewayJson flags.JsonFlag
// TODO: short flags
cmd.Flags().Var(&putAiGatewayJson, "json", `either inline JSON string or @path/to/file.json with request body`)
// TODO: complex arg: guardrails
// TODO: complex arg: inference_table_config
// TODO: array: rate_limits
// TODO: complex arg: usage_tracking_config
cmd.Use = "put-ai-gateway NAME"
cmd.Short = `Update AI Gateway of a serving endpoint.`
cmd.Long = `Update AI Gateway of a serving endpoint.
Used to update the AI Gateway of a serving endpoint. NOTE: Only external model
endpoints are currently supported.
Arguments:
NAME: The name of the serving endpoint whose AI Gateway is being updated. This
field is required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
if cmd.Flags().Changed("json") {
diags := putAiGatewayJson.Unmarshal(&putAiGatewayReq)
if diags.HasError() {
return diags.Error()
}
if len(diags) > 0 {
err := cmdio.RenderDiagnosticsToErrorOut(ctx, diags)
if err != nil {
return err
}
}
}
putAiGatewayReq.Name = args[0]
response, err := w.ServingEndpoints.PutAiGateway(ctx, putAiGatewayReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range putAiGatewayOverrides {
fn(cmd, &putAiGatewayReq)
}
return cmd
}
// start query command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var queryOverrides []func(
*cobra.Command,
*serving.QueryEndpointInput,
)
func newQuery() *cobra.Command {
cmd := &cobra.Command{}
var queryReq serving.QueryEndpointInput
var queryJson flags.JsonFlag
// TODO: short flags
cmd.Flags().Var(&queryJson, "json", `either inline JSON string or @path/to/file.json with request body`)
// TODO: array: dataframe_records
// TODO: complex arg: dataframe_split
// TODO: map via StringToStringVar: extra_params
// TODO: any: input
// TODO: any: inputs
// TODO: array: instances
cmd.Flags().IntVar(&queryReq.MaxTokens, "max-tokens", queryReq.MaxTokens, `The max tokens field used ONLY for __completions__ and __chat external & foundation model__ serving endpoints.`)
// TODO: array: messages
cmd.Flags().IntVar(&queryReq.N, "n", queryReq.N, `The n (number of candidates) field used ONLY for __completions__ and __chat external & foundation model__ serving endpoints.`)
// TODO: any: prompt
// TODO: array: stop
cmd.Flags().BoolVar(&queryReq.Stream, "stream", queryReq.Stream, `The stream field used ONLY for __completions__ and __chat external & foundation model__ serving endpoints.`)
cmd.Flags().Float64Var(&queryReq.Temperature, "temperature", queryReq.Temperature, `The temperature field used ONLY for __completions__ and __chat external & foundation model__ serving endpoints.`)
cmd.Use = "query NAME"
cmd.Short = `Query a serving endpoint.`
cmd.Long = `Query a serving endpoint.
Arguments:
NAME: The name of the serving endpoint. This field is required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
if cmd.Flags().Changed("json") {
diags := queryJson.Unmarshal(&queryReq)
if diags.HasError() {
return diags.Error()
}
if len(diags) > 0 {
err := cmdio.RenderDiagnosticsToErrorOut(ctx, diags)
if err != nil {
return err
}
}
}
queryReq.Name = args[0]
response, err := w.ServingEndpoints.Query(ctx, queryReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range queryOverrides {
fn(cmd, &queryReq)
}
return cmd
}
// start set-permissions command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var setPermissionsOverrides []func(
*cobra.Command,
*serving.ServingEndpointPermissionsRequest,
)
func newSetPermissions() *cobra.Command {
cmd := &cobra.Command{}
var setPermissionsReq serving.ServingEndpointPermissionsRequest
var setPermissionsJson flags.JsonFlag
// TODO: short flags
cmd.Flags().Var(&setPermissionsJson, "json", `either inline JSON string or @path/to/file.json with request body`)
// TODO: array: access_control_list
cmd.Use = "set-permissions SERVING_ENDPOINT_ID"
cmd.Short = `Set serving endpoint permissions.`
cmd.Long = `Set serving endpoint permissions.
Sets permissions on a serving endpoint. Serving endpoints can inherit
permissions from their root object.
Arguments:
SERVING_ENDPOINT_ID: The serving endpoint for which to get or manage permissions.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
if cmd.Flags().Changed("json") {
diags := setPermissionsJson.Unmarshal(&setPermissionsReq)
if diags.HasError() {
return diags.Error()
}
if len(diags) > 0 {
err := cmdio.RenderDiagnosticsToErrorOut(ctx, diags)
if err != nil {
return err
}
}
}
setPermissionsReq.ServingEndpointId = args[0]
response, err := w.ServingEndpoints.SetPermissions(ctx, setPermissionsReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range setPermissionsOverrides {
fn(cmd, &setPermissionsReq)
}
return cmd
}
// start update-config command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var updateConfigOverrides []func(
*cobra.Command,
*serving.EndpointCoreConfigInput,
)
func newUpdateConfig() *cobra.Command {
cmd := &cobra.Command{}
var updateConfigReq serving.EndpointCoreConfigInput
var updateConfigJson flags.JsonFlag
var updateConfigSkipWait bool
var updateConfigTimeout time.Duration
cmd.Flags().BoolVar(&updateConfigSkipWait, "no-wait", updateConfigSkipWait, `do not wait to reach NOT_UPDATING state`)
cmd.Flags().DurationVar(&updateConfigTimeout, "timeout", 20*time.Minute, `maximum amount of time to reach NOT_UPDATING state`)
// TODO: short flags
cmd.Flags().Var(&updateConfigJson, "json", `either inline JSON string or @path/to/file.json with request body`)
// TODO: complex arg: auto_capture_config
// TODO: array: served_entities
// TODO: array: served_models
// TODO: complex arg: traffic_config
cmd.Use = "update-config NAME"
cmd.Short = `Update config of a serving endpoint.`
cmd.Long = `Update config of a serving endpoint.
Updates any combination of the serving endpoint's served entities, the compute
configuration of those served entities, and the endpoint's traffic config. An
endpoint that already has an update in progress can not be updated until the
current update completes or fails.
Arguments:
NAME: The name of the serving endpoint to update. This field is required.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
if cmd.Flags().Changed("json") {
diags := updateConfigJson.Unmarshal(&updateConfigReq)
if diags.HasError() {
return diags.Error()
}
if len(diags) > 0 {
err := cmdio.RenderDiagnosticsToErrorOut(ctx, diags)
if err != nil {
return err
}
}
}
updateConfigReq.Name = args[0]
wait, err := w.ServingEndpoints.UpdateConfig(ctx, updateConfigReq)
if err != nil {
return err
}
if updateConfigSkipWait {
return cmdio.Render(ctx, wait.Response)
}
spinner := cmdio.Spinner(ctx)
info, err := wait.OnProgress(func(i *serving.ServingEndpointDetailed) {
status := i.State.ConfigUpdate
statusMessage := fmt.Sprintf("current status: %s", status)
spinner <- statusMessage
}).GetWithTimeout(updateConfigTimeout)
close(spinner)
if err != nil {
return err
}
return cmdio.Render(ctx, info)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range updateConfigOverrides {
fn(cmd, &updateConfigReq)
}
return cmd
}
// start update-permissions command
// Slice with functions to override default command behavior.
// Functions can be added from the `init()` function in manually curated files in this directory.
var updatePermissionsOverrides []func(
*cobra.Command,
*serving.ServingEndpointPermissionsRequest,
)
func newUpdatePermissions() *cobra.Command {
cmd := &cobra.Command{}
var updatePermissionsReq serving.ServingEndpointPermissionsRequest
var updatePermissionsJson flags.JsonFlag
// TODO: short flags
cmd.Flags().Var(&updatePermissionsJson, "json", `either inline JSON string or @path/to/file.json with request body`)
// TODO: array: access_control_list
cmd.Use = "update-permissions SERVING_ENDPOINT_ID"
cmd.Short = `Update serving endpoint permissions.`
cmd.Long = `Update serving endpoint permissions.
Updates the permissions on a serving endpoint. Serving endpoints can inherit
permissions from their root object.
Arguments:
SERVING_ENDPOINT_ID: The serving endpoint for which to get or manage permissions.`
cmd.Annotations = make(map[string]string)
cmd.Args = func(cmd *cobra.Command, args []string) error {
check := root.ExactArgs(1)
return check(cmd, args)
}
cmd.PreRunE = root.MustWorkspaceClient
cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
ctx := cmd.Context()
w := root.WorkspaceClient(ctx)
if cmd.Flags().Changed("json") {
diags := updatePermissionsJson.Unmarshal(&updatePermissionsReq)
if diags.HasError() {
return diags.Error()
}
if len(diags) > 0 {
err := cmdio.RenderDiagnosticsToErrorOut(ctx, diags)
if err != nil {
return err
}
}
}
updatePermissionsReq.ServingEndpointId = args[0]
response, err := w.ServingEndpoints.UpdatePermissions(ctx, updatePermissionsReq)
if err != nil {
return err
}
return cmdio.Render(ctx, response)
}
// Disable completions since they are not applicable.
// Can be overridden by manual implementation in `override.go`.
cmd.ValidArgsFunction = cobra.NoFileCompletions
// Apply optional overrides to this command.
for _, fn := range updatePermissionsOverrides {
fn(cmd, &updatePermissionsReq)
}
return cmd
}
// end service ServingEndpoints