Upload a Dataset

# Requires Python 3.9 or later and the requests package.
import os
import time

import requests

API = "https://distillation.training.wandb.ai/v1/tasks/ticket-classifier"
HEADERS = {
    "Authorization": f"Bearer {os.environ['WANDB_API_KEY']}",
    "Wandb-Entity": "your-team",
}
FILE = "tickets.jsonl"

# 1. Create the import session.
response = requests.post(
    f"{API}/dataset-imports",
    headers=HEADERS,
    json={
        "idempotency_key": "support-tickets-v1",
        "name": "Support tickets (uploaded)",
        "file": {"name": FILE, "size_bytes": os.path.getsize(FILE)},
        "split_policy": {"mode": "automatic", "val_fraction": 0.1},
        "duplicate_policy": {"split_overlap": "drop_train"},
    },
)
response.raise_for_status()
session = response.json()
import_id = session["id"]
part_size = session["file"]["part_size_bytes"]
part_count = session["file"]["part_count"]

# 2. Request a signed upload URL for every part.
response = requests.post(
    f"{API}/dataset-imports/{import_id}/parts",
    headers=HEADERS,
    json={"part_numbers": list(range(1, part_count + 1))},
)
response.raise_for_status()
url_by_part = {part["part_number"]: part["url"] for part in response.json()["parts"]}

# 3. Upload the parts. Part numbers start at 1. Signed URLs carry their own authorization.
with open(FILE, "rb") as handle:
    for part_number in range(1, part_count + 1):
        chunk = handle.read(part_size)
        requests.put(url_by_part[part_number], data=chunk).raise_for_status()

# 4. Optional: read the session to see which parts arrived.
progress = requests.get(f"{API}/dataset-imports/{import_id}", headers=HEADERS).json()
uploaded = {part["part_number"] for part in progress["file"]["parts"]}
missing = [n for n in range(1, part_count + 1) if n not in uploaded]

# 5. Complete the upload. This queues validation and dataset creation.
requests.post(f"{API}/dataset-imports/{import_id}/complete", headers=HEADERS).raise_for_status()

# 6. Poll until the session is ready or failed.
while True:
    status = requests.get(f"{API}/dataset-imports/{import_id}", headers=HEADERS).json()
    if status["state"] in ("ready", "failed"):
        break
    time.sleep(10)

print(status["state"], status["dataset_id"], status["error"])

// Requires Node.js 18 or later. No packages to install.
import { open, stat } from "node:fs/promises";

const API = "https://distillation.training.wandb.ai/v1/tasks/ticket-classifier";
const HEADERS = {
  Authorization: `Bearer ${process.env.WANDB_API_KEY}`,
  "Wandb-Entity": "your-team",
};
const FILE = "tickets.jsonl";

async function call(path, init = {}) {
  const response = await fetch(`${API}${path}`, { ...init, headers: { ...HEADERS, ...init.headers } });
  if (!response.ok) throw new Error(`${init.method ?? "GET"} ${path} failed: ${await response.text()}`);
  return response.status === 204 ? null : response.json();
}

// 1. Create the import session.
const { size } = await stat(FILE);
const session = await call("/dataset-imports", {
  method: "POST",
  headers: { "Content-Type": "application/json" },
  body: JSON.stringify({
    idempotency_key: "support-tickets-v1",
    name: "Support tickets (uploaded)",
    file: { name: FILE, size_bytes: size },
    split_policy: { mode: "automatic", val_fraction: 0.1 },
    duplicate_policy: { split_overlap: "drop_train" },
  }),
});
const importId = session.id;
const partSize = session.file.part_size_bytes;
const partCount = session.file.part_count;

// 2. Request a signed upload URL for every part.
const urls = await call(`/dataset-imports/${importId}/parts`, {
  method: "POST",
  headers: { "Content-Type": "application/json" },
  body: JSON.stringify({ part_numbers: Array.from({ length: partCount }, (_, i) => i + 1) }),
});
const urlByPart = new Map(urls.parts.map((part) => [part.part_number, part.url]));

// 3. Upload the parts. Part numbers start at 1. Signed URLs carry their own authorization.
const handle = await open(FILE, "r");
try {
  for (let partNumber = 1; partNumber <= partCount; partNumber++) {
    const offset = (partNumber - 1) * partSize;
    const chunk = Buffer.alloc(Math.min(partSize, size - offset));
    await handle.read(chunk, 0, chunk.length, offset);
    const put = await fetch(urlByPart.get(partNumber), { method: "PUT", body: chunk });
    if (!put.ok) throw new Error(`Part ${partNumber} failed with ${put.status}`);
  }
} finally {
  await handle.close();
}

// 4. Optional: read the session to see which parts arrived.
const progress = await call(`/dataset-imports/${importId}`);
const uploaded = new Set(progress.file.parts.map((part) => part.part_number));
const missing = Array.from({ length: partCount }, (_, i) => i + 1).filter((n) => !uploaded.has(n));

// 5. Complete the upload. This queues validation and dataset creation.
await call(`/dataset-imports/${importId}/complete`, { method: "POST" });

// 6. Poll until the session is ready or failed.
let status;
while (true) {
  status = await call(`/dataset-imports/${importId}`);
  if (status.state === "ready" || status.state === "failed") break;
  await new Promise((resolve) => setTimeout(resolve, 10_000));
}
console.log(status.state, status.dataset_id, status.error);