Merge pull request #2204 from nebius/SCHED-1012/capacity-check

SCHED-1012: Add early capacity check

Author: theyoprst

.github/workflows/e2e_test.yml

@@ -16,6 +16,7 @@ on:
 
 permissions:
   contents: read
+  actions: write
 
 concurrency:
   # Prevent cancelling already-running jobs (avoids resource waste).
@@ -58,6 +59,9 @@ jobs:
           make install-e2e-tools
           echo "${{ github.workspace }}/bin" >> $GITHUB_PATH
 
+      - name: Build e2e binary
+        run: go build -o bin/e2e ./cmd/e2e
+
       - name: Resolve profile
         shell: bash
         run: |
@@ -72,6 +76,21 @@ jobs:
           echo "NEBIUS_REGION=$(echo "$E2E_PROFILE" | yq -r '.nebius_region')" >> "$GITHUB_ENV"
           echo "NEBIUS_TENANT_ID=$(echo "$E2E_PROFILE" | yq -r '.nebius_tenant_id')" >> "$GITHUB_ENV"
 
+      - name: Print run parameters
+        shell: bash
+        run: |
+          {
+            echo "### Run Parameters"
+            echo ""
+            echo "- Branch: \`${{ github.ref_name }}\`"
+            echo "- Terraform branch: \`${{ env.TERRAFORM_REPO_REF }}\`"
+            echo "- Profile variable: \`$PROFILE_ENV_VAR\`"
+            echo ""
+            echo '```yaml'
+            echo "$E2E_PROFILE" | yq .
+            echo '```'
+          } >> "$GITHUB_STEP_SUMMARY"
+
       - name: Install Nebius CLI
         shell: bash
         env:
@@ -187,6 +206,21 @@ jobs:
           path: "${{ github.workspace }}/terraform-repo"
           fetch-depth: 0
 
+      - name: Check capacity
+        shell: bash
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        run: |
+          cd ${{ env.PATH_TO_INSTALLATION }}
+          source .envrc
+          cd -
+
+          echo "### Capacity Check" >> "$GITHUB_STEP_SUMMARY"
+          echo '```' >> "$GITHUB_STEP_SUMMARY"
+          bin/e2e check-capacity 2>&1 | tee -a "$GITHUB_STEP_SUMMARY"
+          echo '```' >> "$GITHUB_STEP_SUMMARY"
+
       - name: Terraform Apply
         timeout-minutes: 120
         run: |
@@ -203,10 +237,10 @@ jobs:
           aws configure set region $NEBIUS_REGION
           aws configure set endpoint_url https://storage.$NEBIUS_REGION.nebius.cloud:443
 
-          go run ./cmd/e2e apply
+          bin/e2e apply
 
       - name: K8s Cluster Info and NodeGroups
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: |
           echo "=== Listing K8s clusters ==="
@@ -224,57 +258,57 @@ jobs:
           done
 
       - name: "K8s Cluster: Pods"
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: kubectl get pods -A -o wide
 
       - name: "K8s Cluster: Events"
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: kubectl get events -A --sort-by='.lastTimestamp'
 
       - name: "K8s Cluster: Nodes"
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: |
           kubectl get nodes
           echo ""
           kubectl get nodes -o yaml
 
       - name: "K8s Cluster: Jobs"
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: |
           kubectl -n soperator get job
           echo ""
           kubectl -n soperator get job -o yaml
 
       - name: "K8s Cluster: Helm Releases"
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: |
           kubectl get helmreleases -n flux-system
           echo ""
           kubectl get helmreleases -n flux-system -o yaml
 
       - name: "K8s Cluster: Slurm Cluster CRs"
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: |
           kubectl get slurmclusters -A
           echo ""
           kubectl get slurmclusters -A -o yaml
 
       - name: "K8s Cluster: Slurm Active Checks CRs"
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: |
           kubectl get activechecks -A
           echo ""
           kubectl get activechecks -A -o yaml
 
       - name: Slurm Cluster State
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: |
           kubectl exec -n soperator controller-0 -- sinfo -N || true
@@ -284,37 +318,37 @@ jobs:
           kubectl exec -n soperator controller-0 -- sacct --parsable2 --allusers --starttime=now-6hours | column -t -s'|'
 
       - name: Collect Full Kubernetes Cluster Info
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: |
           mkdir -p ./cluster-info
           kubectl cluster-info dump --namespaces=kruise-system,soperator-system,soperator,flux-system --output-directory=./cluster-info
 
       - name: Upload Full Kubernetes Cluster Info
-        if: always()
+        if: '!cancelled()'
         uses: actions/upload-artifact@v6
         with:
           name: cluster-info
           path: ./cluster-info
           retention-days: 7
 
       - name: Collect Soperator Outputs
-        if: always()
+        if: '!cancelled()'
         shell: bash
         run: |
           mkdir -p ./soperator-outputs
           kubectl cp soperator/controller-0:/mnt/jail/opt/soperator-outputs ./soperator-outputs
 
       - name: Upload Soperator Outputs
-        if: always()
+        if: '!cancelled()'
         uses: actions/upload-artifact@v6
         with:
           name: soperator-outputs
           path: ./soperator-outputs
           retention-days: 7
 
       - name: Terraform Destroy
-        if: always()
+        if: '!cancelled()'
         timeout-minutes: 30
         run: |
           cd ${{ env.PATH_TO_INSTALLATION }}
@@ -327,7 +361,7 @@ jobs:
           aws configure set region $NEBIUS_REGION
           aws configure set endpoint_url https://storage.$NEBIUS_REGION.nebius.cloud:443
 
-          go run ./cmd/e2e destroy
+          bin/e2e destroy
 
       - name: Force cleanup compute instances on failure
         if: failure()
@@ -414,7 +448,11 @@ jobs:
             # Terraform repo latest commits
             echo ""
             echo "### Terraform Repo Latest Commits on ${{ env.TERRAFORM_REPO_REF }}"
-            git -C "${{ github.workspace }}/terraform-repo" log --format='%h|%s|%ar|%an' -3 2>/dev/null | while IFS='|' read -r hash msg date author; do
-              echo "- [$hash]($terraform_repo_url/commit/$hash): $msg ($date) <$author>"
-            done
+            if [[ -d "${{ github.workspace }}/terraform-repo/.git" ]]; then
+              git -C "${{ github.workspace }}/terraform-repo" log --format='%h|%s|%ar|%an' -3 2>/dev/null | while IFS='|' read -r hash msg date author; do
+                echo "- [$hash]($terraform_repo_url/commit/$hash): $msg ($date) <$author>"
+              done
+            else
+              echo "Terraform repo not checked out."
+            fi
           } >> $GITHUB_STEP_SUMMARY

cmd/e2e/main.go

@@ -2,9 +2,11 @@ package main
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"os"
+	"os/exec"
 	"os/signal"
 	"syscall"
 
@@ -14,41 +16,74 @@ import (
 )
 
 func main() {
+	log.SetFlags(0)
+
 	if len(os.Args) < 2 {
-		_, _ = fmt.Fprintf(os.Stderr, "Usage: e2e <apply|destroy>\n")
+		_, _ = fmt.Fprintf(os.Stderr, "Usage: e2e <apply|destroy|check-capacity>\n")
 		os.Exit(2)
 	}
 
-	var cfg e2e.Config
-	if err := envconfig.Process("", &cfg); err != nil {
-		log.Fatalf("parse config: %v", err)
-	}
-
 	profile, err := e2e.LoadProfile()
 	if err != nil {
-		log.Fatalf("load profile: %v", err)
-	}
-	cfg.Profile = profile
-
-	sshPubKey, err := e2e.GenerateSSHPublicKey()
-	if err != nil {
-		log.Fatalf("generate SSH public key: %v", err)
+		log.Fatalf("Load profile: %v", err)
 	}
-	cfg.SSHPublicKey = sshPubKey
 
 	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
 	defer stop()
 
 	switch os.Args[1] {
+	case "check-capacity":
+		err = runCheckCapacity(ctx, profile)
 	case "apply":
+		cfg := loadFullConfig(profile)
 		err = e2e.Apply(ctx, cfg)
 	case "destroy":
+		cfg := loadFullConfig(profile)
 		err = e2e.Destroy(ctx, cfg)
 	default:
-		_, _ = fmt.Fprintf(os.Stderr, "Unknown command: %s\nUsage: e2e <apply|destroy>\n", os.Args[1])
+		_, _ = fmt.Fprintf(os.Stderr, "Unknown command: %s\nUsage: e2e <apply|destroy|check-capacity>\n", os.Args[1])
 		os.Exit(2)
 	}
 	if err != nil {
 		log.Fatalf("%s: %v", os.Args[1], err)
 	}
 }
+
+func loadFullConfig(profile e2e.Profile) e2e.Config {
+	var cfg e2e.Config
+	if err := envconfig.Process("", &cfg); err != nil {
+		log.Fatalf("Parse config: %v", err)
+	}
+	cfg.Profile = profile
+
+	sshPubKey, err := e2e.GenerateSSHPublicKey()
+	if err != nil {
+		log.Fatalf("Generate SSH public key: %v", err)
+	}
+	cfg.SSHPublicKey = sshPubKey
+
+	return cfg
+}
+
+func runCheckCapacity(ctx context.Context, profile e2e.Profile) error {
+	err := e2e.CheckCapacity(ctx, profile)
+	if !errors.Is(err, e2e.ErrInsufficientCapacity) {
+		return err
+	}
+
+	log.Print("Insufficient capacity detected with cancel strategy, cancelling workflow")
+	runID := os.Getenv("GITHUB_RUN_ID")
+	if runID == "" {
+		return fmt.Errorf("GITHUB_RUN_ID is not set, cannot cancel workflow")
+	}
+
+	cmd := exec.CommandContext(ctx, "gh", "run", "cancel", runID)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	if cancelErr := cmd.Run(); cancelErr != nil {
+		return fmt.Errorf("cancel workflow run %s: %w", runID, cancelErr)
+	}
+
+	log.Printf("Workflow run %s cancelled due to insufficient capacity", runID)
+	return nil
+}

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/kubereboot/kured v0.0.0-20241106074119-94e73465adc3
 	github.com/mackerelio/go-osstat v0.2.6
 	github.com/mariadb-operator/mariadb-operator/v25 v25.10.2
+	github.com/nebius/gosdk v0.0.0-20260224104345-c5b8377510ae
 	github.com/onsi/ginkgo/v2 v2.27.3
 	github.com/onsi/gomega v1.38.3
 	github.com/openkruise/kruise-api v1.8.0
@@ -24,32 +25,36 @@ require (
 	go.uber.org/zap v1.27.1
 	golang.org/x/crypto v0.46.0
 	golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0
+	google.golang.org/grpc v1.72.1
 	istio.io/pkg v0.0.0-20241216214326-d70796207df3
 	k8s.io/api v0.34.3
 	k8s.io/apimachinery v0.34.3
 	k8s.io/client-go v0.34.3
-	k8s.io/component-base v0.34.3
 	k8s.io/utils v0.0.0-20251002143259-bc988d571ff4
 	sigs.k8s.io/controller-runtime v0.22.4
 	sigs.k8s.io/security-profiles-operator v0.8.4
 	sigs.k8s.io/yaml v1.6.0
 )
 
 require (
+	buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.36.4-20250130201111-63bb56e20495.1 // indirect
 	github.com/Masterminds/semver/v3 v3.4.0 // indirect
 	github.com/agext/levenshtein v1.2.3 // indirect
 	github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect
 	github.com/apparentlymart/go-textseg/v15 v15.0.0 // indirect
-	github.com/blang/semver/v4 v4.0.0 // indirect
+	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/cert-manager/cert-manager v1.18.5 // indirect
 	github.com/containers/common v0.60.4 // indirect
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/fsnotify/fsnotify v1.9.0 // indirect
 	github.com/fxamacker/cbor/v2 v2.9.0 // indirect
 	github.com/getkin/kin-openapi v0.132.0 // indirect
 	github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
+	github.com/gofrs/flock v0.12.1 // indirect
+	github.com/golang-jwt/jwt/v4 v4.5.1 // indirect
 	github.com/google/btree v1.1.3 // indirect
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
+	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/go-version v1.7.0 // indirect
 	github.com/mitchellh/go-wordwrap v1.0.1 // indirect
@@ -68,12 +73,11 @@ require (
 	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/ugorji/go/codec v1.2.12 // indirect
 	github.com/x448/float16 v0.8.4 // indirect
-	go.opentelemetry.io/otel v1.35.0 // indirect
-	go.opentelemetry.io/otel/trace v1.35.0 // indirect
 	go.yaml.in/yaml/v2 v2.4.3 // indirect
 	go.yaml.in/yaml/v3 v3.0.4 // indirect
 	golang.org/x/mod v0.30.0 // indirect
 	golang.org/x/sync v0.19.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20260223185530-2f722ef697dc // indirect
 	gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect
 	sigs.k8s.io/randfill v1.0.0 // indirect
 	sigs.k8s.io/release-utils v0.8.1 // indirect
@@ -116,7 +120,7 @@ require (
 	golang.org/x/time v0.11.0 // indirect
 	golang.org/x/tools v0.39.0 // indirect
 	gomodules.xyz/jsonpatch/v2 v2.4.0 // indirect
-	google.golang.org/protobuf v1.36.8 // indirect
+	google.golang.org/protobuf v1.36.11 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	k8s.io/apiextensions-apiserver v0.34.2 // indirect