Add --namespace flag support to workloads

ankitkinra · ankitkinra · commit d27b377aff9d · 2026-07-30T14:18:11.000Z
Allows specifying a target Kubernetes namespace for workload commands (create, list, delete, wait).
This avoids hardcoding the 'default' namespace and enables running workloads in custom namespaces.

Also updates setup_k8s_service_accounts to create service accounts and role bindings in the specified namespace.
diff --git a/recipes/Cluster_delete.md b/recipes/Cluster_delete.md
@@ -17,7 +17,7 @@ gcloud container clusters list --project=golden-project --filter=name=golden-clu
 gcloud container clusters get-credentials golden-cluster --location=us-central1 --dns-endpoint --project=golden-project && kubectl config view && kubectl config set-context --current --namespace=default
 [XPK] Get the name of the workloads in the cluster.
 [XPK] Task: `List Jobs with filter-by-status=EVERYTHING` is implemented by the following command not running since it is a dry run. 
-kubectl get workloads --ignore-not-found -o=json
+kubectl get workloads  --ignore-not-found -o=json
 [XPK] Task: `Cluster Delete` is implemented by the following command not running since it is a dry run. 
 gcloud beta container clusters delete golden-cluster --project=golden-project --location=us-central1 --quiet
 [XPK] Task: `Get All Subnets` is implemented by the following command not running since it is a dry run. 
diff --git a/recipes/Workload_create.md b/recipes/Workload_create.md
@@ -9,7 +9,7 @@ xpk workload create --project=golden-project --zone=us-central1-a --cluster=gold
 $ xpk workload create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --command "bash hello" --tpu-type=v5p-8 --num-slices=1 --script-dir=/tmp
 [XPK] Starting xpk v0.0.0
 [XPK] Task: `Check if Workload Already Exists` is implemented by the following command not running since it is a dry run. 
-kubectl get workloads -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
+kubectl get workloads  -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
 [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. 
 kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true
 [XPK] Skipping workload scheduling validation in dry run.
@@ -139,7 +139,7 @@ spec:
               
 
 [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. 
-kubectl apply -f 39eda1549f4c0d68a4f11e6cbd89ba655d49d2faeef6898a140f476e6e70ae0e
+kubectl apply -f 39eda1549f4c0d68a4f11e6cbd89ba655d49d2faeef6898a140f476e6e70ae0e 
 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. 
 gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error
 [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard.
diff --git a/recipes/Workload_create_pathways.md b/recipes/Workload_create_pathways.md
@@ -9,7 +9,7 @@ xpk workload create-pathways --project=golden-project --zone=us-central1-a --clu
 $ xpk workload create-pathways --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --command "bash hello" --tpu-type=v5p-8 --num-slices=1 --script-dir=/tmp
 [XPK] Starting xpk v0.0.0
 [XPK] Task: `Check if Workload Already Exists` is implemented by the following command not running since it is a dry run. 
-kubectl get workloads -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
+kubectl get workloads  -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
 [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. 
 kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true
 [XPK] Skipping workload scheduling validation in dry run.
@@ -268,7 +268,7 @@ spec:
   suspend: false
 
 [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. 
-kubectl apply -f 0678dfe212682ff519efebdd5b50719a953fd288b50b0dd5e42775d0bb87d332
+kubectl apply -f 0678dfe212682ff519efebdd5b50719a953fd288b50b0dd5e42775d0bb87d332 
 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. 
 gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error
 [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard.
diff --git a/recipes/Workload_create_pathways_super-slicing.md b/recipes/Workload_create_pathways_super-slicing.md
@@ -9,7 +9,7 @@ DRY_RUN_RESOURCES_CONFIG_MAP="map[tpu7x-128:80]" xpk workload create-pathways --
 $ DRY_RUN_RESOURCES_CONFIG_MAP="map[tpu7x-128:80]" xpk workload create-pathways --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --command "bash hello" --tpu-type=tpu7x-4x4x20 --script-dir=/tmp
 [XPK] Starting xpk v0.0.0
 [XPK] Task: `Check if Workload Already Exists` is implemented by the following command not running since it is a dry run. 
-kubectl get workloads -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
+kubectl get workloads  -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
 [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. 
 kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true
 [XPK] Task: `Get defined topologies` is implemented by the following command not running since it is a dry run. 
@@ -270,7 +270,7 @@ spec:
   suspend: false
 
 [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. 
-kubectl apply -f e7574bba6a3ad51a173d673c90ad230a9ba39c7f46cd000e09a952479d19460b
+kubectl apply -f e7574bba6a3ad51a173d673c90ad230a9ba39c7f46cd000e09a952479d19460b 
 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. 
 gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error
 [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard.
diff --git a/recipes/Workload_create_sub-slicing.md b/recipes/Workload_create_sub-slicing.md
@@ -9,7 +9,7 @@ SUB_SLICING_ENABLED=true DRY_RUN_RESOURCES_CONFIG_MAP="map[v6e-16:8]" xpk worklo
 $ SUB_SLICING_ENABLED=true DRY_RUN_RESOURCES_CONFIG_MAP="map[v6e-16:8]" xpk workload create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --command "bash hello" --tpu-type=v6e-2x4 --script-dir=/tmp
 [XPK] Starting xpk v0.0.0
 [XPK] Task: `Check if Workload Already Exists` is implemented by the following command not running since it is a dry run. 
-kubectl get workloads -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
+kubectl get workloads  -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
 [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. 
 kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true
 [XPK] Task: `Get defined topologies` is implemented by the following command not running since it is a dry run. 
@@ -144,7 +144,7 @@ spec:
               
 
 [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. 
-kubectl apply -f 2018fe16498f36301979a10667302a0aff6beb09956705b64ff396373af777ba
+kubectl apply -f 2018fe16498f36301979a10667302a0aff6beb09956705b64ff396373af777ba 
 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. 
 gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error
 [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard.
diff --git a/recipes/Workload_create_super-slicing.md b/recipes/Workload_create_super-slicing.md
@@ -9,7 +9,7 @@ DRY_RUN_RESOURCES_CONFIG_MAP="map[tpu7x-128:80]" xpk workload create --project=g
 $ DRY_RUN_RESOURCES_CONFIG_MAP="map[tpu7x-128:80]" xpk workload create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --command "bash hello" --tpu-type=tpu7x-4x4x20 --script-dir=/tmp
 [XPK] Starting xpk v0.0.0
 [XPK] Task: `Check if Workload Already Exists` is implemented by the following command not running since it is a dry run. 
-kubectl get workloads -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
+kubectl get workloads  -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
 [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. 
 kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true
 [XPK] Task: `Get defined topologies` is implemented by the following command not running since it is a dry run. 
@@ -143,7 +143,7 @@ spec:
               
 
 [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. 
-kubectl apply -f 2c5ab381c0d643f8512a07d296d411413080ec652c15e8c676fd58435de5a327
+kubectl apply -f 2c5ab381c0d643f8512a07d296d411413080ec652c15e8c676fd58435de5a327 
 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. 
 gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error
 [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard.
diff --git a/recipes/Workload_create_with_output-manifest-file.md b/recipes/Workload_create_with_output-manifest-file.md
@@ -9,7 +9,7 @@ xpk workload create --project=golden-project --zone=us-central1-a --cluster=gold
 $ xpk workload create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --command "bash hello" --tpu-type=v5p-8 --num-slices=1 --script-dir=/tmp --output-manifest-file=/var/tmp/manifest.yaml
 [XPK] Starting xpk v0.0.0
 [XPK] Task: `Check if Workload Already Exists` is implemented by the following command not running since it is a dry run. 
-kubectl get workloads -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
+kubectl get workloads  -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
 [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. 
 kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true
 [XPK] Skipping workload scheduling validation in dry run.
@@ -140,7 +140,7 @@ spec:
               
 
 [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. 
-kubectl apply -f 39eda1549f4c0d68a4f11e6cbd89ba655d49d2faeef6898a140f476e6e70ae0e
+kubectl apply -f 39eda1549f4c0d68a4f11e6cbd89ba655d49d2faeef6898a140f476e6e70ae0e 
 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. 
 gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error
 [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard.
diff --git a/recipes/Workload_list.md b/recipes/Workload_list.md
@@ -18,7 +18,7 @@ gcloud container clusters get-credentials golden-cluster --location=us-central1
 kubectl get pods
 [XPK] Finished get-credentials and kubectl setup.
 [XPK] Task: `List Jobs with filter-by-status=EVERYTHING` is implemented by the following command not running since it is a dry run. 
-kubectl get workloads --ignore-not-found -o=json
+kubectl get workloads  --ignore-not-found -o=json
 [XPK] Workload List Output:
 Jobset Name   Created Time   Priority   TPU/GPU VMs Needed   TPU/GPU VMs Running/Ran   TPU/GPU VMs Done   Status   Status Message   Status Time
 [XPK] See your workloads in Cloud Console: https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project=golden-project
diff --git a/recipes/comprehensive-demo.md b/recipes/comprehensive-demo.md
@@ -54,7 +54,7 @@ gcloud container clusters get-credentials foo --location=us-central1 --dns-endpo
 kubectl get pods
 [XPK] Finished get-credentials and kubectl setup.
 [XPK] Task: `List Jobs with filter-by-status=EVERYTHING` is implemented by the following command not running since it is a dry run. 
-kubectl get workloads --ignore-not-found -o=json
+kubectl get workloads  --ignore-not-found -o=json
 [XPK] Workload List Output:
 Jobset Name   Created Time   Priority   TPU/GPU VMs Needed   TPU/GPU VMs Running/Ran   TPU/GPU VMs Done   Status   Status Message   Status Time
 [XPK] See your workloads in Cloud Console: https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project=bar
diff --git a/src/xpk/commands/cluster_test.py b/src/xpk/commands/cluster_test.py
@@ -89,6 +89,7 @@ def construct_args(**kwargs: Any) -> Namespace:
       project='project',
       zone='us-central1-a',
       reservation='',
+      namespace='',
       on_demand=False,
       tpu_type=None,
       device_type=None,
diff --git a/src/xpk/commands/workload.py b/src/xpk/commands/workload.py
@@ -393,7 +393,7 @@ def workload_create(args) -> None:
   k8s_api_client = None
   if not is_dry_run():
     k8s_api_client = setup_k8s_env(args)
-    setup_k8s_service_accounts()
+    setup_k8s_service_accounts(args.namespace)
 
   workload_exists = check_if_workload_exists(args)
 
@@ -786,7 +786,8 @@ def workload_create(args) -> None:
     )
 
   tmp = write_tmp_file(yml_string)
-  command = f'kubectl apply -f {str(tmp)}'
+  ns_arg = f'-n {args.namespace}' if args.namespace else ''
+  command = f'kubectl apply -f {str(tmp)} {ns_arg}'
   return_code = run_command_with_updates(command, 'Creating Workload')
 
   if return_code != 0:
@@ -835,7 +836,7 @@ def workload_create(args) -> None:
           " python -c 'import pathwaysutils; import jax; print(jax.devices())'"
       )
       pathways_proxy_link = (
-          f'https://console.cloud.google.com/kubernetes/job/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/default/{args.workload}-proxy-0/details?project={args.project}'
+          f'https://console.cloud.google.com/kubernetes/job/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/{args.namespace or "default"}/{args.workload}-proxy-0/details?project={args.project}'
       )
       xpk_print(
           'Follow the proxy here:'
@@ -850,15 +851,18 @@ def workload_create(args) -> None:
     xpk_print(
         'Follow your workload here:'
         # pylint: disable=line-too-long
-        f' https://console.cloud.google.com/kubernetes/service/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}'
+        f' https://console.cloud.google.com/kubernetes/service/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/{args.namespace or "default"}/{args.workload}/details?project={args.project}'
     )
     duration_of_logs = 'P1D'  # Past 1 Day
+    ns_log_filter = (
+        f'resource.labels.namespace_name="{args.namespace or "default"}"\n'
+    )
     log_filter = (
         'resource.type="k8s_container"\n'
         f'resource.labels.project_id="{args.project}"\n'
         f'resource.labels.location="{get_cluster_location(args.project, args.cluster, args.zone)}"\n'
         f'resource.labels.cluster_name="{args.cluster}"\n'
-        'resource.labels.namespace_name="default"\n'
+        f'{ns_log_filter}'
         f'resource.labels.pod_name:"{args.workload}-slice-job-0-0-"\n'
         'severity>=DEFAULT'
     )
@@ -916,7 +920,8 @@ def delete_workloads(args, workloads: list[str]) -> int:
   task_names = []
   for workload in workloads:
     args.workload = workload
-    command = f'kubectl delete jobset {workload} -n default'
+    ns_arg = f'-n {args.namespace}' if args.namespace else '-n default'
+    command = f'kubectl delete jobset {workload} {ns_arg}'
     task_name = f'WorkloadDelete-{workload}'
     commands.append(command)
     task_names.append(task_name)
diff --git a/src/xpk/core/cluster.py b/src/xpk/core/cluster.py
diff --git a/src/xpk/core/pathways.py b/src/xpk/core/pathways.py
diff --git a/src/xpk/core/workload.py b/src/xpk/core/workload.py
diff --git a/src/xpk/parser/common.py b/src/xpk/parser/common.py
diff --git a/src/xpk/parser/info.py b/src/xpk/parser/info.py
