import coiled, torch 

@coiled.function(
    vm_type="g5.xlarge",
    region="us-west-1",
)
def train(trial: int) -> pytorch.Model:
    torch.device("cuda:0")
    
    model = ...

    return model.to_device("cpu")

model = train(123)

import xarray as xr
import coiled
  
cluster = coiled.Cluster(
    n_workers=500,
    region="us-west-2",
    worker_memory="64 GiB",
    spot_policy="spot",
)
client = cluster.get_client()

ds = xr.open_mfdataset(
    "s3://mybucket/all-data.zarr",
    engine="zarr", parallel=True,
)

import coiled
import s3fs

@coiled.function(region="us-west-2", worker_memory="128 GiB")
def process(filename: str) -> float:
    with s3.open(filename, mode="rb") as f:
        data = f.read()
        
    output = transform(data)
    
    with s3.open(filename + ".out", mode="wb") as f:
        f.write(output)

filenames = s3.ls("s3://mybucket/*.png")
process.map(filenames)