Remote Inference with @lium.machine

This example runs a local Python function on a GPU pod. The decorator provisions the pod, uploads the function, installs Python dependencies, runs inference, returns the result, and tears the pod down. For how it works, what it costs and the keep_warm / .map() / local options, see Serverless functions with @lium.machine.

Authenticate before running the script:

lium init

For CI, scripts, or temporary overrides, set LIUM_API_KEY instead.

The example is written for the lium package since release 0.0.40 (lium#208), where machine="RTX4090" rents the cheapest node with exactly one RTX 4090 (an A100 here would fail: every A100 node listed today is an 8×A100, and "A100" means exactly one) and requirements are installed on top of the image's torch. The decorator in 0.0.37–0.0.39 matches the string against the node name as a substring — RTX4090 matches nothing (the names read NVIDIA GeForce RTX 4090, with a space) and the call stops with No node found matching machine type; machine="RTX 4090" there rents the first such node whatever its GPU count, billed as the whole node — and installs requirements into a venv that does not see the image's packages: before 0.0.40 there is no way to ask for a single GPU here, and "torch" has to be in requirements.

#!/usr/bin/env python3
"""Run a small instruct model on a remote Lium GPU."""

import lium

@lium.machine(machine="RTX4090", requirements=["transformers", "accelerate"])  # cheapest node with one RTX 4090; torch comes with the image
def infer(prompt: str) -> str:
    from transformers import AutoModelForCausalLM, AutoTokenizer

    model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        dtype="auto",
        device_map="cuda",
    )

    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=64,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
    )

    response_tokens = outputs[0][inputs["input_ids"].shape[-1]:]
    return tokenizer.decode(response_tokens, skip_special_tokens=True).strip()

print(infer("In one sentence, who discovered penicillin?"))

The exact generated text can vary by model version, but it should identify Alexander Fleming.