grailrobotics 0.1.0

Python SDK for the Grail Robotics Finetuning API

Grail Robotics SDK

grailrobotics is a Python SDK for interacting with the Grail Robotics job runner API.
It provisions Lambda Cloud instances, performs setup, and runs training jobs for fine tuning.

---

Client Setup

from grailrobotics import Client

c = Client()

---

Health Check

GET /health

print(c.health())

---

Start a Job

POST /jobs

import base64

with open("/path/to/ssh_private_key.pem", "rb") as f:
    ssh_key_b64 = base64.b64encode(f.read()).decode("utf-8")

job = c.start_job({
    "lambda_api_key": "YOUR_LAMBDA_API_KEY",
    "hf_token": "YOUR_HF_TOKEN", 

    "ssh_private_key_b64": ssh_key_b64,
    "ssh_key_name": "your-lambda-ssh-key-name",
    "ssh_username": "ubuntu",

    "instance_type": "gpu_1x_a10",
    "region": "us-west-1",

    "reuse_existing": False,

    "train_flags": {
        "steps": 200,

        "policy.type": "pi0",
        "dataset.repo_id": "<YOUR_HUGGINGFACE_REPO_ID>",
        "policy.repo_id": "<YOUR_HUGGINGFACE_POLICY_NAME>", #Name you want your policy to be
        "job_name": "",
        "output_dir": "outputs/pi0_training",

        "wandb.enable": False,

        "policy.pretrained_path": "lerobot/pi0_base",
        "policy.compile_model": False,
        "policy.gradient_checkpointing": True,
        "policy.dtype": "bfloat16",
        "policy.empty_cameras": 1,
        "policy.device": "cuda",

        "batch_size": 32,

        "rename_map": {
            "observation.images.front": "observation.images.base_0_rgb",
            "observation.images.side": "observation.images.left_wrist_0_rgb",
        },
    },
})

job_id = job["job_id"]
print("Started job:", job_id)

---

Get Job Status

GET /jobs/{job_id}

status = c.job_status(job_id)
print(status)

---

Stream Job Progress

GET /jobs/{job_id}/progress

progress = c.job_progress(job_id)

print("Status:", progress["status"])
print("".join(progress["tail"]))

---

Fetch Job Logs

GET /jobs/{job_id}/logs

logs = c.job_logs(job_id, tail=200)
print("".join(logs["lines"]))

---

Cancel a Job

POST /jobs/{job_id}/cancel

Requires the Lambda API key to terminate the instance.

c.cancel_job(
    job_id,
    lambda_api_key="YOUR_LAMBDA_API_KEY",
)

---

List Lambda Instance Types

GET /lambda/instance-types

types = c.list_instance_types(
    lambda_api_key="YOUR_LAMBDA_API_KEY"
)

print(types["instance_type_names"])

---

Notes

• Secrets (Lambda API key, SSH key, HF token) are never stored.
• Cancelling a job terminates the backing Lambda instance.
• Logs and progress endpoints are safe to poll.