Package: xpk
Overviewxpk (Accelerated Processing Kit, pronounced x-p-k,) is a software tool to help
Cloud developers to orchestrate training jobs on accelerators such as TPUs and
GPUs on GKE. xpk handles the "multihost pods" of TPUs, GPUs (HGX H100) and CPUs
(n2-standard-32) as first class citizens.xpk decouples provisioning capacity from running jobs. There are two structures:
clusters (provisioned VMs) and workloads (training jobs). Clusters represent the
physical resources you have available. Workloads represent training jobs -- at
any time some of these will be completed, others will be running and some will
be queued, waiting for cluster resources to become available.The ideal workflow starts by provisioning the clusters for all of the ML
hardware you have reserved. Then, without re-provisioning, submit jobs as
needed. By eliminating the need for re-provisioning between jobs, using Docker
containers with pre-installed dependencies and cross-ahead of time compilation,
these queued jobs run with minimal start times. Further, because workloads
return the hardware back to the shared pool when they complete, developers can
achieve better use of finite hardware resources. And automated tests can run
overnight while resources tend to be underutilized.xpk supports the following TPU types:v4v5ev5pand the following GPU types:a100h100and the following CPU types:n2-standard-32InstallationTo install xpk, run the following command:pipinstallxpkXPK for Large Scale (>1k VMs)Follow user instructions inxpk-large-scale-guide.shto use xpk for a GKE cluster greater than 1000 VMs. Run these steps to set up a
GKE cluster with large scale training and high throughput support with XPK, and
run jobs with XPK. We recommend you manually copy commands per step and verify
the outputs of each step.Example usages:To get started, be sure to set your GCP Project and Zone as usual viagcloud config set.Below are reference commands. A typical journey starts with aCluster Createfollowed by manyWorkload Creates. To understand the state of the system you
might want to useCluster ListorWorkload Listcommands. Finally, you can
cleanup with aCluster Delete.Cluster CreateFirst set the project and zone through gcloud config or xpk arguments.PROJECT_ID=my-project-idZONE=us-east5-b# gcloud config:gcloudconfigsetproject$PROJECT_IDgcloudconfigsetcompute/zone$ZONE# xpk argumentsxpk..--zone$ZONE--project$PROJECT_IDThe cluster created is a regional cluster to enable the GKE control plane across
all zones.Cluster Create (provision reserved capacity):# Find your reservationsgcloudcomputereservationslist--project=$PROJECT_ID# Run cluster create with reservation.python3xpk.pyclustercreate\--clusterxpk-test--tpu-type=v5litepod-256\--num-slices=2\--reservation=$RESERVATION_IDCluster Create (provision on-demand capacity):python3xpk.pyclustercreate\--clusterxpk-test--tpu-type=v5litepod-16\--num-slices=4--on-demandCluster Create (provision spot / preemptable capacity):python3xpk.pyclustercreate\--clusterxpk-test--tpu-type=v5litepod-16\--num-slices=4--spotCluster Create can be called again with the same--cluster nameto modify
the number of slices or retry failed steps.For example, if a user creates a cluster with 4 slices:python3xpk.pyclustercreate\--clusterxpk-test--tpu-type=v5litepod-16\--num-slices=4--reservation=$RESERVATION_IDand recreates the cluster with 8 slices. The command will rerun to create 4
new slices:python3xpk.pyclustercreate\--clusterxpk-test--tpu-type=v5litepod-16\--num-slices=8--reservation=$RESERVATION_IDand recreates the cluster with 6 slices. The command will rerun to delete 2
slices. The command will warn the user when deleting slices.
Use--forceto skip prompts.python3xpk.pyclustercreate\--clusterxpk-test--tpu-type=v5litepod-16\--num-slices=6--reservation=$RESERVATION_ID# Skip delete prompts using --force.python3xpk.pyclustercreate--force\--clusterxpk-test--tpu-type=v5litepod-16\--num-slices=6--reservation=$RESERVATION_IDCluster DeleteCluster Delete (deprovision capacity):python3xpk.pyclusterdelete\--clusterxpk-testCluster ListCluster List (see provisioned capacity):python3xpk.pyclusterlistCluster DescribeCluster Describe (see capacity):python3xpk.pyclusterdescribe\--clusterxpk-testCluster CacheimageCluster Cacheimage (enables faster start times):python3xpk.pyclustercacheimage\--clusterxpk-test--docker-imagegcr.io/your_docker_image\--tpu-type=v5litepod-16Workload CreateWorkload Create (submit training job):python3xpk.pyworkloadcreate\--workloadxpk-test-workload--command"echo goodbye"--cluster\xpk-test--tpu-type=v5litepod-16Setmax-restartsfor production jobs--max-restarts <value>: By default, this is 0. This will restart the job ""
times when the job terminates. For production jobs, it is recommended to
increase this to a large number, say 50. Real jobs can be interrupted due to
hardware failures and software updates. We assume your job has implemented
checkpointing so the job restarts near where it was interrupted.Workload Priority and PreemptionSet the priority level of your workload with--priority=LEVELWe have five priorities defined: [very-low,low,medium,high,very-high].
The default priority ismedium.Priority determines:Order of queued jobs.Queued jobs are ordered byvery-low<low<medium<high<very-highPreemption of lower priority workloads.A higher priority job willevictlower priority jobs.
Evicted jobs are brought back to the queue and will re-hydrate appropriately.General Example:python3xpk.pyworkloadcreate\--workloadxpk-test-medium-workload--command"echo goodbye"--cluster\xpk-test--tpu-type=v5litepod-16--priority=mediumWorkload DeleteWorkload Delete (delete training job):python3xpk.pyworkloaddelete\--workloadxpk-test-workload--clusterxpk-testThis will only deletexpk-test-workloadworkload inxpk-testcluster.Workload Delete (delete all training jobs in the cluster):python3xpk.pyworkloaddelete\--clusterxpk-testThis will delete all the workloads inxpk-testcluster. Deletion will only begin if you typeyoryesat the prompt. Multiple workload deletions are processed in batches for optimized processing.Workload Delete supports filtering. Delete a portion of jobs that match user criteria. Multiple workload deletions are processed in batches for optimized processing.Filter by Job:filter-by-jobpython3xpk.pyworkloaddelete\--clusterxpk-test--filter-by-job=$USERThis will delete all the workloads inxpk-testcluster whose names start with$USER. Deletion will only begin if you typeyoryesat the prompt.Filter by Status:filter-by-statuspython3xpk.pyworkloaddelete\--clusterxpk-test--filter-by-status=QUEUEDThis will delete all the workloads inxpk-testcluster that have the status as Admitted or Evicted, and the number of running VMs is 0. Deletion will only begin if you typeyoryesat the prompt. Status can be:EVERYTHING,FINISHED,RUNNING,QUEUED,FAILED,SUCCESSFUL.Workload ListWorkload List (see training jobs):python3xpk.pyworkloadlist\--clusterxpk-testExample Workload List Output:The below example shows four jobs of different statuses:user-first-job-failed:filter-statusisFINISHEDandFAILED.user-second-job-success:filter-statusisFINISHEDandSUCCESSFUL.user-third-job-running:filter-statusisRUNNING.user-forth-job-in-queue:filter-statusisQUEUED.user-fifth-job-in-queue-preempted:filter-statusisQUEUED.Jobset Name                     Created Time           Priority   TPU VMs Needed   TPU VMs Running/Ran   TPU VMs Done      Status     Status Message                                                  Status Time
user-first-job-failed           2023-1-1T1:00:00Z      medium     4                4                     <none>            Finished   JobSet failed                                                   2023-1-1T1:05:00Z
user-second-job-success         2023-1-1T1:10:00Z      medium     4                4                     4                 Finished   JobSet finished successfully                                    2023-1-1T1:14:00Z
user-third-job-running          2023-1-1T1:15:00Z      medium     4                4                     <none>            Admitted   Admitted by ClusterQueue cluster-queue                          2023-1-1T1:16:00Z
user-forth-job-in-queue         2023-1-1T1:16:05Z      medium     4                <none>                <none>            Admitted   couldn't assign flavors to pod set slice-job: insufficient unused quota for google.com/tpu in flavor 2xv4-8, 4 more need   2023-1-1T1:16:10Z
user-fifth-job-preempted        2023-1-1T1:10:05Z      low        4                <none>                <none>            Evicted    Preempted to accommodate a higher priority Workload             2023-1-1T1:10:00ZWorkload List supports filtering. Observe a portion of jobs that match user criteria.Filter by Status:filter-by-statusFilter the workload list by the status of respective jobs.
Status can be:EVERYTHING,FINISHED,RUNNING,QUEUED,FAILED,SUCCESSFULFilter by Job:filter-by-jobFilter the workload list by the name of a job.python3xpk.pyworkloadlist\--clusterxpk-test--filter-by-job=$USERGPU usageIn order to use XPK for GPU, you can do so by usingdevice-typeflag.Cluster Create (provision reserved capacity):# Find your reservationsgcloudcomputereservationslist--project=$PROJECT_ID# Run cluster create with reservation.python3xpk.pyclustercreate\--clusterxpk-test--device-type=h100-80gb-8\--num-slices=2\--reservation=$RESERVATION_IDInstall NVIDIA GPU device drivers# List available driver versionsgcloudcomputessh$NODE_NAME--command"sudo cos-extensions list"# Install the default drivergcloudcomputessh$NODE_NAME--command"sudo cos-extensions install gpu"# OR install a specific version of the drivergcloudcomputessh$NODE_NAME--command"sudo cos-extensions install gpu -- -version=DRIVER_VERSION"Run a workload:# Submit a workloadpython3xpk.pyworkloadcreate\--clusterxpk-test--device-typeh100-80gb-8\--workloadxpk-test-workload\--command="echo hello world"CPU usageIn order to use XPK for CPU, you can do so by usingdevice-typeflag.Cluster Create (provision on-demand capacity):# Run cluster create with on demand capacity.python3xpk.pyclustercreate\--clusterxpk-test\--device-type=n2-standard-32-256\--num-slices=1\--default-pool-cpu-machine-type=n2-standard-32\--on-demandNote thatdevice-typefor CPUs is of the format -, thus in the above example, user requests for 256 VMs of type n2-standard-32.
Currently workloads using < 1000 VMs are supported.Run a workload:# Submit a workloadpython3xpk.pyworkloadcreate\--clusterxpk-test\--num-slices=1\--device-type=n2-standard-32-256\--workloadxpk-test-workload\--command="echo hello world"How to add docker images to a xpk workloadThe default behavior isxpk workload createwill layer the local directory (--script-dir) into
the base docker image (--base-docker-image) and run the workload command.
If you don't want this layering behavior, you can directly use--docker-image. Do not mix arguments from the two flows in the same command.Recommended / Default Docker Flow:--base-docker-imageand--script-dirThis flow pulls the--script-dirinto the--base-docker-imageand runs the new docker image.The below arguments are optional by default. xpk will pull the local
directory with a generic base docker image.--base-docker-imagesets the base image that xpk will start with.--script-dirsets which directory to pull into the image. This defaults to the current working directory.Seepython3 xpk.py workload create --helpfor more info.Example with defaults which pulls the local directory into the base image:echo-e'#!/bin/bash \n echo "Hello world from a test script!"'>test.sh
python3xpk.pyworkloadcreate--clusterxpk-test\--workloadxpk-test-workload-base-image--command"bash test.sh"\--tpu-type=v5litepod-16--num-slices=1Recommended Flow For Normal Sized Jobs (fewer than 10k accelerators):python3xpk.pyworkloadcreate--clusterxpk-test\--workloadxpk-test-workload-base-image--command"bash custom_script.sh"\--base-docker-image=gcr.io/your_dependencies_docker_image\--tpu-type=v5litepod-16--num-slices=1Optional Direct Docker Image Configuration:--docker-imageIf a user wants to directly set the docker image used and not layer in the
current working directory, set--docker-imageto the image to be use in the
workload.Running with--docker-image:python3xpk.pyworkloadcreate--clusterxpk-test\--workloadxpk-test-workload-base-image--command"bash test.sh"\--tpu-type=v5litepod-16--num-slices=1--docker-image=gcr.io/your_docker_imageRecommended Flow For Large Sized Jobs (more than 10k accelerators):python3xpk.pyclustercacheimage\--clusterxpk-test--docker-imagegcr.io/your_docker_image# Run workload create with the same image.python3xpk.pyworkloadcreate--clusterxpk-test\--workloadxpk-test-workload-base-image--command"bash test.sh"\--tpu-type=v5litepod-16--num-slices=1--docker-image=gcr.io/your_docker_imageMore advanced facts:Workload create accepts a --env-file flag to allow specifying the container's
environment from a file. Usage is the same as Docker's--env-file flagExample File:LIBTPU_INIT_ARGS=--my-flag=true--performance=highMY_ENV_VAR=helloWorkload create accepts a --debug-dump-gcs flag which is a path to GCS bucket.
Passing this flag sets the XLA_FLAGS='--xla_dump_to=/tmp/xla_dump/' and uploads
hlo dumps to the specified GCS bucket for each worker.TroubleshootingInvalid machine typefor CPUs.XPK will create a regional GKE cluster. If you see issues likeInvalidmachinetypee2-standard-32inzone$ZONE_NAMEPlease select a CPU type that exists in all zones in the region.# Find CPU Types supported in zones.gcloudcomputemachine-typeslist--zones=$ZONE_LIST# Adjust default cpu machine type.python3xpk.pyclustercreate--default-pool-cpu-machine-type=CPU_TYPE...Permission Issues:requires one of ["permission_name"] permission(s).Determine the role needed based on the permission error:# For example: `requires one of ["container.*"] permission(s)`# Add [Kubernetes Engine Admin](https://cloud.google.com/iam/docs/understanding-roles#kubernetes-engine-roles) to your user.Add the role to the user in your project.Go toiam-adminor use gcloud cli:PROJECT_ID=my-project-idCURRENT_GKE_USER=$(gcloudconfiggetaccount)ROLE=roles/container.admin# container.admin is the role needed for Kubernetes Engine Admingcloudprojectsadd-iam-policy-binding$PROJECT_ID--memberuser:$CURRENT_GKE_USER--role=$ROLECheck the permissions are correct for the users.Go toiam-adminor use gcloud cli:PROJECT_ID=my-project-idCURRENT_GKE_USER=$(gcloudconfiggetaccount)gcloudprojectsget-iam-policy$PROJECT_ID--filter="bindings.members:$CURRENT_GKE_USER"--flatten="bindings[].members"Confirm you have logged in locally with the correct user.gcloudauthloginRoles needed based on permission errors:requires one of ["container.*"] permission(s)AddKubernetes Engine Adminto your user.ERROR: (gcloud.monitoring.dashboards.list) User does not have permission to access projects instance (or it may not exist)AddMonitoring Viewerto your user.Reservation Troubleshooting:How to determine your reservation and its size / utilization:PROJECT_ID=my-projectZONE=us-east5-bRESERVATION=my-reservation-name# Find the reservations in your projectgcloudbetacomputereservationslist--project=$PROJECT_ID# Find the tpu machine type and current utilization of a reservation.gcloudbetacomputereservationsdescribe$RESERVATION--project=$PROJECT_ID--zone=$ZONETPU Workload DebuggingVerbose LoggingIf you are having trouble with your workload, try setting the--enable-debug-logswhen you schedule it. This will give you more detailed logs to help pinpoint the issue. For example:python3xpk.pyworkloadcreate\--cluster--workloadxpk-test-workload\--command="echo hello world"--enable-debug-logsPlease checklibtpu loggingandTensorflow loggingfor more information about the flags that are enabled to get the logs.Collect Stack Tracescloud-tpu-diagnosticsPyPI package can be used to generate stack traces for workloads running in GKE. This package dumps the Python traces when a fault such as segmentation fault, floating-point exception, or illegal operation exception occurs in the program. Additionally, it will also periodically collect stack traces to help you debug situations when the program is unresponsive. You must make the following changes in the docker image running in a Kubernetes main container to enable periodic stack trace collection.# main.pyfromcloud_tpu_diagnosticsimportdiagnostic
fromcloud_tpu_diagnostics.configurationimportdebug_configuration
fromcloud_tpu_diagnostics.configurationimportdiagnostic_configuration
fromcloud_tpu_diagnostics.configurationimportstack_trace_configurationstack_trace_config=stack_trace_configuration.StackTraceConfig(collect_stack_trace=True,stack_trace_to_cloud=True)debug_config=debug_configuration.DebugConfig(stack_trace_config=stack_trace_config)diagnostic_config=diagnostic_configuration.DiagnosticConfig(debug_config=debug_config)withdiagnostic.diagnose(diagnostic_config):main_method()# this is the main method to runThis configuration will start collecting stack traces inside the/tmp/debuggingdirectory on each Kubernetes Pod.Explore Stack TracesTo explore the stack traces collected in a temporary directory in Kubernetes Pod, you can run the following command to configure a sidecar container that will read the traces from/tmp/debuggingdirectory.python3xpk.pyworkloadcreate\--workloadxpk-test-workload--command"python3 main.py"--cluster\xpk-test--tpu-type=v5litepod-16--deploy-stacktrace-sidecar