retrieve proper instance IDs in GCP

Signed-off-by: Dmitry Shmulevich <[email protected]>

Author: dmitsh

pkg/providers/gcp/imds.go

@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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 gcp
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"strings"
+
+	"k8s.io/klog/v2"
+
+	"github.com/NVIDIA/topograph/internal/exec"
+)
+
+const (
+	IMDSURL = "http://metadata.google.internal/computeMetadata/v1"
+)
+
+func instanceToNodeMap(ctx context.Context, nodes []string) (map[string]string, error) {
+	url := fmt.Sprintf("%s/instance/id", IMDSURL)
+	args := []string{"-w", strings.Join(nodes, ","), fmt.Sprintf("echo $(curl -s  -H \"Metadata-Flavor: Google\" %s)", url)}
+	stdout, err := exec.Exec(ctx, "pdsh", args, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	i2n := map[string]string{}
+	scanner := bufio.NewScanner(stdout)
+	for scanner.Scan() {
+		arr := strings.Split(scanner.Text(), ": ")
+		if len(arr) == 2 {
+			node, instance := arr[0], arr[1]
+			klog.V(4).Infoln("Node name: ", node, "Instance ID: ", instance)
+			i2n[instance] = node
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return nil, err
+	}
+
+	return i2n, nil
+}
+
+func getRegion(ctx context.Context) (string, error) {
+	url := fmt.Sprintf("%s/zone", IMDSURL)
+	args := []string{"-s", "-H", "Metadata-Flavor: Google", url}
+	stdout, err := exec.Exec(ctx, "curl", args, nil)
+	if err != nil {
+		return "", err
+	}
+
+	return stdout.String(), nil
+}

pkg/providers/gcp/instance_topology.go

@@ -19,12 +19,16 @@ package gcp
 import (
 	"context"
 	"fmt"
+	"strconv"
 	"strings"
 	"time"
 
+	compute "cloud.google.com/go/compute/apiv1"
 	"cloud.google.com/go/compute/apiv1/computepb"
 	"cloud.google.com/go/compute/metadata"
+	"github.com/agrea/ptr"
 	"google.golang.org/api/iterator"
+	"k8s.io/klog/v2"
 
 	"github.com/NVIDIA/topograph/pkg/topology"
 )
@@ -39,72 +43,63 @@ type InstanceInfo struct {
 	name      string
 }
 
-func (p *Provider) generateInstanceTopology(ctx context.Context, instanceToNodeMap map[string]string) (*InstanceTopology, error) {
+func (p *Provider) generateInstanceTopology(ctx context.Context, pageSize *int, cis []topology.ComputeInstances) (*InstanceTopology, error) {
+	insTop := &InstanceTopology{
+		instances: []*InstanceInfo{},
+	}
+
+	maxRes := castPageSize(pageSize)
+	for _, ci := range cis {
+		err := p.generateRegionInstanceTopology(ctx, insTop, maxRes, &ci)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return insTop, nil
+}
+
+func (p *Provider) generateRegionInstanceTopology(ctx context.Context, insTop *InstanceTopology, maxRes *uint32, ci *topology.ComputeInstances) error {
 	client, err := p.clientFactory()
 	if err != nil {
-		return nil, err
+		return fmt.Errorf("unable to get client: %v", err)
 	}
 
 	projectID, err := metadata.ProjectIDWithContext(ctx)
 	if err != nil {
-		return nil, fmt.Errorf("unable to get project ID: %s", err.Error())
+		return fmt.Errorf("unable to get project ID: %v", err)
 	}
-	listZoneRequest := computepb.ListZonesRequest{Project: projectID}
-	zones := make([]string, 0)
 
-	timeNow := time.Now()
-	res := client.Zones.List(ctx, &listZoneRequest)
-	requestLatency.WithLabelValues("ListZones").Observe(time.Since(timeNow).Seconds())
+	klog.InfoS("Getting instance topology", "region", ci.Region, "project", projectID)
 
-	for {
-		zone, err := res.Next()
-		if err == iterator.Done {
-			break
-		}
-		zones = append(zones, *zone.Name)
+	req := computepb.ListInstancesRequest{
+		Project:    projectID,
+		Zone:       ci.Region,
+		MaxResults: maxRes,
+		PageToken:  nil,
 	}
 
-	instanceTopology := &InstanceTopology{instances: make([]*InstanceInfo, 0)}
+	var cycle int
+	for {
+		cycle++
+		klog.V(4).Infof("Starting cycle %d", cycle)
 
-	for _, zone := range zones {
 		timeNow := time.Now()
-		listInstanceRequest := computepb.ListInstancesRequest{Project: projectID, Zone: zone}
+		resp := client.Instances.List(ctx, &req)
 		requestLatency.WithLabelValues("ListInstances").Observe(time.Since(timeNow).Seconds())
 
-		resInstance := client.Instances.List(ctx, &listInstanceRequest)
-		for {
-			instance, err := resInstance.Next()
-			if err == iterator.Done {
-				break
-			}
-			_, isNodeInCluster := instanceToNodeMap[*instance.Name]
+		processInstanceList(insTop, resp, ci)
 
-			if instance.ResourceStatus == nil {
-				resourceStatusNotFound.WithLabelValues(*instance.Name).Set(1)
-				continue
-			}
-			resourceStatusNotFound.WithLabelValues(*instance.Name).Set(0)
+		klog.V(4).Infof("Processed %d nodes", len(insTop.instances))
 
-			if instance.ResourceStatus.PhysicalHost == nil {
-				physicalHostNotFound.WithLabelValues(*instance.Name).Set(1)
-				continue
-			}
-			physicalHostNotFound.WithLabelValues(*instance.Name).Set(0)
-
-			if isNodeInCluster {
-				tokens := strings.Split(*instance.ResourceStatus.PhysicalHost, "/")
-				physicalHostIDChunks.WithLabelValues(*instance.Name).Set(float64(getTokenCount(tokens)))
-				instanceObj := &InstanceInfo{
-					name:      *instance.Name,
-					clusterID: tokens[1],
-					rackID:    tokens[2],
-				}
-				instanceTopology.instances = append(instanceTopology.instances, instanceObj)
-			}
+		if token := resp.PageInfo().Token; token == "" {
+			break
+		} else {
+			req.PageToken = &token
 		}
 	}
 
-	return instanceTopology, nil
+	return nil
 }
 
 func (cfg *InstanceTopology) toGraph() (*topology.Vertex, error) {
@@ -156,6 +151,42 @@ func (cfg *InstanceTopology) toGraph() (*topology.Vertex, error) {
 	return root, nil
 }
 
+func processInstanceList(insTop *InstanceTopology, resp *compute.InstanceIterator, ci *topology.ComputeInstances) {
+	for {
+		instance, err := resp.Next()
+		if err == iterator.Done {
+			return
+		}
+		instanceId := strconv.FormatUint(*instance.Id, 10)
+		klog.Infof("Checking INSTANCE %s", instanceId)
+		if _, ok := ci.Instances[instanceId]; ok {
+			klog.Infof("FOUND INSTANCE %s", instanceId)
+			if instance.ResourceStatus == nil {
+				klog.Infof("ResourceStatus is not set for INSTANCE %s", instanceId)
+				resourceStatusNotFound.WithLabelValues(instanceId).Set(1)
+				continue
+			}
+			resourceStatusNotFound.WithLabelValues(instanceId).Set(0)
+
+			if instance.ResourceStatus.PhysicalHost == nil {
+				klog.Infof("PhysicalHost is not set for INSTANCE %s", instanceId)
+				physicalHostNotFound.WithLabelValues(instanceId).Set(1)
+				continue
+			}
+			physicalHostNotFound.WithLabelValues(instanceId).Set(0)
+
+			tokens := strings.Split(*instance.ResourceStatus.PhysicalHost, "/")
+			physicalHostIDChunks.WithLabelValues(instanceId).Set(float64(getTokenCount(tokens)))
+			instanceObj := &InstanceInfo{
+				name:      instanceId,
+				clusterID: tokens[1],
+				rackID:    tokens[2],
+			}
+			insTop.instances = append(insTop.instances, instanceObj)
+		}
+	}
+}
+
 func getTokenCount(tokens []string) int {
 	c := 0
 	for _, q := range tokens {
@@ -165,3 +196,11 @@ func getTokenCount(tokens []string) int {
 	}
 	return c
 }
+
+func castPageSize(val *int) *uint32 {
+	if val == nil {
+		return nil
+	}
+
+	return ptr.Uint32(uint32(*val))
+}

pkg/providers/gcp/instance_topology_test.go

@@ -20,6 +20,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/agrea/ptr"
 	"github.com/stretchr/testify/require"
 )
 
@@ -58,3 +59,25 @@ func TestGetTokenCount(t *testing.T) {
 		})
 	}
 }
+
+func TestCastPageSize(t *testing.T) {
+	testCases := []struct {
+		name     string
+		input    *int
+		pageSize *uint32
+	}{
+		{
+			name: "Case 1: nil",
+		},
+		{
+			name:     "Case 2: not nil",
+			input:    ptr.Int(5),
+			pageSize: ptr.Uint32(5),
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.pageSize, castPageSize(tc.input))
+		})
+	}
+}

pkg/providers/gcp/provider.go

@@ -81,17 +81,8 @@ func New(clientFactory ClientFactory) (*Provider, error) {
 	}, nil
 }
 
-func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ *int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
-	if len(instances) > 1 {
-		return nil, fmt.Errorf("GCP does not support mult-region topology requests")
-	}
-
-	var instanceToNode map[string]string
-	if len(instances) == 1 {
-		instanceToNode = instances[0].Instances
-	}
-
-	cfg, err := p.generateInstanceTopology(ctx, instanceToNode)
+func (p *Provider) GenerateTopologyConfig(ctx context.Context, pageSize *int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
+	cfg, err := p.generateInstanceTopology(ctx, pageSize, instances)
 	if err != nil {
 		return nil, err
 	}
@@ -103,17 +94,12 @@ func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ *int, instances
 
 // Instances2NodeMap implements slurm.instanceMapper
 func (p *Provider) Instances2NodeMap(ctx context.Context, nodes []string) (map[string]string, error) {
-	i2n := make(map[string]string)
-	for _, node := range nodes {
-		i2n[node] = node
-	}
-
-	return i2n, nil
+	return instanceToNodeMap(ctx, nodes)
 }
 
 // GetComputeInstancesRegion implements slurm.instanceMapper
-func (p *Provider) GetComputeInstancesRegion(_ context.Context) (string, error) {
-	return "", nil
+func (p *Provider) GetComputeInstancesRegion(ctx context.Context) (string, error) {
+	return getRegion(ctx)
 }
 
 // GetNodeRegion implements k8s.k8sNodeInfo
@@ -123,5 +109,5 @@ func (p *Provider) GetNodeRegion(node *v1.Node) (string, error) {
 
 // GetNodeInstance implements k8s.k8sNodeInfo
 func (p *Provider) GetNodeInstance(node *v1.Node) (string, error) {
-	return node.Labels["kubernetes.io/hostname"], nil
+	return node.Annotations["container.googleapis.com/instance_id"], nil
 }