Add local testing support for job command

README.md

@@ -1246,7 +1246,7 @@ xpk interfaces seamlessly with kind to manage Kubernetes clusters locally, facil
 
 ## Local Testing Basics
 
-Local testing is available exclusively through the `batch` command of xpk with the `--kind-cluster` flag. This allows you to simulate training jobs locally:
+Local testing is available exclusively through the `batch` and `job` commands of xpk with the `--kind-cluster` flag. This allows you to simulate training jobs locally:
 
 ```shell
 python xpk.py batch [other-options] --kind-cluster script

src/xpk/commands/job.py

@@ -21,6 +21,7 @@
     run_command_with_updates,
 )
 from .cluster import set_cluster_command
+from .batch import set_local_cluster_command
 
 
 def job_list(args) -> None:
@@ -32,15 +33,17 @@ def job_list(args) -> None:
   Returns:
     None
   """
-  add_zone_and_project(args)
-  set_cluster_command_code = set_cluster_command(args)
+  if not args.kind_cluster:
+    add_zone_and_project(args)
+    set_cluster_command_code = set_cluster_command(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(
-      f'Listing jobs for project {args.project} and zone {args.zone}:',
-      flush=True,
-  )
+  xpk_print(msg, flush=True)
 
   return_code = run_slurm_job_list_command(args)
   xpk_exit(return_code)
@@ -67,8 +70,12 @@ def job_cancel(args) -> None:
     None
   """
   xpk_print(f'Starting job cancel for job: {args.name}', flush=True)
-  add_zone_and_project(args)
-  set_cluster_command_code = set_cluster_command(args)
+  if not args.kind_cluster:
+    add_zone_and_project(args)
+    set_cluster_command_code = set_cluster_command(args)
+  else:
+    set_cluster_command_code = set_local_cluster_command(args)
+
   if set_cluster_command_code != 0:
     xpk_exit(set_cluster_command_code)
 

src/xpk/parser/job.py

@@ -14,6 +14,8 @@
 limitations under the License.
 """
 
+import argparse
+
 from .common import add_shared_arguments
 from ..commands.job import job_list, job_cancel
 
@@ -48,6 +50,14 @@ def set_job_parser(job_parser):
       required=True,
   )
 
+  job_list_optional_arguments.add_argument(
+      '--kind-cluster',
+      type=bool,
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help='Apply command to a local test cluster.',
+  )
+
   add_shared_arguments(job_list_optional_arguments)
   job_list_parser.set_defaults(func=job_list)
 
@@ -81,5 +91,13 @@ def set_job_parser(job_parser):
       required=True,
   )
 
+  job_cancel_optional_arguments.add_argument(
+      '--kind-cluster',
+      type=bool,
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help='Apply command to a local test cluster.',
+  )
+
   add_shared_arguments(job_cancel_optional_arguments)
   job_cancel_parser.set_defaults(func=job_cancel)