diff --git a/saxml/client/go/BUILD b/saxml/client/go/BUILD index 0dd1cbb..55b15fc 100644 --- a/saxml/client/go/BUILD +++ b/saxml/client/go/BUILD @@ -87,6 +87,7 @@ go_library( name = "saxadmin", srcs = ["admin.go"], deps = [ + "//saxml/admin:validator", "//saxml/common:addr", "//saxml/common:config", "//saxml/common:errors", @@ -101,6 +102,7 @@ go_library( "@org_golang_google_grpc//:go_default_library", "@org_golang_google_grpc//codes:go_default_library", "@org_golang_google_grpc//status:go_default_library", + "@org_golang_google_protobuf//proto", ], ) diff --git a/saxml/client/go/admin.go b/saxml/client/go/admin.go index c641f74..486445e 100644 --- a/saxml/client/go/admin.go +++ b/saxml/client/go/admin.go @@ -28,6 +28,8 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/grpc" "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + "saxml/admin/validator" "saxml/common/addr" "saxml/common/config" "saxml/common/errors" @@ -267,6 +269,36 @@ func (a *Admin) GetSaxCellACL(ctx context.Context, saxCell string) (string, erro return cfg.GetAdminAcl(), nil } +// SetSaxCellACL sets the ACL for the sax cell. +func (a *Admin) SetSaxCellACL(ctx context.Context, saxCell string, acl string) error { + _, err := naming.NewCellFullName(saxCell) + if err != nil { + log.ErrorContextf(ctx, "Invalid sax cell: %v", err) + return err + } + + cfg, err := config.Load(ctx, saxCell) + if err != nil { + log.ErrorContextf(ctx, "Failed to load config: %v", err) + return err + } + log.InfoContextf(ctx, "Current config definition:\n%v", cfg) + + change := proto.Clone(cfg).(*pb.Config) + change.AdminAcl = acl + log.InfoContextf(ctx, "Updated config definition:\n%v", change) + + if err := validator.ValidateConfigUpdate(cfg, change); err != nil { + log.ErrorContextf(ctx, "Invalid config update: %v", err) + return err + } + if err := config.Save(ctx, change, saxCell, acl); err != nil { + log.ErrorContextf(ctx, "Failed to save config: %v", err) + return err + } + return nil +} + // addrReplica maintains a set of server addresses for a model. type addrReplica struct { modelID string