xpk 0.14.4__py3-none-any.whl → 0.16.0__py3-none-any.whl

integration/README.md

@@ -0,0 +1,19 @@
+This folder contains integration tests.
+
+To run them env variables are needed:
+
+```bash
+export PROJECT_ID=...
+export REGION=...
+export ZONE=...
+export AUTH_CIDR=...
+export DEPLOYMENT_DIR=...
+export CLUSTER_NAME=...
+export GCLOUD_CFG_PATH=...
+```
+
+To run tests:
+
+```bash
+pytest src/integration
+```

integration/gcluster_a3mega_test.py

@@ -18,6 +18,7 @@ from xpk.commands.cluster_gcluster import get_unique_name
 from xpk.core.docker_manager import DockerManager
 from xpk.core.gcluster_manager import GclusterManager
 from xpk.core.blueprint.blueprint_generator import BlueprintGenerator
+from xpk.utils.versions import ReleaseChannel
 import pytest
 import os
 import shutil
@@ -28,6 +29,8 @@ region = os.getenv("REGION")
 zone = os.getenv("ZONE")
 auth_cidr = os.getenv("AUTH_CIDR")
 cluster_name = os.getenv("A3_MEGA_TEST_CLUSTER_NAME")
+release_channel = os.getenv("RELEASE_CHANNEL")
+cluster_version = os.getenv("CLUSTER_VERSION")
 
 uploads_dir = "uploads"
 
@@ -87,6 +90,8 @@ def test_create_a3_mega_deployment_files(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
 
   blueprint_name = f"{cluster_name}-a3-mega-xpk"
@@ -107,6 +112,8 @@ def test_create_a3_mega_deployment_files(setup_tests):
       auth_cidr=auth_cidr,
       zone=zone,
       system_node_pool_min_node_count=3,
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, prefix, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, prefix, blueprint_name)
@@ -164,6 +171,8 @@ def create_test_a3_mega_deployment(docker_path: str, bp_path: str):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
 
   blueprint_name = f"{cluster_name}-a3-mega-xpk"
   prefix = "prefix"
@@ -183,6 +192,8 @@ def create_test_a3_mega_deployment(docker_path: str, bp_path: str):
       auth_cidr=auth_cidr,
       zone=zone,
       system_node_pool_min_node_count=3,
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
 
   gcluster_manager = GclusterManager(

integration/gcluster_a3ultra_test.py

@@ -24,6 +24,7 @@ from xpk.core.blueprint.blueprint_generator import BlueprintGenerator
 from xpk.core.capacity import CapacityType
 from xpk.core.docker_manager import DockerManager
 from xpk.core.gcluster_manager import GclusterManager
+from xpk.utils.versions import ReleaseChannel
 
 ctk_gcloud_cfg = os.getenv("GCLOUD_CFG_PATH")
 project_id = os.getenv("PROJECT_ID")
@@ -31,6 +32,8 @@ region = os.getenv("REGION")
 zone = os.getenv("ZONE")
 auth_cidr = os.getenv("AUTH_CIDR")
 cluster_name = os.getenv("A3_ULTRA_TEST_CLUSTER_NAME")
+release_channel = os.getenv("RELEASE_CHANNEL")
+cluster_version = os.getenv("CLUSTER_VERSION")
 
 
 @pytest.fixture(name="setup_tests")
@@ -60,6 +63,8 @@ def test_create_a3_ultra_deployment_files(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
   blueprint_name = f"{cluster_name}-a3-ultra-xpk"
 
@@ -80,6 +85,8 @@ def test_create_a3_ultra_deployment_files(setup_tests):
       num_nodes=1,
       system_node_pool_machine_type="e2-standard-16",
       prefix=prefix,
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, prefix, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)
@@ -125,6 +132,8 @@ def test_create_a3_ultra_deployment(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
   blueprint_name = f"{cluster_name}-a3-ultra-xpk"
 
@@ -144,6 +153,8 @@ def test_create_a3_ultra_deployment(setup_tests):
       capacity_type=CapacityType.SPOT,
       num_nodes=1,
       system_node_pool_machine_type="e2-standard-16",
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)

integration/gcluster_a4_test.py

@@ -24,6 +24,7 @@ from xpk.core.blueprint.blueprint_generator import BlueprintGenerator
 from xpk.core.capacity import CapacityType
 from xpk.core.docker_manager import DockerManager
 from xpk.core.gcluster_manager import GclusterManager
+from xpk.utils.versions import ReleaseChannel
 
 ctk_gcloud_cfg = os.getenv("GCLOUD_CFG_PATH")
 project_id = os.getenv("PROJECT_ID")
@@ -31,6 +32,8 @@ region = os.getenv("REGION")
 zone = os.getenv("ZONE")
 auth_cidr = os.getenv("AUTH_CIDR")
 cluster_name = os.getenv("A4_TEST_CLUSTER_NAME")
+release_channel = os.getenv("RELEASE_CHANNEL")
+cluster_version = os.getenv("CLUSTER_VERSION")
 
 
 @pytest.fixture(name="setup_tests")
@@ -60,6 +63,8 @@ def test_create_a4_deployment_files(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
   blueprint_name = f"{cluster_name}-a4-xpk"
 
@@ -80,6 +85,8 @@ def test_create_a4_deployment_files(setup_tests):
       num_nodes=1,
       system_node_pool_machine_type="e2-standard-16",
       prefix=prefix,
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, prefix, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)
@@ -125,6 +132,8 @@ def test_create_a4_deployment(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
   blueprint_name = f"{cluster_name}-a4-xpk"
 
@@ -144,6 +153,8 @@ def test_create_a4_deployment(setup_tests):
       capacity_type=CapacityType.SPOT,
       num_nodes=1,
       system_node_pool_machine_type="e2-standard-16",
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)

xpk/blueprints/a3mega/config-map.yaml.tftpl

@@ -0,0 +1,15 @@
+kind: ConfigMap
+apiVersion: v1
+metadata:
+  name: ${resource_config_name}
+data:
+  h100-mega-80gb-8: "${num_nodes}"
+---
+kind: ConfigMap
+apiVersion: v1
+metadata:
+  name: ${cluster_config_name}
+data:
+  capacity_type: "${capacity_type}"
+  reservation_id: "${reservation}"
+  provisioner: gcluster

xpk/blueprints/a3mega/storage_crd.yaml

@@ -0,0 +1,52 @@
+apiVersion: apiextensions.k8s.io/v1
+kind: CustomResourceDefinition
+metadata:
+  name: storages.xpk.x-k8s.io
+spec:
+  group: xpk.x-k8s.io
+  versions:
+    - name: v1
+      served: true
+      storage: true
+      schema:
+        openAPIV3Schema:
+          type: object
+          properties:
+            spec:
+              type: object
+              properties:
+                type:
+                  type: string
+                cluster:
+                  type: string
+                auto_mount:
+                  type: boolean
+                mount_point:
+                  type: string
+                readonly:
+                  type: boolean
+                manifest:
+                  type: string
+                pv:
+                  type: string
+                pvc:
+                  type: string
+              required:
+                - type
+                - cluster
+                - auto_mount
+                - mount_point
+                - readonly
+                - manifest
+                - pvc
+                - pv
+          x-kubernetes-validations:
+            - message: Value is immutable
+              rule: self == oldSelf
+  scope: Cluster
+  names:
+    plural: storages
+    singular: storage
+    kind: Storage
+    shortNames:
+      - stg

xpk/blueprints/a3ultra/config-map.yaml.tftpl

@@ -0,0 +1,15 @@
+kind: ConfigMap
+apiVersion: v1
+metadata:
+  name: ${resource_config_name}
+data:
+  h200-141gb-8: "${num_nodes}"
+---
+kind: ConfigMap
+apiVersion: v1
+metadata:
+  name: ${cluster_config_name}
+data:
+  capacity_type: "${capacity_type}"
+  reservation_id: "${reservation}"
+  provisioner: gcluster

xpk/blueprints/a3ultra/mlgru-disable.yaml

@@ -0,0 +1,59 @@
+# Copyright 2024 Google Inc. 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.
+
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: disable-mglru
+  namespace: kube-system
+spec:
+  selector:
+    matchLabels:
+      app: disable-mglru
+  template:
+    metadata:
+      labels:
+        app: disable-mglru
+    spec:
+      hostNetwork: true
+      tolerations:
+      - operator: "Exists"
+        key: nvidia.com/gpu
+      containers:
+      - name: disable-mglru
+        image: alpine:latest
+        command: ["/bin/sh"]
+        securityContext:
+          privileged: true
+        args:
+        - -c
+        - |
+          echo n | tee /sys/kernel/mm/lru_gen/enabled
+          sysctl -w net.ipv4.conf.gpu0rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu1rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu2rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu3rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu4rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu5rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu6rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu7rdma0.log_martians=0
+          sleep infinity
+        volumeMounts:
+        - name: sys-kernel-mm-lru-gen
+          mountPath: /sys/kernel/mm/lru_gen
+      # Remount sysfs so that it will be writable.
+      volumes:
+      - name: sys-kernel-mm-lru-gen
+        hostPath:
+          path: /sys/kernel/mm/lru_gen

xpk/blueprints/a3ultra/nccl-installer.yaml

@@ -0,0 +1,95 @@
+# Copyright 2024 Google Inc. 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.
+
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: nccl-rdma-installer
+  namespace: kube-system
+  labels:
+    k8s-app: nccl-rdma-installer
+spec:
+  selector:
+    matchLabels:
+      k8s-app: nccl-rdma-installer
+  updateStrategy:
+    type: RollingUpdate
+  template:
+    metadata:
+      labels:
+        name: nccl-rdma-installer
+        k8s-app: nccl-rdma-installer
+    spec:
+      priorityClassName: system-node-critical
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+            - matchExpressions:
+              - key: cloud.google.com/gke-accelerator
+                operator: In
+                values:
+                - nvidia-h200-141gb
+      tolerations:
+      - operator: "Exists"
+      hostNetwork: true
+      hostPID: true
+      volumes:
+      - name: library-dir-host
+        hostPath:
+          path: /home/kubernetes/bin/nvidia/lib64
+          type: DirectoryOrCreate
+      - name: gib
+        hostPath:
+          path: /home/kubernetes/bin/gib
+      initContainers:
+      - name: disable-log-martian
+        image: alpine:latest
+        command: ["/bin/sh"]
+        securityContext:
+          privileged: true
+        args:
+        - -c
+        - |
+          sysctl -w net.ipv4.conf.gpu0rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu1rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu2rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu3rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu4rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu5rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu6rdma0.log_martians=0
+          sysctl -w net.ipv4.conf.gpu7rdma0.log_martians=0
+      - image: us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib:v1.0.3
+        name: nccl-rdma-installer
+        resources:
+          requests:
+            cpu: 150m
+        securityContext:
+          privileged: true
+        volumeMounts:
+        - name: library-dir-host
+          mountPath: /usr/local/home/kubernetes/bin/nvidia/lib64
+        - name: gib
+          mountPath: /usr/local/home/kubernetes/bin/gib
+        command: ["/bin/sh", "-c"]
+        args:
+        - |
+          set -ex
+          /scripts/container_entry.sh install --install-nccl
+          cp -r /var/lib/gib/lib64/. /usr/local/home/kubernetes/bin/nvidia/lib64
+          cp -r /var/lib/gib/. /usr/local/home/kubernetes/bin/gib
+          echo "installation finishes"
+      containers:
+      - image: "gke.gcr.io/pause:3.8@sha256:880e63f94b145e46f1b1082bb71b85e21f16b99b180b9996407d61240ceb9830"
+        name: pause

xpk/blueprints/a3ultra/storage_crd.yaml

@@ -0,0 +1,52 @@
+apiVersion: apiextensions.k8s.io/v1
+kind: CustomResourceDefinition
+metadata:
+  name: storages.xpk.x-k8s.io
+spec:
+  group: xpk.x-k8s.io
+  versions:
+    - name: v1
+      served: true
+      storage: true
+      schema:
+        openAPIV3Schema:
+          type: object
+          properties:
+            spec:
+              type: object
+              properties:
+                type:
+                  type: string
+                cluster:
+                  type: string
+                auto_mount:
+                  type: boolean
+                mount_point:
+                  type: string
+                readonly:
+                  type: boolean
+                manifest:
+                  type: string
+                pv:
+                  type: string
+                pvc:
+                  type: string
+              required:
+                - type
+                - cluster
+                - auto_mount
+                - mount_point
+                - readonly
+                - manifest
+                - pvc
+                - pv
+          x-kubernetes-validations:
+            - message: Value is immutable
+              rule: self == oldSelf
+  scope: Cluster
+  names:
+    plural: storages
+    singular: storage
+    kind: Storage
+    shortNames:
+      - stg

xpk/blueprints/a4/config-map.yaml.tftpl

@@ -0,0 +1,15 @@
+kind: ConfigMap
+apiVersion: v1
+metadata:
+  name: ${resource_config_name}
+data:
+  b200-8: "${num_nodes}"
+---
+kind: ConfigMap
+apiVersion: v1
+metadata:
+  name: ${cluster_config_name}
+data:
+  capacity_type: "${capacity_type}"
+  reservation_id: "${reservation}"
+  provisioner: gcluster

xpk/blueprints/a4/nccl-rdma-installer-a4.yaml

@@ -0,0 +1,66 @@
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: nccl-rdma-installer
+  namespace: kube-system
+  labels:
+    k8s-app: nccl-rdma-installer
+spec:
+  selector:
+    matchLabels:
+      k8s-app: nccl-rdma-installer
+  updateStrategy:
+    type: RollingUpdate
+  template:
+    metadata:
+      labels:
+        name: nccl-rdma-installer
+        k8s-app: nccl-rdma-installer
+    spec:
+      priorityClassName: system-node-critical
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+              - matchExpressions:
+                  - key: cloud.google.com/gke-accelerator
+                    operator: In
+                    values:
+                      - nvidia-b200
+      tolerations:
+        - operator: "Exists"
+      hostNetwork: true
+      hostPID: true
+      volumes:
+        - name: library-dir-host
+          hostPath:
+            path: /home/kubernetes/bin/nvidia/lib64
+            type: DirectoryOrCreate
+        - name: gib
+          hostPath:
+            path: /home/kubernetes/bin/gib
+      initContainers:
+        - image: us-docker.pkg.dev/kernel-net-team/clouda4-nccl-dev/nccl-plugin-gib-diagnostic:v1.0.3-b200
+          name: nccl-rdma-installer
+          resources:
+            requests:
+              cpu: 150m
+          securityContext:
+            privileged: true
+          volumeMounts:
+            - name: library-dir-host
+              mountPath: /usr/local/home/kubernetes/bin/nvidia/lib64
+            - name: gib
+              mountPath: /usr/local/home/kubernetes/bin/gib
+          command: ["/bin/sh", "-c"]
+          args:
+            - |
+              set -ex
+              /scripts/container_entry.sh install --install-nccl
+              cp -r /var/lib/gib/lib64/. /usr/local/home/kubernetes/bin/nvidia/lib64
+              cp -r /var/lib/gib/. /usr/local/home/kubernetes/bin/gib
+              # ibv_devinfo || exit 1
+              echo "installation finishes"
+      containers:
+        - image: "gke.gcr.io/pause:3.8@sha256:880e63f94b145e46f1b1082bb71b85e21f16b99b180b9996407d61240ceb9830"
+          name: pause

xpk/blueprints/a4/storage_crd.yaml

@@ -0,0 +1,52 @@
+apiVersion: apiextensions.k8s.io/v1
+kind: CustomResourceDefinition
+metadata:
+  name: storages.xpk.x-k8s.io
+spec:
+  group: xpk.x-k8s.io
+  versions:
+    - name: v1
+      served: true
+      storage: true
+      schema:
+        openAPIV3Schema:
+          type: object
+          properties:
+            spec:
+              type: object
+              properties:
+                type:
+                  type: string
+                cluster:
+                  type: string
+                auto_mount:
+                  type: boolean
+                mount_point:
+                  type: string
+                readonly:
+                  type: boolean
+                manifest:
+                  type: string
+                pv:
+                  type: string
+                pvc:
+                  type: string
+              required:
+                - type
+                - cluster
+                - auto_mount
+                - mount_point
+                - readonly
+                - manifest
+                - pvc
+                - pv
+          x-kubernetes-validations:
+            - message: Value is immutable
+              rule: self == oldSelf
+  scope: Cluster
+  names:
+    plural: storages
+    singular: storage
+    kind: Storage
+    shortNames:
+      - stg