xpk 0.0.1__py3-none-any.whl

xpk/core/docker_manager.py

@@ -0,0 +1,308 @@
+"""
+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 abc import ABC, abstractmethod
+import docker
+from docker.errors import ContainerError, APIError, ImageNotFound, BuildError
+from ..utils.console import xpk_print, xpk_exit
+from ..utils.file import ensure_directory_exists
+from ..utils.objects import hash_string
+from shutil import copytree, copy
+import requests
+import os
+import tempfile
+import time
+
+
+DockerRunCommandExitCode = 135
+dockerBuildErrorCode = 134
+ctk_dockerfile_path = "Dockerfile"
+ctk_build_ref = "v1.45.1"
+ctk_docker_image = "xpk-ctk"
+ctk_container_name = "xpk-ctk-container"
+gcloud_cfg_mount_path = "/root/.config/gcloud"
+working_dir_mount_path = "/out"
+dockerfile_gh_path = f"https://raw.githubusercontent.com/GoogleCloudPlatform/cluster-toolkit/refs/tags/{ctk_build_ref}/tools/cloud-build/images/cluster-toolkit-dockerfile/Dockerfile"
+upload_dir_name = "uploads"
+
+
+class CommandRunner(ABC):
+  """This is a base class that defines methods a class for running cluster toolkit command should implement."""
+
+  @abstractmethod
+  def initialize(self) -> None:
+    """initialize is a method that should implement all steps neccessary to run command.
+
+    Returns:
+        None
+    """
+    return None
+
+  @abstractmethod
+  def run_command(self, cmd: str) -> None:
+    """run_command implements executing command. If command execution fails, exception should be raised.
+
+    Args:
+        cmd (str): command to run
+
+    Returns:
+        None:
+    """
+    return None
+
+  @abstractmethod
+  def upload_file_to_working_dir(self, path: str, prefix: str = "") -> str:
+    """Uploads single file to working directory.
+
+    Args:
+        path (str): path to file to upload
+
+    Returns:
+        str: path to a destination file
+    """
+    return ""
+
+  @abstractmethod
+  def upload_directory_to_working_dir(self, path: str, prefix: str = "") -> str:
+    """upload directory and its content to working directory.
+
+    Args:
+        path (str): path pointing to directory that will be uploaded.
+
+    Returns:
+        str: path to a target directory.
+    """
+    return ""
+
+
+class DockerManager(CommandRunner):
+  """DockerManager is a class for managing gcluster execution in docker container.
+  Attributes:
+    - dockerfile_path (str) : path to dockerfile defining gcluster execution image
+    - gcloud_cfg_path (str) : path to directory containing gcloud configuration
+    - working_dir (str) : path to directory in which gcluster deployment directory will be saved
+    - client (DockerClient) : docker client
+    - nocache (bool) : wheter to use docker cache when building image
+    - img_name (str) : name of docker image to create
+    - container_name (str) : name of the container that will be created from img_name
+    - rm_container_after (bool) : if set to True, docker container in which command is executed will be removed after each execution.
+  """
+
+  def __init__(
+      self,
+      gcloud_cfg_path: str,
+      working_dir: str,
+      nocache: bool = False,
+      img_name: str = ctk_docker_image,
+      container_name: str = ctk_container_name,
+      remove_container: bool = True,
+  ) -> None:
+    self.dockerfile_path = ""
+    self.client = docker.from_env()
+    self.gcloud_cfg_path = gcloud_cfg_path
+    self.working_dir = working_dir
+    self.nocache = nocache
+    self.img_name = f"{img_name}:{ctk_build_ref}"
+    self.container_name = container_name
+    self.remove_container = remove_container
+
+  def initialize(self):
+    """Build image from dockerfile pointed by _img_name. This method
+    uses python docker client to build cloud toolkit execution image.
+    Arguments:
+    Returns:
+      - None
+    Raises:
+      - docker.errors.BuildError – If there is an error during the build.
+      - docker.errors.APIError – If the server returns any other error.
+      - TypeError - otherwise
+
+    """
+    self._is_docker_installed()
+    xpk_print("Docker found!")
+
+    if not self._docker_image_exists():
+      xpk_print(f"Docker image {self.img_name} not found.")
+      self._build_image()
+    else:
+      xpk_print(f"Docker image {self.img_name} found!")
+
+  def run_command(
+      self,
+      cmd: str,
+  ) -> None:
+    """Run container from _img_name and mount directories:
+        - gcloud config
+        - deployment directory
+    Arguments:
+    Returns:
+      - bytes
+    Raises:
+      - docker.errors.ContainerError,
+      - docker.errors.ImageNotFound,
+      - docker.errors.APIError
+    """
+    xpk_print(f"Running command: {cmd} ...")
+    xpk_print(
+        f"volumes: {self.gcloud_cfg_path}:{gcloud_cfg_mount_path},"
+        f" {self.working_dir}:{working_dir_mount_path}"
+    )
+    try:
+      container = self.client.containers.run(
+          image=self.img_name,
+          entrypoint=cmd,
+          remove=self.remove_container,
+          name=self._get_container_unique_name(
+              cmd
+          ),  # To allow multiple xpk commands run in one machine.
+          detach=True,
+          volumes=[
+              f"{self.gcloud_cfg_path}:{gcloud_cfg_mount_path}",
+              f"{self.working_dir}:{working_dir_mount_path}",
+          ],
+          environment={
+              "GOOGLE_APPLICATION_CREDENTIALS": (
+                  "/root/.config/gcloud/application_default_credentials.json"
+              )
+          },
+      )
+      self._print_logs_from_container(container)
+      result = container.wait()
+      if result["StatusCode"] != 0:
+        xpk_print(f"Running gcluster command: {cmd} failed.")
+        xpk_exit(result["StatusCode"])
+    except ContainerError as e:
+      xpk_print(
+          "Running command failed due to ContainerError with exit status:"
+          f" {e.exit_status} and stderr: {e.stderr}"
+      )
+      xpk_exit(DockerRunCommandExitCode)
+    except ImageNotFound as _:
+      xpk_print(f"Image {ctk_docker_image} not found. Deploying cluster failed")
+      xpk_exit(DockerRunCommandExitCode)
+    except APIError as e:
+      xpk_print(f"Deploying cluster toolkit failed due to {e.explanation}")
+      xpk_exit(DockerRunCommandExitCode)
+
+  def _print_logs_from_container(self, container):
+    output = container.attach(stdout=True, stream=True, logs=True)
+    for line in output:
+      xpk_print(f"[gcluster] {line.decode('utf-8').strip()}")
+
+  def upload_directory_to_working_dir(self, path: str, prefix: str = "") -> str:
+    """Move file or directory from specified path to directory containing deployment files
+
+    Args:
+        path (str): path of directory/file that will be moved to deployment directory
+    """
+    name = path.split("/")[-1]
+    target_path = os.path.join(self._get_upload_directory(prefix), name)
+    uploaded_path = os.path.join(
+        self._get_upload_directory_mounted(prefix), name
+    )
+    xpk_print(
+        f"Copying directory from {path} to {target_path}. Path in docker:"
+        f" {uploaded_path}"
+    )
+    copytree(path, target_path, dirs_exist_ok=True)
+    return uploaded_path
+
+  def upload_file_to_working_dir(self, path: str, prefix: str = "") -> str:
+    """Move file or directory from specified path to directory containing deployment files
+
+    Args:
+        path (str): path of directory/file that will be moved to deployment directory
+    """
+    name = path.split("/")[-1]
+    target_path = os.path.join(self._get_upload_directory(prefix), name)
+    uploaded_path = os.path.join(
+        self._get_upload_directory_mounted(prefix), name
+    )
+    xpk_print(
+        f"Copying a file from {path} to {target_path}. Path in docker:"
+        f" {uploaded_path}"
+    )
+    copy(path, target_path)
+    return uploaded_path
+
+  def _get_upload_directory(self, prefix: str = "") -> str:
+    upload_dir = os.path.join(self.working_dir, upload_dir_name, prefix)
+    ensure_directory_exists(upload_dir)
+    return upload_dir
+
+  def _get_upload_directory_mounted(self, prefix: str = "") -> str:
+    return os.path.join(working_dir_mount_path, upload_dir_name, prefix)
+
+  def _create_tmp_for_dockerfile(self) -> str:
+    tmp_dir = os.path.join(tempfile.gettempdir(), "xpkutils")
+    ensure_directory_exists(tmp_dir)
+    tmp_path = os.path.join(tmp_dir, "Dockerfile")
+    return tmp_path
+
+  def _is_docker_installed(self) -> None:
+    self.client.info()
+
+  def _docker_image_exists(self) -> bool:
+    try:
+      self.client.images.get(f"{self.img_name}")
+    except ImageNotFound as _:
+      return False
+    return True
+
+  def _download_ctk_dockerfile(self) -> None:
+    """Downloads cluster toolkit dockerfile to dockerfile_path
+
+    Returns:
+        None
+    """
+    xpk_print(f"Downloading Dockerfile from {dockerfile_gh_path} ...")
+    self.dockerfile_path = self._create_tmp_for_dockerfile()
+    r = requests.get(dockerfile_gh_path, timeout=100)
+    with open(self.dockerfile_path, "w+", encoding="utf8") as dockerfile:
+      dockerfile.write(r.text)
+    xpk_print("Downloading Dockerfile completed!")
+
+  def _build_image(self):
+    try:
+      self._download_ctk_dockerfile()
+      dir_path = "/".join(self.dockerfile_path.split("/")[:-1])
+      xpk_print(
+          f"Building  {self.img_name} docker image from dockerfile:"
+          f" {self.dockerfile_path}. It may take a while..."
+      )
+      self.client.images.build(
+          nocache=self.nocache,
+          path=dir_path,
+          tag=f"{self.img_name}",
+          rm=True,
+          buildargs={"CLUSTER_TOOLKIT_REF": ctk_build_ref},
+      )
+    except BuildError as e:
+      xpk_print(f"error while building image {self.img_name}: {e.msg}")
+      xpk_exit(dockerBuildErrorCode)
+    except APIError as e:
+      xpk_print(f"erro while building image {self.img_name}: {e.explanation}")
+      xpk_exit(dockerBuildErrorCode)
+    except TypeError as e:
+      xpk_print(f"TypeError while building image {self.img_name}: {e.args}")
+      xpk_exit(dockerBuildErrorCode)
+    xpk_print("Docker image build succesfully.")
+    os.remove(self.dockerfile_path)
+    tmp_dockerfile_dir = "/".join(self.dockerfile_path.split("/")[:-1])
+    os.rmdir(tmp_dockerfile_dir)
+
+  def _get_container_unique_name(self, cmd):
+    return f"{self.container_name}_{hash_string(cmd + str(time.time_ns()))}"

xpk/core/docker_resources.py

@@ -0,0 +1,350 @@
+"""
+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 .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
+from .cluster import setup_k8s_env
+from .storage import GCS_FUSE_TYPE, GCP_FILESTORE_TYPE, Storage, get_storages_to_mount
+from .system_characteristics import AcceleratorType, SystemCharacteristics
+
+
+def get_main_container_resources(
+    args, system: SystemCharacteristics, resource_type
+) -> str:
+  """Resources for the main container.
+  Args:
+    args: user provided args.
+    system: system characteristics.
+    resource_type: TPU / GPU / CPU
+
+  Returns:
+    str:
+      Workload resources port as a YAML string
+  """
+  # Resources requirements for Pathways workload containers are known.
+  resources_yaml = """cpu: "24"
+                    memory: 100G"""
+  if args.use_pathways:
+    return resources_yaml
+
+  gpu_resources_yaml = """nvidia.com/gpu: {system.chips_per_vm}"""
+  if system.accelerator_type == AcceleratorType['GPU']:
+    return gpu_resources_yaml.format(system=system)
+
+  if system.accelerator_type == AcceleratorType['CPU']:
+    # CPUs don't have chips, but have a subresource called vCPUs.
+    # system.chips_per_vm is used as a proxy for vCPUs.
+    # Some vCPUs get used in hosting system pods of the workloads,
+    # hence an offset of 0.95 is introduced.
+    offset_vCPUs = int(system.chips_per_vm) * 0.95
+    return f'{resource_type}: {offset_vCPUs}'
+
+  return f'{resource_type}: {system.chips_per_vm}'
+
+
+def get_env_container(args, system: SystemCharacteristics) -> str:
+  """Environment configuration for the main container.
+  Args:
+    args: user provided args.
+    system: system characteristics.
+
+  Returns:
+    str:
+      YAML with the env config for the main container, as a YAML string.
+  """
+  pw_env_yaml = """
+                - name: XCLOUD_ENVIRONMENT
+                  value: GCP
+                - name: JAX_PLATFORMS
+                  value: proxy
+                - name: JAX_BACKEND_TARGET
+                  value: {proxy_address}
+                - name: JOBSET_NAME
+                  valueFrom:
+                    fieldRef:
+                      fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name']"""
+  if args.use_pathways:
+    return pw_env_yaml.format(
+        args=args, proxy_address=args.pathways_proxy_address
+    )
+
+  gpu_env_yaml = """
+                  - name: REPLICATED_JOB_NAME
+                    valueFrom:
+                      fieldRef:
+                        fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name']
+                  - name: JOBSET_NAME
+                    valueFrom:
+                      fieldRef:
+                        fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name']
+                  - name: JAX_COORDINATOR_ADDRESS
+                    value: "$(JOBSET_NAME)-$(REPLICATED_JOB_NAME)-0-0.$(JOBSET_NAME)"
+                  - name: NNODES
+                    value: "{args.num_nodes}"
+                  - name: NODE_RANK
+                    valueFrom:
+                      fieldRef:
+                        fieldPath: metadata.annotations['batch.kubernetes.io/job-completion-index']
+                  - name: USE_GPUDIRECT
+                    value: {gpu_direct_name}
+                  - name: GPUS_PER_NODE
+                    value: "{system.chips_per_vm}"
+                  - name: JAX_COORDINATOR_PORT
+                    value: "6002"
+                  - name: COMMAND
+                    value: "{args.command}"
+                  {args.env}"""
+
+  if system.accelerator_type == AcceleratorType['GPU']:
+    gpu_direct_name = 'fastrak'
+    if args.device_type == H100_DEVICE_TYPE:
+      gpu_direct_name = 'tcpx'
+      gpu_env_yaml += """
+                  - name: LD_LIBRARY_PATH
+                    value: /usr/local/nvidia/lib64
+"""
+    elif args.device_type == H100_MEGA_DEVICE_TYPE:
+      gpu_direct_name = 'tcpxo'
+    elif args.device_type == H200_DEVICE_TYPE:
+      gpu_direct_name = 'rdma'
+    return gpu_env_yaml.format(
+        args=args, system=system, gpu_direct_name=gpu_direct_name
+    )
+
+  if system.accelerator_type == AcceleratorType['CPU']:
+    return get_cpu_env(args.num_slices, args.env, system)
+
+  return args.env  # pytype: disable=bad-return-type
+
+
+def get_cpu_env(num_slices, env_vars, system) -> str:
+  """Generate environment variables for CPU nodepools
+  Args:
+    num_slices: Number of slices to be used in the workload.
+    env_vars: Environment variables, processed from user args.
+    system: system characteristics
+
+  Returns:
+    str: yaml containing env variables
+  """
+  yaml = """
+                - name: REPLICATED_JOB_NAME
+                  valueFrom:
+                    fieldRef:
+                      fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name']
+                - name: JOB_INDEX
+                  valueFrom:
+                    fieldRef:
+                      fieldPath: metadata.annotations['jobset.sigs.k8s.io/job-index']
+                - name: JOB_COMPLETION_INDEX
+                  valueFrom:
+                    fieldRef:
+                      fieldPath: metadata.annotations['batch.kubernetes.io/job-completion-index']
+                - name: PROCESSES_IN_JOB
+                  value: "{processes_in_job}"
+                - name: JAX_PROCESS_COUNT
+                  value: "{process_count}"
+                {env_vars}
+                - name: JAX_COORDINATOR_ADDRESS
+                  value: "$(JOBSET_NAME)-$(REPLICATED_JOB_NAME)-0-0.$(JOBSET_NAME)"
+  """
+  return yaml.format(
+      processes_in_job=system.vms_per_slice,
+      process_count=calculate_process_count(num_slices, system.vms_per_slice),
+      env_vars=env_vars,
+  )
+
+
+def get_volumes(args, system: SystemCharacteristics) -> str:
+  """Get volumes accessible to the containers in the pod.
+  Args:
+    args: user provided args.
+    system: system characteristics.
+
+  Returns:
+    str:
+      YAML for the volumes.
+  """
+  volumes = """- emptyDir:
+                  medium: Memory
+                name: dshm-2
+              """
+
+  if args.ramdisk_directory != '':
+    volumes += """
+              - name: cache
+                csi:
+                  driver: phase1-checkpoint.csi.storage.gke.io"""
+
+  if (
+      system.accelerator_type == AcceleratorType['TPU']
+      and args.deploy_stacktrace_sidecar
+  ):
+    volumes += """
+              - name: tpu-stack-trace
+              - name: shared-data
+              """
+
+  storages: list[Storage] = get_storages_to_mount(
+      setup_k8s_env(args), args.storage
+  )
+  for storage in storages:
+    if storage.type == GCS_FUSE_TYPE:
+      volumes += f"""- name: {storage.pv}
+                persistentVolumeClaim:
+                  claimName: {storage.pvc}
+                  readOnly: {storage.readonly}
+              """
+    if storage.type == GCP_FILESTORE_TYPE:
+      volumes += f"""- name: {storage.pv}
+                persistentVolumeClaim:
+                  claimName: {storage.pvc}
+                  readOnly: {storage.readonly}
+              """
+  return volumes
+
+
+def get_volume_mounts(args, system: SystemCharacteristics) -> str:
+  """Resources for the main container.
+  Args:
+    args: user provided args.
+
+  Returns:
+    str:
+      YAML for the volumes mounted within a Pathways container or GPU container as a YAML string.
+  """
+  volume_mount_yaml = """- mountPath: /dev/shm
+                  name: dshm-2
+                """
+
+  if args.ramdisk_directory != '':
+    volume_mount_yaml += f"""
+                - mountPath: /{args.ramdisk_directory}
+                  name: cache"""
+
+  if args.use_pathways:
+    volume_mount_yaml = """- mountPath: /tmp
+                  name: shared-tmp
+                """
+  elif (
+      system.accelerator_type == AcceleratorType['TPU']
+      and args.deploy_stacktrace_sidecar
+  ):
+    volume_mount_yaml += """- name: tpu-stack-trace
+                  mountPath: /tmp/debugging
+                - name: shared-data
+                  mountPath: /shared-volume
+                """
+  elif system.accelerator_type == AcceleratorType['GPU']:
+    if system.device_type == H100_DEVICE_TYPE:
+      volume_mount_yaml = """- name: nvidia-install-dir-host
+                  mountPath: /usr/local/nvidia/lib64
+                - name: tcpx-nccl-plugin-volume
+                  mountPath: /usr/local/tcpx
+                - name: tcpd-socket
+                  mountPath: /tmp
+                - name: shared-memory
+                  mountPath: /dev/shm
+                - name: workload-terminated-volume
+                  mountPath: /usr/share/workload"""
+    elif (
+        system.device_type == H100_MEGA_DEVICE_TYPE
+        or system.device_type == H200_DEVICE_TYPE
+    ):
+      volume_mount_yaml = ''
+
+  storages: list[Storage] = get_storages_to_mount(
+      setup_k8s_env(args), args.storage
+  )
+  for storage in storages:
+    if storage.type == GCS_FUSE_TYPE:
+      volume_mount_yaml += f"""- name: {storage.pv}
+                  mountPath: {storage.mount_point}
+                  readOnly: {storage.readonly}
+                """
+    if storage.type == GCP_FILESTORE_TYPE:
+      volume_mount_yaml += f"""- name: {storage.pv}
+                  mountPath: {storage.mount_point}
+                  readOnly: {storage.readonly}
+                """
+  return volume_mount_yaml
+
+
+def calculate_process_count(num_slices, vms_per_slice) -> str:
+  """Calculates the total number of processes in the workload.
+  Args:
+    num_slices: Number of slices to be used in the workload.
+    vms_per_slice: number of VMs in each slice.
+
+  Returns:
+    str: total number of processes.
+  """
+  num_processes = int(num_slices) * int(vms_per_slice)
+
+  return f'{num_processes}'
+
+
+def add_container_ports(args, system: SystemCharacteristics) -> str:
+  """Add slice builder and megascale container ports,
+  for non-pathways workloads.
+
+  Args:
+    args: user provided args.
+
+  Returns:
+    str:
+      Pathways server port as a YAML string
+  """
+  port_yaml = """- containerPort: 8471
+                - containerPort: 8080"""
+  if args.use_pathways:
+    return ''
+
+  gpu_port_yaml = """- containerPort: 6002"""
+  if system.accelerator_type == AcceleratorType['GPU']:
+    return gpu_port_yaml
+  return port_yaml
+
+
+def add_jax_coordinator_port(system) -> str:
+  """Add jax coordinator port only for CPUs
+
+  Args:
+    system: system characteristics.
+
+  Returns:
+    str:
+      jax coordinator port as a YAML string
+  """
+  if system.accelerator_type == AcceleratorType['CPU']:
+    return '- containerPort: 1234'
+  return ''
+
+
+def add_image_pull_policy_for_pw_or_gpu(args, system: SystemCharacteristics):
+  """Add image pull policy only for Pathways containers.
+  Args:
+    args: user provided args.
+    system: system characteristics
+
+  Returns:
+    str:
+      YAML stating that the image will be pulled fro GCR every time.
+  """
+  yaml = """imagePullPolicy: Always"""
+
+  if args.use_pathways or system.accelerator_type == AcceleratorType['GPU']:
+    return yaml.format(args=args)
+  return ''