Exemplo n.º 1
0
func testSupervisor(t *testing.T, nWorkers int, test func(*Supervisor)) {
	tempdir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Error(err)
		return
	}
	defer os.RemoveAll(tempdir)

	client, err := tlsconfig.NewClient(tempdir)
	if err != nil {
		t.Error(err)
		return
	}

	testStorage = NewTestStorage()
	c := SupervisorConfig{
		Storage:       testStorage,
		TLSHandshaker: client,
		// Discard logs for tests
		Logger: log.New(ioutil.Discard, "", 0),
	}
	supervisor, err := NewSupervisor(c)
	if err != nil {
		t.Errorf("failed to create supervisor: %v", err)
		return
	}
	// discard logs for tests
	supervisor.Logger = log.New(ioutil.Discard, "", 0)

	for i := 0; i < nWorkers; i++ {
		mps, err := mps.NewMPS()
		if err != nil {
			t.Error(err)
			return
		}
		defer mps.Reset()
		// configure server's TLS config directly rather than with a handshake
		server := httptest.NewUnstartedServer(mps)
		server.TLS, err = client.ServerConfig(net.ParseIP("127.0.0.1"))
		if err != nil {
			t.Error(err)
			return
		}
		server.StartTLS()
		defer server.Close()

		addr := strings.TrimPrefix(server.URL, "https://")

		if err := supervisor.AddWorker(addr, int64(i)); err != nil {
			t.Errorf("could not add worker: %v", err)
			return
		}
	}
	test(supervisor)
}
Exemplo n.º 2
0
func TestSupervisorReboot(t *testing.T) {

	tempdir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Error(err)
		return
	}
	defer os.RemoveAll(tempdir)

	client, err := tlsconfig.NewClient(tempdir)
	if err != nil {
		t.Error(err)
		return
	}

	testStorage = NewTestStorage()
	c := SupervisorConfig{
		Storage:       testStorage,
		TLSHandshaker: client,
		// Discard logs for tests
		Logger: log.New(ioutil.Discard, "", 0),
	}
	supervisor1, err := NewSupervisor(c)
	if err != nil {
		t.Errorf("failed to create supervisor: %v", err)
		return
	}

	addrs := [5]string{"", "", "", "", ""}

	for i := 0; i < 5; i++ {
		mps, err := mps.NewMPS()
		if err != nil {
			t.Error(err)
			return
		}
		defer mps.Reset()
		// configure server's TLS config directly rather than with a handshake
		server := httptest.NewUnstartedServer(mps)
		server.TLS, err = client.ServerConfig(net.ParseIP("127.0.0.1"))
		if err != nil {
			t.Error(err)
			return
		}
		server.StartTLS()
		defer server.Close()

		addrs[i] = strings.TrimPrefix(server.URL, "https://")
	}

	type Model struct {
		User, Model string
		Version     int
		Instances   []int64
	}
	deployments := map[int64]*Model{}

	for i, addr := range addrs {
		if err := supervisor1.AddWorker(addr, int64(i)); err != nil {
			t.Errorf("could not add worker: %v", err)
			return
		}
	}

	user := "******"
	// Deploy 5 different models
	modelnames := []string{
		"hellor_0", "hellor_1", "hellor_2", "hellor_3", "hellor_4",
	}

	// Override the storage's NewDeployment function so we can track which
	// deployment is associated with which model version.
	// This is normally done by the database.
	nInstances := 2
	var nextId int64 = 1
	var nextDeployId int64 = 1
	testStorage.newDeployment = func(user, model string, version int) (deployId int64, instIds []int64, err error) {
		deployId = nextDeployId
		nextDeployId++

		ids := make([]int64, nInstances)
		instIds = make([]int64, nInstances)
		for i := range ids {
			ids[i] = nextId
			instIds[i] = nextId
			nextId++
		}

		deployments[deployId] = &Model{user, model, version, ids}
		return deployId, instIds, nil
	}
	for _, model := range modelnames {
		if err := supervisor1.Deploy(user, model, 1); err != nil {
			t.Errorf("could not deploy: %v", err)
			return
		}
	}

	nExp := len(modelnames)

	if n := len(deployments); n != nExp {
		t.Errorf("expected %d deployments, got: %d", nExp, n)
		return
	}
	for _, model := range deployments {
		if n := len(model.Instances); n != nInstances {
			t.Errorf("expected %d deployments per model, got %d", nInstances, n)
			return
		}
	}

	testPred := func(super *Supervisor) {
		for _, model := range modelnames {
			hf := func(w http.ResponseWriter, r *http.Request) {
				super.Predict(user, model, w, r)
			}
			s := httptest.NewServer(http.HandlerFunc(hf))
			defer s.Close()
			for i := 0; i < 100; i++ {
				body := bytes.NewReader([]byte(`{"name":"bigdatabob"}`))
				resp, err := http.Post(s.URL, "application/json", body)
				if err != nil {
					t.Errorf("request failed: %v", err)
					return
				}
				respbody, err := ioutil.ReadAll(resp.Body)
				resp.Body.Close()
				if err != nil {
					t.Errorf("could not read body: %v", err)
					return
				}
				if resp.StatusCode != http.StatusOK {
					t.Errorf("expected 200, got %s: %s", resp.Status, respbody)
					return
				}
			}
		}
	}
	testPred(supervisor1)

	supervisor1.shutdown()

	workers := make([]db.Worker, len(addrs))
	for i, addr := range addrs {
		workers[i] = db.Worker{int64(i), addr}
	}

	deploymentReqs := make([]db.DeploymentReq, len(deployments))
	i := 0
	for deployId, model := range deployments {
		deploymentReqs[i] = db.DeploymentReq{
			Username:         model.User,
			Modelname:        model.Model,
			Version:          model.Version,
			LastDeployId:     deployId,
			ValidInstanceIds: model.Instances,
		}
		i++
	}

	c2 := SupervisorConfig{
		Storage:       testStorage,
		TLSHandshaker: client,
		Workers:       workers,
		Deployments:   deploymentReqs,
		// Discard logs for tests
		Logger: log.New(ioutil.Discard, "", 0),
	}
	supervisor2, err := NewSupervisor(c2)
	if err != nil {
		t.Errorf("failed to start second supervisor: %v", err)
		return
	}
	// discard logs for tests
	supervisor2.Logger = log.New(ioutil.Discard, "", 0)

	testPred(supervisor2)
}
Exemplo n.º 3
0
func NewApp(config *AppConfig) (*App, error) {

	if config.ModelReplication < 1 {
		return nil, fmt.Errorf("model replication level cannot be less than 1")
	}
	var accessLog io.Writer
	if config.AccessLog == "" {
		accessLog = os.Stderr
	} else {
		flags := os.O_WRONLY | os.O_APPEND | os.O_CREATE
		file, err := os.OpenFile(config.AccessLog, flags, 0644)
		if err != nil {
			return nil, err
		}
		log.Println("using access log:", config.AccessLog)
		accessLog = file
	}

	appDB, err := sqlutil.NewReconnectingDB("mysql", config.DBConnStr)
	if err != nil {
		return nil, fmt.Errorf("could not connect to database: %v", err)
	}

	// do the static directories exist?
	for _, dir := range []string{"css", "js", "fonts", "img", "model"} {
		d := filepath.Join(config.StaticDir, dir)
		_, err := os.Stat(d)
		if err != nil {
			return nil, fmt.Errorf("directory %s doesn't exist", d)
		}
	}

	serveDir := func(dirName string) http.Handler {
		d := filepath.Join(config.StaticDir, dirName)
		return http.StripPrefix("/"+dirName, http.FileServer(http.Dir(d)))
	}

	store := sessions.NewCookieStore([]byte(config.SecretSessionKey))

	client, err := tlsconfig.NewClient(config.CertDir)
	if err != nil {
		return nil, fmt.Errorf("certificate loading failed: %v", err)
	}

	app := App{
		store:            store,
		db:               appDB,
		staticDir:        config.StaticDir,
		bundleDir:        config.BundleDir,
		modelLogsDir:     config.ModelLogs,
		modelReplication: config.ModelReplication,
		isdev:            config.IsDev,
		serviceName:      config.ServiceName,
		sharingDisabled:  config.DisableSharing,
	}

	tx, err := app.db.Begin()
	if err != nil {
		return nil, err
	}
	workers, err := db.Workers(tx)
	if err != nil {
		return nil, err
	}
	deploymentReqs, err := db.DeploymentRequests(tx)
	if err != nil {
		return nil, err
	}
	tx.Rollback()
	c := alb.SupervisorConfig{
		Storage:       &storage{new(sync.Mutex), &app},
		TLSHandshaker: client,
		Workers:       workers,
		Deployments:   deploymentReqs,
	}
	app.sup, err = alb.NewSupervisor(c)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize communcation with MPSs: %v", err)
	}
	app.sup.Logger = log.New(os.Stderr, "supervisor: ", log.LstdFlags)

	r := mux.NewRouter()
	r.PathPrefix("/css/").Handler(serveDir("css"))
	r.PathPrefix("/js/").Handler(serveDir("js"))
	r.PathPrefix("/fonts/").Handler(serveDir("fonts"))
	r.PathPrefix("/img/").Handler(serveDir("img"))

	r.HandleFunc("/login", app.handleLogin)
	r.HandleFunc("/register", app.handleRegister)
	r.HandleFunc("/verify-password", app.handleVerifyPassword)
	r.HandleFunc("/logout", app.handleLogout)

	r.Handle("/favicon.ico", app.serveFile("favicon.ico"))

	// deployment routes
	r.PathPrefix("/deployer").HandlerFunc(app.handleDeployment)
	r.HandleFunc("/verify", app.handleOldVerify)

	authedRouter := mux.NewRouter()
	authedRouter.NotFoundHandler = app.serveFile("404.html")

	authedRouter.PathPrefix("/model/").Handler(serveDir("model"))
	authedRouter.Handle("/", app.serveFile("index.html"))
	authedRouter.Handle("/account", app.serveFile("account.html"))
	authedRouter.HandleFunc("/user.json", app.handleUser)
	authedRouter.HandleFunc("/users/{name}", app.handleUserByName)

	// apikey stuff
	authedRouter.HandleFunc("/apikey", app.handleApikey)

	// model specific pages
	authedRouter.Handle("/models/{name}", app.serveFile("model/index.html"))
	authedRouter.HandleFunc("/models/{name}/json", app.handleModel)
	authedRouter.Handle("/models/{name}/scoring", app.serveFile("model/scoring.html"))
	authedRouter.Handle("/models/{name}/versions", app.serveFile("model/versions.html"))
	authedRouter.HandleFunc("/models/{name}/versions.json", app.handleModelVersions)
	authedRouter.Handle("/models/{name}/logs", app.serveFile("model/logs.html"))
	authedRouter.HandleFunc("/models/{name}/logs/json", app.handleModelLogs)
	authedRouter.Handle("/models/{name}/settings", app.serveFile("model/settings.html"))
	authedRouter.Handle("/models/{name}/form-builder", app.serveFile("model/form-builder.html"))
	authedRouter.HandleFunc("/models/{name}/redeploy/{version}", app.handleModelRedeploy)
	authedRouter.HandleFunc("/models/{name}/shared", app.handleModelSharedUsers)
	authedRouter.HandleFunc("/models/{name}/startshare/{user}", app.handleModelStartSharing)
	authedRouter.HandleFunc("/models/{name}/stopshare/{user}", app.handleModelStopSharing)
	authedRouter.HandleFunc("/models/{name}/example", app.handleModelExample)

	// model actions
	authedRouter.HandleFunc("/models/{name}/action/{action}", app.handleModelStateChange)

	authedRouter.HandleFunc("/models", app.handleUserModels)
	authedRouter.HandleFunc("/shared", app.handleSharedModels)
	authedRouter.HandleFunc("/whoami", app.handleWhoami)

	adminRouter := mux.NewRouter()
	authedRouter.PathPrefix("/admin").Handler(app.restrictAdmin(adminRouter))

	adminRouter.Handle("/admin", app.serveFile("admin/index.html"))
	adminRouter.Handle("/admin/users", app.serveFile("admin/users.html"))
	adminRouter.HandleFunc("/admin/users.json", app.handleUsersData)
	adminRouter.HandleFunc("/admin/users/create", app.handleUsersCreate)
	adminRouter.HandleFunc("/admin/users/delete", app.handleUserDelete)
	adminRouter.HandleFunc("/admin/users/makeadmin", app.handleUserMakeAdmin)
	adminRouter.HandleFunc("/admin/users/unmakeadmin", app.handleUserUnmakeAdmin)
	adminRouter.HandleFunc("/admin/users/setpass", app.handleUserSetPass)

	adminRouter.Handle("/admin/models", app.serveFile("admin/models.html"))
	adminRouter.HandleFunc("/admin/models.json", app.handleAllUserModels)

	adminRouter.HandleFunc("/admin/servers", app.handleServers)
	adminRouter.HandleFunc("/admin/servers.json", app.handleServersData)
	adminRouter.HandleFunc("/admin/servers/remove", app.handleServersRemove)

	// get initial authentication for models
	auth, err := func() (*db.PredictionAuth, error) {
		tx, err := app.db.Begin()
		if err != nil {
			return nil, err
		}
		defer tx.Rollback()
		return db.GetAuth(tx)
	}()
	if err != nil {
		return nil, err
	}

	// must enforce no caching for jsx pages
	noCaching := func(handler http.Handler) http.Handler {
		hf := func(w http.ResponseWriter, r *http.Request) {
			// see: http://goo.gl/itaIDo
			w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
			w.Header().Set("Pragma", "no-cache")
			w.Header().Set("Expires", "0")

			handler.ServeHTTP(w, r)
		}
		return http.HandlerFunc(hf)
	}

	h := handlers.LoggingHandler(accessLog, noCaching(app.restrict(authedRouter)))
	r.NotFoundHandler = h

	predRouter := alb.NewPredictionRouter(r, app.sup, auth)

	// TODO: add api routes

	app.router = predRouter

	return &app, nil
}