Add a skeleton for the foundation model class

[shchur]

Issue #, if available:

• This PR outlines the user-facing API for the new FoundationModel class in AutoGluon Cloud. This class allows users to deploy and run inference with pretrained foundation models (Chronos, Mitra, etc.) on AWS.
• How is this different from sagemaker.jumpstart.JumpStartModel?
  • DataFrame in / DataFrame out (no manual JSON payload construction)
  • Supports batch prediction and fine-tuning out of the box
  • Model onboarding is much easier for us in the future - any model available in AutoGluon works as soon as we add it to the config
• How is this different from TimeSeriesCloudPredictor?
  • No fit() required before deploy() or predict() — foundation models work out of the box
  • Cleaner API: the user does need to provide nested dictionaries of predictor_init_kwargs and predictor_fit_kwargs
  • The FoundationModel class is stateless — no need to worry about the underlying state (attached artifacts, endpoints, etc)

Example usage

# ============================================================
# Time Series — Chronos
# ============================================================
from autogluon.cloud import FoundationModel

model = FoundationModel("chronos-bolt-base", role_arn="arn:aws:iam::123456789:role/SageMakerRole")

# --- Deploy an endpoint ---
endpoint = model.deploy(instance_type="ml.g5.xlarge")
predictions = endpoint.predict(df, prediction_length=24)
endpoint.delete_endpoint()

# --- Batch predict (large dataset, uses SageMaker Training job under the hood) ---
predictions = model.predict(df, prediction_length=24, quantile_levels=[0.1, 0.5, 0.9])

# --- Batch predict (async) ---
job = model.predict(df, prediction_length=24, wait=False)
# ... later ...
predictions = job.result()

# --- Fine-tune, then deploy ---
fine_tuned = model.fit(train_data=df, instance_type="ml.g5.xlarge")
endpoint = fine_tuned.deploy()

# --- Use custom/cached weights (VPC-safe) ---
model = FoundationModel(
    "chronos-bolt-base",
    hyperparameters={"model_path": "s3://my-bucket/cached-weights/"},
)
endpoint = model.deploy()

# --- Override hyperparameters at call time ---
predictions = model.predict(df, prediction_length=24, hyperparameters={"cross_learning": False})

# ============================================================
# Tabular — Mitra / TabICL
# ============================================================
model = FoundationModel("mitra-classification", role_arn="arn:aws:iam::123456789:role/SageMakerRole")

predictions = model.predict(train_data=train_df, test_data=test_df, label="class")

To Do:

• Wiring the user-facing API to the internal backend that actually runs stuff on AWS

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[AnirudhDagar]

Awesome design and I like this API! Just one thing to note, #213 introduces autogluon.cloud.init() which persists the user's role_arn, region, and S3 bucket to ~/.autogluon/cloud.yaml at first-time setup. If FoundationModel falls back to that config when those args aren't passed, usage collapses to:

# one time per account
import autogluon.cloud as agc
agc.init()

# Every session after that
from autogluon.cloud import FoundationModel

model = FoundationModel("chronos-bolt-base")            # no role_arn needed it will automatically be picked
endpoint = model.deploy(instance_type="ml.g5.xlarge")
predictions = endpoint.predict(data=df, prediction_length=24)

I can add this, once both the PRs are shipped.

[prateekdesai04]

AG-Cloud requires SageMaker as a package, it would be good to assume that customers may have v3 as well.
v3 has quite some API changes, do these changes fit v3 and v2 both (should be the case ideally)?
Here is also a migration by the Sagemaker team - https://github.com/nargokul/sagemaker-python-sdk/blob/master/migration.md#migration-tool-mcp-server

[abdulfatir]

Thanks @shchur! This looks great and would really simplify the user experience with FMs. A few comments:

• Is FoundationModel supposed to be a magic class which would automatically infer the model/task type? Maybe we can add it more concrete in some way, e.g., by adding a task_type kwarg or subclassing the base FoundationModel class?
• Related to above: how do we plan to handle model-specific hyperparameters (e.g., cross_learning, n_estimators, etc.)? What exactly would model_config map to internally?
• For the batch prediction mode, we will have SM job start overhead. Is there a way to reduce this by using another AWS service here (Lambda?)?

[shchur]

Thanks @abdulfatir,

1. The FoundationModel class automatically resolves to TimeSeriesFoundationModel, TabularFoundationModel, etc based on the task associated with each model_id. The subclasses only differ in their API for predict: e.g. TimeSeriesFoundationModel takes historical data + known covariates as input, and TabularFoundationModel takes train & test DFs as input

[shchur]

@abdulfatir
2. model_config = hyperparameters provided to the TimeSeriesPredictor/TabularPredictor under the hood. Now that you pointed this out, I think that we should separate the training and inference configs, e.g.

"chronos-2": {
    "task": "forecasting",
    "model_name": "Chronos2",
    "inference_config": {"model_path": "amazon/chronos-2"},
    "training_config": {"model_path": "amazon/chronos-2", "fine_tune": True},
    ...
},

The user can provide kwargs to deploy(), predict(), fit() that updates the default inference/train configs stored in the registry.

Do you think we should allow the user to configure these hyperparameters inside the payload (e.g. passing cross_learning: True inside the call to endpoint.predict), or is it okay if this is fixed at the creation of the endpoint?

[shchur]

@abdulfatir
3. We can add support for other backends like Lambda in the future - but I would decouple this from the FoundationModel API so that the user doesn't need to worry about it

[melopeo]

Question:

Are TabularFoundationModel and TimeSeriesFoundationModel part of the public API, or internal to the FoundationModel factory? They're in __all__, but examples only use FoundationModel(...). For instance, should users be able to write:

def forecast(model: TimeSeriesFoundationModel, df: pd.DataFrame) -> pd.DataFrame:
    return model.predict(df, prediction_length=24)

[shchur]

Yes, with the current design the user should be able to directly create a TimeSeriesFoundationModel object bypassing the FoundationModel base class, e.g.

model = TimeSeriesFoundationModel(model_id="chronos-bolt-base")
predictions = model.predict(...)

The FoundationModel base class only exists for convenience.

[melopeo]

Putting fit() on FoundationModel (the base class) implies every foundation model can be fine-tuned.

In general it makes sense to support fine-tuning of a pretrained model. My question is if we are assuming all FoundationModels under consideration will have a fine-tuning capability. If a foundation model can't be fine-tuned, how are we planning to handle that?

[shchur]

Good question, I added a boolean flag to the ModelConfig indicating whether fine-tuning is supported.

[melopeo]

Hey @shchur , thanks for this. It looks very nice.

Question: The example calls job.result(), but RemoteJob doesn't define it. Is the plan to add it in a follow-up?

[shchur]

@melopeo I missed the top-level comment. Here is a temporary solution, I will figure out the best path when implementing the backend.