mirror of
https://github.com/google/ax.git
synced 2026-10-20 05:39:15 +08:00
JBD
864488c67d
Transition the AX architecture from asynchronous event queues (Redis Streams and the `ax-controller` worker pool) to direct, synchronous reconciliation in `ax-server` guarded by fine-grained distributed locks. Key changes: - Implement per-resource locking (`lock:task`, `lock:workspace`, `lock:model`) using Redis distributed locks (`SET NX PX` with token validation) and an in-memory locker for testing. - Move Substrate reconciliation directly into `ax-server` for Task, Workspace, and Model lifecycle methods (`Create`, `Resume`, `Suspend`, `Delete`). - Decommission `ax-controller`, `deploy/ax-controller.yaml`, and Redis Stream event publishing/consuming interfaces. - Simplify `ax delete` and `ax watch` in the CLI to eliminate two-phase polling loops now that operations complete synchronously. - Update `demo.sh`, build configurations (`Makefile`, `.ko.yaml`), and documentation to reflect the single-binary control plane architecture.
603 lines
18 KiB
Go
603 lines
18 KiB
Go
// Copyright 2026 Google LLC
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package server
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log/slog"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/google/ax/internal/lock"
|
|
"github.com/google/ax/internal/store"
|
|
"github.com/google/ax/pkg/apis/v1alpha1"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
// Reconciler coordinates sandbox/actor lifecycles on Agent Substrate directly.
|
|
type Reconciler interface {
|
|
Reconcile(ctx context.Context, task *v1alpha1.Task, workspaces ...*v1alpha1.Workspace) (*v1alpha1.Task, error)
|
|
ReconcileDelete(ctx context.Context, atespace, taskName string) error
|
|
}
|
|
|
|
// Options configures the AX API Server.
|
|
type Options struct {
|
|
Locker lock.Locker
|
|
Reconciler Reconciler
|
|
}
|
|
|
|
// Server provides the gRPC API for AX.
|
|
type Server struct {
|
|
v1alpha1.UnimplementedAXServer
|
|
store store.Store
|
|
locker lock.Locker
|
|
reconciler Reconciler
|
|
grpcServer *grpc.Server
|
|
}
|
|
|
|
// NewServer creates a new AX API server.
|
|
func NewServer(s store.Store, opts ...Options) *Server {
|
|
var opt Options
|
|
if len(opts) > 0 {
|
|
opt = opts[0]
|
|
}
|
|
locker := opt.Locker
|
|
if locker == nil {
|
|
locker = lock.NewMemoryLocker()
|
|
}
|
|
|
|
srv := &Server{
|
|
store: s,
|
|
locker: locker,
|
|
reconciler: opt.Reconciler,
|
|
grpcServer: grpc.NewServer(),
|
|
}
|
|
v1alpha1.RegisterAXServer(srv.grpcServer, srv)
|
|
return srv
|
|
}
|
|
|
|
// GRPCServer returns the underlying gRPC server instance.
|
|
func (s *Server) GRPCServer() *grpc.Server {
|
|
return s.grpcServer
|
|
}
|
|
|
|
// Handler returns the HTTP handler for the server, routing gRPC and HTTP health checks.
|
|
func (s *Server) Handler() http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") {
|
|
s.grpcServer.ServeHTTP(w, r)
|
|
return
|
|
}
|
|
if r.URL.Path == "/healthz" {
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte("ok\n"))
|
|
return
|
|
}
|
|
http.NotFound(w, r)
|
|
})
|
|
}
|
|
|
|
// --- gRPC AXServer implementation ---
|
|
|
|
// --- Tasks ---
|
|
|
|
func (s *Server) GetTask(ctx context.Context, req *v1alpha1.GetTaskRequest) (*v1alpha1.Task, error) {
|
|
if req == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "missing request")
|
|
}
|
|
atespace := req.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
}
|
|
task, err := s.store.GetTask(ctx, atespace, req.Name)
|
|
if err != nil {
|
|
if errors.Is(err, store.ErrNotFound) {
|
|
return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", req.Name, atespace)
|
|
}
|
|
return nil, status.Errorf(codes.Internal, "getting task: %v", err)
|
|
}
|
|
return task, nil
|
|
}
|
|
|
|
func (s *Server) ListTasks(ctx context.Context, req *v1alpha1.ListTasksRequest) (*v1alpha1.ListTasksResponse, error) {
|
|
atespace := ""
|
|
limit := int64(50)
|
|
offset := int64(0)
|
|
if req != nil {
|
|
atespace = req.Atespace
|
|
if req.Limit > 0 {
|
|
limit = req.Limit
|
|
}
|
|
if req.Offset >= 0 {
|
|
offset = req.Offset
|
|
}
|
|
}
|
|
tasks, err := s.store.ListTasks(ctx, atespace, limit, offset)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Internal, "listing tasks: %v", err)
|
|
}
|
|
return &v1alpha1.ListTasksResponse{Tasks: tasks}, nil
|
|
}
|
|
|
|
func (s *Server) CreateTask(ctx context.Context, req *v1alpha1.CreateTaskRequest) (*v1alpha1.Task, error) {
|
|
if req == nil || req.Task == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "task required")
|
|
}
|
|
task := req.Task
|
|
if err := v1alpha1.ValidateTask(task); err != nil {
|
|
return nil, status.Error(codes.InvalidArgument, err.Error())
|
|
}
|
|
if task.Metadata == nil {
|
|
task.Metadata = &v1alpha1.ObjectMeta{}
|
|
}
|
|
atespace := task.Metadata.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
task.Metadata.Atespace = atespace
|
|
}
|
|
taskName := task.Metadata.GetName()
|
|
|
|
// Acquire exclusive lock for this task
|
|
unlock, err := s.locker.Lock(ctx, "task", atespace, taskName)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Aborted, "locking task %s/%s: %v", atespace, taskName, err)
|
|
}
|
|
defer unlock()
|
|
|
|
_, err = s.store.GetTask(ctx, atespace, taskName)
|
|
if err == nil {
|
|
return nil, status.Errorf(codes.FailedPrecondition, "task %s/%s already exists and is immutable", atespace, taskName)
|
|
}
|
|
if !errors.Is(err, store.ErrNotFound) {
|
|
return nil, status.Errorf(codes.Internal, "checking existing task: %v", err)
|
|
}
|
|
|
|
if task.Metadata.CreationTimestamp == nil {
|
|
task.Metadata.CreationTimestamp = timestamppb.Now()
|
|
}
|
|
if task.Status == nil {
|
|
task.Status = &v1alpha1.TaskStatus{}
|
|
}
|
|
task.Status.Phase = "Suspended"
|
|
if err := s.store.SaveTask(ctx, task); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "saving task: %v", err)
|
|
}
|
|
|
|
// Directly reconcile with Substrate
|
|
if s.reconciler != nil {
|
|
workspaces := s.fetchWorkspaces(ctx, atespace, task)
|
|
reconciled, err := s.reconciler.Reconcile(ctx, task, workspaces...)
|
|
if err != nil {
|
|
slog.Error("direct reconcile error on create task", "task", taskName, "error", err)
|
|
task.Status.Phase = "Failed"
|
|
_ = s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status)
|
|
return nil, status.Errorf(codes.Internal, "provisioning task on substrate: %v", err)
|
|
}
|
|
task.Status = reconciled.Status
|
|
if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "updating task status: %v", err)
|
|
}
|
|
}
|
|
|
|
return task, nil
|
|
}
|
|
|
|
func (s *Server) DeleteTask(ctx context.Context, req *v1alpha1.DeleteTaskRequest) (*v1alpha1.DeleteTaskResponse, error) {
|
|
if req == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "missing request")
|
|
}
|
|
atespace := req.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
}
|
|
taskName := req.Name
|
|
|
|
// Acquire exclusive lock for this task
|
|
unlock, err := s.locker.Lock(ctx, "task", atespace, taskName)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Aborted, "locking task %s/%s: %v", atespace, taskName, err)
|
|
}
|
|
defer unlock()
|
|
|
|
task, err := s.store.GetTask(ctx, atespace, taskName)
|
|
if err != nil {
|
|
if errors.Is(err, store.ErrNotFound) {
|
|
return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", taskName, atespace)
|
|
}
|
|
return nil, status.Errorf(codes.Internal, "getting task: %v", err)
|
|
}
|
|
|
|
// Mark the task as terminating while actor cleanup is in progress.
|
|
if task.Status == nil {
|
|
task.Status = &v1alpha1.TaskStatus{}
|
|
}
|
|
task.Status.Phase = v1alpha1.PhaseTerminating
|
|
if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "updating task status: %v", err)
|
|
}
|
|
|
|
// Directly clean up Substrate actor and templates
|
|
if s.reconciler != nil {
|
|
if err := s.reconciler.ReconcileDelete(ctx, atespace, taskName); err != nil {
|
|
if errors.Is(err, context.Canceled) {
|
|
return nil, status.Error(codes.Canceled, err.Error())
|
|
}
|
|
if errors.Is(err, context.DeadlineExceeded) {
|
|
return nil, status.Error(codes.DeadlineExceeded, err.Error())
|
|
}
|
|
if st, ok := status.FromError(err); ok && st.Code() != codes.Unknown {
|
|
return nil, err
|
|
}
|
|
return nil, status.Errorf(codes.Internal, "cleaning up substrate resources for task %s/%s: %v", atespace, taskName, err)
|
|
}
|
|
}
|
|
|
|
if err := s.store.DeleteTask(ctx, atespace, taskName); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "deleting task record: %v", err)
|
|
}
|
|
return &v1alpha1.DeleteTaskResponse{}, nil
|
|
}
|
|
|
|
func (s *Server) SuspendTask(ctx context.Context, req *v1alpha1.SuspendTaskRequest) (*v1alpha1.Task, error) {
|
|
if req == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "missing request")
|
|
}
|
|
atespace := req.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
}
|
|
taskName := req.Name
|
|
|
|
// Acquire exclusive lock for this task
|
|
unlock, err := s.locker.Lock(ctx, "task", atespace, taskName)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Aborted, "locking task %s/%s: %v", atespace, taskName, err)
|
|
}
|
|
defer unlock()
|
|
|
|
task, err := s.store.GetTask(ctx, atespace, taskName)
|
|
if err != nil {
|
|
if errors.Is(err, store.ErrNotFound) {
|
|
return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", taskName, atespace)
|
|
}
|
|
return nil, status.Errorf(codes.Internal, "getting task: %v", err)
|
|
}
|
|
if task.Status == nil {
|
|
task.Status = &v1alpha1.TaskStatus{}
|
|
}
|
|
task.Status.Phase = "Suspended"
|
|
|
|
if s.reconciler != nil {
|
|
workspaces := s.fetchWorkspaces(ctx, atespace, task)
|
|
reconciled, err := s.reconciler.Reconcile(ctx, task, workspaces...)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Internal, "suspending task on substrate: %v", err)
|
|
}
|
|
task.Status = reconciled.Status
|
|
if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "updating task status: %v", err)
|
|
}
|
|
} else {
|
|
if err := s.store.SaveTask(ctx, task); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "suspending task: %v", err)
|
|
}
|
|
}
|
|
|
|
return task, nil
|
|
}
|
|
|
|
func (s *Server) ResumeTask(ctx context.Context, req *v1alpha1.ResumeTaskRequest) (*v1alpha1.Task, error) {
|
|
if req == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "missing request")
|
|
}
|
|
atespace := req.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
}
|
|
taskName := req.Name
|
|
|
|
// Acquire exclusive lock for this task
|
|
unlock, err := s.locker.Lock(ctx, "task", atespace, taskName)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Aborted, "locking task %s/%s: %v", atespace, taskName, err)
|
|
}
|
|
defer unlock()
|
|
|
|
task, err := s.store.GetTask(ctx, atespace, taskName)
|
|
if err != nil {
|
|
if errors.Is(err, store.ErrNotFound) {
|
|
return nil, status.Errorf(codes.NotFound, "task %q not found in atespace %q", taskName, atespace)
|
|
}
|
|
return nil, status.Errorf(codes.Internal, "getting task: %v", err)
|
|
}
|
|
if task.Status == nil {
|
|
task.Status = &v1alpha1.TaskStatus{}
|
|
}
|
|
task.Status.Phase = "Running"
|
|
|
|
if s.reconciler != nil {
|
|
workspaces := s.fetchWorkspaces(ctx, atespace, task)
|
|
reconciled, err := s.reconciler.Reconcile(ctx, task, workspaces...)
|
|
if err != nil {
|
|
task.Status.Phase = "Failed"
|
|
_ = s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status)
|
|
return nil, status.Errorf(codes.Internal, "resuming task on substrate: %v", err)
|
|
}
|
|
task.Status = reconciled.Status
|
|
if err := s.store.UpdateTaskStatus(ctx, atespace, taskName, task.Status); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "updating task status: %v", err)
|
|
}
|
|
} else {
|
|
if err := s.store.SaveTask(ctx, task); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "resuming task: %v", err)
|
|
}
|
|
}
|
|
|
|
return task, nil
|
|
}
|
|
|
|
func (s *Server) fetchWorkspaces(ctx context.Context, atespace string, task *v1alpha1.Task) []*v1alpha1.Workspace {
|
|
var workspaces []*v1alpha1.Workspace
|
|
if task.Spec == nil {
|
|
return workspaces
|
|
}
|
|
for _, ref := range task.Spec.WorkspaceRefs() {
|
|
if ref.Name == "" {
|
|
continue
|
|
}
|
|
if wsp, err := s.store.GetWorkspace(ctx, atespace, ref.Name); err == nil {
|
|
workspaces = append(workspaces, wsp)
|
|
}
|
|
}
|
|
return workspaces
|
|
}
|
|
|
|
func (s *Server) WatchTask(req *v1alpha1.WatchTaskRequest, stream grpc.ServerStreamingServer[v1alpha1.WatchTaskResponse]) error {
|
|
if req == nil {
|
|
return status.Error(codes.InvalidArgument, "missing request")
|
|
}
|
|
atespace := req.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
}
|
|
ctx := stream.Context()
|
|
ch, closer, err := s.store.WatchTask(ctx, atespace, req.Name)
|
|
if err != nil {
|
|
return status.Errorf(codes.Internal, "watching task: %v", err)
|
|
}
|
|
defer closer.Close()
|
|
|
|
if initial, err := s.store.GetTask(ctx, atespace, req.Name); err == nil {
|
|
if err := stream.Send(&v1alpha1.WatchTaskResponse{Task: initial, Action: "INITIAL"}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case task, ok := <-ch:
|
|
if !ok {
|
|
return nil
|
|
}
|
|
if err := stream.Send(&v1alpha1.WatchTaskResponse{Task: task, Action: "MODIFIED"}); err != nil {
|
|
return err
|
|
}
|
|
if task.Status != nil && (task.Status.Phase == "Failed" || task.Status.Phase == "Completed") {
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// --- Workspaces ---
|
|
|
|
func (s *Server) GetWorkspace(ctx context.Context, req *v1alpha1.GetWorkspaceRequest) (*v1alpha1.Workspace, error) {
|
|
if req == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "missing request")
|
|
}
|
|
atespace := req.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
}
|
|
ws, err := s.store.GetWorkspace(ctx, atespace, req.Name)
|
|
if err != nil {
|
|
if errors.Is(err, store.ErrNotFound) {
|
|
return nil, status.Errorf(codes.NotFound, "workspace %q not found in atespace %q", req.Name, atespace)
|
|
}
|
|
return nil, status.Errorf(codes.Internal, "getting workspace: %v", err)
|
|
}
|
|
return ws, nil
|
|
}
|
|
|
|
func (s *Server) ListWorkspaces(ctx context.Context, req *v1alpha1.ListWorkspacesRequest) (*v1alpha1.ListWorkspacesResponse, error) {
|
|
atespace := ""
|
|
if req != nil {
|
|
atespace = req.Atespace
|
|
}
|
|
workspaces, err := s.store.ListWorkspaces(ctx, atespace)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Internal, "listing workspaces: %v", err)
|
|
}
|
|
return &v1alpha1.ListWorkspacesResponse{Workspaces: workspaces}, nil
|
|
}
|
|
|
|
func (s *Server) UpdateWorkspace(ctx context.Context, req *v1alpha1.UpdateWorkspaceRequest) (*v1alpha1.Workspace, error) {
|
|
if req == nil || req.Workspace == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "workspace required")
|
|
}
|
|
if err := v1alpha1.ValidateWorkspace(req.Workspace); err != nil {
|
|
return nil, status.Error(codes.InvalidArgument, err.Error())
|
|
}
|
|
req.Workspace.Metadata = defaultMetadata(req.Workspace.Metadata, func(atespace, name string) *v1alpha1.ObjectMeta {
|
|
existing, err := s.store.GetWorkspace(ctx, atespace, name)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
return existing.GetMetadata()
|
|
})
|
|
|
|
atespace := req.Workspace.Metadata.Atespace
|
|
wsName := req.Workspace.Metadata.Name
|
|
|
|
// Acquire exclusive lock for this workspace
|
|
unlock, err := s.locker.Lock(ctx, "workspace", atespace, wsName)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Aborted, "locking workspace %s/%s: %v", atespace, wsName, err)
|
|
}
|
|
defer unlock()
|
|
|
|
if err := s.store.SaveWorkspace(ctx, req.Workspace); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "saving workspace: %v", err)
|
|
}
|
|
return req.Workspace, nil
|
|
}
|
|
|
|
func (s *Server) DeleteWorkspace(ctx context.Context, req *v1alpha1.DeleteWorkspaceRequest) (*v1alpha1.DeleteWorkspaceResponse, error) {
|
|
if req == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "missing request")
|
|
}
|
|
atespace := req.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
}
|
|
wsName := req.Name
|
|
|
|
// Acquire exclusive lock for this workspace
|
|
unlock, err := s.locker.Lock(ctx, "workspace", atespace, wsName)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Aborted, "locking workspace %s/%s: %v", atespace, wsName, err)
|
|
}
|
|
defer unlock()
|
|
|
|
if err := s.store.DeleteWorkspace(ctx, atespace, wsName); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "deleting workspace: %v", err)
|
|
}
|
|
return &v1alpha1.DeleteWorkspaceResponse{}, nil
|
|
}
|
|
|
|
// --- Models ---
|
|
|
|
func (s *Server) GetModel(ctx context.Context, req *v1alpha1.GetModelRequest) (*v1alpha1.Model, error) {
|
|
if req == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "missing request")
|
|
}
|
|
atespace := req.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
}
|
|
model, err := s.store.GetModel(ctx, atespace, req.Name)
|
|
if err != nil {
|
|
if errors.Is(err, store.ErrNotFound) {
|
|
return nil, status.Errorf(codes.NotFound, "model %q not found in atespace %q", req.Name, atespace)
|
|
}
|
|
return nil, status.Errorf(codes.Internal, "getting model: %v", err)
|
|
}
|
|
return model, nil
|
|
}
|
|
|
|
func (s *Server) ListModels(ctx context.Context, req *v1alpha1.ListModelsRequest) (*v1alpha1.ListModelsResponse, error) {
|
|
atespace := ""
|
|
if req != nil {
|
|
atespace = req.Atespace
|
|
}
|
|
models, err := s.store.ListModels(ctx, atespace)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Internal, "listing models: %v", err)
|
|
}
|
|
return &v1alpha1.ListModelsResponse{Models: models}, nil
|
|
}
|
|
|
|
func (s *Server) UpdateModel(ctx context.Context, req *v1alpha1.UpdateModelRequest) (*v1alpha1.Model, error) {
|
|
if req == nil || req.Model == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "model required")
|
|
}
|
|
if err := v1alpha1.ValidateModel(req.Model); err != nil {
|
|
return nil, status.Error(codes.InvalidArgument, err.Error())
|
|
}
|
|
req.Model.Metadata = defaultMetadata(req.Model.Metadata, func(atespace, name string) *v1alpha1.ObjectMeta {
|
|
existing, err := s.store.GetModel(ctx, atespace, name)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
return existing.GetMetadata()
|
|
})
|
|
|
|
atespace := req.Model.Metadata.Atespace
|
|
modelName := req.Model.Metadata.Name
|
|
|
|
// Acquire exclusive lock for this model
|
|
unlock, err := s.locker.Lock(ctx, "model", atespace, modelName)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Aborted, "locking model %s/%s: %v", atespace, modelName, err)
|
|
}
|
|
defer unlock()
|
|
|
|
if err := s.store.SaveModel(ctx, req.Model); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "saving model: %v", err)
|
|
}
|
|
return req.Model, nil
|
|
}
|
|
|
|
func (s *Server) DeleteModel(ctx context.Context, req *v1alpha1.DeleteModelRequest) (*v1alpha1.DeleteModelResponse, error) {
|
|
if req == nil {
|
|
return nil, status.Error(codes.InvalidArgument, "missing request")
|
|
}
|
|
atespace := req.Atespace
|
|
if atespace == "" {
|
|
atespace = "default"
|
|
}
|
|
modelName := req.Name
|
|
|
|
// Acquire exclusive lock for this model
|
|
unlock, err := s.locker.Lock(ctx, "model", atespace, modelName)
|
|
if err != nil {
|
|
return nil, status.Errorf(codes.Aborted, "locking model %s/%s: %v", atespace, modelName, err)
|
|
}
|
|
defer unlock()
|
|
|
|
if err := s.store.DeleteModel(ctx, atespace, modelName); err != nil {
|
|
return nil, status.Errorf(codes.Internal, "deleting model: %v", err)
|
|
}
|
|
return &v1alpha1.DeleteModelResponse{}, nil
|
|
}
|
|
|
|
func defaultMetadata(meta *v1alpha1.ObjectMeta, existing func(atespace, name string) *v1alpha1.ObjectMeta) *v1alpha1.ObjectMeta {
|
|
if meta == nil {
|
|
meta = &v1alpha1.ObjectMeta{}
|
|
}
|
|
if meta.Atespace == "" {
|
|
meta.Atespace = "default"
|
|
}
|
|
if meta.CreationTimestamp == nil {
|
|
var prev *v1alpha1.ObjectMeta
|
|
if existing != nil {
|
|
prev = existing(meta.Atespace, meta.Name)
|
|
}
|
|
if prev.GetCreationTimestamp() != nil {
|
|
meta.CreationTimestamp = prev.GetCreationTimestamp()
|
|
} else {
|
|
meta.CreationTimestamp = timestamppb.Now()
|
|
}
|
|
}
|
|
return meta
|
|
}
|
|
|