Using TPUs

name: enumerate-tpu-devices
type: job
image:
  type: image
  image_uri: python:3.11
  command: bash /run.sh
mounts:
  - type: string
    mount_path: /run.sh
    data: |
      #!/bin/bash
      pip install --quiet --upgrade 'jax[tpu]>0.3.0' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print(jax.devices())'
retries: 0
trigger:
  type: manual
resources:
  node:
    type: node_selector
    capacity_type: spot_fallback_on_demand
  devices:
    - name: tpu-v5-lite-device
      type: gcp_tpu
      topology: 1x1
  cpu_request: 21
  cpu_limit: 23
  memory_request: 39100
  memory_limit: 46000
  ephemeral_storage_request: 20000
  ephemeral_storage_limit: 100000
trigger_on_deploy: true

#!/bin/bash

gcloud config set project your-project
PROJECT_ID=$(gcloud config get project)
BUCKET_NAME=tpu-models-bucket
REGION=us-central1
LOCATION=us-central1-a
WORKSPACE_NAME=tpu-models-workspace
SA_NAME=tpu-models-sa


gcloud iam service-accounts create "${SA_NAME}"

gcloud storage buckets add-iam-policy-binding \
    "gs://${BUCKET_NAME}" \
    --member "serviceAccount:${SA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com"  \
    --role roles/storage.objectUser \
    --project "${PROJECT_ID}"
gcloud storage buckets add-iam-policy-binding \
  "gs://${BUCKET_NAME}" \
  --member "serviceAccount:${SA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com" \
  --role roles/storage.insightsCollectorService \
  --project "${PROJECT_ID}"

gcloud iam service-accounts add-iam-policy-binding \
    "${SA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role roles/iam.workloadIdentityUser \
    --member "serviceAccount:${PROJECT_ID}.svc.id.goog[${WORKSPACE_NAME}/${SA_NAME}]"

import logging
from truefoundry.deploy import Job, Image, Resources, GcpTPU, TPUType, SecretMount

logging.basicConfig(level=logging.INFO)

job = Job(
    name="maxengine-ckpt-conv",
    # Mention your bucket in command below
    image=Image(
        image_uri="us-docker.pkg.dev/cloud-tpu-images/inference/inference-checkpoint:v0.2.0",
        command="/usr/bin/checkpoint_converter.sh -b=tpu-models-bucket -m=google/gemma/maxtext/2b-it/3",
    ),
    # Service account to use to write to the bucket
    service_account="tpu-models-sa",
    # Secret to use for fetching the model checkpoint from Kaggle
    mounts=[
        SecretMount(
            mount_path="/kaggle/kaggle.json",
            secret_fqn="tfy-secret://your-org:kaggle-secrets:kaggle-json"
        )
    ],
    resources=Resources(
        cpu_request=20,
        cpu_limit=20,
        memory_request=150000,
        memory_limit=190000,
        ephemeral_storage_request=30000,
        ephemeral_storage_limit=40000,
        devices=[
            GcpTPU(
                name=TPUType.V5_LITE_DEVICE,
                topology="2x2"
            )
        ]
    ),
    trigger_on_deploy=True,
)

# Add the fqn of the Workspace we created earlier
job.deploy(workspace_fqn="<your-workspace-fqn>", wait=False)

import logging
from servicefoundry import Service, Image, Port, Resources, GcpTPU, TPUType

logging.basicConfig(level=logging.INFO)

service = Service(
    name="maxengine-server",
    # Model name and bucket are mentioned in the command
    image=Image(
        image_uri="us-docker.pkg.dev/cloud-tpu-images/inference/maxengine-server:v0.2.0",
        command="/usr/bin/maxengine_server_entrypoint.sh model_name=gemma-2b tokenizer_path=assets/tokenizer.gemma per_device_batch_size=4 max_prefill_predict_length=1024 max_target_length=2048 async_checkpointing=false ici_fsdp_parallelism=1 ici_autoregressive_parallelism=-1 ici_tensor_parallelism=1 scan_layers=false weight_dtype=bfloat16 load_parameters_path=gs://tpu-models-bucket/final/unscanned/gemma_2b-it/0/checkpoints/0/items",
    ),
    ports=[
        Port(port=9000, expose=False),
    ],
    # Service account to use to read from the bucket
    service_account="tpu-models-sa",
    resources=Resources(
        cpu_request=20,
        cpu_limit=24,
        memory_request=40000,
        memory_limit=48000,
        ephemeral_storage_request=30000,
        ephemeral_storage_limit=40000,
        devices=[
            GcpTPU(
                name=TPUType.V5_LITE_DEVICE,
                topology="1x1"
            )
        ]
    ),
)

# Add the fqn of the Workspace we created earlier
service.deploy(workspace_fqn="<your-workspace-fqn>", wait=False)

import logging
from servicefoundry import Service, Image, LocalSource, DockerFileBuild, Port, Resources

logging.basicConfig(level=logging.INFO)

service = Service(
    name="jetstream-http",
    image=Image(image_uri="truefoundrycloud/jetstream-http-patched:v0.2.0"),
    ports=[
        Port(
            port=8000,
            host="<Your Host>"
        ),
    ],
    resources=Resources(
        cpu_request=0.5,
        cpu_limit=1,
        memory_request=500,
        memory_limit=1000,
        ephemeral_storage_request=1000,
        ephemeral_storage_limit=4000,
    ),
)

# Add the fqn of the Workspace we created earlier
service.deploy(workspace_fqn="<your-workspace-fqn>", wait=False)