AIMP SDK (Python)

This package provides the class-based SDK used by model authors.

Install

python -m pip install ai-marketplace-sdk

Install (local source)

cd packages/python/sdk
uv pip install .

Development (editable)

cd packages/python/sdk
uv pip install -e .

Model Runtime Entrypoint

Model containers must use the canonical executable module:

python -m aimp_sdk.model_runner --module /app/main.py

aimp_sdk.model_runner is the only supported public model runtime entrypoint. It starts the long-lived stdin/stdout runner that emits the initial {"type":"ready"} protocol message and serves mode requests.

Streaming and Playground (authoring)

• Streaming — use Stream.text(text=Schema.Outputs.<field>, status=Schema.Outputs.<json>) with StreamState for status(...) updates; done() returns validated outputs. The execution sink remains internal.
• Playground — declare PlaygroundSpec in playground.py only: When, SectionRole, Markdown(..., stream=True), sources=PlaygroundSources(history=source.chat_history()), computed=PlaygroundComputed(...) with computed.asset_ref / computed.asset_refs, and typed item fields via JsonField(schema=..., multiple=True) plus Image.asset_ref(...). Only outputs referenced by views are UI-rendered; other outputs stay contract-only.

Fine-Tune From Scratch

The fine-tune flow has two sides:

• model authoring: define the schema and training runtime inside the model project
• consumer dataset creation: build one local bundle that matches the published schema

1. Define the model-scoped tuning contract

Model authors declare fine-tune inputs by subclassing FineTuneSchema in tuning/dataset.py.

Available field types:

• ImageField
• AudioField
• VideoField
• FileField
• TextField
• JsonField

Example:

from aimp_sdk import FineTuneSchema, ImageField


def validate_png(value: object) -> None:
    if not str(value).lower().endswith(".png"):
        raise ValueError("Only PNG files are allowed")


class ModelFineTuneSchema(FineTuneSchema):
    query_image = ImageField(required=True, validators=[validate_png])
    candidate_image = ImageField(required=True, validators=[validate_png])

Use field-level validators=[...] and preprocessors=[...] when one field needs custom cleanup or validation. These hooks run during local bundle build, before upload.

2. Implement the tuning runner

Model authors implement tuning behavior in tuning/runner.py and bind it from main.py using tuning = ....

The runner declares strategy plus tuning contracts and implements run(session) -> TuneResult.

from aimp_sdk import Model, TuneResult


class CoreTuning:
    strategy = "lora"
    dataset_schema = CoreDataset
    hyperparams_schema = CoreHyperparams

    async def run(self, session) -> TuneResult:
        learning_rate = session.suggest_float("learning_rate", 1e-5, 1e-3, default=5e-4)
        epochs = session.suggest_int("epochs", 1, 5, default=3)
        for record in session.train_dataset:
            train_step(
                query_image=session.resolve_path(record.query_image),
                candidate_image=session.resolve_path(record.candidate_image),
                learning_rate=learning_rate,
            )

        _ = epochs
        validation_recall_at_10 = evaluate(session.validation_dataset)
        return TuneResult(metrics={"validation_recall_at_10": validation_recall_at_10})


class CoreModel(Model):
    engine = CoreEngine
    tuning = CoreTuning

The runner owns only model-authored training logic and typed schemas. Dataset bundle format, validation split policy, optimization targets, and hardware selection belong to job/platform inputs, not runner class attributes.

session exposes only training concerns:

• session.train_dataset
• session.validation_dataset
• session.params
• session.resolve_path(...)
• session.suggest_float(...)
• session.suggest_int(...)
• session.suggest_categorical(...)

In the current SDK, authors implement build(self, ctx) to load runtime state once. The separate load(...) method is framework-internal cache/lifecycle plumbing and is not the author-facing hook.

session.params comes from the model's published tuning.parameters contract. Defaults and search space metadata are contract-owned by the platform. Use session.suggest_* helpers for explicit default handling; the SDK does not add random or hidden fallbacks.

Model code does not handle dataset files, S3 uploads, adapter URIs, or trial orchestration.

3. Build the consumer dataset bundle

Consumers do not write model-side tuning code. They use the published schema through the CLI:

1. ai fine-tune schema --model <model-slug>
2. ai fine-tune scaffold --model <model-slug>
3. edit prepare_bundle.py and add any local assets/
4. python prepare_bundle.py
5. ai fine-tune create --model <model-slug> --output-dir .

During bundle build the SDK:

• validates required fields and field types
• runs field preprocessors
• validates the processed values
• copies binary assets into the bundle
• writes data.parquet and manifest.json

That means field preprocessors and validators run before upload, not during remote execution.

4. Use the derived model directly

Successful LoRA fine-tunes produce a derived executable model with its own asset ID and slug. Clients should execute that derived model directly, as if it were a new model:

job = client.fine_tune_job(job_id="ft-123").wait_for_completion()

result = client.models.execute(
    model=job.tuned_model_slug,
    mode="generate",
    input={"prompt": "hello"},
)

Client code should not pass adapter_id to executions. Adapter resolution is platform-internal.