xpk 0.17.1__py3-none-any.whl → 0.17.3__py3-none-any.whl

xpk/parser/core.py

@@ -23,13 +23,9 @@ from .cluster import set_cluster_parser
 from .inspector import set_inspector_parser
 from .storage import set_storage_parser
 from .workload import set_workload_parsers
-from .batch import set_batch_parser
-from .job import set_job_parser
 from .info import set_info_parser
 from .kind import set_kind_parser
-from .shell import set_shell_parser
 from .version import set_version_parser
-from .run import set_run_parser
 
 
 def set_parser(parser: argparse.ArgumentParser):
@@ -54,20 +50,10 @@ def set_parser(parser: argparse.ArgumentParser):
       "info",
       help="Commands around listing kueue clusterqueues and localqueues.",
   )
-  batch_parser = xpk_subcommands.add_parser(
-      "batch",
-      help="commands around running batch job",
-  )
-  job_parser = xpk_subcommands.add_parser(
-      "job", help="commands around listing, cancelling and investigating jobs"
-  )
   kind_parser = xpk_subcommands.add_parser(
       "kind",
       help="commands around Kind cluster management",
   )
-  shell_parser = xpk_subcommands.add_parser(
-      "shell", help="Commands around configuring and using interactive shell."
-  )
   version_parser = xpk_subcommands.add_parser(
       "version", help="Command to get xpk version"
   )
@@ -76,11 +62,6 @@ def set_parser(parser: argparse.ArgumentParser):
       "config", help="Commands to set and retrieve values from xpk config."
   )
 
-  run_parser = xpk_subcommands.add_parser(
-      "run",
-      help="Command to run parallel jobs",
-  )
-
   def default_subcommand_function(
       _args,
   ) -> int:  # args is unused, so pylint: disable=invalid-name
@@ -96,14 +77,10 @@ def set_parser(parser: argparse.ArgumentParser):
     parser.print_help()
     cluster_parser.print_help()
     workload_parser.print_help()
-    batch_parser.print_help()
     info_parser.print_help()
-    job_parser.print_help()
-    shell_parser.print_help()
     version_parser.print_help()
     kind_parser.print_help()
     config_parser.print_help()
-    run_parser.print_help()
 
     storage_parser.print_help()
     return 0
@@ -111,25 +88,17 @@ def set_parser(parser: argparse.ArgumentParser):
   parser.set_defaults(func=default_subcommand_function)
   workload_parser.set_defaults(func=default_subcommand_function)
   cluster_parser.set_defaults(func=default_subcommand_function)
-  batch_parser.set_defaults(func=default_subcommand_function)
   info_parser.set_defaults(func=default_subcommand_function)
-  job_parser.set_defaults(func=default_subcommand_function)
   kind_parser.set_defaults(func=default_subcommand_function)
-  shell_parser.set_defaults(func=default_subcommand_function)
   storage_parser.set_defaults(func=default_subcommand_function)
   version_parser.set_defaults(func=default_subcommand_function)
   config_parser.set_defaults(func=default_subcommand_function)
-  run_parser.set_defaults(func=default_subcommand_function)
 
   set_workload_parsers(workload_parser=workload_parser)
   set_cluster_parser(cluster_parser=cluster_parser)
   set_inspector_parser(inspector_parser=inspector_parser)
-  set_batch_parser(batch_parser=batch_parser)
   set_info_parser(info_parser=info_parser)
-  set_job_parser(job_parser=job_parser)
   set_kind_parser(kind_parser=kind_parser)
-  set_shell_parser(shell_parser=shell_parser)
   set_storage_parser(storage_parser=storage_parser)
   set_version_parser(version_parser=version_parser)
   set_config_parsers(config_parser=config_parser)
-  set_run_parser(run_parser=run_parser)

xpk/utils/validation.py

@@ -37,14 +37,6 @@ class SystemDependency(Enum):
           ' to install xpk prerequisites.'
       ),
   )
-  KJOB = _SystemDependency(
-      command='kubectl kjob --help',
-      message=(
-          '`kjobctl` not installed. Please follow'
-          ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
-          ' to install xpk prerequisites.'
-      ),
-  )
   GCLOUD = _SystemDependency(
       command='gcloud version',
       message=(

{xpk-0.17.1.dist-info → xpk-0.17.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 0.17.1
+Version: 0.17.3
 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.17.1.dist-info → xpk-0.17.3.dist-info}/RECORD

@@ -20,34 +20,29 @@ xpk/blueprints/a4/config-map.yaml.tftpl,sha256=o6LeGIYUfFGyj3vj-8ztV5ildQ46QZVl7
 xpk/blueprints/a4/nccl-rdma-installer-a4.yaml,sha256=if3WOmNLVGTJIJHU76EWC1FyiIXDTRIXcwo4OsBxarQ,2113
 xpk/blueprints/a4/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=DtMiIYdYsciXldoWqAfxPIxl9Hc9kbYIj2LsdBER0PI,46172
-xpk/commands/cluster_gcluster.py,sha256=x26UqoT8RFX5T9ftQXPEL12HMnMFTi8lret16dnZCms,13970
-xpk/commands/cluster_gcluster_test.py,sha256=UcqTTkrQv-R753AtsQvinwgI2vqI6lMHPPEfHPS5e-4,6655
-xpk/commands/cluster_test.py,sha256=-7EjuOoGSZhdnLBNBNCMKM6laDYy02aPncbSfUYcrUs,24147
+xpk/commands/cluster.py,sha256=5ebvHXe8Bi4haMd1RokJbjP0LemqXxC1EseqOLWIkGw,45659
+xpk/commands/cluster_gcluster.py,sha256=Ig8jLjsiyFgw9U4BBEzDK2diA9m0STKQgz-uUTG_vYE,13731
+xpk/commands/cluster_gcluster_test.py,sha256=s1wwkcdY4LTxmk_Tx5PKdh9pZmEoo1n8XlzybWalc0M,6165
+xpk/commands/cluster_test.py,sha256=aMkwKrhoEuqElME16ztx5lwv4zT0z_xV0L3in1RaW6M,24017
 xpk/commands/common.py,sha256=p43sspD5RfYRj3Se_b-X0s0dbBs1PMI1qtySg6zZKKg,2706
 xpk/commands/config.py,sha256=L_zRpQTxMcSh6rxOT8gG263V6YGqzVoz4UxdWywTFdA,850
 xpk/commands/info.py,sha256=uhv5mPfgg9N-5JhQw4dT2jujL9ZC5kzGA18h9NFfm5A,7429
 xpk/commands/inspector.py,sha256=FPasKtGuEZKNXIQin4AG49clfD4b53NxXpWqBPZIIoE,12955
-xpk/commands/job.py,sha256=rPIfWvgm5mLz7K7YDLK721ZcUcg5OEmYVAPAtRtB5Ag,6718
-xpk/commands/kind.py,sha256=GNqsaoLInifFQ_ZGpbN_3xA8ExyeyOqBMdnoPV-PqYI,7813
-xpk/commands/kjob_common.py,sha256=bRaORiGVjPAdN0T3aRmbcQgXYe-EtjoVKePdWzQ5xU4,1928
+xpk/commands/kind.py,sha256=ck4zaJh9kaNluEdBew9OVj4cZXjUYk_ElycqGLo6f7g,7355
 xpk/commands/managed_ml_diagnostics.py,sha256=87wmFbnYQY-kEpJfPo1Up53xM5P_P5wOlXczxHzxJjQ,6984
 xpk/commands/managed_ml_diagnostics_test.py,sha256=pQ1YUGMGRQFJYTS_1o9YyGUzYdLaBdA84LjbnncaeEo,3828
-xpk/commands/run.py,sha256=D0zgmnGeBLATphYhzQj29EScxrMmAKqPRhP6nfWuYcY,4085
-xpk/commands/shell.py,sha256=mRHMwm3Izzsue4bocekm82Rg_cPUaGMClSlvNzNXQ-o,4467
-xpk/commands/storage.py,sha256=kPViq6mrfGeAJwScdMs_kUJg-QxEO6SrEvyBbXhCzEI,11439
+xpk/commands/storage.py,sha256=cSTJN9Mjvdsvk_Nk43kVdQFhp89nxWbanDsTOGZCkpQ,10708
 xpk/commands/version.py,sha256=k30rdLP9clUM8eeSwRFhpfzSb1qwcQImTfuC59Ed6CA,771
 xpk/commands/workload.py,sha256=l99NRFLs7pXuaLdn5d-Pid-cZulKpB3FNus-HdNDtZw,31513
 xpk/commands/workload_test.py,sha256=iXTY7VR1KrlPZZyh1Zm0N946kIP1iV2Fnqx1NtOYDJU,7274
 xpk/core/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
 xpk/core/capacity.py,sha256=_TyWayBkNU8fBpz1LTbCddEFZiZW5Qz-xmJnQMsXh0c,10534
 xpk/core/capacity_test.py,sha256=jZjMHTYlFLdAmBN1t9k29iABCSE5hlW0--q7QLDQpfQ,4330
-xpk/core/cluster.py,sha256=zAv46s-UB2r-I6cSkH7YzwAkGgD8Vxn7dJsXb_CMEQI,24062
+xpk/core/cluster.py,sha256=3nl77I_MgQpBZsZSzsiQ_7IyFRzfLrYNRUL1gsSNhKU,24036
 xpk/core/cluster_private.py,sha256=RLi0C7bV0NEUXl6QKQzvUT0weN9EdqPvjuuOQsNO0DY,6868
 xpk/core/cluster_test.py,sha256=J4Wk7E--ik_IsWWzL_iWGWbx99Ih03m-0bs-uU7gGDg,5853
 xpk/core/commands.py,sha256=at73VJHdZ4rVA8uvW997tNrvnCjP9v6zaw96bU0kd74,10841
-xpk/core/config.py,sha256=L3iPFvzFCpW8IEAvlbkuEHYBYXmRTC0BAaR7I_5_Peo,5146
+xpk/core/config.py,sha256=7U8jI5oZcgV_UnOHSS3huUIlDmPNREM-ml0N1Y9IvGM,4612
 xpk/core/config_test.py,sha256=POSuofK0LFbNNygDAo2fjtKY4NMrRjUFeGcpBh9JOS4,3569
 xpk/core/docker_container.py,sha256=8hqWWNKtjf6dqCFRpfndTMGvN_NS6zhfBr7YuKfh7qo,7626
 xpk/core/docker_image.py,sha256=9vwqbb6Mc3C5ZEOph03WS-EWI5hxMYGGigqzIMkDTjE,6909
@@ -59,7 +54,6 @@ xpk/core/gcloud_context_test.py,sha256=M8rp6S1zaEcAI7u4Bt8ukWKzv82HH5h9oYVojBcKg
 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=Ustta_ygXaacmgb1Av6QW4Epw0S_r-b-tjrMA6uNVj0,14240
 xpk/core/kueue_manager.py,sha256=JB8DcD-RFvBdC9Mk_DDCAkI2Km8W5-KMTRMVec06LlM,20010
 xpk/core/kueue_manager_test.py,sha256=FfBd1vninU_fcJ9wZev45-vpEsH12a9-XKysk_h4auo,22008
 xpk/core/monitoring.py,sha256=__bzTq_DIDAK8yIaN4F3MJh-yjYw5X1OlxmRgYOpf1g,4332
@@ -75,7 +69,7 @@ xpk/core/resources.py,sha256=dDsG_LOtcU17p1UKgOYyjdPxbMfqcb7pJ4SjfLDA6Os,9389
 xpk/core/scheduling.py,sha256=RMoei_HUs03rfrEC-HYk7ONzg9BRKwr59-KljCR2TMo,11560
 xpk/core/scheduling_test.py,sha256=iYnzXv_MjN743pa4zYAgRqb-6dB9nVPpLI7JP5S8M2I,14463
 xpk/core/storage.py,sha256=NILvVAcLNMLmp4wKx_TEKbMMF5X1oL-FrQV46PT0_ds,16902
-xpk/core/system_characteristics.py,sha256=Tam8wjUz77E6jAJib-r0GsTBmdjo9uaEkXmIdWuzGO8,32844
+xpk/core/system_characteristics.py,sha256=ZQbTbjaeT3Q12kmobz14U878w3FWnXDCetiLZQlVAdY,32127
 xpk/core/system_characteristics_test.py,sha256=sREN8u8bC0ze_q9hY3v-ZxC7so-_Ox1mt_DkIbUgHJ4,7477
 xpk/core/telemetry.py,sha256=R7IONNl5heMoNcOurfT3I34XJrBEODKVY88ONiDGuqE,7512
 xpk/core/telemetry_test.py,sha256=ll-B1ut9X-por17fpQnNb6hKrfyoZanMWRPbvqWrXss,8261
@@ -100,25 +94,21 @@ xpk/core/testing/__init__.py,sha256=PkV8D9WOtlJHH5AIxsQaKeIBcmupT_Ol_bwJgN6G2I8,
 xpk/core/testing/commands_tester.py,sha256=mQOSFggESeTdzqG4srAPV9ezmoeT90r22K58yAty9sE,4445
 xpk/core/testing/commands_tester_test.py,sha256=NnLWh7TJ9rKtb-DtB-vwkxvCe5wNtvUJ0f6sOa87Ht4,4023
 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/rdma_decorator.py,sha256=02HVA_jSyzlVtSQnQj7aPdK03h7v5YyioBqEen6pbj0,3636
 xpk/core/workload_decorators/storage_decorator.py,sha256=DDYQVO1OKTLhveDOA4V6b2RWr4n0fbwHdnoFFmW7iaQ,2000
-xpk/core/workload_decorators/tcpx_decorator.py,sha256=6yvofTv6_XmRfI-nESZjGYeLmGrza1rWxeJGET0TqXU,6182
-xpk/core/workload_decorators/tcpx_decorator_test.py,sha256=iTBS3X_-VwA2oveNDjscduLtll0VOJyFRCp4xmsjg7w,8515
-xpk/core/workload_decorators/tcpxo_decorator.py,sha256=_nLX7tbnxhnS-xv4Jijd1JOP76V4LpNCfW3Np404Cqw,6537
+xpk/core/workload_decorators/tcpx_decorator.py,sha256=cLOntH2ekBcPeiPW0sU3TRozSCpcTxgxpzncrMbRj44,5962
+xpk/core/workload_decorators/tcpx_decorator_test.py,sha256=BmTWsFoBeLb9xhQh3kpqSiarkYax4bj2wLeZ9GrQzag,6089
+xpk/core/workload_decorators/tcpxo_decorator.py,sha256=5SgL-7aTHclN7rvCGvEOjZoUixBmyjfuhVIUBFmneug,6124
 xpk/parser/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
-xpk/parser/batch.py,sha256=mJU-Cp1yTLje59vD-B1IiBcUeD-ZmEsoeB4xhj9cflc,1406
 xpk/parser/cluster.py,sha256=U2T-Q4yS86PWeFLNfknYWDDzZfubCKqIhqasxKLmErI,31342
 xpk/parser/cluster_test.py,sha256=xzQEC3IeAMpwsbNbHLuaNKxR3iaZcm3z4m3i61G62d4,6581
-xpk/parser/common.py,sha256=w6u6rqCOO23572C99PV1N8Fsp-vTP0C7Kv5tdWGEQO8,11691
+xpk/parser/common.py,sha256=sJYGjrn2YgFxelDCYB18s1R8Md8GpDcMQNoAezxDDIs,7257
 xpk/parser/common_test.py,sha256=_6Fm2pUF7h4K0G5qxGabXSYr4ng9ihOzlViE6oLQwQs,1557
 xpk/parser/config.py,sha256=-XnWx9aFsBW4Uzo_hpOMD2ZQ0bdZLvq1ksv83_5jqSM,1633
-xpk/parser/core.py,sha256=VRJerlS92ufoQbG1mZv7B04DAP4qGkBHa4pRXgcbAs0,4761
+xpk/parser/core.py,sha256=P2Dx3AbTlDoWnCCrMhVdr3Fs5FEzYFlmSiugkun1GL0,3623
 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=0V1d1htsjoa-SuxOX_vNxz2Lg4Nue9CBe_H0bNS2Hv0,10270
 xpk/parser/storage_test.py,sha256=i_F9cuQXHRvUy4RJwbfuuI8ZVpTpkkY96sZ1GZ4dLPw,1494
 xpk/parser/validators.py,sha256=-NBZelvfwZRzjz-YUCreD8EzMLHll8PZM-d-MVm2PG4,1192
@@ -139,7 +129,6 @@ xpk/templates/kueue_sub_slicing_topology.yaml.j2,sha256=UXjpRFqCIcoebwcMeD9Lo4fe
 xpk/templates/kueue_super_slicing_topology.yaml.j2,sha256=4WkSfQ2A5-jnKWiHWj2WXlv4sQmAcfxzbJCW-cWUE8E,264
 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=AJWSyjuWyLjb7SYt8kPb0gw9N84EN9LbLxYCXjC-6Ds,2464
 xpk/utils/console_test.py,sha256=x1v7v9VrIZwAKH-eOzj1lAY4EsHxJ6ruhfEOzpssO6o,2944
@@ -158,13 +147,13 @@ xpk/utils/user_agent.py,sha256=1NMtixC1RIr_MwM5pJ0THQ0x1-fCQA92TFHjWAVZldw,1083
 xpk/utils/user_agent_test.py,sha256=lkv8LqzhlA1gXFVeBzoLwE1_iGnm8G9LzkkElMrIrx0,1774
 xpk/utils/user_input.py,sha256=kMdCcPWdkI31f1mJcMsNGda-xKyKxEerpSLpCqIWYPc,1503
 xpk/utils/user_input_test.py,sha256=xO34jkMoTAk5Cmw7yHTk-7YexzC2UZ6ajihV8lnlAyI,2666
-xpk/utils/validation.py,sha256=irL9579RbvwxiGn1t3zhhPo-0oHgdUPOSYsUuFqsDSM,3039
+xpk/utils/validation.py,sha256=rE9LTkXJT7jIesodFb9pONL7ixhLqiQleyoaz7N39Dw,2765
 xpk/utils/validation_test.py,sha256=PEDSMUqZdt_Lx1FSR-LOTXKKtsJ47JH1fxugM0Gfz6Y,1168
 xpk/utils/versions.py,sha256=_Ep68W70a9605XjiaOOpBa9Is9jXlsoOiwL8v5Xt-WA,897
 xpk/utils/yaml.py,sha256=j8xuAJ9yAAwnQi6ozwZ-nMnDyDnc3xWkeBZMtSuP4RU,844
-xpk-0.17.1.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-xpk-0.17.1.dist-info/METADATA,sha256=7g7GYCSOZ1MjXNwXaMHto0yLF-MV2A2zwdfm98Qw2Eo,7930
-xpk-0.17.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xpk-0.17.1.dist-info/entry_points.txt,sha256=mzEtiIesFkT1kmcTUVDA1o3uOhiniX6tIz2wmOlMu1M,38
-xpk-0.17.1.dist-info/top_level.txt,sha256=TQKZWgV7LSElvmunYT9V_627qOMoxq3qYzWAFzKudB8,16
-xpk-0.17.1.dist-info/RECORD,,
+xpk-0.17.3.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+xpk-0.17.3.dist-info/METADATA,sha256=ONK-6JpzJboT0wF60svxA4SKJTSqeQ4KNfgSvZ_kkDY,7930
+xpk-0.17.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xpk-0.17.3.dist-info/entry_points.txt,sha256=mzEtiIesFkT1kmcTUVDA1o3uOhiniX6tIz2wmOlMu1M,38
+xpk-0.17.3.dist-info/top_level.txt,sha256=TQKZWgV7LSElvmunYT9V_627qOMoxq3qYzWAFzKudB8,16
+xpk-0.17.3.dist-info/RECORD,,

xpk/commands/batch.py

@@ -1,144 +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.
-"""
-
-import re
-from argparse import Namespace
-
-from ..core.cluster import (
-    setup_k8s_service_accounts,
-    get_cluster_credentials,
-)
-from ..core.commands import run_command_for_value
-from ..core.gcloud_context import add_zone_and_project
-from ..core.kjob import (
-    AppProfileDefaults,
-    JobTemplateDefaults,
-    get_storage_annotations,
-    prepare_kjob,
-)
-from ..core.kueue_manager import LOCAL_QUEUE_NAME
-from ..utils.console import xpk_exit, xpk_print
-from ..utils.execution_context import is_dry_run
-from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
-from .kind import set_local_cluster_command
-from .kjob_common import add_gpu_networking_annotations_to_command, add_TAS_annotations_to_command
-
-
-def batch(args: Namespace) -> None:
-  """Run batch task.
-     This function runs passed script in non-blocking manner.
-  Args:
-    args: user provided arguments for running the command.
-  Returns:
-    None
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  if not args.kind_cluster:
-    add_zone_and_project(args)
-    get_cluster_credentials(args)
-  else:
-    set_cluster_command_code = set_local_cluster_command(args)
-    if set_cluster_command_code != 0:
-      xpk_exit(set_cluster_command_code)
-
-  if not is_dry_run():
-    err_code = prepare_kjob(args)
-    if err_code > 0:
-      xpk_exit(err_code)
-    setup_k8s_service_accounts()
-
-  submit_job(args)
-
-
-def submit_job(args: Namespace) -> None:
-  cmd = (
-      'kubectl kjob create slurm'
-      f' --profile {AppProfileDefaults.NAME.value}'
-      f' --localqueue {LOCAL_QUEUE_NAME}'
-      f' --worker-container {JobTemplateDefaults.CONTAINER_NAME.value}'
-      ' --first-node-ip'
-  )
-  cmd = add_gpu_networking_annotations_to_command(args, cmd)
-  cmd = add_TAS_annotations_to_command(args, cmd)
-
-  annotations = [] if is_dry_run() else get_storage_annotations(args)
-  for annotation in annotations:
-    cmd += f' --pod-template-annotation {annotation}'
-
-  if args.ignore_unknown_flags:
-    cmd += ' --ignore-unknown-flags'
-
-  cmd += f' -- {args.script} --partition {LOCAL_QUEUE_NAME}'
-
-  if args.array is not None:
-    cmd += f' --array {args.array}'
-
-  if args.cpus_per_task is not None:
-    cmd += f' --cpus-per-task {args.cpus_per_task}'
-
-  if args.gpus_per_task is not None:
-    cmd += f' --gpus-per-task {args.gpus_per_task}'
-
-  if args.mem is not None:
-    cmd += f' --mem {args.mem}'
-
-  if args.mem_per_task is not None:
-    cmd += f' --mem-per-task {args.mem_per_task}'
-
-  if args.mem_per_cpu is not None:
-    cmd += f' --mem-per-cpu {args.mem_per_cpu}'
-
-  if args.mem_per_gpu is not None:
-    cmd += f' --mem-per-gpu {args.mem_per_gpu}'
-
-  if args.nodes is not None:
-    cmd += f' --nodes {args.nodes}'
-
-  if args.ntasks is not None:
-    cmd += f' --ntasks {args.ntasks}'
-
-  if args.output is not None:
-    cmd += f' --output {args.output}'
-
-  if args.error is not None:
-    cmd += f' --error {args.error}'
-
-  if args.input is not None:
-    cmd += f' --input {args.input}'
-
-  if args.job_name is not None:
-    cmd += f' --job-name {args.job_name}'
-
-  if args.chdir is not None:
-    cmd += f' --chdir {args.chdir}'
-
-  if args.time is not None:
-    cmd += f' --time {args.time}'
-
-  return_code, return_value = run_command_for_value(cmd, 'submit job')
-
-  if return_code != 0:
-    xpk_print(f'Running batch job returned ERROR {return_code}')
-    xpk_exit(return_code)
-
-  m = re.match(r'job\.batch/([-a-z0-9]+)', return_value)
-  if m:
-    xpk_print(f'Job name: {m.group(1)}')

xpk/commands/job.py

@@ -1,244 +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.
-"""
-
-import re
-import sys
-
-from ruamel.yaml import YAML
-from typing import cast
-
-from ..core.commands import run_command_for_value, run_command_with_updates
-from ..core.cluster import get_cluster_credentials
-from ..core.gcloud_context import add_zone_and_project
-from ..core.kjob import AppProfileDefaults
-from ..utils.console import xpk_exit, xpk_print
-from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
-from .kind import set_local_cluster_command
-
-
-JOBS_DRY_RUN_YAML = """
-items:
-- apiVersion: slurm.k8s.io/v1alpha1
-  kind: SlurmJob
-  metadata:
-    annotations:
-      kjobctl.x-k8s.io/script: echo hello
-    creationTimestamp: '2024-04-29T12:00:00Z'
-    labels:
-      kjobctl.x-k8s.io/app-profile: default
-    name: golden-job
-    namespace: default
-  spec:
-    script: echo hello
-"""
-
-PODS_DRY_RUN_RESULT = """
-foo-pod 2/2 Running 0 2d
-bar-pod 1/1 Evicted 0 1d
-"""
-
-
-def job_info(args):
-  """Run commands obtaining information about a job given by name.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    None
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  job_name = args.name
-
-  desc_command = f'kubectl-kjob describe slurm {job_name}'
-  desc_code, desc_text = run_command_for_value(desc_command, 'Getting job data')
-  if desc_code != 0:
-    xpk_print(f'Data info request returned ERROR {desc_code}')
-    xpk_exit(desc_code)
-
-  job_command = (
-      'kubectl-kjob list slurm -o yaml --field-selector'
-      f' metadata.name=={job_name}'
-  )
-  job_code, job_text = run_command_for_value(
-      job_command,
-      'Getting job info',
-      dry_run_return_val=JOBS_DRY_RUN_YAML,
-  )
-  if job_code != 0:
-    xpk_print(f'Job info request returned ERROR {job_code}')
-    xpk_exit(job_code)
-
-  pods_command = f'kubectl get pods -l=job-name={job_name} --no-headers'
-  pods_code, pods_text = run_command_for_value(
-      pods_command,
-      'Getting pods list',
-      dry_run_return_val=PODS_DRY_RUN_RESULT,
-  )
-  if pods_code != 0:
-    xpk_print(f'Pods list request returned ERROR {pods_code}')
-    xpk_exit(pods_code)
-
-  yaml = YAML(typ='safe')
-  job_yaml = yaml.load(job_text)['items'][0]
-
-  output = {
-      'Job name': job_name,
-      'Script name': get_script_name(job_yaml),
-      'Profile': get_profile(job_yaml),
-      'Labels': job_yaml.get('metadata').get('labels', []),
-      'Mounts': get_mounts(job_yaml),
-      'Pods': get_pods(pods_text),
-      'Entrypoint environment variables template': get_kjob_env_vars(desc_text),
-  }
-
-  yaml.default_flow_style = False
-  yaml.sort_base_mapping_type_on_output = False
-  yaml.dump(output, sys.stdout)
-
-
-def get_profile(job_yaml: dict) -> str:
-  containers: list[dict] = (
-      job_yaml.get('spec', {})
-      .get('template', {})
-      .get('spec', {})
-      .get('containers', [])
-  )
-  env_vars = next(iter(containers), {}).get('env', [])
-  profile = next((x['value'] for x in env_vars if x['name'] == 'PROFILE'), '')
-  return profile
-
-
-def get_mounts(job_yaml: dict) -> list[dict]:
-  containers: list[dict] = (
-      job_yaml.get('spec', {})
-      .get('template', {})
-      .get('spec', {})
-      .get('containers', [])
-  )
-  mounts: list[dict] = next(iter(containers), {}).get('volumeMounts', [])
-  return mounts
-
-
-def get_kjob_env_vars(job_desc_text: str) -> list[tuple[str, str]]:
-  regex = r'(SLURM_[A-Z_]*=.*)'
-  search_res = re.findall(regex, job_desc_text)
-  return search_res
-
-
-def get_pods(pods_text: str) -> list[dict[str, str]]:
-  pods_lines = pods_text.strip().split('\n')
-  pods_lines_tokenized = [line.split() for line in pods_lines]
-  return [
-      {
-          'Name': tokens[0],
-          'Status': tokens[2],
-      }
-      for tokens in pods_lines_tokenized
-  ]
-
-
-def get_script_name(job_yaml: dict) -> str | None:
-  return cast(
-      str | None,
-      job_yaml.get('metadata', {})
-      .get('annotations', {})
-      .get('kjobctl.x-k8s.io/script', ''),
-  )
-
-
-def job_list(args) -> None:
-  """Function around job list.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    None
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  if not args.kind_cluster:
-    add_zone_and_project(args)
-    get_cluster_credentials(args)
-    msg = f'Listing jobs for project {args.project} and zone {args.zone}:'
-  else:
-    set_cluster_command_code = set_local_cluster_command(args)
-    msg = 'Listing jobs:'
-    if set_cluster_command_code != 0:
-      xpk_exit(set_cluster_command_code)
-
-  xpk_print(msg, flush=True)
-
-  return_code = run_slurm_job_list_command()
-  xpk_exit(return_code)
-
-
-def run_slurm_job_list_command() -> int:
-  cmd = f'kubectl-kjob list slurm  --profile {AppProfileDefaults.NAME.value}'
-
-  return_code = run_command_with_updates(cmd, 'list jobs')
-  if return_code != 0:
-    xpk_print(f'Listing jobs returned ERROR {return_code}')
-  return return_code
-
-
-def job_cancel(args) -> None:
-  """Function around job cancel.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    None
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-
-  xpk_print(f'Starting job cancel for job: {args.name}', flush=True)
-  if not args.kind_cluster:
-    add_zone_and_project(args)
-    get_cluster_credentials(args)
-  else:
-    set_cluster_command_code = set_local_cluster_command(args)
-    if set_cluster_command_code != 0:
-      xpk_exit(set_cluster_command_code)
-
-  return_code = run_slurm_job_delete_command(args)
-  xpk_exit(return_code)
-
-
-def run_slurm_job_delete_command(args) -> int:
-  list_of_jobs = ' '.join(args.name)
-  cmd = f'kubectl-kjob delete slurm {list_of_jobs}'
-
-  return_code = run_command_with_updates(cmd, 'delete job')
-  if return_code != 0:
-    xpk_print(f'Delete job request returned ERROR {return_code}')
-  return return_code

xpk/commands/kjob_common.py

@@ -1,60 +0,0 @@
-"""
-Copyright 2025 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 ..core.capacity import (
-    B200_DEVICE_TYPE,
-    H100_MEGA_DEVICE_TYPE,
-    H200_DEVICE_TYPE,
-)
-from ..core.cluster import get_gpu_type_from_cluster
-from ..core.kjob import (
-    get_a3mega_pod_template_annotations,
-    get_a3ultra_pod_template_annotations,
-    get_a4_pod_template_annotations,
-    Kueue_TAS_annotation,
-)
-from .common import is_TAS_possible
-from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
-
-
-def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
-  gpu_type = get_gpu_type_from_cluster(args)
-
-  annotations: tuple
-  if gpu_type == H100_MEGA_DEVICE_TYPE:
-    annotations = get_a3mega_pod_template_annotations()
-  elif gpu_type == H200_DEVICE_TYPE:
-    annotations = get_a3ultra_pod_template_annotations()
-  elif gpu_type == B200_DEVICE_TYPE:
-    annotations = get_a4_pod_template_annotations()
-  else:
-    annotations = tuple()
-
-  flags = [
-      f" --pod-template-annotation {annotation} " for annotation in annotations
-  ]
-  cmd += "\\\n".join(flags)
-
-  return cmd
-
-
-def add_TAS_annotations_to_command(args, cmd: str) -> str:
-  system_characteristics = get_cluster_system_characteristics(args)
-  capacity_type = get_cluster_capacity_type(args)
-  if is_TAS_possible(system_characteristics, capacity_type):
-    cmd += f" --pod-template-annotation {Kueue_TAS_annotation}"
-
-  return cmd