xpk 0.5.0__py3-none-any.whl → 0.6.0__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/gcluster_manager.py

@@ -0,0 +1,158 @@
+"""
+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 .docker_manager import CommandRunner
+from ..utils.console import xpk_print
+
+
+xpk_gcloud_cfg_path = '~/gcloud/cfg'
+xpk_deployment_dir = '/deployment'
+gcluster_deploy_command = 'gcluster deploy'
+gcluster_create_command = 'gcluster create'
+gcluster_destroy_command = 'gcluster destroy'
+blueprint_file_name = 'xpk_blueprint.yaml'
+deployment_module = '/out/xpk-deployment'
+a3_utils_dir_name = 'a3-mega-xpk'
+config_map_repo_path = 'src/xpk/blueprints/a3-mega-xpk/config-map.yaml.tftpl'
+kueue_config_repo_path = (
+    'src/xpk/blueprints/a3-mega-xpk/kueue-xpk-configuration.yaml.tftpl'
+)
+
+
+class GclusterManager:
+  """Manager is a class responsible for running cluster toolkit commands.
+  Attributes:
+    - gcluster_command_runner (CommandRunner) : instance of class implementing CommandRunner abstract methods.
+  Methods:
+    - deploy : run a deployment process of cluster toolkit. This method will invoke gcluster create and than gcluster deploy commands.
+    - destroy_deployment : run gcluster command to destroy existing deployment.
+  """
+
+  def __init__(
+      self,
+      gcluster_command_runner: CommandRunner,
+  ) -> None:
+    self.gcluster_command_runner = gcluster_command_runner
+
+  def _run_create_deployment_cmd(
+      self, blueprint_container_path: str, prefix: str = ''
+  ):
+    xpk_print('Creating deployment resources...')
+    cluster_create_cmd = (
+        f'{gcluster_create_command} -o {self._get_deployment_path(prefix)}'
+        f' {blueprint_container_path} -w --force'
+    )
+    self.gcluster_command_runner.run_command(cluster_create_cmd)
+    xpk_print('Creating deployment resources completed.')
+
+  def _run_deploy_cmd(
+      self,
+      deployment_name: str,
+      auto_approve: bool,
+      dry_run: bool,
+      prefix: str = '',
+  ):
+    xpk_print('Deploying resources...')
+    deploy_cmd = (
+        f'{gcluster_deploy_command} {self._get_deployment_path(prefix)}/{deployment_name}'
+    )
+    if auto_approve is True:
+      deploy_cmd += ' --auto-approve'
+    if dry_run is True:
+      return
+    self.gcluster_command_runner.run_command(deploy_cmd)
+    xpk_print('Deployment completed.')
+
+  def deploy(
+      self,
+      blueprint_path: str,
+      deployment_name: str,
+      prefix: str = '',
+      auto_approve: bool = True,
+      dry_run: bool = False,
+  ) -> None:
+    """ "deploy method provisions a new cluster using Cluster Toolkit.
+    It will invoke gcluster create and then gcluster deploy commands.
+    The files staged or created during running gcluster command will be managed by gcluster_command_runner in its working directory."
+
+    Args:
+        blueprint_path (str): path pointing to blueprint which will be deployed.
+        deployment_name (str): name of the deployment.
+        auto_approve (bool, optional): If set to true deployment command will be auto approved. Currently only True is supported. Defaults to True.
+        dry_run (bool, optional): If set to True gcluster will not deploy. Defaults to False.
+    Returns:
+      None
+    """
+    xpk_print(f'Deploying blueprint from path {blueprint_path} ...')
+    self._run_create_deployment_cmd(
+        blueprint_container_path=blueprint_path, prefix=prefix
+    )
+    self._run_deploy_cmd(
+        deployment_name=deployment_name,
+        prefix=prefix,
+        auto_approve=auto_approve,
+        dry_run=dry_run,
+    )
+    xpk_print('Deploying blueprint completed!')
+
+  def _run_destroy_command(
+      self,
+      deployment_name: str,
+      prefix: str = '',
+      auto_approve: bool = True,
+      dry_run: bool = False,
+  ):
+    destroy_cmd = (
+        f'{gcluster_destroy_command} {self._get_deployment_path(prefix)}/{deployment_name}'
+    )
+    if auto_approve is True:
+      destroy_cmd += ' --auto-approve'
+    if dry_run is True:
+      xpk_print(f'executing command {destroy_cmd}')
+      return
+    self.gcluster_command_runner.run_command(destroy_cmd)
+
+  def _get_deployment_path(self, prefix: str = '') -> str:
+    prefix = f'/{prefix}' if prefix != '' else ''
+    return f'deployments{prefix}'
+
+  def destroy_deployment(self, deployment_name: str, prefix: str = '') -> None:
+    """Destroy deployment.
+
+    Args:
+        deployment_name (str): name of deployment to destroy.
+    """
+    xpk_print(f'Destroying {deployment_name} started...')
+    self._run_destroy_command(deployment_name, prefix=prefix)
+    xpk_print(f'Destroying {deployment_name} completed!')
+
+  def stage_files(
+      self, blueprint_file: str, blueprint_dependencies: str, prefix: str = ''
+  ) -> str:
+    """Uploads blueprint file and directory to gcluster working directory."""
+    xpk_print(
+        "Staging (sending) blueprint file to gcluster's working directory..."
+    )
+    staged_blueprint = self.gcluster_command_runner.upload_file_to_working_dir(
+        blueprint_file, prefix
+    )
+    if len(blueprint_dependencies) > 0:
+      self.gcluster_command_runner.upload_directory_to_working_dir(
+          blueprint_dependencies, prefix
+      )
+    xpk_print('Staging blueprint completed!')
+    xpk_print(f"File path in gcluster's working directory: {staged_blueprint}")
+    return staged_blueprint

xpk/core/kjob.py

@@ -0,0 +1,205 @@
+"""
+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 argparse import Namespace
+from ..utils.console import xpk_print
+from .commands import run_command_for_value, run_kubectl_apply, run_command_with_updates
+from enum import Enum
+
+
+class AppProfileDefaults(Enum):
+  NAME = "xpk-def-app-profile"
+
+
+class JobTemplateDefaults(Enum):
+  NAME = "xpk-def-batch"
+  PARALLELISM = 1
+  COMPLETIONS = 1
+  CONTAINER_NAME = "xpk-batch-container"
+  IMAGE = "ubuntu:22.04"
+
+
+class PodTemplateDefaults(Enum):
+  NAME = "xpk-def-pod"
+  CONTAINER_NAME = "xpk-interactive-container"
+  IMAGE = "busybox:1.28"
+  INTERACTIVE_COMMAND = "/bin/sh"
+
+
+job_template_yaml = """
+  apiVersion: kjobctl.x-k8s.io/v1alpha1
+  kind: JobTemplate
+  metadata:
+    name: {name}
+    namespace: default
+  template:
+    spec:
+      parallelism: {parallelism}
+      completions: {completions}
+      completionMode: Indexed
+      template:
+        spec:
+          containers:
+            - name: {container_name}
+              image: {image}
+          restartPolicy: OnFailure"""
+
+app_profile_yaml = """
+apiVersion: kjobctl.x-k8s.io/v1alpha1
+kind: ApplicationProfile
+metadata:
+  name: {name}
+  namespace: default
+spec:
+  supportedModes:
+    - name: Slurm
+      template: {batch_template}
+      requiredFlags: []
+    - name: Interactive
+      template: {interactive_template}
+"""
+
+pod_template_yaml = """
+apiVersion: v1
+kind: PodTemplate
+metadata:
+  name: {name}
+  namespace: default
+template:
+  spec:
+    containers:
+      - name: {container_name}
+        image: {image}
+        command: [{interactive_command}]
+"""
+
+
+def verify_kjob_installed(args: Namespace) -> int:
+  """Check if kjob is installed. If not provide user with proper communicate and exit.
+  Args:
+    args - user provided arguments.
+  Returns:
+    error code > if kjob not installed, otherwise 0
+  """
+  command = "kubectl-kjob help"
+  task = "Verify kjob installation "
+  verify_kjob_installed_code, _ = run_command_for_value(command, task, args)
+
+  if verify_kjob_installed_code == 0:
+    xpk_print("kjob found")
+    return 0
+
+  if verify_kjob_installed_code != 0:
+    xpk_print(
+        " kjob not found. Please follow"
+        " https://github.com/kubernetes-sigs/kjob/blob/main/docs/installation.md"
+        " to install kjob."
+    )
+    return verify_kjob_installed_code
+  return 0
+
+
+def create_app_profile_instance(args: Namespace) -> int:
+  """Create new AppProfile instance on cluster with default settings.
+
+  Args:
+    args - user provided arguments
+  Returns:
+    exit_code > 0 if creating AppProfile fails, 0 otherwise
+  """
+  return run_kubectl_apply(
+      yml_string=app_profile_yaml.format(
+          name=AppProfileDefaults.NAME.value,
+          batch_template=JobTemplateDefaults.NAME.value,
+          interactive_template=PodTemplateDefaults.NAME.value,
+      ),
+      task="Creating AppProfile",
+      args=args,
+  )
+
+
+def create_job_template_instance(args: Namespace) -> int:
+  """Create new JobTemplate instance on cluster with default settings.
+
+  Args:
+    args - user provided arguments
+  Returns:
+    exit_code > 0 if creating JobTemplate fails, 0 otherwise
+  """
+  return run_kubectl_apply(
+      yml_string=job_template_yaml.format(
+          name=JobTemplateDefaults.NAME.value,
+          parallelism=JobTemplateDefaults.PARALLELISM.value,
+          completions=JobTemplateDefaults.COMPLETIONS.value,
+          container_name=JobTemplateDefaults.CONTAINER_NAME.value,
+          image=JobTemplateDefaults.IMAGE.value,
+      ),
+      task="Creating JobTemplate",
+      args=args,
+  )
+
+
+def create_pod_template_instance(args: Namespace) -> int:
+  """Create new PodTemplate instance on cluster with default settings.
+
+  Args:
+    args - user provided arguments
+  Returns:
+    exit_code > 0 if creating PodTemplate fails, 0 otherwise
+  """
+  return run_kubectl_apply(
+      yml_string=pod_template_yaml.format(
+          name=PodTemplateDefaults.NAME.value,
+          container_name=PodTemplateDefaults.CONTAINER_NAME.value,
+          image=PodTemplateDefaults.IMAGE.value,
+          interactive_command=PodTemplateDefaults.INTERACTIVE_COMMAND.value,
+      ),
+      task="Creating PodTemplate",
+      args=args,
+  )
+
+
+def prepare_kjob(args) -> int:
+  job_err_code = create_job_template_instance(args)
+  if job_err_code > 0:
+    return job_err_code
+
+  pod_err_code = create_pod_template_instance(args)
+  if pod_err_code > 0:
+    return pod_err_code
+
+  return create_app_profile_instance(args)
+
+
+def apply_kjob_crds(args: Namespace) -> int:
+  """Apply kjob CRDs on cluster.
+
+  This function install kjob CRDs files from kjobctl printcrds.
+  It creates all neccessary kjob CRDs.
+
+  Args:
+    args - user provided arguments
+  Returns:
+    None
+  """
+  command = "kubectl kjob printcrds | kubectl apply --server-side -f -"
+  task = "Create kjob CRDs on cluster"
+  return_code = run_command_with_updates(command, task, args)
+  if return_code != 0:
+    xpk_print(f"{task} returned ERROR {return_code}")
+    return return_code
+  xpk_print("Creating kjob CRDs succeeded")
+  return 0