LeRobot.js / packages /web /src /runpodTrainer.ts
Aditya Shankar
feat: added dataset recording; hf uploader, s3 uploader; runpod trainer (#6)
4384839 unverified
/**
* Class to train a model using runpod
*/
class LeRobotRunpodTrainer extends EventTarget {
api_token: string;
constructor(api_token: string) {
super();
this.api_token = api_token;
}
async _startPod(api_token: string) {
const podOptions = {
allowedCudaVersions: ["12.8"],
cloudType: "SECURE",
computeType: "GPU",
containerDiskInGb: 50,
countryCodes: ["<string>"],
cpuFlavorIds: ["cpu3c"],
cpuFlavorPriority: "availability",
dataCenterIds: ["EU-RO-1","CA-MTL-1"],
dataCenterPriority: "availability",
dockerEntrypoint: [],
dockerStartCmd: [],
env: {"ENV_VAR":"value"},
globalNetworking: true,
gpuCount: 1,
gpuTypeIds: ["NVIDIA GeForce RTX 4090"],
gpuTypePriority: "availability",
imageName: "runpod/pytorch:2.1.0-py3.10-cuda11.8.0-devel-ubuntu22.04",
interruptible: false,
locked: false,
minDiskBandwidthMBps: 123,
minDownloadMbps: 123,
minRAMPerGPU: 8,
minUploadMbps: 123,
minVCPUPerGPU: 2,
name: "my pod",
ports: ["8888/http","22/tcp"],
supportPublicIp: true,
templateId: null,
vcpuCount: 2,
volumeInGb: 20,
volumeMountPath: "/workspace"
}
const options = {
method: 'POST',
headers: {Authorization: 'Bearer ' + api_token, 'Content-Type': 'application/json'},
body: JSON.stringify(podOptions)
};
const response = await fetch('https://rest.runpod.io/v1/pods', options);
const data = await response.json();
return data;
}
async start() {
await this._startPod(this.api_token);
this.dispatchEvent(new CustomEvent("deployed_training_pod"));
}
}