diff --git a/models/sys/gpu/client.go b/models/sys/gpu/client.go index e202c3fa..1c04b862 100644 --- a/models/sys/gpu/client.go +++ b/models/sys/gpu/client.go @@ -9,16 +9,11 @@ import ( ) type Client struct { - ExcludedHosts []string - ExcludedGPUs []string - resource.ResourceClient[GPU] } func New() *Client { return &Client{ - ExcludedHosts: make([]string, 0), - ExcludedGPUs: make([]string, 0), ResourceClient: resource.ResourceClient[GPU]{ Collection: db.DB.GetCollection("gpus"), }, @@ -43,8 +38,12 @@ func (client *Client) WithExclusion(excludedHosts []string, excludedGPUs []strin excludedGPUs = make([]string, 0) } - client.ExcludedHosts = excludedHosts - client.ExcludedGPUs = excludedGPUs + filter := bson.D{ + {"host", bson.M{"$nin": excludedHosts}}, + {"id", bson.M{"$nin": excludedGPUs}}, + } + + client.ResourceClient.AddExtraFilter(filter) return client } @@ -56,8 +55,6 @@ func (client *Client) OnlyAvailable() *Client { bson.M{"lease.vmId": ""}, bson.M{"lease.end": bson.M{"$lte": time.Now()}}, }}, - {"host", bson.M{"$nin": client.ExcludedHosts}}, - {"id", bson.M{"$nin": client.ExcludedGPUs}}, } client.ResourceClient.AddExtraFilter(filter) diff --git a/models/sys/gpu/dto.go b/models/sys/gpu/dto.go new file mode 100644 index 00000000..f0494caf --- /dev/null +++ b/models/sys/gpu/dto.go @@ -0,0 +1,30 @@ +package gpu + +import ( + "encoding/base64" + "go-deploy/models/dto/body" +) + +func (gpu *GPU) ToDTO(addUserInfo bool) body.GpuRead { + id := base64.StdEncoding.EncodeToString([]byte(gpu.ID)) + + var lease *body.GpuLease + + if gpu.Lease.VmID != "" { + lease = &body.GpuLease{ + End: gpu.Lease.End, + Expired: gpu.Lease.IsExpired(), + } + + if addUserInfo { + lease.User = &gpu.Lease.UserID + lease.VmID = &gpu.Lease.VmID + } + } + + return body.GpuRead{ + ID: id, + Name: gpu.Data.Name, + Lease: lease, + } +} diff --git a/models/sys/gpu/helpers.go b/models/sys/gpu/helpers.go index f0316fca..558baaed 100644 --- a/models/sys/gpu/helpers.go +++ b/models/sys/gpu/helpers.go @@ -2,40 +2,14 @@ package gpu import ( "context" - "encoding/base64" "errors" "fmt" - "go-deploy/models/dto/body" "go-deploy/utils" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "time" ) -func (gpu *GPU) ToDTO(addUserInfo bool) body.GpuRead { - id := base64.StdEncoding.EncodeToString([]byte(gpu.ID)) - - var lease *body.GpuLease - - if gpu.Lease.VmID != "" { - lease = &body.GpuLease{ - End: gpu.Lease.End, - Expired: gpu.Lease.IsExpired(), - } - - if addUserInfo { - lease.User = &gpu.Lease.UserID - lease.VmID = &gpu.Lease.VmID - } - } - - return body.GpuRead{ - ID: id, - Name: gpu.Data.Name, - Lease: lease, - } -} - func (client *Client) Create(id, host string, data GpuData, zone string) error { currentGPU, err := client.GetByID(id) if err != nil { @@ -67,66 +41,6 @@ func (client *Client) Create(id, host string, data GpuData, zone string) error { return nil } -func (client *Client) GetByID(id string) (*GPU, error) { - var gpu GPU - err := client.Collection.FindOne(context.TODO(), bson.D{{"id", id}}).Decode(&gpu) - if err != nil { - if errors.Is(err, mongo.ErrNoDocuments) { - return nil, nil - } - - err = fmt.Errorf("failed to fetch gpu. details: %w", err) - return nil, err - } - - return &gpu, err -} - -func (client *Client) List() ([]GPU, error) { - filter := bson.D{ - {"host", bson.M{"$nin": client.ExcludedHosts}}, - {"data.name", bson.M{"$nin": client.ExcludedGPUs}}, - } - - var gpus []GPU - cursor, err := client.Collection.Find(context.Background(), filter) - if err != nil { - return nil, err - } - - err = cursor.All(context.Background(), &gpus) - if err != nil { - return nil, err - } - - return gpus, nil -} - -func (client *Client) GetAllLeased() ([]GPU, error) { - // filter lease exist and vmId is not empty - filter := bson.D{ - {"$and", []interface{}{ - bson.M{"lease.vmId": bson.M{"$ne": ""}}, - bson.M{"lease": bson.M{"$exists": true}}, - }}, - {"host", bson.M{"$nin": client.ExcludedHosts}}, - {"id", bson.M{"$nin": client.ExcludedGPUs}}, - } - - var gpus []GPU - cursor, err := client.Collection.Find(context.Background(), filter) - if err != nil { - return nil, err - } - - err = cursor.All(context.Background(), &gpus) - if err != nil { - return nil, err - } - - return gpus, nil -} - func (client *Client) Delete(gpuID string) error { err := client.Collection.FindOneAndDelete(context.Background(), bson.D{{"id", gpuID}}).Err() if err != nil { diff --git a/service/vm_service/cs_service/client.go b/service/vm_service/cs_service/client.go index f5bc275e..05d18a77 100644 --- a/service/vm_service/cs_service/client.go +++ b/service/vm_service/cs_service/client.go @@ -82,9 +82,10 @@ func (c *Client) Get(opts *client.Opts) (*vmModels.VM, *cs.Client, *resources.Cs var cc *cs.Client if opts.Client { - // If creating a client and a VM, use the VM's zone. var zone *configModels.VmZone - if vm != nil { + if opts.ExtraOpts.Zone != nil { + zone = opts.ExtraOpts.Zone + } else if vm != nil { zone = config.Config.VM.GetZone(vm.Zone) } diff --git a/service/vm_service/gpu_service.go b/service/vm_service/gpu_service.go index 9a103bb7..b47b1ccf 100644 --- a/service/vm_service/gpu_service.go +++ b/service/vm_service/gpu_service.go @@ -128,8 +128,12 @@ func (c *Client) AttachGPU(vmID string, gpuIDs []string, leaseDuration float64) continue } - err = c.CheckGpuHardwareAvailable(vmID) + err = c.CheckGpuHardwareAvailable(gpuID) if err != nil { + if errors.Is(err, sErrors.GpuNotFoundErr) { + continue + } + return makeError(err) }