Run a GPU job and keep the result

You will upscale an image on your own Replicate account without holding a request open, with a cost limit, and store the result in files. It takes about 20 minutes.

You need

• A MainSDK app with files set up (see file uploads).
• A Replicate API token in secret:REPLICATE_API_TOKEN, and a webhook secret in secret:REPLICATE_WEBHOOK_SECRET.

1. Declare the target

compute:
  poll: 15s
  limits: { per-tenant: $50/month }
  targets:
    upscale:
      driver: replicate
      model: nightmareai/real-esrgan:f121d640bd286e1fdc67f9799164c1d5be36ff74576ee11c803ae5b665dd46aa
      region: us
      timeout: 30m
      price: $0.0014/s
      max-cost: $2

A target with no price or timeout, a timeout over 24 hours, or no version hash stops mainsdk check. Use the local driver with a program in tests; it runs through system.run.

2. Start a job

import { compute, type ComputeRunOptions } from "mainsdk/experimental/compute";

type StoredFile = NonNullable<ComputeRunOptions["files"]>[number];   // a file stored with `files`

export async function upscale(image: StoredFile) {
  const { jobId } = await compute.run("upscale", { scale: 4 }, { files: [image] });
  return jobId;
}

The job is queued when the call returns. If price times timeout could pass maxCost, the call is refused first with MAIN-COMPUTE-0002 and nothing is sent.

3. Watch, cancel and read

import { compute } from "mainsdk/experimental/compute";

declare const jobId: string;

const job = await compute.get(jobId); // state, progress, cost, result
await compute.cancel(jobId, { reason: "user changed their mind" });

Progress is also published on the realtime channel compute:<jobId>. A job with no sign of life for three heartbeats becomes lost and is cancelled at the provider.

4. Use it inside a workflow

import { compute } from "mainsdk/experimental/compute";

const result = await compute.step("upscale-step", "upscale", { scale: 4 });

The step waits for the job and returns its result, and the workflow survives a restart in between.

Test it

Use FakeComputeProvider and ParkingFlows from mainsdk/testing to run a job without a provider account, and computeConformance to check a driver of your own.