xpk 0.4.0__tar.gz → 0.6.0__tar.gz

{xpk-0.4.0 → xpk-0.6.0}/PKG-INFO

@@ -1,8 +1,8 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: xpk
-Version: 0.4.0
+Version: 0.6.0
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
-Author-email: Cloud TPU Team <cloud-tpu-eng@google.com>
+Author-email: XPK team <xpk-code-reviewers@google.com>
 License: Apache-2.0
 Project-URL: Homepage, https://github.com/google/xpk
 Project-URL: Bug Tracker, https://github.com/google/xpk/issues
@@ -12,10 +12,17 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: cloud-accelerator-diagnostics
+Requires-Dist: tabulate
+Requires-Dist: ruamel.yaml
+Requires-Dist: pyyaml
+Requires-Dist: docker
+Requires-Dist: packaging
 Provides-Extra: dev
 Requires-Dist: pyink==24.3.0; extra == "dev"
 Requires-Dist: pylint>=2.6.0; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: docker; extra == "dev"
 
 <!--
  Copyright 2023 Google LLC
@@ -62,31 +69,73 @@ xpk supports the following TPU types:
 * v4
 * v5e
 * v5p
+* Trillium (v6e)
 
 and the following GPU types:
-* a100
-* h100
+* A100
+* A3-Highgpu (h100)
+* A3-Mega (h100-mega) - [Create cluster](#provisioning-a3-ultra-and-a3-mega-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-and-a3-mega-clusters-gpu-machines)
+* A3-Ultra (h200) - [Create cluster](#provisioning-a3-ultra-and-a3-mega-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-and-a3-mega-clusters-gpu-machines)
 
 and the following CPU types:
 * n2-standard-32
 
+# Cloud Console Permissions on the user or service account needed to run XPK:
+
+* Artifact Registry Writer
+* Compute Admin
+* Kubernetes Engine Admin
+* Logging Admin
+* Monitoring Admin
+* Service Account User
+* Storage Admin
+* Vertex AI Administrator
+
+# Prerequisites
+
+Following tools must be installed:
+
+- python >= 3.10 (download from [here](https://www.python.org/downloads/))
+- pip ([installation instruction](https://pip.pypa.io/en/stable/installation/))
+- python venv ([installation instruction](https://virtualenv.pypa.io/en/latest/installation.html))
+(all three of above can be installed at once from [here](https://packaging.python.org/en/latest/guides/installing-using-linux-tools/#installing-pip-setuptools-wheel-with-linux-package-managers))
+- gcloud (install from [here](https://cloud.google.com/sdk/gcloud#download_and_install_the))
+  - Run `gcloud init` 
+  - [Authenticate](https://cloud.google.com/sdk/gcloud/reference/auth/application-default/login) to Google Cloud
+- kubectl (install from [here](https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl#install_kubectl))
+  - Install `gke-gcloud-auth-plugin` from [here](https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl#install_plugin)
+- docker ([installation instruction](https://docs.docker.com/engine/install/))
+  - Run `gcloud auth configure-docker` to ensure images can be uploaded to registry 
+- make - please run below command.
+```shell
+# sudo may be required
+apt-get -y install make
+```
+In addition, below dependencies will be installed with `make install` command:
+- kueuectl (install from [here](https://kueue.sigs.k8s.io/docs/reference/kubectl-kueue/installation/))
+- kjob (installation instructions [here](https://github.com/kubernetes-sigs/kjob/blob/main/docs/installation.md))
+
 # Installation
-To install xpk, run the following command:
+To install xpk, run the following command and install additional tools, mentioned in [prerequisites](#prerequisites). [Makefile](https://github.com/AI-Hypercomputer/xpk/blob/main/Makefile) provides a way to install all neccessary tools:
 
 ```shell
 pip install xpk
 ```
 
+
 If you are running XPK by cloning GitHub repository, first run the
 following commands to begin using XPK commands:
 
 ```shell
 git clone https://github.com/google/xpk.git
 cd xpk
-# Install dependencies such as cloud-accelerator-diagnostics
-pip install .
+# Install required dependencies with make
+make install && export PATH=$PATH:$PWD/bin
 ```
 
+If you want to have installed dependecies persist in your PATH please run:
+`echo $PWD/bin` and add its value to `PATH` in .bashrc  or .zshrc
+
 If you see an error saying: `This environment is externally managed`, please use a virtual environment.
 
 Example:
@@ -100,8 +149,8 @@ Example:
   ## Clone the repository and installing dependencies.
   git clone https://github.com/google/xpk.git
   cd xpk
-  # Install dependencies such as cloud-accelerator-diagnostics
-  pip install .
+  # Install required dependencies with make
+  make install && export PATH=$PATH:$PWD/bin
 ```
 
 # XPK for Large Scale (>1k VMs)
@@ -139,14 +188,6 @@ gcloud config set compute/zone $ZONE
 xpk .. --zone $ZONE --project $PROJECT_ID
 ```
 
-`Cluster Create` command will create a project-specific Service Account. Note that only one service 
-account will be created per project. This service account will be attached to the node pools instead of default 
-[Compute Engine Service Account](https://cloud.google.com/compute/docs/access/service-accounts#default_service_account). 
-All the required permissions will be assigned to this service account by XPK. Make sure you have 
-[Service Account Admin](https://cloud.google.com/iam/docs/understanding-roles#iam.serviceAccountAdmin) and 
-[Project IAM Admin](https://cloud.google.com/iam/docs/understanding-roles#resourcemanager.projectIamAdmin) 
-roles assigned to your user account.
-
 The cluster created is a regional cluster to enable the GKE control plane across
 all zones.
 
@@ -179,13 +220,22 @@ all zones.
     ```
 
 * Cluster Create for Pathways:
-    Pathways compatible cluster can be created using `--enable-pathways`
+    Pathways compatible cluster can be created using `cluster create-pathways`.
     ```shell
-    python3 xpk.py cluster create \
+    python3 xpk.py cluster create-pathways \
     --cluster xpk-pw-test \
     --num-slices=4 --on-demand \
-    --tpu-type=v5litepod-16 \
-    --enable-pathways
+    --tpu-type=v5litepod-16
+    ```
+
+*   Cluster Create for Ray:
+    A cluster with KubeRay enabled and a RayCluster can be created using `cluster create-ray`.
+    ```shell
+    python3 xpk.py cluster create-ray \
+    --cluster xpk-rc-test \
+    --ray-version=2.39.0 \
+    --num-slices=4 --on-demand \
+    --tpu-type=v5litepod-8
     ```
 
 *   Cluster Create can be called again with the same `--cluster name` to modify
@@ -222,11 +272,77 @@ all zones.
     python3 xpk.py cluster create --force \
     --cluster xpk-test --tpu-type=v5litepod-16 \
     --num-slices=6  --reservation=$RESERVATION_ID
+    ```
+
+    and recreates the cluster with 4 slices of v4-8. The command will rerun to delete
+    6 slices of v5litepod-16 and create 4 slices of v4-8. The command will warn the
+    user when deleting slices. Use `--force` to skip prompts.
 
+    ```shell
+    python3 xpk.py cluster create \
+    --cluster xpk-test --tpu-type=v4-8 \
+    --num-slices=4  --reservation=$RESERVATION_ID
+
+    # Skip delete prompts using --force.
+
+    python3 xpk.py cluster create --force \
+    --cluster xpk-test --tpu-type=v4-8 \
+    --num-slices=4  --reservation=$RESERVATION_ID
     ```
 
+### Create Private Cluster
+
+XPK allows you to create a private GKE cluster for enhanced security. In a private cluster, nodes and pods are isolated from the public internet, providing an additional layer of protection for your workloads.
+
+To create a private cluster, use the following arguments:
+
+**`--private`**
+
+This flag enables the creation of a private GKE cluster. When this flag is set:
+
+*  Nodes and pods are isolated from the direct internet access.
+*  `master_authorized_networks` is automatically enabled.
+*  Access to the cluster's control plane is restricted to your current machine's IP address by default.
+
+**`--authorized-networks`**
+
+This argument allows you to specify additional IP ranges (in CIDR notation) that are authorized to access the private cluster's control plane and perform `kubectl` commands. 
+
+*  Even if this argument is not set when you have `--private`, your current machine's IP address will always be given access to the control plane.
+*  If this argument is used with an existing private cluster, it will replace the existing authorized networks.
+
+**Example Usage:**
+
+* To create a private cluster and allow access to Control Plane only to your current machine:
+
+  ```shell
+  python3 xpk.py cluster create \
+    --cluster=xpk-private-cluster \
+    --tpu-type=v4-8 --num-slices=2 \
+    --private
+  ```
+
+* To create a private cluster and allow access to Control Plane only to your current machine and the IP ranges `1.2.3.0/24` and `1.2.4.5/32`:
+
+  ```shell
+  python3 xpk.py cluster create \
+    --cluster=xpk-private-cluster \
+    --tpu-type=v4-8 --num-slices=2 \
+    --authorized-networks 1.2.3.0/24 1.2.4.5/32
+
+    # --private is optional when you set --authorized-networks
+  ```
+
+> **Important Notes:** 
+> * The argument `--private` is only applicable when creating new clusters. You cannot convert an existing public cluster to a private cluster using these flags.
+> * The argument `--authorized-networks` is applicable when creating new clusters or using an existing _*private*_ cluster. You cannot convert an existing public cluster to a private cluster using these flags.
+> * You need to [set up a Cluster NAT for your VPC network](https://cloud.google.com/nat/docs/set-up-manage-network-address-translation#creating_nat) so that the Nodes and Pods have outbound access to the internet. This is required because XPK installs and configures components such as kueue that need access to external sources like `registry.k8.io`.
+
+
 ### Create Vertex AI Tensorboard
-*Note: This feature is available in XPK >= 0.4.0. Enable [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis) in your Google Cloud console to use this feature.*
+*Note: This feature is available in XPK >= 0.4.0. Enable [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis) in your Google Cloud console to use this feature. Make sure you have
+[Vertex AI Administrator](https://cloud.google.com/vertex-ai/docs/general/access-control#aiplatform.admin) role
+assigned to your user account.*
 
 Vertex AI Tensorboard is a fully managed version of open-source Tensorboard. To learn more about Vertex AI Tensorboard, visit [this](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-introduction). Note that Vertex AI Tensorboard is only available in [these](https://cloud.google.com/vertex-ai/docs/general/locations#available-regions) regions.
 
@@ -312,6 +428,26 @@ will fail the cluster creation process because Vertex AI Tensorboard is not supp
     --tpu-type=v5litepod-16
     ```
 
+## Provisioning A3-Ultra and A3-Mega clusters (GPU machines)
+To create a cluster with A3 machines, run the below command. To create workloads on these clusters see [here](#workloads-for-a3-ultra-and-a3-mega-clusters-gpu-machines).
+  * For A3-Ultra: --device-type=h200-141gb-8
+  * For A3-Mega: --device-type=h100-mega-80gb-8
+
+  ```shell
+  python3 xpk.py cluster create \
+  --cluster CLUSTER_NAME --device-type=h200-141gb-8 \
+  --zone=$COMPUTE_ZONE  --project=$PROJECT_ID \
+  --num-nodes=4 --reservation=$RESERVATION_ID
+  ```
+Currently, the below flags/arguments are supported for A3-Mega and A3-Ultra machines:
+  * --num-nodes
+  * --default-pool-cpu-machine-type
+  * --default-pool-cpu-num-nodes
+  * --reservation
+  * --spot
+  * --on-demand (only A3-Mega)
+
+
 ## Workload Create
 *   Workload Create (submit training job):
 
@@ -323,26 +459,25 @@ will fail the cluster creation process because Vertex AI Tensorboard is not supp
     ```
 
 *   Workload Create for Pathways:
-    Pathways workload can be submitted using `--use-pathways` on a Pathways enabled cluster (created with `--enable-pathways`)
+    Pathways workload can be submitted using `workload create-pathways` on a Pathways enabled cluster (created with `cluster create-pathways`)
 
     Pathways workload example:
     ```shell
-    python3 xpk.py workload create \
+    python3 xpk.py workload create-pathways \
     --workload xpk-pw-test \
     --num-slices=1 \
     --tpu-type=v5litepod-16 \
-    --use-pathways \
     --cluster xpk-pw-test \
     --docker-name='user-workload' \
     --docker-image=<maxtext docker image> \
     --command='python3 MaxText/train.py MaxText/configs/base.yml base_output_directory=<output directory> dataset_path=<dataset path> per_device_batch_size=1 enable_checkpointing=false enable_profiler=false remat_policy=full global_parameter_scale=4 steps=300 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True run_name=$(USER)-pw-xpk-test-1'
     ```
 
-    Regular workload can also be submitted on a Pathways enabled cluster (created with `--enable-pathways`)
+    Regular workload can also be submitted on a Pathways enabled cluster (created with `cluster create-pathways`)
 
     Pathways workload example:
     ```shell
-    python3 xpk.py workload create \
+    python3 xpk.py workload create-pathways \
     --workload xpk-regular-test \
     --num-slices=1 \
     --tpu-type=v5litepod-16 \
@@ -352,6 +487,25 @@ will fail the cluster creation process because Vertex AI Tensorboard is not supp
     --command='python3 MaxText/train.py MaxText/configs/base.yml base_output_directory=<output directory> dataset_path=<dataset path> per_device_batch_size=1 enable_checkpointing=false enable_profiler=false remat_policy=full global_parameter_scale=4 steps=300 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True run_name=$(USER)-pw-xpk-test-1'
     ```
 
+    Pathways in headless mode - Pathways now offers the capability to run JAX workloads in Vertex AI notebooks or in GCE VMs!
+    Specify `--headless` with `workload create-pathways` when the user workload is not provided in a docker container.
+    ```shell
+    python3 xpk.py workload create-pathways --headless \
+    --workload xpk-pw-headless \
+    --num-slices=1 \
+    --tpu-type=v5litepod-16 \
+    --cluster xpk-pw-test
+    ```
+    Executing the command above would provide the address of the proxy that the user job should connect to.
+    ```shell
+    kubectl get pods
+    kubectl port-forward pod/<proxy-pod-name> 29000:29000
+    ```
+    ```shell
+    JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python -c 'import pathwaysutils; import jax; print(jax.devices())'
+    ```
+    Specify `JAX_PLATFORMS=proxy` and `JAX_BACKEND_TARGET=<proxy address from above>` and `import pathwaysutils` to establish this connection between the user's JAX code and the Pathways proxy. Execute Pathways workloads interactively on Vertex AI notebooks!
+
 ### Set `max-restarts` for production jobs
 
 * `--max-restarts <value>`: By default, this is 0. This will restart the job ""
@@ -360,6 +514,20 @@ 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.
 
+### Workloads for A3-Ultra and A3-Mega clusters (GPU machines)
+To submit jobs on a cluster with A3 machines, run the below command. To create a cluster with A3 machines see [here](#provisioning-a3-ultra-and-a3-mega-clusters-gpu-machines).
+  * For A3-Ultra: --device-type=h200-141gb-8
+  * For A3-Mega: --device-type=h100-mega-80gb-8
+
+  ```shell
+  python3 xpk.py workload create \
+  --workload=$WORKLOAD_NAME --command="echo goodbye" \
+  --cluster=$CLUSTER_NAME --device-type=h200-141gb-8 \
+  --zone=$COMPUTE_ZONE  --project=$PROJECT_ID \
+  --num-nodes=$WOKRKLOAD_NUM_NODES
+  ```
+> The docker image flags/arguments introduced in [workloads section](#workload-create) can be used with A3 machines as well.
+
 ### Workload Priority and Preemption
 * Set the priority level of your workload with `--priority=LEVEL`
 
@@ -386,7 +554,9 @@ checkpointing so the job restarts near where it was interrupted.
   ```
 
 ### Create Vertex AI Experiment to upload data to Vertex AI Tensorboard
-*Note: This feature is available in XPK >= 0.4.0. Enable [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis) in your Google Cloud console to use this feature.*
+*Note: This feature is available in XPK >= 0.4.0. Enable [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis) in your Google Cloud console to use this feature. Make sure you have
+[Vertex AI Administrator](https://cloud.google.com/vertex-ai/docs/general/access-control#aiplatform.admin) role
+assigned to your user account and to the [Compute Engine Service account](https://cloud.google.com/compute/docs/access/service-accounts#default_service_account) attached to the node pools in the cluster.*
 
 Vertex AI Experiment is a tool that helps to track and analyze an experiment run on Vertex AI Tensorboard. To learn more about Vertex AI Experiments, visit [this](https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments).
 
@@ -495,7 +665,7 @@ Check out [MaxText example](https://github.com/google/maxtext/pull/570) on how t
     --cluster xpk-test --filter-by-job=$USER
     ```
 
-* Workload List supports waiting for the completion of a specific job. XPK will follow an existing job until it has finished or the `timeout`, if provided, has been reached  and then list the job. If no `timeout` is specified, the default value is set to the max value, 1 week. You may also set `timeout=0` to poll the job once.  
+* Workload List supports waiting for the completion of a specific job. XPK will follow an existing job until it has finished or the `timeout`, if provided, has been reached  and then list the job. If no `timeout` is specified, the default value is set to the max value, 1 week. You may also set `timeout=0` to poll the job once.
 (Note: `restart-on-user-code-failure` must be set
 when creating the workload otherwise the workload will always finish with `Completed` status.)
 
@@ -514,11 +684,37 @@ when creating the workload otherwise the workload will always finish with `Compl
     --timeout=300
     ```
 
-  Return codes  
-    `0`: Workload finished and completed successfully.  
-    `124`: Timeout was reached before workload finished.  
-    `125`: Workload finished but did not complete successfully.  
-    `1`: Other failure.  
+  Return codes
+    `0`: Workload finished and completed successfully.
+    `124`: Timeout was reached before workload finished.
+    `125`: Workload finished but did not complete successfully.
+    `1`: Other failure.
+
+## Job List
+
+*   Job List (see jobs submitted via batch command):
+
+    ```shell
+    python3 xpk.py job ls --cluster xpk-test
+    ```
+
+* Example Job List Output:
+
+  ```
+    NAME                              PROFILE               LOCAL QUEUE   COMPLETIONS   DURATION   AGE
+    xpk-def-app-profile-slurm-74kbv   xpk-def-app-profile                 1/1           15s        17h
+    xpk-def-app-profile-slurm-brcsg   xpk-def-app-profile                 1/1           9s         3h56m
+    xpk-def-app-profile-slurm-kw99l   xpk-def-app-profile                 1/1           5s         3h54m
+    xpk-def-app-profile-slurm-x99nx   xpk-def-app-profile                 3/3           29s        17h
+  ```
+
+## Job Cancel
+
+*   Job Cancel (delete job submitted via batch command):
+
+    ```shell
+    python3 xpk.py job cancel xpk-def-app-profile-slurm-74kbv --cluster xpk-test
+    ```
 
 ## Inspector
 * Inspector provides debug info to understand cluster health, and why workloads are not running.
@@ -975,6 +1171,14 @@ gcloud compute machine-types list --zones=$ZONE_LIST
 python3 xpk.py cluster create --default-pool-cpu-machine-type=CPU_TYPE ...
 ```
 
+## Workload creation fails
+
+Some XPK cluster configuration might be missing, if workload creation fails with the below error.
+
+`[XPK] b'error: the server doesn\'t have a resource type "workloads"\n'`
+
+Mitigate this error by re-running your `xpk.py cluster create ...` command, to refresh the cluster configurations.
+
 ## Permission Issues: `requires one of ["permission_name"] permission(s)`.
 
 1) Determine the role needed based on the permission error:
@@ -1035,6 +1239,9 @@ gcloud beta compute reservations list --project=$PROJECT_ID
 gcloud beta compute reservations describe $RESERVATION --project=$PROJECT_ID --zone=$ZONE
 ```
 
+## 403 error on workload create when using `--base-docker-image` flag
+You need authority to push to the registry from your local machine. Try running `gcloud auth configure-docker`.
+
 # TPU Workload Debugging
 
 ## Verbose Logging
@@ -1076,3 +1283,65 @@ To explore the stack traces collected in a temporary directory in Kubernetes Pod
   --workload xpk-test-workload --command "python3 main.py" --cluster \
   xpk-test --tpu-type=v5litepod-16 --deploy-stacktrace-sidecar
  ```
+
+### Get information about jobs, queues and resources.
+
+To list available resources and queues use ```xpk info``` command. It allows to see localqueues and clusterqueues and check for available resources.
+
+To see queues with usage and workload info use:
+```shell
+python3 xpk.py info --cluster my-cluster
+```
+
+You can specify what kind of resources(clusterqueue or localqueue) you want to see using flags --clusterqueue or --localqueue.
+```shell
+python3 xpk.py info --cluster my-cluster --localqueue
+```
+
+# Local testing with Kind
+
+To facilitate development and testing locally, we have integrated support for testing with `kind`. This enables you to simulate a Kubernetes environment on your local machine.
+
+## Prerequisites
+
+- Install kind on your local machine. Follow the official documentation here: [Kind Installation Guide.](https://kind.sigs.k8s.io/docs/user/quick-start#installation)
+
+## Usage
+
+xpk interfaces seamlessly with kind to manage Kubernetes clusters locally, facilitating the orchestration and management of workloads. Below are the commands for managing clusters:
+
+### Cluster Create
+*   Cluster create:
+
+    ```shell
+    python3 xpk.py kind create \
+    --cluster xpk-test
+    ```
+
+### Cluster Delete
+*   Cluster Delete:
+
+    ```shell
+    python3 xpk.py kind delete \
+    --cluster xpk-test
+    ```
+
+### Cluster List
+*   Cluster List:
+
+    ```shell
+    python3 xpk.py kind list
+    ```
+
+## Local Testing Basics
+
+Local testing is available exclusively through the `batch` and `job` commands of xpk with the `--kind-cluster` flag. This allows you to simulate training jobs locally:
+
+```shell
+python xpk.py batch [other-options] --kind-cluster script
+```
+
+Please note that all other xpk subcommands are intended for use with cloud systems on Google Cloud Engine (GCE) and don't support local testing. This includes commands like cluster, info, inspector, etc.
+
+# Other advanced usage
+[Use a Jupyter notebook to interact with a Cloud TPU cluster](xpk-notebooks.md)