Skip to content

Commit

Permalink
Make saxutil create require an admin ACL, check cell existence befo…
Browse files Browse the repository at this point in the history
…re creation/deletion, and write a dummy location.proto for newly created cells.

PiperOrigin-RevId: 586095002
Change-Id: Idba060c7a239214a641ab0305db01a6294a73666
  • Loading branch information
jiawenhao authored and copybara-github committed Nov 28, 2023
1 parent 9dbd276 commit 62f1aa4
Show file tree
Hide file tree
Showing 7 changed files with 66 additions and 38 deletions.
2 changes: 2 additions & 0 deletions saxml/bin/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -55,10 +55,12 @@ go_library(
"//saxml/admin:validator",
"//saxml/client/go:sax",
"//saxml/client/go:saxadmin",
"//saxml/common:addr",
"//saxml/common:cell",
"//saxml/common:config",
"//saxml/common:naming",
"//saxml/common:watchable",
"//saxml/common/platform:env",
"//saxml/protobuf:admin_go_proto_grpc",
"//saxml/protobuf:common_go_proto",
"@com_github_golang_glog//:go_default_library",
Expand Down
64 changes: 47 additions & 17 deletions saxml/bin/saxutil_cmd_admin.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
"io"
"math/rand"
"os"
"path/filepath"
"sort"
"strconv"
"strings"
Expand All @@ -35,17 +36,20 @@ import (
"github.com/olekukonko/tablewriter"
"saxml/admin/validator"
"saxml/client/go/saxadmin"
"saxml/common/addr"
"saxml/common/cell"
"saxml/common/config"
"saxml/common/naming"
"saxml/common/platform/env"
"saxml/common/watchable"

apb "saxml/protobuf/admin_go_proto_grpc"
cpb "saxml/protobuf/common_go_proto"
)

var (
cmdTimeout = flag.Duration("sax_timeout", 60*time.Second, "How many seconds to wait for command completion.")
cmdTimeout = flag.Duration("sax_timeout", 60*time.Second, "How many seconds to wait for command completion.")
dummyLocation = "No admin server has been started for this Sax cell."
)

// CreateCmd creates a new Sax cell.
Expand All @@ -60,7 +64,7 @@ func (*CreateCmd) Synopsis() string { return "Create a Sax cell." }
// Usage returns the full usage of CreateCmd.
func (*CreateCmd) Usage() string {
return `create <cell name> <file system path> [admin ACL string]:
Create a SAX cell and initialize its state.
Create a Sax cell and initialize its state.
`
}

Expand All @@ -69,30 +73,51 @@ func (c *CreateCmd) SetFlags(f *flag.FlagSet) {}

// Execute executes CreateCmd.
func (c *CreateCmd) Execute(ctx context.Context, f *flag.FlagSet, args ...any) subcommands.ExitStatus {
if len(f.Args()) != 2 && len(f.Args()) != 3 {
log.Errorf("Provide a SAX cell name (e.g. /sax/bar), a file system path for persistence (e.g. gs://bucket/path), and an optional admin ACL string.")
if len(f.Args()) != 3 {
log.Errorf("Provide a Sax cell name (e.g. /sax/bar), a file system path for persistence (e.g. gs://bucket/path), and an admin ACL string.")
return subcommands.ExitUsageError
}
saxCell := f.Args()[0]
fsRoot := f.Args()[1]
var adminACL string
if len(f.Args()) == 3 {
adminACL = f.Args()[2]
adminACL := f.Args()[2]
if adminACL == "" {
log.Errorf("Provide a non-empty admin ACL string.")
return subcommands.ExitFailure
}

if err := cell.Exists(ctx, saxCell); err == nil {
log.Errorf("Sax cell %s already exists.", saxCell)
return subcommands.ExitFailure
}

// If adminACL is not empty, it is set as the write ACL on the created cell subdirectory.
// All files created within, including config.proto and location.proto will inherit this ACL
// as writer.
// adminACL is set as the write ACL on the created cell subdirectory.
if err := cell.Create(ctx, saxCell, adminACL); err != nil {
log.Errorf("Failed to create SAX cell %s: %v", saxCell, err)
log.Errorf("Failed to create Sax cell %s: %v", saxCell, err)
return subcommands.ExitFailure
}
// If adminACL is empty, the config.proto created here will have the Sax dev group as the
// write ACL, to prevent regular users from accidentally deleting others' Sax cells.
// adminACL is set as the write ACL and the AdminAcl field content of config.proto.
if err := config.Create(ctx, saxCell, fsRoot, adminACL); err != nil {
log.Errorf("Failed to create config %s: %v", saxCell, err)
return subcommands.ExitFailure
}
// Write a dummy message to location.proto for users who never ran an admin cell in a newly
// created cell. adminACL is set as the write ACL on location.proto.
path, err := cell.Path(ctx, saxCell)
if err != nil {
log.Errorf("Failed to get path for Sax cell %s: %v", saxCell, err)
return subcommands.ExitFailure
}
fname := filepath.Join(path, addr.LocationFile)
location := &apb.Location{Location: dummyLocation}
content, err := proto.Marshal(location)
if err != nil {
log.Errorf("Failed to marshal location: %v", err)
return subcommands.ExitFailure
}
if err := env.Get().WriteFile(ctx, fname, adminACL, content); err != nil {
log.Errorf("Failed to write location file: %v", err)
return subcommands.ExitFailure
}

return subcommands.ExitSuccess
}
Expand All @@ -109,7 +134,7 @@ func (*DeleteCmd) Synopsis() string { return "Delete a Sax cell." }
// Usage returns the full usage of DeleteCmd.
func (*DeleteCmd) Usage() string {
return `delete <cell name>:
Delete a SAX cell.
Delete a Sax cell.
`
}

Expand All @@ -119,7 +144,7 @@ func (c *DeleteCmd) SetFlags(f *flag.FlagSet) {}
// Execute executes DeleteCmd.
func (c *DeleteCmd) Execute(ctx context.Context, f *flag.FlagSet, args ...any) subcommands.ExitStatus {
if len(f.Args()) != 1 {
log.Errorf("Provide a SAX cell name (e.g. /sax/bar).")
log.Errorf("Provide a Sax cell name (e.g. /sax/bar).")
return subcommands.ExitUsageError
}
saxCell := f.Args()[0]
Expand All @@ -137,8 +162,13 @@ func (c *DeleteCmd) Execute(ctx context.Context, f *flag.FlagSet, args ...any) s
return subcommands.ExitFailure
}

if err := cell.Exists(ctx, saxCell); err != nil {
log.Errorf("Sax cell %s does not exist: %v.", saxCell, err)
return subcommands.ExitFailure
}

if err := cell.Delete(ctx, saxCell); err != nil {
log.Errorf("Failed to delete SAX cell %s: %v", saxCell, err)
log.Errorf("Failed to delete Sax cell %s: %v", saxCell, err)
return subcommands.ExitFailure
}

Expand Down Expand Up @@ -214,7 +244,7 @@ func (*ListCmd) Synopsis() string { return "List published models." }
// Usage returns the full usage of ListCmd.
func (*ListCmd) Usage() string {
return `ls /sax[/cell[/model]]:
List all sax cells, models in a cell, or information of a model.
List all Sax cells, models in a cell, or information of a model.
`
}

Expand Down
1 change: 1 addition & 0 deletions saxml/common/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@ go_library(
srcs = ["config.go"],
deps = [
":cell",
":errors",
"//saxml/common/platform:env",
"//saxml/protobuf:admin_go_proto_grpc",
"@com_github_golang_glog//:go_default_library",
Expand Down
7 changes: 5 additions & 2 deletions saxml/common/cell.go
Original file line number Diff line number Diff line change
Expand Up @@ -71,14 +71,17 @@ func ListAll(ctx context.Context) ([]string, error) {
}

// Create creates a Sax cell.
func Create(ctx context.Context, saxCell, adminACL string) error {
func Create(ctx context.Context, saxCell, writeACL string) error {
path, err := Path(ctx, saxCell)
if err != nil {
return err
}
if !env.Get().InTest(ctx) && writeACL == "" {
return fmt.Errorf("cell must be created with a non-empty write ACL: %w", errors.ErrInvalidArgument)
}

log.Infof("Creating directory %q", path)
if err := env.Get().CreateDir(ctx, path, adminACL); err != nil {
if err := env.Get().CreateDir(ctx, path, writeACL); err != nil {
return err
}

Expand Down
11 changes: 5 additions & 6 deletions saxml/common/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,14 @@ package config

import (
"context"
"fmt"
"path/filepath"

log "github.com/golang/glog"
"google.golang.org/protobuf/encoding/prototext"
"google.golang.org/protobuf/proto"
"saxml/common/cell"
"saxml/common/errors"
"saxml/common/platform/env"
pb "saxml/protobuf/admin_go_proto_grpc"
)
Expand Down Expand Up @@ -87,16 +89,13 @@ func Watch(ctx context.Context, saxCell string) (<-chan *pb.Config, error) {
}

// Save saves the server config.
func Save(ctx context.Context, config *pb.Config, saxCell string, adminACL string) error {
func Save(ctx context.Context, config *pb.Config, saxCell string, writeACL string) error {
fname, err := configFileName(ctx, saxCell)
if err != nil {
return err
}
// If adminACL is not set, we should at least make the Sax dev group the write ACL on config.proto
// in prod environment to prevent regular users from accidentally deleting others' Sax cells.
writeACL := adminACL
if writeACL == "" {
writeACL = env.Get().DefaultWriteACL(ctx)
if !env.Get().InTest(ctx) && writeACL == "" {
return fmt.Errorf("config must be saved with a non-empty write ACL: %w", errors.ErrInvalidArgument)
}
in, err := proto.Marshal(config)
if err != nil {
Expand Down
15 changes: 5 additions & 10 deletions saxml/common/platform/cloud.go
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,7 @@ func (e *Env) ReadCachedFile(ctx context.Context, path string) ([]byte, error) {
// WriteFile writes the content of a file.
func (e *Env) WriteFile(ctx context.Context, path, writeACL string, data []byte) error {
if writeACL != "" {
return fmt.Errorf("WriteFile with ACL is not supported: %w", errors.ErrUnimplemented)
log.Warningf("Ignoring file write ACL: %s", writeACL)
}
if strings.HasPrefix(path, gcsPathPrefix) {
_, object, err := gcsBucketAndObject(ctx, path)
Expand Down Expand Up @@ -330,9 +330,9 @@ func (e *Env) FsRootDir(fsRoot string) string {
}

// CreateDir creates a directory.
func (e *Env) CreateDir(ctx context.Context, path, acl string) error {
if acl != "" {
return fmt.Errorf("CreateDir with ACL is not supported: %w", errors.ErrUnimplemented)
func (e *Env) CreateDir(ctx context.Context, path, writeACL string) error {
if writeACL != "" {
log.Warningf("Ignoring directory write ACL: %s", writeACL)
}
if strings.HasPrefix(path, gcsPathPrefix) {
_, object, err := gcsBucketAndObject(ctx, filepath.Join(path, metadataFile))
Expand Down Expand Up @@ -418,16 +418,11 @@ func (e *Env) DirExists(ctx context.Context, path string) (bool, error) {
return false, err
}

// DefaultWriteACL returns the default write ACL.
func (e *Env) DefaultWriteACL(ctx context.Context) string {
return ""
}

// CheckACLs returns nil iff the given principal passes an ACL check.
func (e *Env) CheckACLs(principal string, acls []string) error {
for _, acl := range acls {
if acl != "" {
return fmt.Errorf("ACL check is not supported: %w", errors.ErrUnimplemented)
log.Warningf("Ignoring ACL: %s", acl)
}
}
return nil
Expand Down
4 changes: 1 addition & 3 deletions saxml/common/platform/env.go
Original file line number Diff line number Diff line change
Expand Up @@ -87,16 +87,14 @@ type Env interface {
// FsRootDir reformats the directory path where the admin server periodically dumps its state.
FsRootDir(fsRoot string) string
// CreateDir creates a directory.
CreateDir(ctx context.Context, path, acl string) error
CreateDir(ctx context.Context, path, writeACL string) error
// DeleteDir deletes a directory.
DeleteDir(ctx context.Context, path string) error
// ListSubdirs lists subdirectories in a directory.
ListSubdirs(ctx context.Context, path string) ([]string, error)
// DirExists checks the existence of a directory.
DirExists(ctx context.Context, path string) (bool, error)

// DefaultWriteACL returns the default write ACL.
DefaultWriteACL(ctx context.Context) string
// CheckACLs returns nil iff the given principal passes an ACL check.
CheckACLs(principal string, acls []string) error
// ValidateACLName returns nil iff the given aclname is valid and exists.
Expand Down

0 comments on commit 62f1aa4

Please sign in to comment.