xpk 0.11.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

xpk/commands/batch.py

@@ -31,6 +31,7 @@ from ..core.kjob import (
 )
 from ..core.kueue import LOCAL_QUEUE_NAME
 from ..utils.console import xpk_exit, xpk_print
+from ..utils.execution_context import is_dry_run
 from .kind import set_local_cluster_command
 from .kjob_common import add_gpu_networking_annotations_to_command, add_TAS_annotations_to_command
 
@@ -51,18 +52,16 @@ def batch(args: Namespace) -> None:
     if set_cluster_command_code != 0:
       xpk_exit(set_cluster_command_code)
 
-  err_code = prepare_kjob(args)
-  if err_code > 0:
-    xpk_exit(err_code)
-  setup_k8s_service_accounts()
+  if not is_dry_run():
+    err_code = prepare_kjob(args)
+    if err_code > 0:
+      xpk_exit(err_code)
+    setup_k8s_service_accounts()
 
   submit_job(args)
 
 
 def submit_job(args: Namespace) -> None:
-
-  setup_k8s_service_accounts()
-
   cmd = (
       'kubectl kjob create slurm'
       f' --profile {AppProfileDefaults.NAME.value}'
@@ -73,7 +72,8 @@ def submit_job(args: Namespace) -> None:
   cmd = add_gpu_networking_annotations_to_command(args, cmd)
   cmd = add_TAS_annotations_to_command(args, cmd)
 
-  for annotation in get_storage_annotations(args):
+  annotations = [] if is_dry_run() else get_storage_annotations(args)
+  for annotation in annotations:
     cmd += f' --pod-template-annotation {annotation}'
 
   if args.ignore_unknown_flags:

xpk/commands/cluster.py

@@ -76,6 +76,7 @@ from ..core.vertex import create_vertex_tensorboard
 from ..core.workload import get_workload_list
 from ..utils.console import get_user_input, xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
+from ..utils.execution_context import is_dry_run
 from . import cluster_gcluster
 from .common import set_cluster_command
 import shutil
@@ -92,7 +93,7 @@ def cluster_adapt(args) -> None:
 
   system, return_code = get_system_characteristics(args)
 
-  if return_code > 0:
+  if return_code > 0 or system is None:
     xpk_print('Fetching system characteristics failed!')
     xpk_exit(return_code)
 
@@ -128,9 +129,10 @@ def cluster_adapt(args) -> None:
 
   get_cluster_credentials(args)
 
-  k8s_client = setup_k8s_env(args)
+  if not is_dry_run():
+    k8s_client = setup_k8s_env(args)
+    install_storage_crd(k8s_client)
 
-  install_storage_crd(k8s_client)
   install_storage_csis(args)
 
   # create Vertex Tensorboard for new and existing clusters if create-vertex-tensorboard is set
@@ -141,8 +143,6 @@ def cluster_adapt(args) -> None:
     if not tensorboard_config:
       xpk_exit(1)
 
-  # Provision node pools dynamically based on incoming workloads:
-  # Currently autoprovisioning is not supported with Pathways.
   autoprovisioning_config = None
   if args.enable_autoprovisioning:
     xpk_print('Enabling Autoprovisioning')
@@ -201,7 +201,7 @@ def cluster_create(args) -> None:
   """
   system, return_code = get_system_characteristics(args)
 
-  if return_code > 0:
+  if return_code > 0 or system is None:
     xpk_print('Fetching system characteristics failed!')
     xpk_exit(return_code)
 
@@ -217,13 +217,13 @@ def cluster_create(args) -> None:
     xpk_exit(0)
 
   return_code, gke_server_config = get_gke_server_config(args)
-  if return_code != 0:
+  if return_code != 0 or gke_server_config is None:
     xpk_exit(return_code)
 
   return_code, gke_control_plane_version = get_gke_control_plane_version(
       args, gke_server_config
   )
-  if return_code != 0:
+  if return_code != 0 or gke_control_plane_version is None:
     xpk_exit(return_code)
 
   create_cluster_command_code = create_cluster_if_necessary(
@@ -253,9 +253,10 @@ def cluster_create(args) -> None:
   if update_coredns_command_code != 0:
     xpk_exit(update_cluster_command_code)
 
-  k8s_client = setup_k8s_env(args)
+  if not is_dry_run():
+    k8s_client = setup_k8s_env(args)
+    install_storage_crd(k8s_client)
 
-  install_storage_crd(k8s_client)
   install_storage_csis(args)
 
   # create Vertex Tensorboard for new and existing clusters if create-vertex-tensorboard is set
@@ -294,7 +295,7 @@ def cluster_create(args) -> None:
   # Provision node pools dynamically based on incoming workloads:
   # Currently autoprovisioning is not supported with Pathways.
   autoprovisioning_config = None
-  if not args.enable_pathways and args.enable_autoprovisioning:
+  if args.enable_autoprovisioning:
     xpk_print('Enabling Autoprovisioning')
     autoprovisioning_config, return_code = enable_autoprovisioning_on_cluster(
         args, system
@@ -398,7 +399,7 @@ def cluster_cacheimage(args) -> None:
   get_cluster_credentials(args)
   system, return_code = get_system_characteristics(args)
 
-  if return_code > 0:
+  if return_code > 0 or system is None:
     xpk_print('Fetching system characteristics failed!')
     xpk_exit(return_code)
 
@@ -411,10 +412,8 @@ def cluster_cacheimage(args) -> None:
       nodeSelectorKey=node_selector_key,
   )
   tmp = write_tmp_file(yml_string)
-  command_apply = f'kubectl apply -f {str(tmp.file.name)}'
-  command_delete = (
-      f'kubectl delete -f {str(tmp.file.name)} --ignore-not-found=true'
-  )
+  command_apply = f'kubectl apply -f {str(tmp)}'
+  command_delete = f'kubectl delete -f {str(tmp)} --ignore-not-found=true'
 
   return_code = run_command_with_updates(
       command_delete, 'Deleting Cached Image', args
@@ -808,6 +807,7 @@ def scale_up_coredns(args, replicas: int = 15, namespace: str = 'kube-system'):
 
 def check_deployment_exists(args, deployment_name: str, namespace: str) -> bool:
   """Check for the existence of a specific Deployment in a given namespace."""
+  # TODO: rewrite this to be more obvious, check if it is correct
   command = (
       f'kubectl get deployment {deployment_name} -n'
       f' {namespace} --ignore-not-found'
@@ -815,11 +815,11 @@ def check_deployment_exists(args, deployment_name: str, namespace: str) -> bool:
   result = run_command_with_updates(
       command, 'Waiting for kubeDNS to be checked.', args
   )
-  return result
+  return result != 0
 
 
 def verify_coredns_readiness(
-    args, timeout: int = 120, namespace: str = 'kube-system'
+    args, timeout: int = 240, namespace: str = 'kube-system'
 ):
   """Verifies CoreDNS readiness using kubectl wait commands."""
   xpk_print('Now verifying CoreDNS readiness...')
@@ -874,7 +874,7 @@ def cleanup_coredns_repo(coredns_repo_full_path: str):
     xpk_print(f'Error deleting directory {coredns_repo_full_path}: {e}')
 
 
-def update_coredns(args):
+def update_coredns(args) -> int:
   """Updates and deploys CoreDNS within a cluster.
 
   Args:

xpk/commands/cluster_gcluster.py

@@ -310,4 +310,5 @@ def generate_blueprint(
           system_node_pool_machine_type=args.default_pool_cpu_machine_type,
           system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
       )
-  return None
+  xpk_print('Device type is not supported.')
+  xpk_exit(1)

xpk/commands/common.py

@@ -18,6 +18,7 @@ from ..core.commands import run_command_with_updates_retry
 from ..core.capacity import H100_MEGA_DEVICE_TYPE, CapacityType
 from ..core.gcloud_context import zone_to_region
 from ..utils.console import xpk_print, xpk_exit
+from ..utils.execution_context import is_dry_run
 from ..core.system_characteristics import (
     SystemCharacteristics,
 )
@@ -50,8 +51,8 @@ def set_cluster_command(args) -> int:
 
 
 def is_TAS_possible(
-    system_characteristics: SystemCharacteristics,
-    capacity_type: CapacityType,
+    system_characteristics: SystemCharacteristics | None,
+    capacity_type: CapacityType | None,
     flex: bool,
 ) -> bool:
   """Check cluster's machine_type and capacity type to determine if Kueue TAS is possible
@@ -63,6 +64,9 @@ def is_TAS_possible(
     True if possible and False otherwise.
   """
 
+  if is_dry_run():
+    return True
+
   if system_characteristics is None:
     xpk_print('system_characteristics data was not found in configmaps.')
     xpk_exit(1)
@@ -71,7 +75,7 @@ def is_TAS_possible(
     xpk_print('capacity_type data was not found in configmaps.')
     xpk_exit(1)
 
-  if flex:
+  if not flex:
     return False
 
   if (

xpk/commands/info.py

@@ -51,19 +51,19 @@ def info(args: Namespace) -> None:
   cqs = run_kueuectl_list_clusterqueue(args)
   quotas = get_nominal_quotas(cqs)
 
-  if lq:
+  if lq and lqs is not None:
     print_formatted_lqs(lqs, quotas)
 
   if cq:
     print_formatted_cqs(cqs, quotas)
 
 
-def get_nominal_quotas(cqs: list[dict]) -> dict[str, dict[str, str]]:
+def get_nominal_quotas(cqs: str) -> dict[str, dict[str, str]]:
   """Get quotas from clusterqueues.
   This function retrieves how much of resource in each flavor is assigned to cluster queue.
   It parses flavors of passed cluster queues.
   Args:
-    - cqs - list of cluster queues.
+    - cqs - string containing a list of cluster queues in JSON format.
   Returns:
     - dictionary of cluster queues resources quotas in format:
     {cq_name:{"flavorName:resourceName":quota}}
@@ -75,7 +75,7 @@ def get_nominal_quotas(cqs: list[dict]) -> dict[str, dict[str, str]]:
     xpk_print(cqs)
     xpk_exit(1)
 
-  quotas = {}
+  quotas: dict[str, dict] = {}
   for cq in cq_list:
     spec = cq['spec']
     cq_name = cq['metadata']['name']
@@ -89,7 +89,7 @@ def get_nominal_quotas(cqs: list[dict]) -> dict[str, dict[str, str]]:
   return quotas
 
 
-def print_formatted_cqs(cqs: list[dict], nominalQuotas) -> None:
+def print_formatted_cqs(cqs: str, nominalQuotas) -> None:
   try:
     cq_list = json.loads(cqs)['items']
   except ValueError:
@@ -105,7 +105,7 @@ def print_formatted_cqs(cqs: list[dict], nominalQuotas) -> None:
   )
 
 
-def print_formatted_lqs(lqs: list[dict], nominalQuotas) -> None:
+def print_formatted_lqs(lqs: str, nominalQuotas) -> None:
   try:
     lq_list = json.loads(lqs)['items']
   except ValueError:
@@ -143,18 +143,18 @@ def parse_queue_lists(
 
 
 def get_flavors_resources_reservations(
-    cq_name: str, flavors_res: list[dict]
+    cq_name: str, flavors_res: dict
 ) -> dict[str, dict[str, str]]:
   """Get usage of flavors resources.
   This function parser flavorsReservation section of clusterQueue of LocalQueue.
   Args:
     - cq_name - name of ClusterQueue to which flavors belong.
-    - flavors_res - list of reservations made by flavors
+    - flavors_res - dict of reservations made by flavors
   Returns:
     Dict containing usage of each resource in flavor for each flavor in cluster or local queue.
     Dict format: {cq_name: {{flavor:resource}:reservation}}
   """
-  reservations = {}
+  reservations: dict[str, dict] = {}
   reservations[cq_name] = {}
   for flavor_name, flavor_resources_reservation_list in flavors_res.items():
     for resource in flavor_resources_reservation_list:
@@ -167,15 +167,15 @@ def get_flavors_resources_reservations(
 
 def get_flavors_usage(
     q_entry: dict, res_field: str, flavor_resource_quotas: dict
-) -> list[dict]:
+) -> dict[str, str]:
   """Parse q_entry to retrieve list of each resource usage in flavour.
   Args:
     q_entry - single entry into either LocalQueue or ClusterQueue structured as json
     flavor_resource_quotas - nominalQuota of flavors resource usage for each clusterqueue
   Returns:
-    list of dicts where each list entry is in format (key, entry) where:
+    Dict where for each (key, value):
     - key is flavorName:resourceName
-    - entry is flavorResourceReservation/flavorResourceQuota
+    - value is string formatted as 'flavorResourceReservation/flavorResourceQuota'
   """
   status = q_entry['status']
   flavors_res = status[res_field]

xpk/commands/inspector.py

@@ -346,7 +346,7 @@ def inspector(args) -> None:
       )
 
   # Summarize inspector:
-  xpk_print(f'Find xpk inspector output file: {inspector_file.name}')
+  xpk_print(f'Find xpk inspector output file: {inspector_file}')
 
   if final_return_code != 0:
     xpk_print(

xpk/commands/job.py

@@ -18,6 +18,7 @@ import re
 import sys
 
 from ruamel.yaml import YAML
+from typing import cast
 
 from ..core.commands import run_command_for_value, run_command_with_updates
 from ..core.cluster import get_cluster_credentials
@@ -27,6 +28,28 @@ from ..utils.console import xpk_exit, xpk_print
 from .kind import set_local_cluster_command
 
 
+JOBS_DRY_RUN_YAML = """
+items:
+- apiVersion: slurm.k8s.io/v1alpha1
+  kind: SlurmJob
+  metadata:
+    annotations:
+      kjobctl.x-k8s.io/script: echo hello
+    creationTimestamp: '2024-04-29T12:00:00Z'
+    labels:
+      kjobctl.x-k8s.io/app-profile: default
+    name: golden-job
+    namespace: default
+  spec:
+    script: echo hello
+"""
+
+PODS_DRY_RUN_RESULT = """
+foo-pod 2/2 Running 0 2d
+bar-pod 1/1 Evicted 0 1d
+"""
+
+
 def job_info(args):
   """Run commands obtaining information about a job given by name.
 
@@ -51,7 +74,10 @@ def job_info(args):
       f' metadata.name=={job_name}'
   )
   job_code, job_text = run_command_for_value(
-      job_command, 'Getting job info', args
+      job_command,
+      'Getting job info',
+      args,
+      dry_run_return_val=JOBS_DRY_RUN_YAML,
   )
   if job_code != 0:
     xpk_print(f'Job info request returned ERROR {job_code}')
@@ -59,7 +85,10 @@ def job_info(args):
 
   pods_command = f'kubectl get pods -l=job-name={job_name} --no-headers'
   pods_code, pods_text = run_command_for_value(
-      pods_command, 'Getting pods list', args
+      pods_command,
+      'Getting pods list',
+      args,
+      dry_run_return_val=PODS_DRY_RUN_RESULT,
   )
   if pods_code != 0:
     xpk_print(f'Pods list request returned ERROR {pods_code}')
@@ -84,7 +113,7 @@ def job_info(args):
 
 
 def get_profile(job_yaml: dict) -> str:
-  containers = (
+  containers: list[dict] = (
       job_yaml.get('spec', {})
       .get('template', {})
       .get('spec', {})
@@ -96,13 +125,13 @@ def get_profile(job_yaml: dict) -> str:
 
 
 def get_mounts(job_yaml: dict) -> list[dict]:
-  containers = (
+  containers: list[dict] = (
       job_yaml.get('spec', {})
       .get('template', {})
       .get('spec', {})
       .get('containers', [])
   )
-  mounts = next(iter(containers), {}).get('volumeMounts', [])
+  mounts: list[dict] = next(iter(containers), {}).get('volumeMounts', [])
   return mounts
 
 
@@ -112,23 +141,24 @@ def get_kjob_env_vars(job_desc_text: str) -> list[tuple[str, str]]:
   return search_res
 
 
-def get_pods(pods_text: str) -> list[str]:
+def get_pods(pods_text: str) -> list[dict[str, str]]:
   pods_lines = pods_text.strip().split('\n')
-  pods_lines = [line.split() for line in pods_lines]
+  pods_lines_tokenized = [line.split() for line in pods_lines]
   return [
       {
-          'Name': line[0],
-          'Status': line[2],
+          'Name': tokens[0],
+          'Status': tokens[2],
       }
-      for line in pods_lines
+      for tokens in pods_lines_tokenized
   ]
 
 
 def get_script_name(job_yaml: dict) -> str | None:
-  return (
+  return cast(
+      str | None,
       job_yaml.get('metadata', {})
       .get('annotations', {})
-      .get('kjobctl.x-k8s.io/script', '')
+      .get('kjobctl.x-k8s.io/script', ''),
   )
 
 

xpk/commands/kjob_common.py

@@ -33,6 +33,7 @@ from ..core.resources import get_cluster_capacity_type, get_cluster_system_chara
 def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
   gpu_type = get_gpu_type_from_cluster(args)
 
+  annotations: tuple
   if gpu_type == H100_MEGA_DEVICE_TYPE:
     annotations = get_a3mega_pod_template_annotations(args)
   elif gpu_type == H200_DEVICE_TYPE:
@@ -40,7 +41,7 @@ def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
   elif gpu_type == B200_DEVICE_TYPE:
     annotations = get_a4_pod_template_annotations(args)
   else:
-    annotations = []
+    annotations = tuple()
 
   flags = [
       f" --pod-template-annotation {annotation} " for annotation in annotations

xpk/commands/storage.py

@@ -58,6 +58,7 @@ from ..core.storage import (
 )
 from ..utils.console import get_user_input, xpk_exit, xpk_print
 from ..utils.kubectl import apply_kubectl_manifest
+from ..utils.execution_context import is_dry_run
 
 
 def storage_create(args: Namespace) -> None:
@@ -141,7 +142,7 @@ def storage_delete(args: Namespace) -> None:
 
 def storage_attach(args: Namespace) -> None:
   add_zone_and_project(args)
-  manifest = [{}]
+  manifest: list[dict] = [{}]
   if args.type == GCP_FILESTORE_TYPE:
     if args.instance is None:
       args.instance = args.name
@@ -243,8 +244,10 @@ def enable_csi_drivers_if_necessary(args: Namespace) -> None:
 
 
 def storage_list(args: Namespace) -> None:
-  k8s_api_client = setup_k8s_env(args)
-  storages = list_storages(k8s_api_client)
+  storages = []
+  if not is_dry_run():
+    k8s_api_client = setup_k8s_env(args)
+    storages = list_storages(k8s_api_client)
   print_storages_for_cluster(storages)
 
 

xpk/commands/workload.py

@@ -84,6 +84,7 @@ from ..core.system_characteristics import (
 from ..core.vertex import create_vertex_experiment
 from ..core.workload import (
     check_if_workload_exists,
+    get_jobsets_list_gcp_link,
     get_workload_list,
     wait_for_job_completion,
     zone_to_region,
@@ -96,6 +97,7 @@ from ..core.workload_decorators import (
 )
 from ..utils.console import get_user_input, xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
+from ..utils.execution_context import is_dry_run
 from . import cluster_gcluster
 from .common import is_TAS_possible
 
@@ -226,7 +228,8 @@ spec:
             metadata:
               labels:
                 xpk.google.com/workload: {args.workload}
-              annotations: {annotations}
+              annotations:
+                {annotations}
             spec:
               priorityClassName: {args.priority}
               restartPolicy: Never
@@ -304,8 +307,10 @@ def workload_create(args) -> None:
   Returns:
     0 if successful and 1 otherwise.
   """
-  k8s_api_client = setup_k8s_env(args)
-  setup_k8s_service_accounts()
+  k8s_api_client = None
+  if not is_dry_run():
+    k8s_api_client = setup_k8s_env(args)
+    setup_k8s_service_accounts()
 
   workload_exists = check_if_workload_exists(args)
 
@@ -319,7 +324,7 @@ def workload_create(args) -> None:
   xpk_print('Starting workload create', flush=True)
   system, return_code = get_system_characteristics(args)
 
-  if return_code > 0:
+  if return_code > 0 or system is None:
     xpk_print('Fetching system characteristics failed!')
     xpk_exit(return_code)
 
@@ -345,7 +350,7 @@ def workload_create(args) -> None:
   ):
     xpk_print(
         'Warning: Cluster has been created using XPK version:'
-        f' {cluster_config_map["xpk_version"]} but the XPK version you are'
+        f' {cluster_xpk_version} but the XPK version you are'
         f' using to schedule workload is: {XPK_CURRENT_VERSION}. Some features'
         ' might not be available for this cluster. We recommend to'
         ' upgrade/downgrade your XPK version or cluster by running `xpk'
@@ -354,7 +359,7 @@ def workload_create(args) -> None:
 
   debugging_dashboard_id = None
 
-  tensorboard_config = {}
+  tensorboard_config: dict | None = {}
   if VERTEX_TENSORBOARD_FEATURE_FLAG and args.use_vertex_tensorboard:
     tensorboard_config = create_vertex_experiment(args)
     # exit if failed to create Experiment in Vertex AI
@@ -381,8 +386,10 @@ def workload_create(args) -> None:
   all_storages = []
   # Currently storage customization is not supported for Pathways workloads. b/408468941
   if not args.use_pathways:
-    storages: list[Storage] = get_storages_to_mount(
-        k8s_api_client, args.storage
+    storages: list[Storage] = (
+        []
+        if k8s_api_client is None
+        else get_storages_to_mount(k8s_api_client, args.storage)
     )
     gcs_fuse_storages = list(
         filter(lambda storage: storage.type == GCS_FUSE_TYPE, storages)
@@ -450,8 +457,8 @@ def workload_create(args) -> None:
       - action: FailJobSet
         onJobFailureReasons:
         - PodFailurePolicy"""
-    restart_on_exit_codes = get_restart_exit_codes(args)
-    restart_on_exit_codes = ','.join(map(str, restart_on_exit_codes))
+    restart_on_exit_codes_list = get_restart_exit_codes(args)
+    restart_on_exit_codes = ','.join(map(str, restart_on_exit_codes_list))
     pod_failure_policy = f"""
           podFailurePolicy:
             rules:
@@ -567,14 +574,14 @@ def workload_create(args) -> None:
         pod_failure_policy=pod_failure_policy,
     )
   tmp = write_tmp_file(yml_string)
-  command = f'kubectl apply -f {str(tmp.file.name)}'
+  command = f'kubectl apply -f {str(tmp)}'
   return_code = run_command_with_updates(command, 'Creating Workload', args)
 
   if return_code != 0:
     xpk_print(f'Create Workload request returned ERROR {return_code}')
     xpk_exit(return_code)
 
-  if not args.use_pathways:
+  if not args.use_pathways and not is_dry_run():
     add_bucket_iam_members(args, storages)
 
   # Get GKE outlier dashboard for TPU
@@ -723,7 +730,11 @@ def workload_delete(args) -> None:
       )
     else:
       return_code = run_commands(
-          commands, 'Delete Workload', task_names, batch=100
+          commands,
+          'Delete Workload',
+          task_names,
+          batch=100,
+          dry_run=args.dry_run,
       )
 
     if return_code != 0:
@@ -741,8 +752,6 @@ def workload_list(args) -> None:
   Returns:
     0 if successful and 1 otherwise.
   """
-  xpk_print(args)
-
   xpk_print('Starting workload list', flush=True)
   add_zone_and_project(args)
   get_cluster_credentials(args)
@@ -760,4 +769,8 @@ def workload_list(args) -> None:
     xpk_print(f'List Job request returned ERROR {return_code}')
     xpk_exit(return_code)
   xpk_print(f'Workload List Output:\n{return_value}')
+
+  workload_list_gcp_link = get_jobsets_list_gcp_link(project=args.project)
+  xpk_print(f'See your workloads in Cloud Console: {workload_list_gcp_link}')
+
   xpk_exit(0)

xpk/core/blueprint/blueprint_generator.py

@@ -34,7 +34,7 @@ from ..system_characteristics import get_system_characteristics_by_device_type
 from .blueprint_definitions import Blueprint, DeploymentGroup, DeploymentModule
 from ..kueue import KUEUE_VERSION
 
-yaml = yaml.YAML()
+yaml_parser = yaml.YAML()
 
 a3high_device_type = H100_DEVICE_TYPE
 a3mega_device_type = H100_MEGA_DEVICE_TYPE
@@ -52,7 +52,7 @@ blueprint_dependencies_dir = {
 }
 
 cluster_toolkit_url = "github.com/GoogleCloudPlatform/cluster-toolkit"
-cluster_toolkit_version = "v1.57.1"
+cluster_toolkit_version = "v1.62.2"
 
 
 class BlueprintGeneratorOutput:
@@ -1019,7 +1019,7 @@ class BlueprintGenerator:
   ) -> str:
     blueprint_path = self._get_blueprint_path(blueprint_name, prefix)
     with open(blueprint_path, "w+", encoding="utf-8") as blueprint_file:
-      yaml.dump(xpk_blueprint, blueprint_file)
+      yaml_parser.dump(xpk_blueprint, blueprint_file)
     return blueprint_path
 
   def _get_blueprint_path(self, blueprint_name, prefix: str = ""):
@@ -1033,7 +1033,7 @@ class BlueprintGenerator:
     ensure_directory_exists(storage_path_with_prefix)
     return storage_path_with_prefix
 
-  def blueprint_exists(self, blueprint_name, prefix: str = ""):
+  def blueprint_exists(self, blueprint_name, prefix: str = "") -> bool:
     blueprint_path = self._get_blueprint_path(blueprint_name, prefix)
     return os.path.exists(blueprint_path)
 
@@ -1061,6 +1061,6 @@ class BlueprintGenerator:
     }
 
 
-yaml.register_class(Blueprint)
-yaml.register_class(DeploymentGroup)
-yaml.register_class(DeploymentModule)
+yaml_parser.register_class(Blueprint)
+yaml_parser.register_class(DeploymentGroup)
+yaml_parser.register_class(DeploymentModule)