Merge pull request #340 from kthcloud/dev

fix gpu not attaching, and clean up gpu model client

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)

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,
+	}
+}

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 {

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)
 		}
 

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)
 		}