From 03e0d4e01c4f9e2ed9a4357023a5b1e128d487ca Mon Sep 17 00:00:00 2001 From: Ankit Kinra Date: Wed, 29 Jul 2026 22:44:32 +0000 Subject: [PATCH] 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. --- src/xpk/commands/cluster_test.py | 1 + src/xpk/commands/workload.py | 17 +++++++---- src/xpk/core/cluster.py | 48 +++++++++++++++----------------- src/xpk/core/pathways.py | 9 +++++- src/xpk/core/workload.py | 32 ++++++++++++++------- src/xpk/parser/common.py | 6 ++++ src/xpk/parser/info.py | 7 ----- 7 files changed, 70 insertions(+), 50 deletions(-) diff --git a/src/xpk/commands/cluster_test.py b/src/xpk/commands/cluster_test.py index fc88eaaea..a4a51e3e8 100644 --- 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 index 19a7d87a1..7b6d21f0a 100644 --- 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 index 920678a8d..9698b148f 100644 --- a/src/xpk/core/cluster.py +++ b/src/xpk/core/cluster.py @@ -428,27 +428,28 @@ 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 """ + namespace = namespace or 'default' 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 +462,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,7 +475,7 @@ 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. """ @@ -482,7 +483,7 @@ def create_pod_reader_role() -> str: 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: @@ -491,9 +492,7 @@ def create_pod_reader_role() -> str: xpk_exit(1) role = k8s_client.V1Role( - metadata=k8s_client.V1ObjectMeta( - name=role_name, namespace=DEFAULT_NAMESPACE - ), + metadata=k8s_client.V1ObjectMeta(name=role_name, namespace=namespace), rules=[ k8s_client.V1PolicyRule( api_groups=[''], @@ -508,12 +507,9 @@ def create_pod_reader_role() -> str: ], ) - xpk_print( - f'Attempting to create Role: {role_name} in namespace:' - f' {DEFAULT_NAMESPACE}' - ) + xpk_print(f'Attempting to create Role: {role_name} in namespace: {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 +521,9 @@ 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. @@ -535,9 +533,7 @@ def create_role_binding(sa: str, role_name: str) -> None: role_binding_name = f'{sa}-{role_name}-binding' try: - k8s_rbac_client.read_namespaced_role_binding( - role_binding_name, DEFAULT_NAMESPACE - ) + k8s_rbac_client.read_namespaced_role_binding(role_binding_name, namespace) xpk_print( f'RoleBinding: {role_binding_name} already exists. Skipping its' ' creation.' @@ -550,11 +546,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 +561,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: diff --git a/src/xpk/core/pathways.py b/src/xpk/core/pathways.py index 9f1677597..ae7d84e1a 100644 --- a/src/xpk/core/pathways.py +++ b/src/xpk/core/pathways.py @@ -101,11 +101,17 @@ def check_if_pathways_job_is_installed(args) -> bool: def get_pathways_unified_query_link(args) -> str: """Get the unified query link for the pathways workload.""" + ns_log_filter = ( + f'resource.labels.namespace_name="{args.namespace}"\n' + if args.namespace + else '' + ) 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' + f'{ns_log_filter}' f'resource.labels.pod_name:"{args.workload}-"\n' 'severity>=DEFAULT' ) @@ -143,7 +149,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) diff --git a/src/xpk/core/workload.py b/src/xpk/core/workload.py index b75966281..9d1e1ae8d 100644 --- a/src/xpk/core/workload.py +++ b/src/xpk/core/workload.py @@ -240,9 +240,11 @@ def _parse_workload_item(item: dict[str, Any]) -> _WorkloadListRow: def _fetch_workloads( filter_by_status: _StatusFilter, filter_by_job: Optional[str] = None, + namespace: str = '', ) -> tuple[int, list[_WorkloadListRow]]: """Fetches and parses the raw workload list from the cluster.""" - command = 'kubectl get workloads --ignore-not-found -o=json' + ns_arg = f' -n {namespace}' if namespace else '' + command = f'kubectl get workloads{ns_arg} --ignore-not-found -o=json' task = f'List Jobs with filter-by-status={filter_by_status.value}' if filter_by_job: @@ -399,7 +401,9 @@ 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 +429,8 @@ 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}'" + ns_arg = f' -n {args.namespace}' if args.namespace else '' + command = f"kubectl get workloads{ns_arg} -o=custom-columns='{s}'" return_code, return_msg = run_command_for_value( command, 'Check if Workload Already Exists' ) @@ -442,16 +447,20 @@ 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 = '' +) -> 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' + ns_arg = f' -n {namespace}' if namespace else '' + status_cmd = f'kubectl get jobset {workload_name}{ns_arg} -o json' return_code, return_value = run_command_for_value( status_cmd, 'Get jobset status' ) @@ -497,7 +506,10 @@ 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}' + ns_arg = f' -n {args.namespace}' if args.namespace else '' + get_workload_name_cmd = ( + f'kubectl get workloads{ns_arg} | grep jobset-{args.workload}' + ) return_code, return_value = run_command_for_value( get_workload_name_cmd, 'Get full workload name' ) @@ -512,7 +524,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{ns_arg} --for=condition=Finished' f' workload {full_workload_name} --timeout={timeout_val}s' ) return_code, return_value = run_command_for_value( @@ -526,7 +538,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 or "default"}/{args.workload}/details?project={args.project}' ) return 124 else: @@ -536,9 +548,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 or "default"}/{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 diff --git a/src/xpk/parser/common.py b/src/xpk/parser/common.py index aba699a0a..b416de17d 100644 --- a/src/xpk/parser/common.py +++ b/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, diff --git a/src/xpk/parser/info.py b/src/xpk/parser/info.py index 0bf122a16..9cbf7c133 100644 --- a/src/xpk/parser/info.py +++ b/src/xpk/parser/info.py @@ -36,13 +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() )