xpk 0.13.0__py3-none-any.whl → 0.14.1__py3-none-any.whl

xpk-0.14.1.dist-info/RECORD

@@ -0,0 +1,133 @@
+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=CG4Ieg3QADNUAwyPzZh0P3mop4Pv1DILrNMJVLvgbi8,41682
+xpk/commands/cluster_gcluster.py,sha256=MOxQfQ19sxaDtDBfIzUdrxw4FpboDIiGFKEtFLwfEgQ,13080
+xpk/commands/cluster_gcluster_test.py,sha256=MoMlHbHnCI4ZrnMai1Zw71IP1ORxVVh23cMD5qgXOtQ,6136
+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=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=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_private.py,sha256=RLi0C7bV0NEUXl6QKQzvUT0weN9EdqPvjuuOQsNO0DY,6868
+xpk/core/commands.py,sha256=mtkT_ZsuHj5HTvFAj7mMd8kMXsWoWSA4zqHjlRiFL4o,10369
+xpk/core/config.py,sha256=G7MKXARftcuoKSV90F40HXh7kYj1ePde8O8dCp95QoI,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
+xpk/core/docker_manager.py,sha256=JBFgyD6O7LKwEHJC7YuSoCDZqrFRtb-LjgWNqkfAbR0,10566
+xpk/core/docker_resources.py,sha256=_aKgpUjyJB2krQ1PkHrotB7K4kByLmPLbuvl_UVvuX8,12843
+xpk/core/filestore.py,sha256=mcuUzsAPARbnrBG4fIGsEoN8NmzjaQ6k0tvIwMtjO9k,8068
+xpk/core/gcloud_context.py,sha256=xZdVoRNLlE-kwXY5djoyQ0I0-KEh5nAohrVql7Jl42k,6649
+xpk/core/gcloud_context_test.py,sha256=YY0R6j-m62coVK2MAjWXDIdxdP6J5yn6R1RiTDkuExQ,2719
+xpk/core/gcluster_manager.py,sha256=lyv_MvdnkByy9_PEBj_ugAEBwnCbFNiWTSrEFjrMlPc,6236
+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=TFdl33N-b-eSrOrsnJueP1H4qYWpkX_lFDhLSVvfeGM,13292
+xpk/core/kueue_manager_test.py,sha256=PtsSSfwtdk-2wK_dr_OQb8oYWVce5nmDN6iowND_3b4,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
+xpk/core/network.py,sha256=Oulb7U69lWkpOKxOC1C7ekJDpC51TLwd7XdZA3NQ7E0,10505
+xpk/core/nodepool.py,sha256=LfS_RvKHmS7f97hc_UidUC13FFEyGvfageSlt8d-5hw,23227
+xpk/core/nodepool_test.py,sha256=YNvp8WXznAI8DscrN9-BSnqII2AtRju2guxrobvCH8A,8805
+xpk/core/pathways.py,sha256=s-h_ofMrbFn3J6NFmT5OMe_HiUQIkI90ty7xbS05iA8,10710
+xpk/core/ray.py,sha256=JWhc_ToRHpF4_URGnuE_47FMgamaRsA4KVUMpqThWzw,6145
+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=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
+xpk/core/blueprint/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
+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
+xpk/core/workload_decorators/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
+xpk/core/workload_decorators/rdma_decorator.py,sha256=isbgPnjdu2AT_Da1nVUIRoGE_qZ7jMDOKCgZOLq5r2A,4006
+xpk/core/workload_decorators/storage_decorator.py,sha256=DDYQVO1OKTLhveDOA4V6b2RWr4n0fbwHdnoFFmW7iaQ,2000
+xpk/core/workload_decorators/tcpx_decorator.py,sha256=m5EgzEHjbcOD13ygY91mQdhwQt4Gr5PyalVkKcHyeV8,5975
+xpk/core/workload_decorators/tcpx_decorator_test.py,sha256=iTBS3X_-VwA2oveNDjscduLtll0VOJyFRCp4xmsjg7w,8515
+xpk/core/workload_decorators/tcpxo_decorator.py,sha256=_nLX7tbnxhnS-xv4Jijd1JOP76V4LpNCfW3Np404Cqw,6537
+xpk/parser/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
+xpk/parser/batch.py,sha256=mJU-Cp1yTLje59vD-B1IiBcUeD-ZmEsoeB4xhj9cflc,1406
+xpk/parser/cluster.py,sha256=1cc2pCMbceF6Lx6MXINlgEfiIRUYMITha3eCE6R8SrI,30708
+xpk/parser/cluster_test.py,sha256=h7lILF2V37WzQSZjGD2jCqI1iWJNY7BLODZf1uhK9QA,1806
+xpk/parser/common.py,sha256=1wihGYBe01ci4quZnkeR60-u_lapSOG4WhRFrE3yFR8,8108
+xpk/parser/config.py,sha256=-XnWx9aFsBW4Uzo_hpOMD2ZQ0bdZLvq1ksv83_5jqSM,1633
+xpk/parser/core.py,sha256=VRJerlS92ufoQbG1mZv7B04DAP4qGkBHa4pRXgcbAs0,4761
+xpk/parser/info.py,sha256=UJohxVVWdt9IgUXoPsrVae2DN1BjAVGWrSN2ajrB8RQ,1860
+xpk/parser/inspector.py,sha256=hAPAZ2k9iSJgC1mjnz3rMleInsAQ8PmkyyUKFyBmsgY,1997
+xpk/parser/job.py,sha256=5RdE70rucGfrsn65l7Ho6RmO06mag1S0AO-3saVuXyw,4328
+xpk/parser/kind.py,sha256=sgPCqNVrgmFLcOBEbhlaphwVXxMh_opP9ntCq4KPePE,2682
+xpk/parser/run.py,sha256=oi_ksSyJ8Ooffe2EgoV_ecpmXEmNGVotjpIQH-HjufE,1481
+xpk/parser/shell.py,sha256=VC8p-kz9XjJZW9DXZ-rnv41XnRDRpQRFywHpB5j7tfc,1970
+xpk/parser/storage.py,sha256=XNynqulEzTmT8_G6wkeBwfXX0XQ1lsd6BFcx0H6rGfU,9971
+xpk/parser/validators.py,sha256=-NBZelvfwZRzjz-YUCreD8EzMLHll8PZM-d-MVm2PG4,1192
+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=wW3qt6p3VDPgFVX7Ozw4-O4QgQ3mhH8U3osKnmuOFaE,299
+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
+xpk/utils/feature_flags.py,sha256=b8zB_zEiKcoOnyGHc7vuQN0ruRTdY8ixqUjWT7Ilp-M,824
+xpk/utils/file.py,sha256=hi9v4gfwiB3JHi3tnelPbm_dlTUt47U0wvvWKQqMjiQ,2500
+xpk/utils/gcs_utils.py,sha256=zg-XSTv4G4TFjeT2bNBm2WLdDXPrOZi0rNv_JdppNg4,4113
+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=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.1.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+xpk-0.14.1.dist-info/METADATA,sha256=41WVkaR-Bgj64aC9vi26hgI9g6CpnDtEF1G3yXhbpzI,71934
+xpk-0.14.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xpk-0.14.1.dist-info/entry_points.txt,sha256=mzEtiIesFkT1kmcTUVDA1o3uOhiniX6tIz2wmOlMu1M,38
+xpk-0.14.1.dist-info/top_level.txt,sha256=TQKZWgV7LSElvmunYT9V_627qOMoxq3qYzWAFzKudB8,16
+xpk-0.14.1.dist-info/RECORD,,

xpk-0.14.1.dist-info/top_level.txt

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

xpk/core/kueue.py

@@ -1,561 +0,0 @@
-"""
-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 argparse import Namespace
-
-import math
-import packaging
-from packaging.version import Version
-
-from ..utils.console import xpk_exit, xpk_print
-from ..utils.file import write_tmp_file
-from .capacity import B200_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
-from .commands import (
-    run_command_for_value,
-    run_command_with_updates,
-    run_command_with_updates_retry,
-)
-from .pathways import add_pw_resource_flavors, add_pw_resources_to_kueue
-from .resources import AutoprovisioningConfig
-from .scheduling import (
-    create_accelerator_label,
-    create_machine_label,
-    get_total_chips_requested_from_args,
-)
-from .system_characteristics import (
-    AcceleratorTypeToAcceleratorCharacteristics,
-    SystemCharacteristics,
-)
-
-KUEUE_VERSION = 'v0.12.2'
-CLUSTER_QUEUE_NAME = 'cluster-queue'
-LOCAL_QUEUE_NAME = 'multislice-queue'
-WAIT_FOR_KUEUE_TIMEOUT = '10m'
-MEMORY_SIZE_PER_VM = 1.2
-MIN_MEMORY_LIMIT_SIZE = 4096
-
-packaging.version.VERSION_PATTERN = r'^v\d+\.\d+\.\d+$'
-
-topology_yaml = """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"
----
-"""
-
-cluster_set_crd_yaml = """apiVersion: kueue.x-k8s.io/v1beta1
-kind: ResourceFlavor
-metadata:
-  name: {cluster_hardware_name}
-spec:
-  nodeLabels:
-    {accelerator_label}
-    {machine_label}
-  {topology_label}
----
-apiVersion: kueue.x-k8s.io/v1beta1
-kind: AdmissionCheck
-metadata:
-  name: dws-prov
-spec:
-  controllerName: kueue.x-k8s.io/provisioning-request
-  parameters:
-    apiGroup: kueue.x-k8s.io
-    kind: ProvisioningRequestConfig
-    name: dws-config
----
-apiVersion: kueue.x-k8s.io/v1beta1
-kind: ProvisioningRequestConfig
-metadata:
-  name: dws-config
-spec:
-  provisioningClassName: queued-provisioning.gke.io
-  podSetUpdates:
-    nodeSelector:
-    - key: autoscaling.gke.io/provisioning-request
-      valueFromProvisioningClassDetail: ResizeRequestName
-  managedResources:
-  - {managed_resource}
----
-{pw_resource_flavors}
-apiVersion: kueue.x-k8s.io/v1beta1
-kind: ClusterQueue
-metadata:
-  name: {cluster_queue_name}
-spec:
-  preemption:
-      reclaimWithinCohort: Never # Don't preempt other queues in the cohort.
-      withinClusterQueue: LowerPriority
-  namespaceSelector: {{}} # match all.
-  resourceGroups:
-  {covered_resources_config}
-  {pw_resources_kueue}
-  {admission_checks}
----
-apiVersion: kueue.x-k8s.io/v1beta1
-kind: LocalQueue
-metadata:
-  namespace: default
-  name: {local_queue_name}
-spec:
-  clusterQueue: {cluster_queue_name}
----
-apiVersion: scheduling.k8s.io/v1
-kind: PriorityClass
-metadata:
-  name: very-low
-value: 100
-globalDefault: false
-description: "Very Low"
----
-apiVersion: scheduling.k8s.io/v1
-kind: PriorityClass
-metadata:
-  name: low
-value: 250
-globalDefault: false
-description: "Low"
----
-apiVersion: scheduling.k8s.io/v1
-kind: PriorityClass
-metadata:
-  name: medium
-value: 500
-globalDefault: false
-description: "Medium"
----
-apiVersion: scheduling.k8s.io/v1
-kind: PriorityClass
-metadata:
-  name: high
-value: 750
-globalDefault: false
-description: "High"
----
-apiVersion: scheduling.k8s.io/v1
-kind: PriorityClass
-metadata:
-  name: very-high
-value: 1000
-globalDefault: false
-description: "Very High"
-"""
-
-cluster_preheat_yml = """
-apiVersion: apps/v1
-kind: DaemonSet
-metadata:
-  name: {cachekey}
-  labels:
-    k8s-app: {cachekey}
-spec:
-  selector:
-    matchLabels:
-      k8s-app: {cachekey}
-  updateStrategy:
-    type: RollingUpdate
-  template:
-    metadata:
-      labels:
-        name: {cachekey}
-        k8s-app: {cachekey}
-    spec:
-      affinity:
-        nodeAffinity:
-          requiredDuringSchedulingIgnoredDuringExecution:
-            nodeSelectorTerms:
-            - matchExpressions:
-              - key: {nodeSelectorKey}
-                operator: Exists
-      tolerations:
-      - operator: "Exists"
-      containers:
-      - image: {image_name}
-        name: {cachekey}
-        command: [ "sleep", "inf" ]
-"""
-
-kueue_controller_manager_yml = """
-apiVersion: apps/v1
-kind: Deployment
-metadata:
-  labels:
-    app.kubernetes.io/component: controller
-    app.kubernetes.io/name: kueue
-    control-plane: controller-manager
-  name: kueue-controller-manager
-  namespace: kueue-system
-spec:
-  replicas: 1
-  selector:
-    matchLabels:
-      control-plane: controller-manager
-  template:
-    metadata:
-      annotations:
-        kubectl.kubernetes.io/default-container: manager
-      labels:
-        app.kubernetes.io/component: controller
-        app.kubernetes.io/name: kueue
-        control-plane: controller-manager
-    spec:
-      containers:
-      - args:
-        - --config=/controller_manager_config.yaml
-        - --zap-log-level=2
-        command:
-        - /manager
-        image: registry.k8s.io/kueue/kueue:{KUEUE_VERSION}
-        imagePullPolicy: Always
-        livenessProbe:
-          httpGet:
-            path: /healthz
-            port: 8081
-          initialDelaySeconds: 15
-          periodSeconds: 20
-        name: manager
-        ports:
-        - containerPort: 8082
-          name: visibility
-          protocol: TCP
-        - containerPort: 9443
-          name: webhook-server
-          protocol: TCP
-        readinessProbe:
-          httpGet:
-            path: /readyz
-            port: 8081
-          initialDelaySeconds: 5
-          periodSeconds: 10
-        resources:
-          limits:
-            cpu: 1000m
-            memory: {memory_limit_size}
-          requests:
-            cpu: 1000m
-            memory: 512Mi
-        securityContext:
-          allowPrivilegeEscalation: false
-        volumeMounts:
-        - mountPath: /visibility
-          name: visibility
-        - mountPath: /tmp/k8s-webhook-server/serving-certs
-          name: cert
-          readOnly: true
-        - mountPath: /controller_manager_config.yaml
-          name: manager-config
-          subPath: controller_manager_config.yaml
-      securityContext:
-        runAsNonRoot: true
-      serviceAccountName: kueue-controller-manager
-      terminationGracePeriodSeconds: 10
-      volumes:
-      - name: visibility
-        emptyDir: {{}}
-      - name: cert
-        secret:
-          defaultMode: 420
-          secretName: kueue-webhook-server-cert
-      - configMap:
-          name: kueue-manager-config
-        name: manager-config
-"""
-
-
-def verify_kueuectl(args: Namespace) -> None:
-  """Verify if kueuectl is installed.
-  Args:
-    args: user provided arguments.
-  Returns:
-    None
-  """
-  xpk_print('Veryfing kueuectl installation')
-
-  command = 'kubectl kueue version'
-  task = 'Verify kueuectl installation on cluster'
-  verify_kueuectl_installed_code, _ = run_command_for_value(command, task, args)
-
-  if verify_kueuectl_installed_code == 0:
-    xpk_print('kueuectl found')
-
-  if verify_kueuectl_installed_code != 0:
-    xpk_print(
-        'kueuectl not found. Please follow'
-        ' https://kueue.sigs.k8s.io/docs/reference/kubectl-kueue/installation/'
-        ' to install kueuectl.'
-    )
-    xpk_exit(verify_kueuectl_installed_code)
-
-
-def delete_multikueueconfigs_definitions(args) -> int:
-  command = 'kubectl delete crd multikueueconfigs.kueue.x-k8s.io'
-  task = 'Delete multikueueconfigs crds'
-  return_code = run_command_with_updates_retry(command, task, args)
-  if return_code != 0:
-    xpk_print(f'{task} returned ERROR {return_code}')
-  return return_code
-
-
-def delete_multikueueclusters_definitions(args) -> int:
-  command = 'kubectl delete crd multikueueclusters.kueue.x-k8s.io'
-  task = 'Delete multikueueclusters crds'
-  return_code = run_command_with_updates_retry(command, task, args)
-  if return_code != 0:
-    xpk_print(f'{task} returned ERROR {return_code}')
-  return return_code
-
-
-def get_kueue_version(args) -> tuple[int, str]:
-  command = 'kubectl kueue version'
-  task = 'Get kueue version on server'
-  return_code, val = run_command_for_value(command, task, args)
-  if return_code != 0:
-    return return_code, ''
-  lines = val.splitlines()
-  if len(lines) == 1:
-    return 1, ''
-  server_version_line = lines[1]
-  manager_image_version = server_version_line.split(':')[-1]
-  return return_code, manager_image_version
-
-
-def install_kueue_on_cluster(args) -> int:
-  """Install Kueue on the cluster.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-
-  err_code, kueue_version_installed = get_kueue_version(args)
-  if err_code == 0:
-    if Version(kueue_version_installed) < Version('v0.9.0') and Version(
-        KUEUE_VERSION
-    ) >= Version('v0.9.0'):
-      xpk_print('Upgrading kueue on cluster from version < 0.9.0.')
-      upgrade_code = delete_multikueueclusters_definitions(args)
-      if upgrade_code != 0:
-        return upgrade_code
-      upgrade_code = delete_multikueueconfigs_definitions(args)
-      if upgrade_code != 0:
-        return upgrade_code
-
-  command = (
-      'kubectl apply --server-side --force-conflicts -f'
-      f' https://github.com/kubernetes-sigs/kueue/releases/download/{KUEUE_VERSION}/manifests.yaml'
-  )
-  task = 'Set Kueue On Cluster'
-  return_code = run_command_with_updates_retry(command, task, args)
-  if return_code != 0:
-    xpk_print(f'{task} returned ERROR {return_code}')
-  return return_code
-
-
-def wait_for_kueue_available(args: Namespace) -> int:
-  """Wait for Kueue to be fully available.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  command = (
-      'kubectl wait deploy/kueue-controller-manager -nkueue-system'
-      f' --for=condition=available --timeout={WAIT_FOR_KUEUE_TIMEOUT}'
-  )
-  task = 'Wait for Kueue to be available'
-  return_code = run_command_with_updates(command, task, args)
-  if return_code != 0:
-    xpk_print(f'{task} returned ERROR {return_code}')
-  return return_code
-
-
-def install_kueue_crs(
-    args,
-    system: SystemCharacteristics,
-    autoprovisioning_config: AutoprovisioningConfig | None,
-    flex_with_tpu=False,
-) -> int:
-  """Install Kueue Custom Resources.
-
-  Args:
-    args: user provided arguments for running the command.
-    system: system level arguments.
-    autoprovisioning_config: Autoprovisioning config to configure kueue with if
-        autoprovisioning is enabled.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  device_type = system.device_type
-  cluster_hardware_name = f'{args.num_slices}x{device_type}'
-  resource_type = AcceleratorTypeToAcceleratorCharacteristics[
-      system.accelerator_type
-  ].resource_type
-
-  autoprovisioning_enabled = False
-  if autoprovisioning_config:
-    # Determine total resources available based on autoprovisioning max chips.
-    autoprovisioning_enabled = True
-    total_chips = autoprovisioning_config.maximum_chips
-    cluster_hardware_name = f'{system.gke_accelerator}'
-  else:
-    # Determine total chips based on user specified topology.
-    total_chips = get_total_chips_requested_from_args(args, system)
-  if args.flex and flex_with_tpu is False:
-    admission_checks = """
-  admissionChecks:
-  - dws-prov
-"""
-  else:
-    admission_checks = ''
-
-  covered_resources_config = get_kueue_covered_resources_config(
-      cluster_hardware_name=cluster_hardware_name,
-      resource_type=resource_type,
-      total_chips=total_chips,
-      cpu_limit=args.cpu_limit,
-      memory_limit=args.memory_limit,
-  )
-  topology_label = ''
-  if system.device_type in [
-      H100_MEGA_DEVICE_TYPE,
-      H200_DEVICE_TYPE,
-      B200_DEVICE_TYPE,
-  ]:
-    topology_label = 'topologyName: "gke-default"'
-  res_type = AcceleratorTypeToAcceleratorCharacteristics[
-      system.accelerator_type
-  ].resource_type
-  yml_string = cluster_set_crd_yaml.format(
-      system=system,
-      cluster_hardware_name=cluster_hardware_name,
-      accelerator_label=create_accelerator_label(
-          system.accelerator_type, system
-      ),
-      machine_label=create_machine_label(
-          system.accelerator_type, system, autoprovisioning_enabled
-      ),
-      topology_label=topology_label,
-      covered_resources_config=covered_resources_config,
-      resource_type=res_type,
-      pw_resource_flavors=add_pw_resource_flavors(args),
-      pw_resources_kueue=add_pw_resources_to_kueue(args),
-      admission_checks=admission_checks,
-      managed_resource=res_type,
-      cluster_queue_name=CLUSTER_QUEUE_NAME,
-      local_queue_name=LOCAL_QUEUE_NAME,
-  )
-  if system.device_type in [
-      H100_MEGA_DEVICE_TYPE,
-      H200_DEVICE_TYPE,
-      B200_DEVICE_TYPE,
-  ]:
-    yml_string = topology_yaml + yml_string
-
-  tmp = write_tmp_file(yml_string)
-  command = f'kubectl apply -f {str(tmp)}'
-
-  task = 'Applying Kueue Custom Resources'
-  return_code = run_command_with_updates_retry(command, task, args)
-  if return_code != 0:
-    xpk_print(f'{task} returned ERROR {return_code}')
-  return return_code
-
-
-def get_kueue_covered_resources_config(
-    cluster_hardware_name, resource_type, total_chips, cpu_limit, memory_limit
-) -> str:
-  """Gets Kueue covered resources configuration.
-
-  Args:
-    cluster_hardware_name: cluster hardware name.
-    resource_type: resource type of tpu or gpu.
-    total_chips: total number of chips for the specific resource type.
-
-  Returns:
-    A string of Kueue covered resources configuration.
-  """
-  config_format = """
-  - coveredResources: {resource_types}
-    flavors:
-    - name: {cluster_hardware_name}
-      resources:
-      - name: "{resource_type}"
-        nominalQuota: {total_chips}"""
-  resource_types = [resource_type]
-  if cpu_limit:
-    config_format = config_format + """
-      - name: "cpu"
-        nominalQuota: {cpu_limit}"""
-    resource_types.append('cpu')
-  if memory_limit:
-    config_format = config_format + """
-      - name: "memory"
-        nominalQuota: {memory_limit}"""
-    resource_types.append('memory')
-
-  config_string = config_format.format(
-      cluster_hardware_name=cluster_hardware_name,
-      resource_types=resource_types,
-      resource_type=resource_type,
-      total_chips=total_chips,
-      cpu_limit=cpu_limit,
-      memory_limit=memory_limit,
-  )
-  return config_string
-
-
-def update_kueue_resources_if_necessary(args):
-  """Update the kueue manifest to increase the resources for the kueue controller manager.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  # Get total number of nodes
-  cmd_total_node_num = 'kubectl get node --no-headers | wc -l'
-  return_code, out = run_command_for_value(
-      cmd_total_node_num, 'Count total nodes', args
-  )
-  if return_code != 0:
-    xpk_exit(1)
-  # 1.2MiB per VM or 4GiB (whichever is greater).
-  new_memory_limit = (
-      f'{max(math.ceil(int(out) * MEMORY_SIZE_PER_VM), MIN_MEMORY_LIMIT_SIZE)}Mi'
-  )
-  yml_string = kueue_controller_manager_yml.format(
-      memory_limit_size=new_memory_limit, KUEUE_VERSION=KUEUE_VERSION
-  )
-  tmp = write_tmp_file(yml_string)
-  command = f'kubectl apply -f {str(tmp)}'
-
-  task = 'Updating Kueue Controller Manager resources'
-  return_code = run_command_with_updates_retry(command, task, args)
-  if return_code != 0:
-    xpk_print(f'{task} returned ERROR {return_code}')
-  return return_code