Add --namespace flag support to workloads

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.

Author: ankitkinra

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,

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}/{args.workload}-proxy-0/details?project={args.project}'
       )
       xpk_print(
           'Follow the proxy here:'
@@ -850,15 +851,15 @@ 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}/{args.workload}/details?project={args.project}'
     )
     duration_of_logs = 'P1D'  # Past 1 Day
     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'resource.labels.namespace_name="{args.namespace}"\n'
         f'resource.labels.pod_name:"{args.workload}-slice-job-0-0-"\n'
         'severity>=DEFAULT'
     )
@@ -916,7 +917,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)

src/xpk/core/cluster.py

@@ -428,27 +428,27 @@ def get_gpu_type_from_cluster(args) -> str:
   return ''
 
 
-def setup_k8s_service_accounts() -> None:
+def setup_k8s_service_accounts(namespace: str = 'default') -> None:
   """
   Creates/sets up SAs and the roles for them
   """
   default_sa = 'default'
 
-  create_xpk_k8s_service_account()
+  create_xpk_k8s_service_account(namespace)
 
-  role_name = create_pod_reader_role()
-  create_role_binding(default_sa, role_name)
-  create_role_binding(XPK_SA, role_name)
+  role_name = create_pod_reader_role(namespace)
+  create_role_binding(default_sa, role_name, namespace)
+  create_role_binding(XPK_SA, role_name, namespace)
 
 
-def create_xpk_k8s_service_account() -> None:
+def create_xpk_k8s_service_account(namespace: str = 'default') -> None:
   k8s_core_client = k8s_client.CoreV1Api()
   sa = k8s_client.V1ServiceAccount(
       metadata=k8s_client.V1ObjectMeta(name=XPK_SA)
   )
 
   try:
-    k8s_core_client.read_namespaced_service_account(XPK_SA, DEFAULT_NAMESPACE)
+    k8s_core_client.read_namespaced_service_account(XPK_SA, namespace)
     xpk_print(
         f'Service account: {XPK_SA} already exists. Skipping its creation.'
     )
@@ -461,7 +461,7 @@ def create_xpk_k8s_service_account() -> None:
   xpk_print(f'Creating a new service account: {XPK_SA}')
   try:
     k8s_core_client.create_namespaced_service_account(
-        DEFAULT_NAMESPACE, sa, pretty=True
+        namespace, sa, pretty=True
     )
     xpk_print(f'Created a new service account: {XPK_SA} successfully')
   except ApiException as e:
@@ -474,15 +474,15 @@ def create_xpk_k8s_service_account() -> None:
       xpk_exit(1)
 
 
-def create_pod_reader_role() -> str:
+def create_pod_reader_role(namespace: str = 'default') -> str:
   """
   Creates the 'pod-reader' Role in the default namespace.
   """
   k8s_rbac_client = k8s_client.RbacAuthorizationV1Api()
   role_name = 'pod-reader'
 
   try:
-    k8s_rbac_client.read_namespaced_role(role_name, DEFAULT_NAMESPACE)
+    k8s_rbac_client.read_namespaced_role(role_name, namespace)
     xpk_print(f'Role: {role_name} already exists. Skipping its creation.')
     return role_name
   except ApiException as e:
@@ -492,7 +492,7 @@ def create_pod_reader_role() -> str:
 
   role = k8s_client.V1Role(
       metadata=k8s_client.V1ObjectMeta(
-          name=role_name, namespace=DEFAULT_NAMESPACE
+          name=role_name, namespace=namespace
       ),
       rules=[
           k8s_client.V1PolicyRule(
@@ -510,10 +510,10 @@ def create_pod_reader_role() -> str:
 
   xpk_print(
       f'Attempting to create Role: {role_name} in namespace:'
-      f' {DEFAULT_NAMESPACE}'
+      f' {namespace}'
   )
   try:
-    k8s_rbac_client.create_namespaced_role(DEFAULT_NAMESPACE, role, pretty=True)
+    k8s_rbac_client.create_namespaced_role(namespace, role, pretty=True)
     xpk_print(f'Successfully created Role: {role_name}')
     return role_name
   except ApiException as e:
@@ -525,7 +525,7 @@ def create_pod_reader_role() -> str:
       xpk_exit(1)
 
 
-def create_role_binding(sa: str, role_name: str) -> None:
+def create_role_binding(sa: str, role_name: str, namespace: str = 'default') -> None:
   """
   Creates a RoleBinding to associate the Service Account
   with the Role in the default namespace.
@@ -536,7 +536,7 @@ def create_role_binding(sa: str, role_name: str) -> None:
 
   try:
     k8s_rbac_client.read_namespaced_role_binding(
-        role_binding_name, DEFAULT_NAMESPACE
+        role_binding_name, namespace
     )
     xpk_print(
         f'RoleBinding: {role_binding_name} already exists. Skipping its'
@@ -550,11 +550,11 @@ def create_role_binding(sa: str, role_name: str) -> None:
 
   role_binding = k8s_client.V1RoleBinding(
       metadata=k8s_client.V1ObjectMeta(
-          name=role_binding_name, namespace=DEFAULT_NAMESPACE
+          name=role_binding_name, namespace=namespace
       ),
       subjects=[
           k8s_client.RbacV1Subject(
-              kind='ServiceAccount', name=sa, namespace=DEFAULT_NAMESPACE
+              kind='ServiceAccount', name=sa, namespace=namespace
           )
       ],
       role_ref=k8s_client.V1RoleRef(
@@ -565,11 +565,11 @@ def create_role_binding(sa: str, role_name: str) -> None:
   xpk_print(
       f'Attempting to create RoleBinding: {role_binding_name} for Service'
       f' Account: {sa} to Role: {role_name} in namespace:'
-      f' {DEFAULT_NAMESPACE}'
+      f' {namespace}'
   )
   try:
     k8s_rbac_client.create_namespaced_role_binding(
-        DEFAULT_NAMESPACE, role_binding, pretty=True
+        namespace, role_binding, pretty=True
     )
     xpk_print(f'Successfully created RoleBinding: {role_binding_name} for {sa}')
   except ApiException as e:

src/xpk/core/pathways.py

@@ -106,6 +106,7 @@ def get_pathways_unified_query_link(args) -> str:
       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'
+      f'resource.labels.namespace_name="{args.namespace}"\n'
       f'resource.labels.pod_name:"{args.workload}-"\n'
       'severity>=DEFAULT'
   )
@@ -143,7 +144,8 @@ def try_to_delete_pathwaysjob_first(args, workloads) -> bool:
   task_names = []
   for workload in workloads:
     args.workload = workload
-    command = f'kubectl delete pathwaysjob {workload} -n default'
+    ns_arg = f'-n {args.namespace}' if args.namespace else '-n default'
+    command = f'kubectl delete pathwaysjob {workload} {ns_arg}'
     task_name = f'PathwaysWorkloadDelete-{workload}'
     commands.append(command)
     task_names.append(task_name)

src/xpk/core/workload.py

@@ -240,9 +240,10 @@ def _parse_workload_item(item: dict[str, Any]) -> _WorkloadListRow:
 def _fetch_workloads(
     filter_by_status: _StatusFilter,
     filter_by_job: Optional[str] = None,
+    namespace: str = 'default',
 ) -> tuple[int, list[_WorkloadListRow]]:
   """Fetches and parses the raw workload list from the cluster."""
-  command = 'kubectl get workloads --ignore-not-found -o=json'
+  command = f'kubectl get workloads -n {namespace} --ignore-not-found -o=json'
 
   task = f'List Jobs with filter-by-status={filter_by_status.value}'
   if filter_by_job:
@@ -399,7 +400,7 @@ def get_workload_list(args: argparse.Namespace) -> tuple[int, str]:
   filter_by_job = getattr(args, 'filter_by_job', None)
   filter_by_status = _get_status_filter(args.filter_by_status)
 
-  return_code, raw_rows = _fetch_workloads(filter_by_status, filter_by_job)
+  return_code, raw_rows = _fetch_workloads(filter_by_status, filter_by_job, args.namespace)
   if return_code != 0:
     return return_code, ''
 
@@ -425,7 +426,7 @@ def check_if_workload_exists(args: argparse.Namespace) -> bool:
 
   s = ','.join([key + ':' + value for key, value in columns.items()])
 
-  command = f"kubectl get workloads -o=custom-columns='{s}'"
+  command = f"kubectl get workloads -n {args.namespace} -o=custom-columns='{s}'"
   return_code, return_msg = run_command_for_value(
       command, 'Check if Workload Already Exists'
   )
@@ -442,16 +443,17 @@ def check_if_workload_exists(args: argparse.Namespace) -> bool:
   return False
 
 
-def _get_jobset_status(workload_name: str) -> tuple[int, str]:
+def _get_jobset_status(workload_name: str, namespace: str = 'default') -> tuple[int, str]:
   """Retrieves the current status of a given jobset workload.
 
   Args:
     workload_name: The name of the workload to retrieve the status for.
+    namespace: The Kubernetes namespace to retrieve the status from.
 
   Returns:
     A tuple containing the return code of the command (0 for success) and the status string.
   """
-  status_cmd = f'kubectl get jobset {workload_name} -o json'
+  status_cmd = f'kubectl get jobset {workload_name} -n {namespace} -o json'
   return_code, return_value = run_command_for_value(
       status_cmd, 'Get jobset status'
   )
@@ -497,7 +499,7 @@ def wait_for_job_completion(args: argparse.Namespace) -> int:
     return 1
 
   # Get the full workload name
-  get_workload_name_cmd = f'kubectl get workloads | grep jobset-{args.workload}'
+  get_workload_name_cmd = f'kubectl get workloads -n {args.namespace} | grep jobset-{args.workload}'
   return_code, return_value = run_command_for_value(
       get_workload_name_cmd, 'Get full workload name'
   )
@@ -512,7 +514,7 @@ def wait_for_job_completion(args: argparse.Namespace) -> int:
       f'{timeout_val}s' if timeout_val != -1 else 'max timeout (1 week)'
   )
   wait_cmd = (
-      'kubectl wait --for=condition=Finished'
+      f'kubectl wait -n {args.namespace} --for=condition=Finished'
       f' workload {full_workload_name} --timeout={timeout_val}s'
   )
   return_code, return_value = run_command_for_value(
@@ -526,7 +528,7 @@ def wait_for_job_completion(args: argparse.Namespace) -> int:
           f'Timed out waiting for your workload after {timeout_msg}, see 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}/{args.workload}/details?project={args.project}'
       )
       return 124
     else:
@@ -536,9 +538,9 @@ def wait_for_job_completion(args: argparse.Namespace) -> int:
   xpk_print(
       'Finished waiting for your workload, see 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}/{args.workload}/details?project={args.project}'
   )
-  return_code, return_value = _get_jobset_status(args.workload)
+  return_code, return_value = _get_jobset_status(args.workload, args.namespace)
   if return_code != 0:
     return return_code
 

src/xpk/parser/common.py

@@ -108,6 +108,12 @@ def add_shared_arguments(
       ),
       required=required,
   )
+  custom_parser_or_group.add_argument(
+      '--namespace',
+      type=str,
+      default='',
+      help='Kubernetes namespace to use. Defaults to active namespace.',
+  )
   custom_parser_or_group.add_argument(
       '--dry-run',
       type=bool,

src/xpk/parser/info.py

@@ -36,12 +36,6 @@ def set_info_parser(info_parser: argparse.ArgumentParser) -> None:
       required=True,
   )
 
-  info_optional_arguments.add_argument(
-      '--namespace',
-      type=str,
-      default='',
-      help='Namespace to which resources and queues belong',
-  )
 
   queues_flitering_group = (
       info_optional_arguments.add_mutually_exclusive_group()