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)