xpk 0.0.1__py3-none-any.whl

xpk/core/pathways.py

@@ -0,0 +1,377 @@
+"""
+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 .cluster import XPK_SA
+from ..core.docker_container import get_user_workload_container
+from ..core.gcloud_context import zone_to_region
+from ..core.nodepool import get_all_nodepools_programmatic
+from ..utils.console import xpk_exit, xpk_print
+from .config import AcceleratorType
+from .storage import Storage, get_storage_volumes_yaml, GCS_FUSE_ANNOTATION
+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:
+    if args.custom_pathways_worker_args:
+      yaml = append_custom_pathways_args(yaml, args.custom_pathways_worker_args)
+    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:
+    if args.custom_pathways_proxy_server_args:
+      yaml = append_custom_pathways_args(
+          yaml, args.custom_pathways_proxy_server_args
+      )
+    return yaml.format(args=args, rm_address=get_rm_address(args))
+  else:
+    return ''
+
+
+def get_pathways_sidecar_container(args) -> str:
+  """This is a sidecar container that runs the remote python server.
+
+      It is a special case of the initContainer (designated by restartPolicy:
+      Always)
+      See https://kubernetes.io/docs/concepts/workloads/pods/sidecar-containers/
+      for more details.
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    str: yaml containing arguments for the Pathways sidecar container.
+  """
+  yaml = """initContainers:
+            - name: remote-python-sidecar
+              image: {args.remote_python_sidecar_image}
+              imagePullPolicy: Always
+              securityContext:
+                privileged: true
+              volumeMounts:
+              - mountPath: /tmp  # Shared volume mount with the main container.
+                name: shared-tmp
+              restartPolicy: Always
+              ports:
+              - containerPort: 50051
+              env:
+              - name: GRPC_SERVER_ADDRESS
+                value: '0.0.0.0:50051'"""
+  if args.use_pathways and args.remote_python_sidecar_image is not None:
+    return yaml.format(args=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: 480
+      - name: "memory"
+        nominalQuota: 2000G
+    - 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'
+
+  return True
+
+
+def get_pathways_unified_query_link(args) -> str:
+  """Get the unified query link for the pathways workload."""
+  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:"{args.workload}-"%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:
+    if args.custom_pathways_server_args:
+      yaml = append_custom_pathways_args(yaml, args.custom_pathways_server_args)
+    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 append_custom_pathways_args(yaml, custom_args) -> str:
+  """Append custom Pathways args to the YAML with proper indentation.
+
+  Args:
+      yaml (string): existing yaml containing args
+
+  Returns:
+      yaml (string): yaml with additional args appended.
+  """
+  second_line = yaml.split('\n')[1]
+  if (
+      not second_line
+  ):  # to cover edge case if only one arg remains, we would have to look at the entire YAML in this case.
+    return yaml
+  # Calculate the indentation based on the second line of existing YAML.
+  indentation = ' ' * (len(second_line) - len(second_line.lstrip()))
+  custom_args = custom_args.split(' ')
+  for arg in custom_args:
+    yaml += '\n' + indentation + '- ' + arg
+  return yaml
+
+
+def get_user_workload_for_pathways(
+    args,
+    system: SystemCharacteristics,
+    pod_failure_policy,
+    storages: list[Storage],
+) -> 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
+          {pod_failure_policy}
+          template:
+            metadata:
+              annotations:
+                {gcs_fuse_annotation}
+            spec:
+              containers:
+              {container}
+              serviceAccountName: {service_account}
+              nodeSelector:
+                cloud.google.com/gke-nodepool: cpu-user-np
+              hostNetwork: true
+              dnsPolicy: ClusterFirstWithHostNet
+              restartPolicy: Never
+              volumes:
+              - hostPath:
+                  path: /tmp
+                  type: DirectoryOrCreate
+                name: shared-tmp
+              {storage_volumes}"""
+  if args.headless:
+    return ''
+  else:
+    container, _ = get_user_workload_container(args, system)
+    storage_volumes = get_storage_volumes_yaml(storages)
+    return user_workload_yaml.format(
+        args=args,
+        container=container,
+        storage_volumes=storage_volumes,
+        pod_failure_policy=pod_failure_policy,
+        service_account=XPK_SA,
+        gcs_fuse_annotation=GCS_FUSE_ANNOTATION,
+    )
+
+
+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

xpk/core/remote_state/__init__.py

@@ -0,0 +1,15 @@
+"""
+Copyright 2025 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.
+"""

xpk/core/remote_state/fuse_remote_state.py

@@ -0,0 +1,99 @@
+"""
+Copyright 2025 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 .remote_state_client import RemoteStateClient
+from ...utils.gcs_utils import upload_directory_to_gcs, check_file_exists, download_bucket_to_dir, upload_file_to_gcs
+from ...utils.console import xpk_print
+from google.cloud.storage import Client
+import os
+
+
+class FuseStateClient(RemoteStateClient):
+  """FuseStateClient is a class for managing remote xpk state stored in GCS Fuse."""
+
+  def __init__(
+      self,
+      bucket: str,
+      state_directory: str,
+      cluster: str,
+      deployment_name: str,
+      prefix: str,
+  ) -> None:
+    self.bucket = bucket
+    self.state_dir = state_directory
+    self.storage_client = Client()
+    self.cluster = cluster
+    self.prefix = prefix
+    self.deployment_name = deployment_name
+
+  def _get_bucket_path(self) -> str:
+    return (
+        f'xpk_terraform_state/{self.prefix}/blueprints/{self.deployment_name}/'
+    )
+
+  def _get_bucket_path_blueprint(self) -> str:
+    return f'xpk_terraform_state/{self.prefix}/blueprints/'
+
+  def _get_deployment_filename(self) -> str:
+    return f'{self.deployment_name}.yaml'
+
+  def _get_blueprint_path(self) -> str:
+    blueprint_dir = '/'.join(self.state_dir.split('/')[:-1])
+    return os.path.join(blueprint_dir, self.deployment_name) + '.yaml'
+
+  def upload_state(self) -> None:
+    xpk_print(
+        f'Uploading dependecies from directory {self.state_dir} to bucket:'
+        f' {self.bucket}. Path within bucket is: {self._get_bucket_path()}'
+    )
+    upload_directory_to_gcs(
+        storage_client=self.storage_client,
+        bucket_name=self.bucket,
+        bucket_path=self._get_bucket_path(),
+        source_directory=self.state_dir,
+    )
+    blueprint_bucket_path = (
+        self._get_bucket_path_blueprint() + self._get_deployment_filename()
+    )
+    xpk_print(
+        f'Uploading blueprint file: {self._get_blueprint_path()} to bucket'
+        f' {self.bucket}. Path within bucket is: {blueprint_bucket_path}'
+    )
+    upload_file_to_gcs(
+        storage_client=self.storage_client,
+        bucket_name=self.bucket,
+        bucket_path=blueprint_bucket_path,
+        file=self._get_blueprint_path(),
+    )
+
+  def download_state(self) -> None:
+    xpk_print(
+        f'Downloading from bucket: {self.bucket}, from path:'
+        f' {self._get_bucket_path()} to directory: {self.state_dir}'
+    )
+    download_bucket_to_dir(
+        self.storage_client,
+        self.bucket,
+        self._get_bucket_path(),
+        destination_directory=self.state_dir,
+    )
+
+  def check_remote_state_exists(self) -> bool:
+    return check_file_exists(
+        self.storage_client,
+        self.bucket,
+        self._get_bucket_path_blueprint() + self._get_deployment_filename(),
+    )