xpk 0.4.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

xpk/core/pathways.py

@@ -0,0 +1,298 @@
+"""
+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.
+"""
+
+from ..utils.console import xpk_exit, xpk_print
+from .core import (
+    AcceleratorType,
+    get_all_nodepools_programmatic,
+    get_user_workload_container,
+    zone_to_region,
+)
+from .system_characteristics import SystemCharacteristics
+
+PathwaysExpectedInstancesMap = {
+    'v6e': 'tpuv6e',
+    'v5p': 'tpuv5',
+    'v5litepod': 'tpuv5e',
+    'v4': 'tpuv4',
+    'v3': 'tpuv3',
+}
+
+
+def get_pathways_worker_args(args) -> str:
+  """Arguments for the Pathways workers.
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    str: yaml containing arguments for the Pathways workers.
+  """
+  yaml = """- --server_port=29001
+              - --resource_manager_address={rm_address}
+              - --gcs_scratch_location={args.pathways_gcs_location}"""
+  if args.use_pathways:
+    return yaml.format(args=args, rm_address=get_rm_address(args))
+  else:
+    return ''
+
+
+def get_pathways_proxy_args(args) -> str:
+  """Arguments for the Pathways proxy.
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    str: yaml containing arguments for the Pathways proxy.
+  """
+  yaml = """- --server_port=29000
+              - --resource_manager_address={rm_address}
+              - --gcs_scratch_location={args.pathways_gcs_location}"""
+
+  if args.use_pathways:
+    return yaml.format(args=args, rm_address=get_rm_address(args))
+  else:
+    return ''
+
+
+def add_pw_resource_flavors(args):
+  """Add resource flavors required for Pathways enabled clusters."""
+  resource_flavor_yaml = """apiVersion: kueue.x-k8s.io/v1beta1
+kind: ResourceFlavor
+metadata:
+  name: cpu-rm
+spec:
+  nodeLabels:
+    cloud.google.com/gke-nodepool: cpu-rm-np
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ResourceFlavor
+metadata:
+  name: cpu-proxy
+spec:
+  nodeLabels:
+    cloud.google.com/gke-nodepool: cpu-proxy-np
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ResourceFlavor
+metadata:
+  name: cpu-user
+spec:
+  nodeLabels:
+    cloud.google.com/gke-nodepool: cpu-user-np
+---"""
+  if args.enable_pathways:
+    return resource_flavor_yaml
+  return ''
+
+
+def add_pw_resources_to_kueue(args):
+  """Add resource flavors required for Pathways, to the cluster queue."""
+  resources_yaml = """- coveredResources: ["cpu", "memory"]
+    flavors:
+    - name: cpu-rm
+      resources:
+      - name: "cpu"
+        nominalQuota: 80
+      - name: "memory"
+        nominalQuota: 160G
+    - name: cpu-proxy
+      resources:
+      - name: "cpu"
+        nominalQuota: 480
+      - name: "memory"
+        nominalQuota: 2000G
+    - name: cpu-user
+      resources:
+      - name: "cpu"
+        nominalQuota: 480
+      - name: "memory"
+        nominalQuota: 2000G"""
+  if args.enable_pathways:
+    return resources_yaml
+  return ''
+
+
+def ensure_pathways_workload_prerequisites(args, system) -> bool:
+  """Check all Pathways workload prerequisites and set necessary args.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    True once conditions satisfy and variables are set. Exits otherwise.
+  """
+  # Ensure command is provided if not using Pathways in headless mode
+  if args.command is None and not args.headless:
+    xpk_print(
+        'Please provide a command using "--command" for the docker container to'
+        ' execute. Command is not required if you wish to run Pathways'
+        ' workloads in headless mode (`xpk workload create-pathways'
+        ' --headless`).'
+    )
+    xpk_exit(1)
+
+  # Ensure the cluster and CPU nodepools were created with create-pathways
+  all_node_pools = get_all_nodepools_programmatic(args)
+  desired_pw_cpu_node_pools = {'cpu-user-np', 'cpu-rm-np', 'cpu-proxy-np'}
+  if not desired_pw_cpu_node_pools.issubset(set(all_node_pools[0])):
+    xpk_print(
+        'Cluster needs to be created with `xpk create-pathways` to run'
+        ' Pathways workloads.'
+    )
+    xpk_exit(1)
+
+  # Ensure device type is TPUs - currently Pathways supports TPUs only.
+  if system.accelerator_type != AcceleratorType['TPU']:
+    xpk_print('Currently, Pathways workloads can only be run on TPUs.')
+    xpk_exit(1)
+
+  # Set proxy address to be consumed in helper methods and displayed to user.
+  args.pathways_proxy_address = get_proxy_address(args)
+
+  # Set the job which determines the life of other Pathways jobs
+  args.targetReplicatedJob = 'proxy' if args.headless else 'main'
+
+  # Always report user code failures back to JobSet.
+  args.restart_on_user_code_failure = True
+
+  return True
+
+
+def get_pathways_unified_query_link(args) -> str:
+  """Get the unified query link for the pathways workload."""
+  pw_suffixes = ['main', 'rm', 'proxy']
+  pw_pod_names = [f'"{args.workload}-{suffix}-0"' for suffix in pw_suffixes]
+  pw_pod_names_query = '%20OR%20'.join(pw_pod_names + ['worker-0-0'])
+  query_params = (
+      'resource.type%3D"k8s_container"%0A'
+      f'resource.labels.project_id%3D"{args.project}"%0A'
+      f'resource.labels.location%3D"{zone_to_region(args.zone)}"%0A'
+      f'resource.labels.cluster_name%3D"{args.cluster}"%0A'
+      f'resource.labels.pod_name:{pw_pod_names_query}%0A'
+      'severity>%3DDEFAULT'
+  )
+
+  return f'https://console.cloud.google.com/logs/query;query={query_params}'
+
+
+def get_pathways_rm_args(args, system: SystemCharacteristics) -> str:
+  """Arguments for the Pathways resource manager.
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    str: yaml containing arguments for the Pathways resource manager.
+  """
+  yaml = """- --server_port=29001
+              - --gcs_scratch_location={args.pathways_gcs_location}
+              - --node_type=resource_manager
+              - --instance_count={instance_count}
+              - --instance_type={instance_type}"""
+  if args.use_pathways:
+    return yaml.format(
+        args=args,
+        instance_count=args.num_slices,
+        instance_type=f'{get_pathways_expected_tpu_type(system.device_type)}:{system.topology}',
+    )
+  else:
+    return ''
+
+
+def get_user_workload_for_pathways(args, system: SystemCharacteristics) -> str:
+  """
+  Create a user workload container for Pathways.
+  Don't create one for Pathways headless mode.
+
+  Args:
+    args: user provided args.
+    system: system characteristics.
+
+
+  Returns:
+    str:
+      Pathways server port as a YAML string
+  """
+  user_workload_yaml = """- name: main
+    replicas: 1
+    template:
+      metadata:
+        labels:
+          xpk.google.com/workload: {args.workload}
+      spec:
+        backoffLimit: 0
+        completions: 1
+        parallelism: 1
+        template:
+          spec:
+            containers:
+              {container}
+            nodeSelector:
+              cloud.google.com/gke-nodepool: cpu-user-np
+            restartPolicy: OnFailure
+            volumes:
+            - hostPath:
+                path: /tmp
+                type: DirectoryOrCreate
+              name: shared-tmp"""
+  if args.headless:
+    return ''
+  else:
+    container, _ = get_user_workload_container(args, system)
+    return user_workload_yaml.format(args=args, container=container)
+
+
+def get_rm_address(args) -> str:
+  """Generates the Pathways resource manager address.
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    str: Fully qualified RM address.
+  """
+  rm_address = f'{args.workload}-rm-0-0.{args.workload}:29001'
+  return rm_address
+
+
+def get_proxy_address(args) -> str:
+  """Generates the Pathways proxy address.
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    str: Fully qualified proxy address.
+  """
+  proxy_address = f'grpc://{args.workload}-proxy-0-0.{args.workload}:29000'
+  return proxy_address
+
+
+def get_pathways_expected_tpu_type(device_type: str) -> str:
+  """Returns the device type expected by Pathways
+  Args:
+    device_type: the system characteristic device type
+
+  Returns:
+    str: the device type expected by pathways.
+  """
+  raw_type = device_type.split('-')[0].lower()
+  pathways_expected_instance = PathwaysExpectedInstancesMap[raw_type]
+  if not pathways_expected_instance:
+    xpk_print(
+        f'Passed in device_type {device_type} is incorrect. Please pass in a'
+        ' valid device type'
+    )
+    xpk_exit(1)
+  return pathways_expected_instance

xpk/core/ray.py

@@ -0,0 +1,222 @@
+"""
+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 ..utils.console import xpk_exit, xpk_print
+from ..utils.file import write_tmp_file
+from .commands import run_command_for_value, run_command_with_updates_retry
+
+
+HEAD_CPU = 0.5
+WORKER_CPU = 0.9
+GCS_SERVER = 6379
+DASHBOARD = 8265
+CLIENT = 10001
+MULTISLICE = 8081
+
+ray_cluster_crd_yaml = """apiVersion: v1
+kind: Namespace
+metadata:
+  name: ray
+---
+apiVersion: ray.io/v1
+kind: RayCluster
+metadata:
+  name: raycluster
+  namespace: ray
+spec:
+  rayVersion: '{version}'
+  headGroupSpec:
+    rayStartParams: {{}}
+    #pod template
+    template:
+      spec:
+        containers:
+        - name: ray-head
+          image: rayproject/ray:{version}
+          resources:
+            limits:
+              cpu: {head_cpu}
+              memory: {head_mem}
+            requests:
+              cpu: {head_cpu}
+              memory: {head_mem}
+          ports:
+          - containerPort: {gcs_server}
+            name: gcs-server
+          - containerPort: {dashboard} # Ray dashboard
+            name: dashboard
+          - containerPort: {client}
+            name: client
+          - containerPort: {multislice}
+            name: multislice
+  workerGroupSpecs:
+    - replicas: {replicas} # TODO: Set min and max replicas
+      numOfHosts: {num_hosts}
+      minReplicas: {replicas}
+      maxReplicas: {replicas}
+      groupName: workergroup0
+      rayStartParams:
+        block: 'true'
+      template:
+        spec:
+          containers:
+            - name: ray-worker
+              image: rayproject/ray:{version}
+              resources:
+                limits:
+                  cpu: {worker_cpu}
+                  google.com/tpu: {chips_per_vm}
+                  memory: {worker_mem}
+                requests:
+                  cpu: {worker_cpu}
+                  google.com/tpu: {chips_per_vm}
+                  memory: {worker_mem}
+          nodeSelector:
+            cloud.google.com/gke-tpu-accelerator: {accelerator}
+            cloud.google.com/gke-tpu-topology: {topology}
+"""
+
+
+def install_ray_cluster(args, system) -> int:
+  """Install a RayCluster on the cluster
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  delete_ray_cluster(args)
+
+  label = 'cloud.google.com/gke-nodepool=default-pool'
+  available_head_cpu, available_head_mem = generate_available_resources(
+      label, args, HEAD_CPU
+  )
+
+  label = f'cloud.google.com/gke-tpu-accelerator={system.gke_accelerator}'
+  available_worker_cpu, available_worker_mem = generate_available_resources(
+      label, args, WORKER_CPU
+  )
+
+  yml_string = ray_cluster_crd_yaml.format(
+      accelerator=system.gke_accelerator,
+      topology=system.topology,
+      chips_per_vm=system.chips_per_vm,
+      num_hosts=system.vms_per_slice,
+      replicas=args.num_slices,
+      version=args.ray_version,
+      worker_cpu=available_worker_cpu,
+      worker_mem=available_worker_mem,
+      head_cpu=available_head_cpu,
+      head_mem=available_head_mem,
+      gcs_server=GCS_SERVER,
+      dashboard=DASHBOARD,
+      client=CLIENT,
+      multislice=MULTISLICE,
+  )
+
+  tmp = write_tmp_file(yml_string)
+  command = f'kubectl apply -f {str(tmp.file.name)}'
+  task = 'Applying RayCluster'
+  retry_attempts = 1
+  return_code = run_command_with_updates_retry(
+      command, task, args, num_retry_attempts=retry_attempts
+  )
+  if return_code != 0:
+    xpk_print(f'{task} not successful.')
+    xpk_exit(return_code)
+  return return_code
+
+
+def delete_ray_cluster(args) -> None:
+  """Delete all RayClusters on the cluster
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    None
+  """
+
+  command = 'kubectl delete rayclusters -n ray --all'
+  task = 'Deleting old RayCluster'
+  retry_attempts = 1
+  return_code = run_command_with_updates_retry(
+      command, task, args, num_retry_attempts=retry_attempts
+  )
+
+  if return_code != 0:
+    xpk_print(f'{task} not successful.')
+    xpk_exit(return_code)
+
+  return
+
+
+def generate_available_resources(label, args, percent) -> tuple:
+  """Generate the available resources for the nodes that match the given label
+
+  Args:
+    label: the label used to match the appropriate nodes
+    args: user provided arguments for running the command
+    percent: the percent of the available resources to use
+
+  Returns:
+    A tuple with the available cpu and memory
+  """
+
+  command = (
+      f"kubectl get nodes -l {label} -o jsonpath='{{.items[0].metadata.name}}'"
+  )
+  task = f'Getting nodes with label {label}'
+  _, node_name = run_command_for_value(command, task, args)
+
+  command = (
+      f"kubectl get node {node_name} -o jsonpath='{{.status.allocatable.cpu}}'"
+  )
+  task = 'Fetching available CPU on node'
+  _, available_cpu = run_command_for_value(command, task, args)
+  match = re.match(r'(\d+)([a-zA-Z]+)', available_cpu)
+  if not match:
+    xpk_print(
+        'Could not find a regex match for allocatable cpu on TPU node'
+        f' {node_name}'
+    )
+    xpk_exit(1)
+  value, units = match.group(1), match.group(2)
+  cpu_value = int(int(value) * percent)
+  adjusted_available_cpu = str(cpu_value) + units
+
+  command = (
+      f'kubectl get node {node_name} -o'
+      " jsonpath='{.status.allocatable.memory}'"
+  )
+  task = 'Fetching available memory on node'
+  _, available_memory = run_command_for_value(command, task, args)
+  match = re.match(r'(\d+)([a-zA-Z]+)', available_memory)
+  if not match:
+    xpk_print(
+        'Could not find a regex match for allocatable memory on TPU node'
+        f' {node_name}'
+    )
+    xpk_exit(1)
+  value, units = match.group(1), match.group(2)
+  memory_value = int(int(value) * percent)
+  adjusted_available_memory = str(memory_value) + units
+
+  return adjusted_available_cpu, adjusted_available_memory