xpk 0.14.0__py3-none-any.whl → 0.14.2__py3-none-any.whl

xpk/templates/kueue_gke_default_topology.yaml.j2

@@ -0,0 +1,10 @@
+apiVersion: kueue.x-k8s.io/v1alpha1
+kind: Topology
+metadata:
+  name: "gke-default"
+spec:
+  levels:
+  - nodeLabel: "cloud.google.com/gce-topology-block"
+  - nodeLabel: "cloud.google.com/gce-topology-subblock"
+  - nodeLabel: "cloud.google.com/gce-topology-host"
+  - nodeLabel: "kubernetes.io/hostname"

xpk/templates/kueue_sub_slicing_topology.yaml.j2

@@ -0,0 +1,14 @@
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: Topology
+metadata:
+  name: {{ sub_slice_topology_name }}
+spec:
+  levels:
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-16x16-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-8x16-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-8x8-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-4x8-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-4x4-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-2x4-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-2x2-id"
+  - nodeLabel: "kubernetes.io/hostname"

xpk/templates/mtc-cpc.yaml

@@ -0,0 +1,15 @@
+apiVersion: checkpointing.gke.io/v1
+kind: CheckpointConfiguration
+metadata:
+  name: my-checkpointconfiguration
+spec:
+  cloudStorageBucketName:
+  # This field is optional
+  nodeSelector:
+    node.kubernetes.io/instance-type:
+  # This field is optional
+  tolerations:
+  - key:
+    operator: Exists
+    effect: NoSchedule
+  inMemoryVolumeSize:

xpk/templates/volume_bundle.yaml

@@ -0,0 +1,7 @@
+apiVersion: kjobctl.x-k8s.io/v1alpha1
+kind: VolumeBundle
+metadata:
+  name: $NAME
+spec:
+  volumes: []
+  containerVolumeMounts: []

xpk/utils/templates.py

@@ -18,7 +18,7 @@ import os
 
 import ruamel.yaml
 
-TEMPLATE_PATH = "src/xpk/templates/"
+TEMPLATE_PATH = "templates"
 
 yaml = ruamel.yaml.YAML()
 
@@ -28,3 +28,16 @@ def load(path: str) -> dict:
   with open(template_path, "r", encoding="utf-8") as file:
     data: dict = yaml.load(file)
   return data
+
+
+def get_templates_absolute_path(templates_path: str = TEMPLATE_PATH) -> str:
+  """
+  Return the absolute path to the templates folder
+
+  Args:
+    templates_path: The path to the templates folder relative to the src/xpk directory
+  """
+  current_file_path = os.path.abspath(__file__)
+  current_dir = os.path.dirname(current_file_path)
+  xpk_package_dir = os.path.dirname(current_dir)
+  return os.path.join(xpk_package_dir, templates_path)

xpk/utils/topology.py

@@ -35,3 +35,12 @@ def parse_topology(topology: str) -> list[int]:
     raise ValueError("Topology is an empty string")
 
   return [int(el) for el in topology.lower().split("x")]
+
+
+def is_topology_contained(contained: str, container: str) -> bool:
+  contained_parsed = parse_topology(contained)
+  container_parsed = parse_topology(container)
+  return len(contained_parsed) == len(container_parsed) and all(
+      contained <= container
+      for contained, container in zip(contained_parsed, container_parsed)
+  )

xpk/utils/topology_test.py

@@ -15,7 +15,7 @@ limitations under the License.
 """
 
 import pytest
-from .topology import is_topology_valid, get_topology_product, parse_topology
+from .topology import is_topology_valid, get_topology_product, parse_topology, is_topology_contained
 
 
 def test_is_topology_valid_with_invalid_topology():
@@ -41,3 +41,23 @@ def test_parse_topology_with_empty_input():
 def test_get_topology_product():
   result = get_topology_product("1x2x3")
   assert result == 6
+
+
+def test_is_topology_contained_with_container_smaller_than_contained_returns_false():
+  result = is_topology_contained(contained="3x3x3", container="2x2x2")
+  assert result is False
+
+
+def test_is_topology_contained_with_container_larger_than_contained_returns_true():
+  result = is_topology_contained(contained="1x1x1", container="2x2x2")
+  assert result is True
+
+
+def test_is_topology_contained_with_container_equal_to_contained_returns_true():
+  result = is_topology_contained(contained="2x2x2", container="2x2x2")
+  assert result is True
+
+
+def test_is_topology_contained_with_different_topologies_dimensions_returns_false():
+  result = is_topology_contained(contained="2x2", container="2x2x2")
+  assert result is False

{xpk-0.14.0.dist-info → xpk-0.14.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 0.14.0
+Version: 0.14.2
 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

{xpk-0.14.0.dist-info → xpk-0.14.2.dist-info}/RECORD

@@ -1,30 +1,38 @@
+integration/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
+integration/docker_manager_test.py,sha256=J2xijy6crRtrwQXrEvtOEY7mo1kEJYhcIYMZ7w0OGa4,2514
+integration/gcluster_a3mega_test.py,sha256=KRxgAsTbo6QBVvhhaEHSis_EvTPi2B-ZIlhEoGClcFs,6046
+integration/gcluster_a3ultra_test.py,sha256=8wEtlQN1_uIBUsidvH_l7Ab-ikDpnABrlu9k1TIBz4Q,5846
+integration/gcluster_a4_test.py,sha256=GCe6BujHCvM62kIGOd-9Wvz-IrR0BY5d83bGD1cmsQ0,5754
+integration/gcluster_test.py,sha256=3GSOMszzNW6Yr4T4PFIpmszonwDAAGpSdKutUA77O-g,3304
 xpk/__init__.py,sha256=7mu-VQDQMyxM5To0KOhuYe4y2TYGsEkfV7hXZmUyih4,561
 xpk/main.py,sha256=7YBpzpHxV61c2Js0-uUHUXxTW7cWhLLdMc263rYRW3E,2416
 xpk/api/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
 xpk/api/storage_crd.yaml,sha256=r4WFXnSJJ25EUF-t4Ljfbl-cJoSaiFiZkP8451eTub4,1260
 xpk/commands/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
 xpk/commands/batch.py,sha256=Cj1bDpzPMoPdhaKKrOJJLJ3JzRvJrCMn8huQoHHIZJI,4192
-xpk/commands/cluster.py,sha256=plS9p9PW-NPdNgvf30K9fdKo-0VyDk5gd6BX8tizcpM,41219
-xpk/commands/cluster_gcluster.py,sha256=YWNC-ljRJfZG-T1ttkPEklplMsYlBXwJ7JpGpJfdKuE,12918
+xpk/commands/cluster.py,sha256=CG4Ieg3QADNUAwyPzZh0P3mop4Pv1DILrNMJVLvgbi8,41682
+xpk/commands/cluster_gcluster.py,sha256=MOxQfQ19sxaDtDBfIzUdrxw4FpboDIiGFKEtFLwfEgQ,13080
 xpk/commands/cluster_gcluster_test.py,sha256=MoMlHbHnCI4ZrnMai1Zw71IP1ORxVVh23cMD5qgXOtQ,6136
-xpk/commands/common.py,sha256=wGsZH9qhV4rgzZ5oMvMgrU7i4Z3orNjFU-OgymLs9Lc,2308
+xpk/commands/cluster_test.py,sha256=tT1oEUXqNkzElgHZ-UXJGkJEGyakuaPb1cLHjV-eeM0,2800
+xpk/commands/common.py,sha256=nxLKPhXuAMW7wq-5xL2YeOgDTMUDcLGMujatOUHJ3s4,2504
 xpk/commands/config.py,sha256=gFNkf3ibsvZmcPpkpKXe-KJmHO5IKucNwLCXNgKvaDc,836
 xpk/commands/info.py,sha256=uhv5mPfgg9N-5JhQw4dT2jujL9ZC5kzGA18h9NFfm5A,7429
 xpk/commands/inspector.py,sha256=JruseZl9ZIlR9-Lv_pn8YHfLdyiMgHHSf7xPGBAtXTM,12616
 xpk/commands/job.py,sha256=rPIfWvgm5mLz7K7YDLK721ZcUcg5OEmYVAPAtRtB5Ag,6718
-xpk/commands/kind.py,sha256=8z9B80Xc-D8bt9iDdWPlqZqzmGFYpqdwokhmySQ5zLY,7619
+xpk/commands/kind.py,sha256=NM2CoKJXrYtY9DHZ2yM097WGcI1kno_G2b23tlLVIEI,7658
 xpk/commands/kjob_common.py,sha256=bRaORiGVjPAdN0T3aRmbcQgXYe-EtjoVKePdWzQ5xU4,1928
 xpk/commands/run.py,sha256=D0zgmnGeBLATphYhzQj29EScxrMmAKqPRhP6nfWuYcY,4085
 xpk/commands/shell.py,sha256=mRHMwm3Izzsue4bocekm82Rg_cPUaGMClSlvNzNXQ-o,4467
 xpk/commands/storage.py,sha256=90nz84_ut-uwdIVrxnLYq0K4uYpZWncKM-ZZ_2bRzcI,11505
 xpk/commands/version.py,sha256=k30rdLP9clUM8eeSwRFhpfzSb1qwcQImTfuC59Ed6CA,771
-xpk/commands/workload.py,sha256=Lx1BZx_PsQ7llML4VomurJRSTO4fC9Ohf_DolotiSoA,27894
+xpk/commands/workload.py,sha256=S6g2B8bzhM_XNGY4NJUYVODcIuvutKQmdl0VV7jwGzw,29239
+xpk/commands/workload_test.py,sha256=e6D1j2Akewq50_76oD84gSCdcW44PEOsIxLPjaL11s8,2344
 xpk/core/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
 xpk/core/capacity.py,sha256=ohrVvguNwSv7HufdhAhp9dEtf49QT9eXPaPuI9BuAck,7401
-xpk/core/cluster.py,sha256=H5-x0xYijQbgINbAIXovu6JoR-9iKnuVQrMVy1xrEP0,26359
+xpk/core/cluster.py,sha256=BoZzQO7j6Q5vfooFPCv72KDg7zYL_RFLLjxVEMKQsus,26359
 xpk/core/cluster_private.py,sha256=RLi0C7bV0NEUXl6QKQzvUT0weN9EdqPvjuuOQsNO0DY,6868
 xpk/core/commands.py,sha256=mtkT_ZsuHj5HTvFAj7mMd8kMXsWoWSA4zqHjlRiFL4o,10369
-xpk/core/config.py,sha256=JCvW7vetC9-VLEOkeGkydiSklupBRl20imP8zqjR5S4,3407
+xpk/core/config.py,sha256=EBGJC2MEmkks-9RpRsW7HFyfa4k-W-8geFIlkm3y44c,3407
 xpk/core/config_test.py,sha256=v1qfyFRzLkYSQ7Wn4nx1N0dBSOFXidLWDfhkeHDZOVM,1847
 xpk/core/docker_container.py,sha256=GvkCJ2S5UKn8uh3pZhRd3X7iS0-PsQpRO8l7QhywVGc,7604
 xpk/core/docker_image.py,sha256=MIU397IGIPwkTZFK-ZGEWuc3RmUIF3sQQZUiUj2gLqA,6775
@@ -37,8 +45,8 @@ xpk/core/gcluster_manager.py,sha256=lyv_MvdnkByy9_PEBj_ugAEBwnCbFNiWTSrEFjrMlPc,
 xpk/core/gcsfuse.py,sha256=kg5pgxdTjgiqquuGjev9fXzJPb8oiWPTK6wzCddzheQ,2125
 xpk/core/jobset.py,sha256=PJ4Fd8TNNLuYKNOMehoMYRIUEXyc5jsbHctJGqfW_8Y,4037
 xpk/core/kjob.py,sha256=symYO3DQfSNP6MAJE54QZuCPDF2kseaxZ-_mmsoBQjo,14478
-xpk/core/kueue_manager.py,sha256=ujEvPLhnnNiHAB0Bosl28OXR3WpWvOOS0K7oP1sVSMc,12129
-xpk/core/kueue_manager_test.py,sha256=BBJLPoTAB0oIA4DGL0MYx08BKcnYXUQMxoQ_8ORQShs,18623
+xpk/core/kueue_manager.py,sha256=PjZ1vjEfKvFdB5Tg5gK2qOm4J6LoHrpaafn5TizUo1A,13292
+xpk/core/kueue_manager_test.py,sha256=xN2viTlwrd5pPwn0OSn8RlD_gy13t3vpQfuN7jYFZWo,20062
 xpk/core/monitoring.py,sha256=__bzTq_DIDAK8yIaN4F3MJh-yjYw5X1OlxmRgYOpf1g,4332
 xpk/core/mtc.py,sha256=pO7p3l-EzLFdTE8MdwWV8i0Zu-7epGql_kPoksVofIU,6259
 xpk/core/nap.py,sha256=uA33XccGjEF5RZRO5IpUMMzCf-u6D73cwwPEj4q1qvc,12820
@@ -51,7 +59,8 @@ xpk/core/resources.py,sha256=HlYNPQlaJa5y-pb70aVJzdiiOLDnoACF-wJXePM4ejs,8077
 xpk/core/scheduling.py,sha256=_o8QkQxVM-8z5K5ATslM_qQ87f5dzYE0ZeBcwQ6Oqic,9702
 xpk/core/scheduling_test.py,sha256=m9KcglAbg0qly095PmrUOZxJYUE2UeQmkBNIWn5nFyk,979
 xpk/core/storage.py,sha256=NILvVAcLNMLmp4wKx_TEKbMMF5X1oL-FrQV46PT0_ds,16902
-xpk/core/system_characteristics.py,sha256=kRNWwQ6vcLyfII-VcWP1Q3Pvgpo5GRUDTcDHq9e_sPQ,22206
+xpk/core/system_characteristics.py,sha256=qHgqyQF6RVjU059YU0eB66_ia74tV-ogDyZbqCTUfw0,22589
+xpk/core/system_characteristics_test.py,sha256=MWvbaEQRoepYNZKOQcs73ppcTKW8plYQ9FQBotfA27I,2250
 xpk/core/vertex.py,sha256=orIZAVwZruRJQ6-vgc1wShuTsiipdH-zHQ9O4ie_HSA,3638
 xpk/core/workload.py,sha256=6TVZM15n8W7046VgmmH9Jv54MrhExtLQH3GaiwlV8Xs,8959
 xpk/core/workload_test.py,sha256=tVTvrwDRXD3O1GCoftgEBWilCYTN74ayP1KRP0vptx0,857
@@ -59,6 +68,7 @@ xpk/core/blueprint/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlL
 xpk/core/blueprint/blueprint_definitions.py,sha256=5i331XA-2yP_ALyB6XU5tP2Tf9iHcIX5g0TilxQi8zE,1800
 xpk/core/blueprint/blueprint_generator.py,sha256=VWhp89upY7vqVFUqcpWR6zwTBoFUSahX2sQxXEdUpSk,36393
 xpk/core/blueprint/blueprint_test.py,sha256=T058Dq-x4wQQqjs33BWhjHdT4qJLDwsIcc7vareHh_c,7204
+xpk/core/blueprint/testing/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
 xpk/core/remote_state/__init__.py,sha256=PkV8D9WOtlJHH5AIxsQaKeIBcmupT_Ol_bwJgN6G2I8,561
 xpk/core/remote_state/fuse_remote_state.py,sha256=3Dx4ZZd0NFF5-MlqGWHzz8H4bjYiPOWdF_YSEnKUPQ8,3246
 xpk/core/remote_state/remote_state_client.py,sha256=6PcR92Xy_RMjlF4AscanQ1jXNHnewLWGNC2v53jbzD4,1077
@@ -87,7 +97,18 @@ xpk/parser/version.py,sha256=eJo4PAbbmRQZulgKBs_ytbVgV9zAaaXeNzMMxmgFMVY,769
 xpk/parser/workload.py,sha256=oI6y66Old_z6PxFPTZ3LvektUsOZ7U8f6g6cscuzh9g,27208
 xpk/parser/workload_test.py,sha256=t0aAqEsAMz1U3xTMrBcm773d9kwYOcvLtQk5n4jqWPw,2075
 xpk/templates/__init__.py,sha256=7mu-VQDQMyxM5To0KOhuYe4y2TYGsEkfV7hXZmUyih4,561
+xpk/templates/cluster_preheat.yaml.j2,sha256=1e8jYagQE6O7BjAfuwmEqGG1b8AOsLRlQm4V68ZnGNs,721
+xpk/templates/filestore-pv.yaml,sha256=FxKZkAXa2czIYblq77iewQjCjOjs-FptuF3YLOByfLo,316
+xpk/templates/filestore-pvc.yaml,sha256=Rf80UNYs3XTUdOJuWCeFq80TKXq5FhfafRE89hq7y9o,161
+xpk/templates/filestore-sc.yaml,sha256=vHzcU7jk0B5z7EgTfQmMM1m2TIzR70ny6MvthyYdqhE,213
+xpk/templates/fuse-pv.yaml,sha256=-CM6AYAy4HKOErd9ogiM-6vWIuWS5yXoXyBAp2EoZsM,321
+xpk/templates/fuse-pvc.yaml,sha256=heGWvRIetukAI9pH9auXxnU2H-G_8pL_wRIzoWDLVH8,218
+xpk/templates/kueue_config.yaml.j2,sha256=ZAZwzZ28piwlXi9Offo2CvrQ3K9gv1r9BJTPavnUEdY,2055
+xpk/templates/kueue_gke_default_topology.yaml.j2,sha256=GDFUf0DUnVTVjjSVRqfyEhQoDj2sEoL2YWhHYJZU9N4,300
+xpk/templates/kueue_sub_slicing_topology.yaml.j2,sha256=Lv54uS3fC7sJNJdk4DFtnOiLSgq4wq38n_BQCx4Tz6Q,542
+xpk/templates/mtc-cpc.yaml,sha256=MPx75tog09kjRAvHoNOPCEobigQ17d7pYCUnZCevSDQ,340
 xpk/templates/storage.yaml,sha256=AykdyMtDnKZF8Y_0BYxoYP03hEIzEk6iNalXAQHgAls,163
+xpk/templates/volume_bundle.yaml,sha256=sqeag7GPWqGNQ5doZtO9IVAX_vKYRO73-aBE7waEtSY,129
 xpk/utils/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
 xpk/utils/console.py,sha256=hRbvtog_VAzuxt5GfwK5GZdd5SWaa7kvWG8zo_qFRQc,1519
 xpk/utils/execution_context.py,sha256=WYxm6NExBIP6iLAWaL5aV858riGJbAHn0Zs6fmKlmzE,784
@@ -98,15 +119,15 @@ xpk/utils/kubectl.py,sha256=WKB9UhpouPN9G4n2ejRi_PgsYLI0R01gzkS1WGU6mJA,1828
 xpk/utils/kueue.py,sha256=P1Pu_crGuOgYxjl8CczTgtQoum0w1sbSLGPOaEZ5180,713
 xpk/utils/network.py,sha256=dGS5rxIm_zaayDElHNlzalaf09M99by5ckL_lGDl_yQ,4293
 xpk/utils/objects.py,sha256=OwMNxB4TGX21qnJPdZo2YBMPMbQPqOtHMh19QhoRNRY,2498
-xpk/utils/templates.py,sha256=mH7up63_-vev7TQF_n2IwhEZHarbzSqCKPRDVP51gVc,845
-xpk/utils/topology.py,sha256=heKgMQxeTEYCD-PVp8m_c7rscdGcGVLx2_kHCDQca50,1039
-xpk/utils/topology_test.py,sha256=UiCurxYGTDrqVlGlAeg-I95-ScOEZQq3roSSwtdG10U,1196
+xpk/utils/templates.py,sha256=5VAUtv-F6ICL5mxZ3Xtzdh8FEc0-86jFbqhgwW-QtcM,1277
+xpk/utils/topology.py,sha256=WcAG3kzA8krlbyiZpTRDI2TteNozn3r1GH914Lykcp8,1374
+xpk/utils/topology_test.py,sha256=jDXCPgBPfByqjhi0W9A5c8uOHOYjRav53nNlg71ipjk,1943
 xpk/utils/validation.py,sha256=-Qd5jqkVzQHkJvmQnGjHjtAcfvz064Vbo6_sl4EnYKw,3497
 xpk/utils/validation_test.py,sha256=PEDSMUqZdt_Lx1FSR-LOTXKKtsJ47JH1fxugM0Gfz6Y,1168
 xpk/utils/yaml.py,sha256=j8xuAJ9yAAwnQi6ozwZ-nMnDyDnc3xWkeBZMtSuP4RU,844
-xpk-0.14.0.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-xpk-0.14.0.dist-info/METADATA,sha256=Y2-MopqRDfoUx-owouglTHeDP-88ruOPr4o6oWIBmSE,71934
-xpk-0.14.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xpk-0.14.0.dist-info/entry_points.txt,sha256=mzEtiIesFkT1kmcTUVDA1o3uOhiniX6tIz2wmOlMu1M,38
-xpk-0.14.0.dist-info/top_level.txt,sha256=aDe4N0jicmuWExx_6w0TxWQJaEuPSs9BnLU-3aF1GLo,4
-xpk-0.14.0.dist-info/RECORD,,
+xpk-0.14.2.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+xpk-0.14.2.dist-info/METADATA,sha256=XP5pIZaZR0oSR9WDn_vbpq5zHRURGHGJ8kRuVw4-s0I,71934
+xpk-0.14.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xpk-0.14.2.dist-info/entry_points.txt,sha256=mzEtiIesFkT1kmcTUVDA1o3uOhiniX6tIz2wmOlMu1M,38
+xpk-0.14.2.dist-info/top_level.txt,sha256=TQKZWgV7LSElvmunYT9V_627qOMoxq3qYzWAFzKudB8,16
+xpk-0.14.2.dist-info/RECORD,,

xpk-0.14.2.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+integration
+xpk

xpk-0.14.0.dist-info/top_level.txt

@@ -1 +0,0 @@
-xpk