feat(orchestrator): remove staticMachineCapacity and integrate local telemetry tracking for internal users

Author: Neelabh94

pkg/orchestrator/gke/gke_job_orchestrator.go

@@ -24,6 +24,7 @@ import (
 	"hpc-toolkit/pkg/orchestrator"
 	"hpc-toolkit/pkg/scheduling"
 	"hpc-toolkit/pkg/shell"
+	"hpc-toolkit/pkg/telemetry"
 	"io"
 	"net/http"
 	"os"
@@ -181,8 +182,21 @@ func NewGKEOrchestrator() (*GKEOrchestrator, error) {
 func (g *GKEOrchestrator) SubmitJob(job orchestrator.JobDefinition) error {
 	logging.Info("Starting gcluster job submit workflow...")
 
+	startTime := time.Now()
+	var success bool
+	defer func() {
+		latencySecs := time.Since(startTime).Seconds()
+		profile := map[string]string{
+			"accelerator_type": job.AcceleratorType,
+			"nodes":           fmt.Sprintf("%d", job.NumSlices),
+		}
+
+		telemetry.RecordLocalMetrics(job.WorkloadName, latencySecs, success, profile)
+	}()
+
 	var err error
 	job, err = g.initializeJobSubmission(job)
+
 	if err != nil {
 		return err
 	}
@@ -231,9 +245,11 @@ func (g *GKEOrchestrator) SubmitJob(job orchestrator.JobDefinition) error {
 	}
 
 	logging.Info("gcluster job submit workflow completed.")
+	success = true
 	return nil
 }
 
+
 func (g *GKEOrchestrator) generatePathwaysManifest(job orchestrator.JobDefinition, fullImageName string) (string, error) {
 	// Set default values for Pathways-specific fields if not provided
 	if job.Pathways.ProxyServerImage == "" {
@@ -1254,77 +1270,8 @@ type MachineTypeCap struct {
 	GuestCpus int `json:"guestCpus"` // Parse vCPUs for CPU-only machines
 }
 
-var staticMachineCapacity = map[string]int{
-	// GPUs - RTX-Pro 6000 (g4 series)
-	"g4-standard-48":  1,
-	"g4-standard-96":  2,
-	"g4-standard-192": 4,
-	"g4-standard-384": 8,
-
-	// GPUs - Tesla A100 (a2 series)
-	"a2-highgpu-1g":  1,
-	"a2-highgpu-2g":  2,
-	"a2-highgpu-4g":  4,
-	"a2-highgpu-8g":  8,
-	"a2-megagpu-16g": 16,
-
-	// GPUs - A100 80GB
-	"a2-ultragpu-1g": 1,
-	"a2-ultragpu-2g": 2,
-	"a2-ultragpu-4g": 4,
-	"a2-ultragpu-8g": 8,
-
-	// GPUs - H100 80GB
-	"a3-highgpu-1g": 1,
-	"a3-highgpu-2g": 2,
-	"a3-highgpu-4g": 4,
-	"a3-highgpu-8g": 8,
-
-	// GPUs - H100 Mega & H200
-	"a3-megagpu-8g":  8,
-	"a3-ultragpu-8g": 8,
-
-	// GPUs - B200 and GB200
-	"a4-highgpu-8g-lowmem":  8,
-	"a4-highgpu-8g":         8,
-	"a4x-highgpu-4g":        4,
-	"a4x-highgpu-4g-nolssd": 4,
-
-	// GPUs - L4 (g2 series)
-	"g2-standard-4":  1,
-	"g2-standard-8":  1,
-	"g2-standard-12": 1,
-	"g2-standard-24": 2,
-	"g2-standard-48": 4,
-	"g2-standard-96": 8,
-
-	// TPUs
-	"ct4p-hightpu-4t":  4,
-	"ct5p-hightpu-1t":  1,
-	"ct5p-hightpu-2t":  2,
-	"ct5p-hightpu-4t":  4,
-	"ct6e-standard-1t": 1,
-	"ct6e-standard-4t": 4,
-	"ct6e-standard-8t": 8,
-}
-
 var acceleratorShorthandMap = map[string]string{
-	"nvidia-l4":             "g2-standard-12",
-	"nvidia-tesla-a100":     "a2-highgpu-1g",
-	"nvidia-a100-80gb":      "a3-highgpu-8g",
-	"nvidia-h100-80gb":      "a3-highgpu-8g",
-	"nvidia-h100-mega-80gb": "a3-megagpu-8g",
-	"nvidia-h200-141gb":     "a3-ultragpu-8g",
-	"nvidia-b200":           "a4-highgpu-8g",
-	"nvidia-gb200":          "a4x-highgpu-4g",
-	"tpu-v7":                "tpu7-standard-1t",
-	"tpu-v7x":               "tpu7x-standard-4t",
-	"tpu-v6e-slice":         "ct6e-standard-4t",
-	"tpu-v5p-slice":         "ct5p-hightpu-4t",
-	"tpu-v5-lite-podslice":  "ct5lp-hightpu-4t",
-	"tpu-v4-podslice":       "ct4p-hightpu-4t",
-
-	// xpk compatibility mappings (GPUs)
+	// GPU mappings
 	"l4-1":             "g2-standard-12",
 	"l4-2":             "g2-standard-24",
 	"l4-4":             "g2-standard-48",
@@ -1351,22 +1298,23 @@ var acceleratorShorthandMap = map[string]string{
 	"b200-8":           "a4-highgpu-8g",
 	"gb200-4":          "a4x-highgpu-4g",
 
-	// xpk compatibility mappings (TPUs - Single VM configurations)
-	"v4-8":  "ct4p-hightpu-4t",
-	"v5p-2": "ct5p-hightpu-1t",
-	"v5p-4": "ct5p-hightpu-2t",
-	"v5p-8": "ct5p-hightpu-4t",
-	"v6e-1": "ct6e-standard-1t",
-	"v6e-4": "ct6e-standard-4t",
-	"v6e-8": "ct6e-standard-8t",
+	// TPUs mappings
+	"v4-8":    "ct4p-hightpu-4t",
+	"v5p-1":   "ct5p-hightpu-1t",
+	"v5p-2":   "ct5p-hightpu-2t",
+	"v5p-4":   "ct5p-hightpu-4t",
+	"v5e-1":   "ct5lp-hightpu-1t",
+	"v5e-4":   "ct5lp-hightpu-4t",
+	"v5e-8":   "ct5lp-hightpu-8t",
+	"v6e-1":   "ct6e-standard-1t",
+	"v6e-4":   "ct6e-standard-4t",
+	"v6e-8":   "ct6e-standard-8t",
+	"tpu-v7":  "tpu7-standard-1t",
+	"tpu-v7x": "tpu7x-standard-4t",
+
 }
 
 func (g *GKEOrchestrator) FetchMachineCapacity(machineType, zone string) (int, error) {
-	if cap, exists := staticMachineCapacity[machineType]; exists {
-		logging.Info("Using hardcoded capacity for machine type %s: %d", machineType, cap)
-		return cap, nil
-	}
-
 	if zone == "" {
 		return 0, fmt.Errorf("zone is required for machine capacity lookup")
 	}
@@ -1504,11 +1452,15 @@ func (g *GKEOrchestrator) calculateGCPMachineResourceLimits(opts ManifestOptions
 		if strings.Contains(strings.ToLower(mapped), "tpu") {
 			return "", "", "", fmt.Sprintf("%d", count), nil
 		}
+		return "", "", "", "", fmt.Errorf("machine type %s resolved to %d capacity but could not be classified as GPU or TPU (mapped label: %s)", machineName, count, mapped)
 	}
-	return "", "", "", "", nil
+	return "", "", "", "", fmt.Errorf("failed to determine capacity for machine type %s", machineName)
 }
 
 
+
+
+
 func (g *GKEOrchestrator) calculateCPUMachineResourceLimits(opts ManifestOptions, profile JobProfile) (string, error) {
 	count := profile.CapacityCount
 	logging.Info("Using cached capacity for CPU machine %s during limits calculation: %d", opts.AcceleratorType, count)

pkg/telemetry/telemetry.go

@@ -0,0 +1,95 @@
+// 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 telemetry
+
+import (
+	"encoding/json"
+	"hpc-toolkit/pkg/logging"
+	"hpc-toolkit/pkg/shell"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+)
+
+// MetricEntry holds telemetry data for internal Google usage.
+type MetricEntry struct {
+	Timestamp             time.Time         `json:"timestamp"`
+	WorkloadName          string            `json:"workload_name"`
+	LatencySeconds        float64           `json:"latency_seconds"`
+	SubmissionSuccess     bool              `json:"submission_success"`
+	StaticResourceProfile map[string]string `json:"static_resource_profile"`
+}
+
+// isInternalUser checks if the active gcloud account is a @google.com domain.
+func isInternalUser() bool {
+	res := shell.ExecuteCommand("gcloud", "auth", "list", "--filter=status:ACTIVE", "--format=value(account)")
+	if res.ExitCode != 0 {
+		return false
+	}
+	account := strings.TrimSpace(res.Stdout)
+	return strings.HasSuffix(account, "@google.com")
+}
+
+// RecordLocalMetrics appends a telemetry metric entry to ~/.gcluster-job/telemetry_metrics.jsonl only for internal Googlers.
+func RecordLocalMetrics(workloadName string, latency float64, success bool, profile map[string]string) {
+	if os.Getenv("GCLUSTER_SKIP_TELEMETRY") == "true" {
+		logging.Info("Skipping telemetry metrics due to GCLUSTER_SKIP_TELEMETRY environment variable.")
+		return
+	}
+
+	if !isInternalUser() {
+		return // Skip quietly if not an internal Googler
+	}
+
+	entry := MetricEntry{
+		Timestamp:             time.Now(),
+		WorkloadName:          workloadName,
+		LatencySeconds:        latency,
+		SubmissionSuccess:     success,
+		StaticResourceProfile: profile,
+	}
+
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		logging.Error("Failed to resolve user home for internal telemetry: %v", err)
+		return
+	}
+
+	metricsDir := filepath.Join(homeDir, ".gcluster-job")
+	if err := os.MkdirAll(metricsDir, 0755); err != nil {
+		logging.Error("Failed to create metrics storage directory %s: %v", metricsDir, err)
+		return
+	}
+
+	metricsFile := filepath.Join(metricsDir, "telemetry_metrics.jsonl")
+
+	data, err := json.Marshal(entry)
+	if err != nil {
+		logging.Error("Failed to marshal telemetry metrics to JSON: %v", err)
+		return
+	}
+
+	f, err := os.OpenFile(metricsFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		logging.Error("Failed to open telemetry metrics file %s: %v", metricsFile, err)
+		return
+	}
+	defer f.Close()
+
+	if _, err := f.Write(append(data, '\n')); err != nil {
+		logging.Error("Failed to write telemetry metrics to %s: %v", metricsFile, err)
+	}
+}