xpk 0.13.0__tar.gz → 0.14.1__tar.gz

{xpk-0.13.0/src/xpk.egg-info → xpk-0.14.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 0.13.0
+Version: 0.14.1
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
 Author-email: XPK team <xpk-code-reviewers@google.com>
 License: Apache-2.0
@@ -22,15 +22,18 @@ Requires-Dist: google-api-core==2.24.1
 Requires-Dist: packaging==24.2
 Requires-Dist: google-cloud-filestore==1.12.0
 Requires-Dist: google-cloud-storage
+Requires-Dist: Jinja2==3.1.6
 Provides-Extra: dev
 Requires-Dist: pyink==24.3.0; extra == "dev"
 Requires-Dist: pylint>=2.6.0; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-mock==3.15.1; extra == "dev"
 Requires-Dist: docker==7.1.0; extra == "dev"
 Requires-Dist: mypy~=1.17; extra == "dev"
 Requires-Dist: types-PyYAML==6.0.2; extra == "dev"
 Requires-Dist: types-docker~=7.1.0.0; extra == "dev"
+Requires-Dist: pylint-per-file-ignores==1.4.0; extra == "dev"
 Dynamic: license-file
 
 <!--
@@ -76,6 +79,7 @@ XPK supports the following TPU types:
 * v5e
 * v5p
 * Trillium (v6e)
+* Ironwood (tpu7x)
 
 and the following GPU types:
 * A100
@@ -83,6 +87,7 @@ and the following GPU types:
 * A3-Mega (h100-mega) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A3-Ultra (h200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A4 (b200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
+* A4X (gb200)
 
 and the following CPU types:
 * n2-standard-32

{xpk-0.13.0 → xpk-0.14.1}/README.md

@@ -41,6 +41,7 @@ XPK supports the following TPU types:
 * v5e
 * v5p
 * Trillium (v6e)
+* Ironwood (tpu7x)
 
 and the following GPU types:
 * A100
@@ -48,6 +49,7 @@ and the following GPU types:
 * A3-Mega (h100-mega) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A3-Ultra (h200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A4 (b200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
+* A4X (gb200)
 
 and the following CPU types:
 * n2-standard-32

{xpk-0.13.0 → xpk-0.14.1}/pyproject.toml

@@ -40,7 +40,8 @@ dependencies = [
   "google-api-core==2.24.1",
   "packaging==24.2",
   "google-cloud-filestore==1.12.0",
-  "google-cloud-storage"
+  "google-cloud-storage",
+  "Jinja2==3.1.6"
 ]
 
 [project.urls]
@@ -62,19 +63,24 @@ dev = [
     "pylint>=2.6.0",
     "pre-commit",
     "pytest",
+    "pytest-mock==3.15.1",
     "docker==7.1.0",
     "mypy ~= 1.17",
     "types-PyYAML == 6.0.2",
     "types-docker ~= 7.1.0.0",
+    "pylint-per-file-ignores == 1.4.0",
 ]
 
 [tool.setuptools.dynamic]
 version = {attr = "xpk.core.config.__version__"}
 
 [tool.setuptools]
-packages = ["xpk", "xpk.parser", "xpk.core", "xpk.commands", "xpk.api", "xpk.templates", "xpk.utils", "xpk.core.blueprint", "xpk.core.remote_state", "xpk.core.workload_decorators"]
 package-dir = {"" = "src"}
-package-data = {"xpk.api" = ["storage_crd.yaml"], "xpk.templates" = ["storage.yaml"]}
+packages = { find = { where = ["src"] } }
+
+[tool.setuptools.package-data]
+"xpk" = ["templates/*"]
+"xpk.api" = ["*.yaml"]
 
 [tool.pyink]
 # Formatting configuration to follow Google style-guide.

xpk-0.14.1/src/integration/docker_manager_test.py

@@ -0,0 +1,102 @@
+"""
+Copyright 2024 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
+
+     https://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.
+"""
+
+import docker
+from docker.errors import APIError
+from xpk.core.docker_manager import DockerManager, ctk_build_ref
+import pytest
+import os
+import time
+
+test_cfg_path = '/tmp/xpk_gcloud_cfg'
+test_deployment_dir = '/tmp/xpk_deployment'
+test_gcluster_cmd = 'gcluster --version'
+test_ctk_xpk_img = 'gcluster-xpk'
+test_ctk_xpk_container = 'xpk-test-container'
+
+
+def remove_img():
+  dc = docker.from_env()
+  try:
+    dc.images.remove(test_ctk_xpk_img, force=True)
+  except APIError as _:
+    pass
+
+
+def remove_container():
+  dc = docker.from_env()
+  try:
+    container = dc.containers.get(test_ctk_xpk_container)
+    container.remove(force=True)
+  except APIError as _:
+    pass
+
+
+def create_tmp_dirs():
+  os.mkdir(test_cfg_path)
+  os.mkdir(test_deployment_dir)
+
+
+def remove_tmp_dirs():
+  os.removedirs(test_cfg_path)
+  os.removedirs(test_deployment_dir)
+
+
+@pytest.fixture(name='setup_img_name')
+def remove_test_ctk_img():
+  create_tmp_dirs()
+  remove_container()
+  remove_img()
+  yield test_ctk_xpk_img
+  remove_container()
+  remove_img()
+  remove_tmp_dirs()
+
+
+def test_docker_build_image(setup_img_name):
+  dm = DockerManager(
+      gcloud_cfg_path=test_cfg_path,
+      working_dir=test_deployment_dir,
+      img_name=setup_img_name,
+  )
+  dm.initialize()
+
+  dc = docker.from_env()
+  containers_before = dc.containers.list(all=True)
+  dc.images.get(f'{setup_img_name}:{ctk_build_ref}')
+  containers_after = dc.containers.list(all=True)
+  assert len(containers_before) == len(containers_after)
+
+
+def test_run_command(setup_img_name):
+
+  dm = DockerManager(
+      gcloud_cfg_path=test_cfg_path,
+      working_dir=test_deployment_dir,
+      img_name=setup_img_name,
+      remove_container=True,
+  )
+  dc = docker.from_env()
+
+  containers_before = dc.containers.list(all=True)
+  dm.initialize()
+  dm.run_command(test_gcluster_cmd)
+
+  time.sleep(2)
+
+  containers_after = dc.containers.list(all=True)
+
+  assert len(containers_after) - len(containers_before) == 0

xpk-0.14.1/src/integration/gcluster_a3mega_test.py

@@ -0,0 +1,204 @@
+"""
+Copyright 2024 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
+
+     https://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.
+"""
+
+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
+import pytest
+import os
+import shutil
+
+ctk_gcloud_cfg = os.getenv("GCLOUD_CFG_PATH")
+project_id = os.getenv("PROJECT_ID")
+region = os.getenv("REGION")
+zone = os.getenv("ZONE")
+auth_cidr = os.getenv("AUTH_CIDR")
+cluster_name = os.getenv("A3_MEGA_TEST_CLUSTER_NAME")
+
+uploads_dir = "uploads"
+
+
+@pytest.fixture(name="setup_tests")
+def prepare_test():
+  pwd = os.getcwd()
+  docker_path = os.path.join(pwd, "xpk_test_docker_dir")
+  bp_path = os.path.join(pwd, "xpk_bp_path")
+  if not os.path.exists(docker_path):
+    os.makedirs(docker_path)
+  if not os.path.exists(bp_path):
+    os.makedirs(bp_path)
+  yield (docker_path, bp_path)
+  shutil.rmtree(docker_path)
+  shutil.rmtree(bp_path)
+
+
+@pytest.mark.skip(
+    reason=(
+        "This test requires A3 capacity, therefore it should not be run on each"
+        " build. Please invoke it manually if needed. "
+    )
+)
+def test_deploy_a3_mega_deployment(setup_tests):
+  docker_path, bp_path = setup_tests[0], setup_tests[1]
+  (
+      blueprint_name,
+      prefix,
+      gcluster_manager,
+      staged_bp_path,
+  ) = create_test_a3_mega_deployment(docker_path, bp_path)
+  gcluster_manager.deploy(
+      blueprint_path=staged_bp_path,
+      deployment_name=blueprint_name,
+      prefix=prefix,
+  )
+
+  #   cleanup part
+  gcluster_manager.destroy_deployment(
+      deployment_name=blueprint_name, prefix=prefix
+  )
+  shutil.rmtree(docker_path)
+  shutil.rmtree(bp_path)
+
+
+@pytest.mark.skip(
+    reason=(
+        "This test requires A3 capacity, therefore it should not be run on each"
+        " build. Please invoke it manually if needed. "
+    )
+)
+def test_create_a3_mega_deployment_files(setup_tests):
+  assert project_id is not None
+  assert region is not None
+  assert zone is not None
+  assert auth_cidr is not None
+  assert ctk_gcloud_cfg is not None
+  assert cluster_name is not None
+  docker_path, bp_path = setup_tests[0], setup_tests[1]
+
+  blueprint_name = f"{cluster_name}-a3-mega-xpk"
+  prefix = "prefix"
+
+  docker_manager = DockerManager(
+      gcloud_cfg_path=ctk_gcloud_cfg, working_dir=docker_path
+  )
+  docker_manager.initialize()
+
+  bpm = BlueprintGenerator(storage_path=bp_path)
+  a3_mega_blueprint = bpm.generate_a3_mega_blueprint(
+      cluster_name=cluster_name,
+      blueprint_name=blueprint_name,
+      prefix=prefix,
+      region=region,
+      project_id=project_id,
+      auth_cidr=auth_cidr,
+      zone=zone,
+      system_node_pool_min_node_count=3,
+  )
+  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)
+
+  assert a3_mega_blueprint.blueprint_file == blueprint_test_path
+  assert a3_mega_blueprint.blueprint_dependencies == blueprint_deps_test_path
+
+  assert os.path.isfile(blueprint_test_path)
+  assert os.path.isdir(blueprint_deps_test_path)
+  assert os.path.isfile(
+      os.path.join(blueprint_deps_test_path, "config-map.yaml.tftpl")
+  )
+
+  gcluster_manager = GclusterManager(
+      gcluster_command_runner=docker_manager, remote_state_client=None
+  )
+
+  staged_bp_path = gcluster_manager.stage_files(
+      blueprint_file=a3_mega_blueprint.blueprint_file,
+      blueprint_dependencies=a3_mega_blueprint.blueprint_dependencies,
+      prefix=prefix,
+  )
+
+  assert staged_bp_path == os.path.join(
+      "/out", uploads_dir, prefix, f"{blueprint_name}.yaml"
+  )
+
+  staged_bp_path_local = os.path.join(
+      docker_path, uploads_dir, prefix, f"{blueprint_name}.yaml"
+  )
+  staged_bp_deps_path_local = os.path.join(
+      docker_path, uploads_dir, prefix, blueprint_name
+  )
+
+  assert os.path.isfile(staged_bp_path_local)
+  assert os.path.isdir(staged_bp_deps_path_local)
+  assert os.path.isfile(
+      os.path.join(staged_bp_deps_path_local, "config-map.yaml.tftpl")
+  )
+  assert os.path.isfile(
+      os.path.join(
+          staged_bp_deps_path_local, "kueue-xpk-configuration.yaml.tftpl"
+      )
+  )
+  unique_name = get_unique_name(project_id, region, zone)
+  gcluster_manager.deploy(
+      blueprint_path=staged_bp_path, deployment_name=unique_name, dry_run=True
+  )
+
+
+def create_test_a3_mega_deployment(docker_path: str, bp_path: str):
+  assert project_id is not None
+  assert region is not None
+  assert zone is not None
+  assert auth_cidr is not None
+  assert ctk_gcloud_cfg is not None
+  assert cluster_name is not None
+
+  blueprint_name = f"{cluster_name}-a3-mega-xpk"
+  prefix = "prefix"
+
+  docker_manager = DockerManager(
+      gcloud_cfg_path=ctk_gcloud_cfg, working_dir=docker_path
+  )
+  docker_manager.initialize()
+
+  bpm = BlueprintGenerator(storage_path=bp_path)
+  a3_mega_blueprint = bpm.generate_a3_mega_blueprint(
+      cluster_name=cluster_name,
+      blueprint_name=blueprint_name,
+      prefix=prefix,
+      region=region,
+      project_id=project_id,
+      auth_cidr=auth_cidr,
+      zone=zone,
+      system_node_pool_min_node_count=3,
+  )
+
+  gcluster_manager = GclusterManager(
+      gcluster_command_runner=docker_manager,
+      remote_state_client=None,
+  )
+
+  staged_bp_path = gcluster_manager.stage_files(
+      blueprint_file=a3_mega_blueprint.blueprint_file,
+      blueprint_dependencies=a3_mega_blueprint.blueprint_dependencies,
+      prefix=prefix,
+  )
+
+  return (
+      blueprint_name,
+      prefix,
+      gcluster_manager,
+      staged_bp_path,
+  )

xpk-0.14.1/src/integration/gcluster_a3ultra_test.py

@@ -0,0 +1,176 @@
+"""
+Copyright 2024 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
+
+     https://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.
+"""
+
+import os
+import shutil
+
+import pytest
+
+from xpk.commands.cluster_gcluster import get_unique_name
+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
+
+ctk_gcloud_cfg = os.getenv("GCLOUD_CFG_PATH")
+project_id = os.getenv("PROJECT_ID")
+region = os.getenv("REGION")
+zone = os.getenv("ZONE")
+auth_cidr = os.getenv("AUTH_CIDR")
+cluster_name = os.getenv("A3_ULTRA_TEST_CLUSTER_NAME")
+
+
+@pytest.fixture(name="setup_tests")
+def prepare_test():
+  pwd = os.getcwd()
+  docker_path = os.path.join(pwd, "xpk_test_docker_dir")
+  bp_path = os.path.join(pwd, "xpk_test_bp_dir")
+  if not os.path.exists(docker_path):
+    os.makedirs(docker_path)
+  if not os.path.exists(bp_path):
+    os.makedirs(bp_path)
+  yield (docker_path, bp_path)
+  shutil.rmtree(docker_path)
+  shutil.rmtree(bp_path)
+
+
+@pytest.mark.skip(
+    reason=(
+        "This test requires A3 capacity, therefore it should not be run on each"
+        " build. Please invoke it manually if needed. "
+    )
+)
+def test_create_a3_ultra_deployment_files(setup_tests):
+  assert project_id is not None
+  assert region is not None
+  assert zone is not None
+  assert auth_cidr is not None
+  assert ctk_gcloud_cfg is not None
+  assert cluster_name is not None
+  docker_path, bp_path = setup_tests[0], setup_tests[1]
+  blueprint_name = f"{cluster_name}-a3-ultra-xpk"
+
+  docker_manager = DockerManager(
+      gcloud_cfg_path=ctk_gcloud_cfg, working_dir=docker_path
+  )
+  docker_manager.initialize()
+  prefix = f"{project_id}-{region}".lower()
+  bpm = BlueprintGenerator(storage_path=bp_path)
+  a3_mega_blueprint = bpm.generate_a3_ultra_blueprint(
+      cluster_name=cluster_name,
+      blueprint_name=blueprint_name,
+      region=region,
+      project_id=project_id,
+      auth_cidr=auth_cidr,
+      zone=zone,
+      reservation="foo",
+      num_nodes=1,
+      system_node_pool_machine_type="e2-standard-16",
+      prefix=prefix,
+  )
+  blueprint_test_path = os.path.join(bp_path, prefix, f"{blueprint_name}.yaml")
+  blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)
+  assert a3_mega_blueprint.blueprint_file == blueprint_test_path
+  assert a3_mega_blueprint.blueprint_dependencies == blueprint_deps_test_path
+
+  assert os.path.isfile(blueprint_test_path)
+  assert os.path.isdir(blueprint_deps_test_path)
+  assert os.path.isfile(
+      os.path.join(blueprint_deps_test_path, "mlgru-disable.yaml")
+  )
+  assert os.path.isfile(
+      os.path.join(blueprint_deps_test_path, "nccl-installer.yaml")
+  )
+  gcluster_manager = GclusterManager(
+      gcluster_command_runner=docker_manager, remote_state_client=None
+  )
+
+  staged_bp_path = gcluster_manager.stage_files(
+      blueprint_file=a3_mega_blueprint.blueprint_file,
+      blueprint_dependencies=a3_mega_blueprint.blueprint_dependencies,
+      prefix=prefix,
+  )
+  assert staged_bp_path == os.path.join(
+      "/out/uploads", prefix, f"{blueprint_name}.yaml"
+  )
+  unique_name = get_unique_name(project_id, region, zone)
+  gcluster_manager.deploy(
+      blueprint_path=staged_bp_path, deployment_name=unique_name, dry_run=True
+  )
+
+
+@pytest.mark.skip(
+    reason=(
+        "This test requires A3 capacity, therefore it should not be run on each"
+        " build. Please invoke it manually if needed. "
+    )
+)
+def test_create_a3_ultra_deployment(setup_tests):
+  assert project_id is not None
+  assert region is not None
+  assert zone is not None
+  assert auth_cidr is not None
+  assert ctk_gcloud_cfg is not None
+  assert cluster_name is not None
+  docker_path, bp_path = setup_tests[0], setup_tests[1]
+  blueprint_name = f"{cluster_name}-a3-ultra-xpk"
+
+  docker_manager = DockerManager(
+      gcloud_cfg_path=ctk_gcloud_cfg, working_dir=docker_path
+  )
+  docker_manager.initialize()
+
+  bpm = BlueprintGenerator(storage_path=bp_path)
+  a3_mega_blueprint = bpm.generate_a3_ultra_blueprint(
+      cluster_name=cluster_name,
+      blueprint_name=blueprint_name,
+      region=region,
+      project_id=project_id,
+      auth_cidr=auth_cidr,
+      zone=zone,
+      capacity_type=CapacityType.SPOT,
+      num_nodes=1,
+      system_node_pool_machine_type="e2-standard-16",
+  )
+  blueprint_test_path = os.path.join(bp_path, f"{blueprint_name}.yaml")
+  blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)
+
+  assert a3_mega_blueprint.blueprint_file == blueprint_test_path
+  assert a3_mega_blueprint.blueprint_dependencies == blueprint_deps_test_path
+
+  assert os.path.isfile(blueprint_test_path)
+  assert os.path.isdir(blueprint_deps_test_path)
+  assert os.path.isfile(
+      os.path.join(blueprint_deps_test_path, "mlgru-disable.yaml")
+  )
+  assert os.path.isfile(
+      os.path.join(blueprint_deps_test_path, "nccl-installer.yaml")
+  )
+  gcluster_manager = GclusterManager(
+      gcluster_command_runner=docker_manager, remote_state_client=None
+  )
+
+  staged_bp_path = gcluster_manager.stage_files(
+      blueprint_file=a3_mega_blueprint.blueprint_file,
+      blueprint_dependencies=a3_mega_blueprint.blueprint_dependencies,
+  )
+
+  gcluster_manager.deploy(
+      blueprint_path=staged_bp_path, deployment_name=blueprint_name
+  )
+
+  #   cleanup part
+  gcluster_manager.destroy_deployment(deployment_name=blueprint_name)