Local Models

No AWS Required

Local artifacts need no AWS account, no config, and no credentials. Paired with PublicData, you can go from install to a trained model without touching AWS at all.

The Local classes mirror the Workbench artifact API against your filesystem. The chain is the same — DataSource → FeatureSet → Model → Endpoint — and training runs the same generated model script that SageMaker runs, as a subprocess. So a script written locally publishes to AWS and produces the same model.

Storage lives under WORKBENCH_LOCAL_PATH (default ~/.workbench/local).

Local is where you iterate: try a feature list, a framework, a set of hyperparameters. There's no build cost and deleting is instant. AWS is where a model lands once it's worth keeping — that's where monitoring, deployed endpoints, and everything else the team consumes live.

Try it

PublicData reads public S3 anonymously, so this runs with no AWS setup:

from workbench.local import LocalDataSource, PublicData, ModelType, ModelFramework

df = PublicData().get("comp_chem/aqsol/aqsol_public_data")

ds = LocalDataSource(df, name="aqsol_local")
fs = ds.to_features("aqsol_local_features", id_column="ID")
model = fs.to_model(
    "aqsol-local-reg",
    model_type=ModelType.REGRESSOR,
    model_framework=ModelFramework.XGBOOST,
    target_column="solubility",
    feature_list=["molwt", "mollogp", "tpsa", "numrotatablebonds"],
)

print(model.get_inference_metrics())
predictions = model.to_endpoint().inference(fs.pull_dataframe().head(10))

to_features() lowercases column names, same as the AWS path, so target_column and feature_list refer to the FeatureSet's names rather than the source frame's. Use fs.columns to see them.

A LocalEndpoint isn't deployed anywhere. inference() loads the model in-process through the same model_fn/predict_fn a real endpoint container uses, so the predictions match what a deployed endpoint returns.

validation_ids, sample_weights, and exclude_ids work as they do in AWS, and they're recorded so publishing can replay them.

The Workbench REPL exposes all of these, so none of the imports are needed there.

Scoring

Inference runs work the same as they do on an AWS Model, so a script that walks them runs against either. Metrics are computed from the run's predictions.

model.list_inference_runs()                     # ["full_cross_fold", ...]
model.get_inference_metrics()                   # defaults to full_cross_fold
model.get_inference_predictions("full_cross_fold")

model.oof_predictions()          # the cross-fold predictions directly
model.validation_predictions()   # held-out rows, when validation_ids were used

# Naming a capture adds it to the run list
model.to_endpoint().inference(eval_df, capture_name="holdout")
model.get_inference_metrics("holdout")

The model_training run an AWS Model carries has no local equivalent — those metrics come from SageMaker scraping the training job's output.

Publishing

model.publish_plan()        # what it would create, creates nothing
aws_model = model.publish() # ds -> fs -> model -> endpoint

Publishing walks up the lineage and creates whatever AWS doesn't already have, then deploys an endpoint (pass endpoint=False to stop at the model). It retrains in AWS from the published FeatureSet rather than uploading local artifacts, so the model lands in the registry like any other — with the row roles replayed, so it trains on the same rows.

That launches a real SageMaker training job, which is why publish_plan() is a separate call: look before you leap.

If a published model disagrees with the local one, model.version_drift() reports package versions that differ between this machine and the training image.

Listing and deleting

from workbench.local import LocalMeta

LocalMeta().models()     # also data_sources(), feature_sets(), endpoints()

The Workbench REPL prints a local summary at startup and on local_summary().

Always delete through the API. LocalModel.delete() takes its endpoints with it; removing directories by hand leaves endpoints pointing at a model that no longer exists. That's the only cascade — deleting a LocalFeatureSet leaves its models alone.

What's not here

Local covers training and scoring. Plots, the inference store, monitoring, contests, and promotion are all properties of published artifacts — when you want those, publish.

LocalDataSource: A DataFrame on local disk, queryable with DuckDB.

LocalDataSource

Bases: LocalArtifact

LocalDataSource: Workbench Local DataSource Class

Common Usage

my_data = LocalDataSource(df, name="my_data")
my_data.query("select * from my_data where height > 0.3")
my_features = my_data.to_features("my_features", id_column="id")

Source code in src/workbench/local/local_data_source.py

class LocalDataSource(LocalArtifact):
    """LocalDataSource: Workbench Local DataSource Class

    Common Usage:
        ```python
        my_data = LocalDataSource(df, name="my_data")
        my_data.query("select * from my_data where height > 0.3")
        my_features = my_data.to_features("my_features", id_column="id")
        ```
    """

    artifact_type = "data_source"

    def __init__(self, source: Union[str, pd.DataFrame] = None, name: str = None, **kwargs):
        """Initialize a LocalDataSource

        Args:
            source (Union[str, pd.DataFrame]): A DataFrame, a CSV/parquet file path, or an
                existing LocalDataSource name. If None, `name` must reference an existing source.
            name (str): The name of the data source (must be lowercase). Required for DataFrames.
        """
        # A bare name refers to an existing local data source
        if isinstance(source, str) and name is None and not os.path.isfile(source):
            name = source
            source = None

        # Derive a name from a file path when one wasn't given
        if name is None and isinstance(source, str):
            name = Artifact.generate_valid_name(os.path.splitext(os.path.basename(source))[0])

        if name is None:
            msg = "Set the 'name' argument: LocalDataSource(df, name='my_data')"
            self.log.critical(msg)
            raise ValueError(msg)
        Artifact.is_name_valid(name)

        # Call superclass init (sets up paths)
        super().__init__(name, **kwargs)
        self.data_path = os.path.join(self.path, "data.parquet")

        # Load the source (if given)
        if source is not None:
            self._load_source(source)

    def query(self, query: str) -> pd.DataFrame:
        """Query this DataSource with DuckDB

        Args:
            query (str): SQL to run; reference this artifact by its name

        Returns:
            pd.DataFrame: The results of the query
        """
        if not self.exists():
            self.log.error(f"Local artifact {self.name} does not exist...")
            return pd.DataFrame()

        with duckdb.connect() as con:
            con.execute(f"CREATE VIEW \"{self.name}\" AS SELECT * FROM read_parquet('{self.data_path}')")
            return con.execute(query).df()

    def pull_dataframe(self, limit: int = None) -> pd.DataFrame:
        """Return a DataFrame of ALL the data from this DataSource

        Args:
            limit (int): Limit the number of rows returned (default: None = all rows)

        Returns:
            pd.DataFrame: A DataFrame of the data from this DataSource
        """
        if not self.exists():
            self.log.error(f"Local artifact {self.name} does not exist...")
            return pd.DataFrame()
        df = pd.read_parquet(self.data_path)
        return df.head(limit) if limit else df

    @property
    def columns(self) -> list[str]:
        """Return the column names for this DataSource"""
        return list(self.workbench_meta().get("columns", []))

    @property
    def column_types(self) -> list[str]:
        """Return the column types for this DataSource"""
        return list(self.workbench_meta().get("column_types", []))

    def num_rows(self) -> int:
        """Return the number of rows for this DataSource"""
        return self.workbench_meta().get("num_rows", 0)

    def num_columns(self) -> int:
        """Return the number of columns for this DataSource"""
        return len(self.columns)

    def details(self, **kwargs) -> dict:
        """LocalDataSource Details

        Returns:
            dict: A dictionary of details about the LocalDataSource
        """
        return {**super().details(), "num_rows": self.num_rows(), "num_columns": self.num_columns()}

    def to_features(
        self,
        name: str,
        id_column: str,
        tags: list = None,
        event_time_column: str = None,
        one_hot_columns: list = None,
    ) -> Union["LocalFeatureSet", None]:  # noqa: F821
        """Convert this LocalDataSource to a LocalFeatureSet

        Args:
            name (str): Set the name for the feature set (must be lowercase).
            id_column (str): The ID column (must be specified, use "auto" for auto-generated IDs).
            tags (list, optional): Set the tags for the feature set (unused, kept for API parity).
            event_time_column (str, optional): The event time column (default: None).
            one_hot_columns (list, optional): Columns to one-hot encode (default: None).

        Returns:
            LocalFeatureSet: The FeatureSet created from this DataSource (or None on invalid name)
        """
        from workbench.local.local_feature_set import LocalFeatureSet

        if not Artifact.is_name_valid(name):
            self.log.critical(f"Invalid FeatureSet name: {name}, not creating FeatureSet!")
            return None

        return LocalFeatureSet.from_dataframe(
            self.pull_dataframe(),
            name=name,
            id_column=id_column,
            event_time_column=event_time_column,
            one_hot_columns=one_hot_columns,
            input_name=self.name,
        )

    def aws_exists(self) -> bool:
        """Does an AWS DataSource by this name already exist?

        Returns:
            bool: True if AWS already has this DataSource
        """
        from workbench.api import DataSource

        return DataSource(self.name).exists()

    def _aws_artifact(self):
        """Internal: The AWS DataSource for this local one"""
        from workbench.api import DataSource

        return DataSource(self.name)

    def _publish_self(self, **kwargs):
        """Internal: Create the AWS DataSource from this local one

        Returns:
            DataSource: The created AWS DataSource
        """
        from workbench.api import DataSource

        return DataSource(self.pull_dataframe(), name=self.name)

    def _load_source(self, source: Union[str, pd.DataFrame]):
        """Internal: Write the source data to local storage

        Args:
            source (Union[str, pd.DataFrame]): A DataFrame or a CSV/parquet file path
        """
        if isinstance(source, pd.DataFrame):
            df = source
        elif source.endswith(".parquet"):
            df = pd.read_parquet(source)
        else:
            df = pd.read_csv(source)

        self.log.important(f"Storing local data source {self.name} ({len(df)} rows)...")
        storage.local_root(create=True)
        os.makedirs(self.path, exist_ok=True)
        df.to_parquet(self.data_path, index=False)
        self._init_storage(input_name="dataframe" if isinstance(source, pd.DataFrame) else str(source))
        self.upsert_workbench_meta(
            {
                "num_rows": len(df),
                "columns": list(df.columns),
                "column_types": [str(dtype) for dtype in df.dtypes],
            }
        )

column_types property

Return the column types for this DataSource

columns property

Return the column names for this DataSource

__init__(source=None, name=None, **kwargs)

Initialize a LocalDataSource

Parameters:

Name	Type	Description	Default
source	Union[str, DataFrame]	A DataFrame, a CSV/parquet file path, or an existing LocalDataSource name. If None, name must reference an existing source.	None
name	str	The name of the data source (must be lowercase). Required for DataFrames.	None

Source code in src/workbench/local/local_data_source.py

def __init__(self, source: Union[str, pd.DataFrame] = None, name: str = None, **kwargs):
    """Initialize a LocalDataSource

    Args:
        source (Union[str, pd.DataFrame]): A DataFrame, a CSV/parquet file path, or an
            existing LocalDataSource name. If None, `name` must reference an existing source.
        name (str): The name of the data source (must be lowercase). Required for DataFrames.
    """
    # A bare name refers to an existing local data source
    if isinstance(source, str) and name is None and not os.path.isfile(source):
        name = source
        source = None

    # Derive a name from a file path when one wasn't given
    if name is None and isinstance(source, str):
        name = Artifact.generate_valid_name(os.path.splitext(os.path.basename(source))[0])

    if name is None:
        msg = "Set the 'name' argument: LocalDataSource(df, name='my_data')"
        self.log.critical(msg)
        raise ValueError(msg)
    Artifact.is_name_valid(name)

    # Call superclass init (sets up paths)
    super().__init__(name, **kwargs)
    self.data_path = os.path.join(self.path, "data.parquet")

    # Load the source (if given)
    if source is not None:
        self._load_source(source)

aws_exists()

Does an AWS DataSource by this name already exist?

Returns:

Name	Type	Description
bool	bool	True if AWS already has this DataSource

Source code in src/workbench/local/local_data_source.py

def aws_exists(self) -> bool:
    """Does an AWS DataSource by this name already exist?

    Returns:
        bool: True if AWS already has this DataSource
    """
    from workbench.api import DataSource

    return DataSource(self.name).exists()

details(**kwargs)

LocalDataSource Details

Returns:

Name	Type	Description
dict	dict	A dictionary of details about the LocalDataSource

Source code in src/workbench/local/local_data_source.py

def details(self, **kwargs) -> dict:
    """LocalDataSource Details

    Returns:
        dict: A dictionary of details about the LocalDataSource
    """
    return {**super().details(), "num_rows": self.num_rows(), "num_columns": self.num_columns()}

num_columns()

Return the number of columns for this DataSource

Source code in src/workbench/local/local_data_source.py

def num_columns(self) -> int:
    """Return the number of columns for this DataSource"""
    return len(self.columns)

num_rows()

Return the number of rows for this DataSource

Source code in src/workbench/local/local_data_source.py

def num_rows(self) -> int:
    """Return the number of rows for this DataSource"""
    return self.workbench_meta().get("num_rows", 0)

pull_dataframe(limit=None)

Return a DataFrame of ALL the data from this DataSource

Parameters:

Name	Type	Description	Default
limit	int	Limit the number of rows returned (default: None = all rows)	None

Returns:

Type	Description
DataFrame	pd.DataFrame: A DataFrame of the data from this DataSource

Source code in src/workbench/local/local_data_source.py

def pull_dataframe(self, limit: int = None) -> pd.DataFrame:
    """Return a DataFrame of ALL the data from this DataSource

    Args:
        limit (int): Limit the number of rows returned (default: None = all rows)

    Returns:
        pd.DataFrame: A DataFrame of the data from this DataSource
    """
    if not self.exists():
        self.log.error(f"Local artifact {self.name} does not exist...")
        return pd.DataFrame()
    df = pd.read_parquet(self.data_path)
    return df.head(limit) if limit else df

query(query)

Query this DataSource with DuckDB

Parameters:

Name	Type	Description	Default
query	str	SQL to run; reference this artifact by its name	required

Returns:

Type	Description
DataFrame	pd.DataFrame: The results of the query

Source code in src/workbench/local/local_data_source.py

def query(self, query: str) -> pd.DataFrame:
    """Query this DataSource with DuckDB

    Args:
        query (str): SQL to run; reference this artifact by its name

    Returns:
        pd.DataFrame: The results of the query
    """
    if not self.exists():
        self.log.error(f"Local artifact {self.name} does not exist...")
        return pd.DataFrame()

    with duckdb.connect() as con:
        con.execute(f"CREATE VIEW \"{self.name}\" AS SELECT * FROM read_parquet('{self.data_path}')")
        return con.execute(query).df()

to_features(name, id_column, tags=None, event_time_column=None, one_hot_columns=None)

Convert this LocalDataSource to a LocalFeatureSet

Parameters:

Name	Type	Description	Default
name	str	Set the name for the feature set (must be lowercase).	required
id_column	str	The ID column (must be specified, use "auto" for auto-generated IDs).	required
tags	list	Set the tags for the feature set (unused, kept for API parity).	None
event_time_column	str	The event time column (default: None).	None
one_hot_columns	list	Columns to one-hot encode (default: None).	None

Returns:

Name	Type	Description
LocalFeatureSet	Union[LocalFeatureSet, None]	The FeatureSet created from this DataSource (or None on invalid name)

Source code in src/workbench/local/local_data_source.py

def to_features(
    self,
    name: str,
    id_column: str,
    tags: list = None,
    event_time_column: str = None,
    one_hot_columns: list = None,
) -> Union["LocalFeatureSet", None]:  # noqa: F821
    """Convert this LocalDataSource to a LocalFeatureSet

    Args:
        name (str): Set the name for the feature set (must be lowercase).
        id_column (str): The ID column (must be specified, use "auto" for auto-generated IDs).
        tags (list, optional): Set the tags for the feature set (unused, kept for API parity).
        event_time_column (str, optional): The event time column (default: None).
        one_hot_columns (list, optional): Columns to one-hot encode (default: None).

    Returns:
        LocalFeatureSet: The FeatureSet created from this DataSource (or None on invalid name)
    """
    from workbench.local.local_feature_set import LocalFeatureSet

    if not Artifact.is_name_valid(name):
        self.log.critical(f"Invalid FeatureSet name: {name}, not creating FeatureSet!")
        return None

    return LocalFeatureSet.from_dataframe(
        self.pull_dataframe(),
        name=name,
        id_column=id_column,
        event_time_column=event_time_column,
        one_hot_columns=one_hot_columns,
        input_name=self.name,
    )

LocalFeatureSet: Engineered features on local disk, queryable with DuckDB.

LocalFeatureSet

Bases: LocalArtifact

LocalFeatureSet: Workbench Local FeatureSet Class

Common Usage

my_features = LocalFeatureSet("my_features")
my_features.query("select * from my_features where solubility < -5")
my_model = my_features.to_model(...)

Source code in src/workbench/local/local_feature_set.py

class LocalFeatureSet(LocalArtifact):
    """LocalFeatureSet: Workbench Local FeatureSet Class

    Common Usage:
        ```python
        my_features = LocalFeatureSet("my_features")
        my_features.query("select * from my_features where solubility < -5")
        my_model = my_features.to_model(...)
        ```
    """

    artifact_type = "feature_set"

    def __init__(self, name: str, **kwargs):
        """Initialize a LocalFeatureSet

        Args:
            name (str): The name of an existing local feature set
        """
        Artifact.is_name_valid(name)
        super().__init__(name, **kwargs)
        self.data_path = os.path.join(self.path, "data.parquet")

    @classmethod
    def from_dataframe(
        cls,
        df: pd.DataFrame,
        name: str,
        id_column: str,
        event_time_column: str = None,
        one_hot_columns: list = None,
        input_name: str = "dataframe",
    ) -> "LocalFeatureSet":
        """Create a LocalFeatureSet from a DataFrame, running the shared column prep.

        Args:
            df (pd.DataFrame): The DataFrame of features
            name (str): The name for the feature set (must be lowercase)
            id_column (str): The ID column (use "auto" for auto-generated IDs)
            event_time_column (str, optional): Event time column (default: None)
            one_hot_columns (list, optional): Columns to one-hot encode (default: None)
            input_name (str): Name of this feature set's input (default: "dataframe")

        Returns:
            LocalFeatureSet: The created feature set
        """
        # Same prep the AWS ingest path runs, so columns/names/dtypes match after publish
        df, id_column = feature_prep.prep_dataframe(
            df.copy(),
            id_column=id_column,
            event_time_column=event_time_column,
            one_hot_columns=one_hot_columns,
        )

        fs = cls(name)
        fs.log.important(f"Storing local feature set {name} ({len(df)} rows, {len(df.columns)} columns)...")
        storage.local_root(create=True)
        os.makedirs(fs.path, exist_ok=True)
        df.to_parquet(fs.data_path, index=False)
        fs._init_storage(input_name=input_name)
        fs.upsert_workbench_meta(
            {
                "id_column": id_column,
                "num_rows": len(df),
                "columns": list(df.columns),
                "column_types": [str(dtype) for dtype in df.dtypes],
            }
        )
        return fs

    @property
    def id_column(self) -> str:
        """The ID column for this FeatureSet"""
        return self.workbench_meta().get("id_column")

    @property
    def columns(self) -> list[str]:
        """Return the column names for this FeatureSet"""
        return list(self.workbench_meta().get("columns", []))

    @property
    def column_types(self) -> list[str]:
        """Return the column types for this FeatureSet"""
        return list(self.workbench_meta().get("column_types", []))

    def num_rows(self) -> int:
        """Return the number of rows for this FeatureSet"""
        return self.workbench_meta().get("num_rows", 0)

    def num_columns(self) -> int:
        """Return the number of columns for this FeatureSet"""
        return len(self.columns)

    def query(self, query: str) -> pd.DataFrame:
        """Query this FeatureSet with DuckDB

        Args:
            query (str): SQL to run; reference this artifact by its name

        Returns:
            pd.DataFrame: The results of the query
        """
        if not self.exists():
            self.log.error(f"Local artifact {self.name} does not exist...")
            return pd.DataFrame()

        with duckdb.connect() as con:
            con.execute(f"CREATE VIEW \"{self.name}\" AS SELECT * FROM read_parquet('{self.data_path}')")
            return con.execute(query).df()

    def pull_dataframe(self, limit: int = None) -> pd.DataFrame:
        """Return a DataFrame of ALL the data from this FeatureSet

        Args:
            limit (int): Limit the number of rows returned (default: None = all rows)

        Returns:
            pd.DataFrame: A DataFrame of the data from this FeatureSet
        """
        if not self.exists():
            self.log.error(f"Local artifact {self.name} does not exist...")
            return pd.DataFrame()
        df = pd.read_parquet(self.data_path)
        return df.head(limit) if limit else df

    def details(self, **kwargs) -> dict:
        """LocalFeatureSet Details

        Returns:
            dict: A dictionary of details about the LocalFeatureSet
        """
        return {
            **super().details(),
            "id_column": self.id_column,
            "num_rows": self.num_rows(),
            "num_columns": self.num_columns(),
        }

    def to_model(self, name: str, model_type, model_framework, **kwargs: Any) -> "LocalModel":  # noqa: F821
        """Train a LocalModel from this FeatureSet.

        Args:
            name (str): The name of the Model to create
            model_type (ModelType): The type of model to create
            model_framework (ModelFramework): The framework to use
            **kwargs: Passed to LocalModel.from_feature_set (target_column, feature_list,
                hyperparameters, sample_weights, validation_ids, exclude_ids, wait)

        Returns:
            LocalModel: The Model created from this FeatureSet
        """
        from workbench.local.local_model import LocalModel

        return LocalModel.from_feature_set(
            self, name=name, model_type=model_type, model_framework=model_framework, **kwargs
        )

    def parent(self):
        """The LocalDataSource this FeatureSet came from, if it still exists locally"""
        from workbench.local.local_data_source import LocalDataSource

        source = LocalDataSource(self.get_input())
        return source if source.exists() else None

    def aws_exists(self) -> bool:
        """Does an AWS FeatureSet by this name already exist?

        Returns:
            bool: True if AWS already has this FeatureSet
        """
        from workbench.api import FeatureSet

        return FeatureSet(self.name).exists()

    def _aws_artifact(self):
        """Internal: The AWS FeatureSet for this local one"""
        from workbench.api import FeatureSet

        return FeatureSet(self.name)

    def _publish_self(self, **kwargs):
        """Internal: Push this FeatureSet's engineered features to AWS.

        The local frame is published as-is rather than recomputing features from the
        AWS DataSource, so the published feature values are exactly the ones trained on.

        Returns:
            FeatureSet: The created AWS FeatureSet
        """
        from workbench.api import FeatureSet
        from workbench.core.transforms.pandas_transforms.pandas_to_features import PandasToFeatures

        to_features = PandasToFeatures(self.name)
        to_features.set_input(self.pull_dataframe(), id_column=self.id_column)
        to_features.set_output_tags([self.name])
        to_features.transform()
        return FeatureSet(self.name)

    def training_view(
        self,
        sample_weights: Union[dict, pd.DataFrame] = None,
        validation_ids: list = None,
        exclude_ids: list = None,
    ) -> pd.DataFrame:
        """Build the training frame: features plus the three role columns.

        Mirrors the AWS model training view: `sample_weight` (default 1.0),
        `validation` (default False), and `exclude` (default False). Excluded rows
        are dropped entirely, and exclude wins over validation.

        Args:
            sample_weights (Union[dict, pd.DataFrame], optional): id -> weight, forwarded as-is
            validation_ids (list, optional): ids held out of training and scored as a holdout
            exclude_ids (list, optional): ids dropped from the training frame entirely

        Returns:
            pd.DataFrame: The feature columns plus sample_weight/validation/exclude
        """
        df = self.pull_dataframe()
        ids = df[self.id_column]

        if isinstance(sample_weights, pd.DataFrame):
            sample_weights = dict(zip(sample_weights[self.id_column], sample_weights["sample_weight"]))
        df["sample_weight"] = ids.map(sample_weights).fillna(1.0) if sample_weights else 1.0
        df["validation"] = ids.isin(validation_ids) if validation_ids else False
        df["exclude"] = ids.isin(exclude_ids) if exclude_ids else False

        # Excluded rows never reach a model (exclude wins over validation)
        return df[~df["exclude"]].reset_index(drop=True)

column_types property

Return the column types for this FeatureSet

columns property

Return the column names for this FeatureSet

id_column property

The ID column for this FeatureSet

__init__(name, **kwargs)

Initialize a LocalFeatureSet

Parameters:

Name	Type	Description	Default
name	str	The name of an existing local feature set	required

Source code in src/workbench/local/local_feature_set.py

def __init__(self, name: str, **kwargs):
    """Initialize a LocalFeatureSet

    Args:
        name (str): The name of an existing local feature set
    """
    Artifact.is_name_valid(name)
    super().__init__(name, **kwargs)
    self.data_path = os.path.join(self.path, "data.parquet")

aws_exists()

Does an AWS FeatureSet by this name already exist?

Returns:

Name	Type	Description
bool	bool	True if AWS already has this FeatureSet

Source code in src/workbench/local/local_feature_set.py

def aws_exists(self) -> bool:
    """Does an AWS FeatureSet by this name already exist?

    Returns:
        bool: True if AWS already has this FeatureSet
    """
    from workbench.api import FeatureSet

    return FeatureSet(self.name).exists()

details(**kwargs)

LocalFeatureSet Details

Returns:

Name	Type	Description
dict	dict	A dictionary of details about the LocalFeatureSet

Source code in src/workbench/local/local_feature_set.py

def details(self, **kwargs) -> dict:
    """LocalFeatureSet Details

    Returns:
        dict: A dictionary of details about the LocalFeatureSet
    """
    return {
        **super().details(),
        "id_column": self.id_column,
        "num_rows": self.num_rows(),
        "num_columns": self.num_columns(),
    }

from_dataframe(df, name, id_column, event_time_column=None, one_hot_columns=None, input_name='dataframe') classmethod

Create a LocalFeatureSet from a DataFrame, running the shared column prep.

Parameters:

Name	Type	Description	Default
df	DataFrame	The DataFrame of features	required
name	str	The name for the feature set (must be lowercase)	required
id_column	str	The ID column (use "auto" for auto-generated IDs)	required
event_time_column	str	Event time column (default: None)	None
one_hot_columns	list	Columns to one-hot encode (default: None)	None
input_name	str	Name of this feature set's input (default: "dataframe")	'dataframe'

Returns:

Name	Type	Description
LocalFeatureSet	LocalFeatureSet	The created feature set

Source code in src/workbench/local/local_feature_set.py

@classmethod
def from_dataframe(
    cls,
    df: pd.DataFrame,
    name: str,
    id_column: str,
    event_time_column: str = None,
    one_hot_columns: list = None,
    input_name: str = "dataframe",
) -> "LocalFeatureSet":
    """Create a LocalFeatureSet from a DataFrame, running the shared column prep.

    Args:
        df (pd.DataFrame): The DataFrame of features
        name (str): The name for the feature set (must be lowercase)
        id_column (str): The ID column (use "auto" for auto-generated IDs)
        event_time_column (str, optional): Event time column (default: None)
        one_hot_columns (list, optional): Columns to one-hot encode (default: None)
        input_name (str): Name of this feature set's input (default: "dataframe")

    Returns:
        LocalFeatureSet: The created feature set
    """
    # Same prep the AWS ingest path runs, so columns/names/dtypes match after publish
    df, id_column = feature_prep.prep_dataframe(
        df.copy(),
        id_column=id_column,
        event_time_column=event_time_column,
        one_hot_columns=one_hot_columns,
    )

    fs = cls(name)
    fs.log.important(f"Storing local feature set {name} ({len(df)} rows, {len(df.columns)} columns)...")
    storage.local_root(create=True)
    os.makedirs(fs.path, exist_ok=True)
    df.to_parquet(fs.data_path, index=False)
    fs._init_storage(input_name=input_name)
    fs.upsert_workbench_meta(
        {
            "id_column": id_column,
            "num_rows": len(df),
            "columns": list(df.columns),
            "column_types": [str(dtype) for dtype in df.dtypes],
        }
    )
    return fs

num_columns()

Return the number of columns for this FeatureSet

Source code in src/workbench/local/local_feature_set.py

def num_columns(self) -> int:
    """Return the number of columns for this FeatureSet"""
    return len(self.columns)

num_rows()

Return the number of rows for this FeatureSet

Source code in src/workbench/local/local_feature_set.py

def num_rows(self) -> int:
    """Return the number of rows for this FeatureSet"""
    return self.workbench_meta().get("num_rows", 0)

parent()

The LocalDataSource this FeatureSet came from, if it still exists locally

Source code in src/workbench/local/local_feature_set.py

def parent(self):
    """The LocalDataSource this FeatureSet came from, if it still exists locally"""
    from workbench.local.local_data_source import LocalDataSource

    source = LocalDataSource(self.get_input())
    return source if source.exists() else None

pull_dataframe(limit=None)

Return a DataFrame of ALL the data from this FeatureSet

Parameters:

Name	Type	Description	Default
limit	int	Limit the number of rows returned (default: None = all rows)	None

Returns:

Type	Description
DataFrame	pd.DataFrame: A DataFrame of the data from this FeatureSet

Source code in src/workbench/local/local_feature_set.py

def pull_dataframe(self, limit: int = None) -> pd.DataFrame:
    """Return a DataFrame of ALL the data from this FeatureSet

    Args:
        limit (int): Limit the number of rows returned (default: None = all rows)

    Returns:
        pd.DataFrame: A DataFrame of the data from this FeatureSet
    """
    if not self.exists():
        self.log.error(f"Local artifact {self.name} does not exist...")
        return pd.DataFrame()
    df = pd.read_parquet(self.data_path)
    return df.head(limit) if limit else df

query(query)

Query this FeatureSet with DuckDB

Parameters:

Name	Type	Description	Default
query	str	SQL to run; reference this artifact by its name	required

Returns:

Type	Description
DataFrame	pd.DataFrame: The results of the query

Source code in src/workbench/local/local_feature_set.py

def query(self, query: str) -> pd.DataFrame:
    """Query this FeatureSet with DuckDB

    Args:
        query (str): SQL to run; reference this artifact by its name

    Returns:
        pd.DataFrame: The results of the query
    """
    if not self.exists():
        self.log.error(f"Local artifact {self.name} does not exist...")
        return pd.DataFrame()

    with duckdb.connect() as con:
        con.execute(f"CREATE VIEW \"{self.name}\" AS SELECT * FROM read_parquet('{self.data_path}')")
        return con.execute(query).df()

to_model(name, model_type, model_framework, **kwargs)

Train a LocalModel from this FeatureSet.

Parameters:

Name	Type	Description	Default
name	str	The name of the Model to create	required
model_type	ModelType	The type of model to create	required
model_framework	ModelFramework	The framework to use	required
**kwargs	Any	Passed to LocalModel.from_feature_set (target_column, feature_list, hyperparameters, sample_weights, validation_ids, exclude_ids, wait)	{}

Returns:

Name	Type	Description
LocalModel	LocalModel	The Model created from this FeatureSet

Source code in src/workbench/local/local_feature_set.py

def to_model(self, name: str, model_type, model_framework, **kwargs: Any) -> "LocalModel":  # noqa: F821
    """Train a LocalModel from this FeatureSet.

    Args:
        name (str): The name of the Model to create
        model_type (ModelType): The type of model to create
        model_framework (ModelFramework): The framework to use
        **kwargs: Passed to LocalModel.from_feature_set (target_column, feature_list,
            hyperparameters, sample_weights, validation_ids, exclude_ids, wait)

    Returns:
        LocalModel: The Model created from this FeatureSet
    """
    from workbench.local.local_model import LocalModel

    return LocalModel.from_feature_set(
        self, name=name, model_type=model_type, model_framework=model_framework, **kwargs
    )

training_view(sample_weights=None, validation_ids=None, exclude_ids=None)

Build the training frame: features plus the three role columns.

Mirrors the AWS model training view: sample_weight (default 1.0), validation (default False), and exclude (default False). Excluded rows are dropped entirely, and exclude wins over validation.

Parameters:

Name	Type	Description	Default
sample_weights	Union[dict, DataFrame]	id -> weight, forwarded as-is	None
validation_ids	list	ids held out of training and scored as a holdout	None
exclude_ids	list	ids dropped from the training frame entirely	None

Returns:

Type	Description
DataFrame	pd.DataFrame: The feature columns plus sample_weight/validation/exclude

Source code in src/workbench/local/local_feature_set.py

def training_view(
    self,
    sample_weights: Union[dict, pd.DataFrame] = None,
    validation_ids: list = None,
    exclude_ids: list = None,
) -> pd.DataFrame:
    """Build the training frame: features plus the three role columns.

    Mirrors the AWS model training view: `sample_weight` (default 1.0),
    `validation` (default False), and `exclude` (default False). Excluded rows
    are dropped entirely, and exclude wins over validation.

    Args:
        sample_weights (Union[dict, pd.DataFrame], optional): id -> weight, forwarded as-is
        validation_ids (list, optional): ids held out of training and scored as a holdout
        exclude_ids (list, optional): ids dropped from the training frame entirely

    Returns:
        pd.DataFrame: The feature columns plus sample_weight/validation/exclude
    """
    df = self.pull_dataframe()
    ids = df[self.id_column]

    if isinstance(sample_weights, pd.DataFrame):
        sample_weights = dict(zip(sample_weights[self.id_column], sample_weights["sample_weight"]))
    df["sample_weight"] = ids.map(sample_weights).fillna(1.0) if sample_weights else 1.0
    df["validation"] = ids.isin(validation_ids) if validation_ids else False
    df["exclude"] = ids.isin(exclude_ids) if exclude_ids else False

    # Excluded rows never reach a model (exclude wins over validation)
    return df[~df["exclude"]].reset_index(drop=True)

LocalModel: A model trained on this machine by the generated model script.

LocalModel

Bases: LocalArtifact

LocalModel: Workbench Local Model Class

Training runs the same generated model script that SageMaker runs, with the same arguments, as a subprocess against local directories.

Common Usage

my_model = LocalModel("my_model")
my_model.training_state()
my_model.oof_predictions()

Source code in src/workbench/local/local_model.py

class LocalModel(LocalArtifact):
    """LocalModel: Workbench Local Model Class

    Training runs the same generated model script that SageMaker runs, with the
    same arguments, as a subprocess against local directories.

    Common Usage:
        ```python
        my_model = LocalModel("my_model")
        my_model.training_state()
        my_model.oof_predictions()
        ```
    """

    artifact_type = "model"
    data_files = ()

    def __init__(self, name: str, **kwargs):
        """Initialize a LocalModel

        Args:
            name (str): The name of the model
        """
        Artifact.is_name_valid(name, delimiter="-", lower_case=False)
        super().__init__(name, **kwargs)
        self.script_dir = os.path.join(self.path, "script")
        self.model_dir = os.path.join(self.path, "model_artifacts")
        self.output_dir = os.path.join(self.path, "output")
        self.train_dir = os.path.join(self.path, "input", "train")
        self.log_path = os.path.join(self.output_dir, "training.log")
        self.status_path = os.path.join(self.path, "status.json")

    @classmethod
    def from_feature_set(
        cls,
        feature_set,
        name: str,
        model_type: ModelType,
        model_framework: ModelFramework,
        target_column: Union[str, list[str]] = None,
        feature_list: list = None,
        model_class: str = None,
        model_import_str: str = None,
        custom_script: str = None,
        hyperparameters: dict = None,
        sample_weights: Union[dict, pd.DataFrame] = None,
        validation_ids: list = None,
        exclude_ids: list = None,
        wait: bool = True,
    ) -> "LocalModel":
        """Train a LocalModel from a LocalFeatureSet.

        Args:
            feature_set (LocalFeatureSet): The feature set to train on
            name (str): The name of the model to create
            model_type (ModelType): The type of model to create
            model_framework (ModelFramework): The framework to use
            target_column (str or list[str], optional): Target column(s), None for unsupervised
            feature_list (list, optional): Feature columns; derived from the FeatureSet if omitted
            model_class (str, optional): Model class for scikit-learn models (e.g. "KMeans")
            model_import_str (str, optional): Import line for the model class
            custom_script (str, optional): Path to a custom model script or template
            hyperparameters (dict, optional): Hyperparameters for the model
            sample_weights (Union[dict, pd.DataFrame], optional): id -> framework weight
            validation_ids (list, optional): ids held out and scored as a validation set
            exclude_ids (list, optional): ids dropped from training entirely
            wait (bool): Block until training finishes (default: True)

        Returns:
            LocalModel: The model (trained if wait=True, still training otherwise)
        """
        supervised = model_type in (
            ModelType.CLASSIFIER,
            ModelType.REGRESSOR,
            ModelType.UQ_REGRESSOR,
            ModelType.ENSEMBLE_REGRESSOR,
        )
        if target_column is None and supervised:
            raise ValueError("target_column is required for supervised models (pass target_column=...)")

        model = cls(name)
        model._init_dirs(input_name=feature_set.name)

        # Stage the training data: features plus the sample_weight/validation/exclude roles
        train_df = feature_set.training_view(
            sample_weights=sample_weights, validation_ids=validation_ids, exclude_ids=exclude_ids
        )
        train_df.to_csv(os.path.join(model.train_dir, "train.csv"), index=False)
        model.log.important(f"Staged {len(train_df)} training rows for {name}...")

        # Derive the feature list the same way the AWS path does when it isn't given
        target_list = [target_column] if isinstance(target_column, str) else (target_column or [])
        if feature_list is None:
            feature_list = model._derive_feature_list(feature_set, target_list)
        model.log.important(f"Feature List for Modeling: {feature_list}")

        # Generate the script, then keep it with the model so a training run is reproducible
        target_for_template = target_list if model_framework == ModelFramework.CHEMPROP else target_column
        template_params = {
            "model_imports": model_import_str,
            "model_type": model_type,
            "model_framework": model_framework,
            "model_class": model_class,
            "target_column": target_for_template,
            "feature_list": feature_list,
            "compressed_features": [],
            "model_metrics_path": model.output_dir,
            "id_column": feature_set.id_column,
            "hyperparameters": hyperparameters or {},
        }
        script_path = model._build_script(template_params, custom_script)
        shutil.copytree(os.path.dirname(script_path), model.script_dir, dirs_exist_ok=True)

        model.upsert_workbench_meta(
            {
                "workbench_model_features": feature_list,
                "workbench_model_target": target_column,
                "model_type": model_type.value,
                "model_framework": model_framework.value,
                "id_column": feature_set.id_column,
                "hyperparameters": hyperparameters or {},
                # All three row roles are kept, so publish() trains in AWS on the same
                # rows with the same weights. Dropping any of them here would make the
                # published model quietly differ from the local one.
                "sample_weights": cls._weights_as_pairs(sample_weights),
                "validation_ids": list(validation_ids) if validation_ids else None,
                "exclude_ids": list(exclude_ids) if exclude_ids else None,
            }
        )

        model._launch_training(wait=wait)
        return model

    @staticmethod
    def _weights_as_pairs(sample_weights: Union[dict, pd.DataFrame, None]) -> Union[list, None]:
        """Internal: Normalize sample weights to JSON-storable [id, weight] pairs.

        Stored as pairs rather than a mapping because JSON object keys are always
        strings: integer ids would come back as strings and fail to join against the
        FeatureSet's id column, silently dropping the weights.

        Args:
            sample_weights (Union[dict, pd.DataFrame, None]): Weights as given by the caller

        Returns:
            Union[list, None]: [[id, weight], ...], or None when there are no weights
        """
        if sample_weights is None:
            return None
        if isinstance(sample_weights, pd.DataFrame):
            if sample_weights.empty:
                return None
            id_column = sample_weights.columns[0]
            sample_weights = dict(zip(sample_weights[id_column], sample_weights["sample_weight"]))
        pairs = [[key, float(value)] for key, value in sample_weights.items()]
        return pairs or None

    def _build_script(self, template_params: dict, custom_script: str = None) -> str:
        """Internal: Produce the model script, from a built-in template or a custom one.

        Args:
            template_params (dict): Parameters filled into the template
            custom_script (str, optional): Path to a custom script or .template

        Returns:
            str: Path to the script to run
        """
        if not custom_script:
            return generate_model_script(template_params)

        if not str(custom_script).endswith(".template"):
            return str(custom_script)

        # A custom template gets the same params, with the enum flattened like the generator does
        template_params = {**template_params, "model_type": template_params["model_type"].value}
        return fill_template(custom_script, template_params, "generated_model_script.py")

    def _derive_feature_list(self, feature_set, target_list: list) -> list:
        """Internal: Guess a feature list from the FeatureSet's numeric columns.

        Args:
            feature_set (LocalFeatureSet): The feature set being trained on
            target_list (list): Target column(s) to exclude

        Returns:
            list: The derived feature list
        """
        self.log.warning("Guessing at the feature list, HIGHLY RECOMMENDED to specify an explicit feature list!")
        skip = {"id", "auto_id", "__index_level_0__", "event_time", "training", feature_set.id_column} | set(
            target_list
        )
        df = feature_set.pull_dataframe(limit=1)
        return [c for c in df.columns if c not in skip and pd.api.types.is_numeric_dtype(df[c])]

    def _init_dirs(self, input_name: str):
        """Internal: Create the model's directory layout

        Args:
            input_name (str): Name of this model's input FeatureSet
        """
        storage.local_root(create=True)
        os.makedirs(self.path, exist_ok=True)

        # Start each run from empty dirs. Leftovers from a previous run would otherwise
        # survive a failure and be served as this run's artifacts and predictions.
        for directory in (self.model_dir, self.output_dir, self.train_dir):
            shutil.rmtree(directory, ignore_errors=True)
            os.makedirs(directory, exist_ok=True)
        shutil.rmtree(self.script_dir, ignore_errors=True)

        self._init_storage(input_name=input_name)

    def _launch_training(self, wait: bool):
        """Internal: Run the generated model script as a subprocess.

        Args:
            wait (bool): Block until training finishes, streaming output
        """
        script = os.path.join(self.script_dir, "generated_model_script.py")
        command = [
            sys.executable,
            script,
            "--model-dir",
            self.model_dir,
            "--train",
            self.train_dir,
            "--output-data-dir",
            self.output_dir,
        ]
        self.log.important(f"Training {self.name} locally: {' '.join(command)}")
        self.upsert_workbench_meta({"workbench_status": "training"})

        if not wait:
            # The child dups the descriptor, so ours is closed as soon as it's launched
            with open(self.log_path, "w") as log_file:
                proc = subprocess.Popen(
                    command, cwd=self.script_dir, stdout=log_file, stderr=subprocess.STDOUT, text=True
                )
            self._write_status(state="training", pid=proc.pid)
            job_tracker.watch_subprocess(
                self.name,
                proc,
                kind="Local training",
                log_path=self.log_path,
                on_finish=self._record_outcome,
            )
            return

        proc = subprocess.Popen(
            command,
            cwd=self.script_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        )
        self._write_status(state="training", pid=proc.pid)

        # Stream the child's output so a long training run isn't a silent block
        with open(self.log_path, "w") as log_file:
            for line in proc.stdout:
                print(line, end="")
                log_file.write(line)
        proc.wait()

        self._finish_training(proc.returncode)

    def _record_outcome(self, returncode: int) -> bool:
        """Internal: Write the durable outcome of a training run.

        Called by the wait=True path directly and by the job watcher for wait=False,
        so a detached run leaves the same on-disk state as a blocking one.

        Args:
            returncode (int): The model script's exit code

        Returns:
            bool: True if training succeeded
        """
        success = returncode == 0
        if success:
            self._bundle_for_inference()
        self.refresh_meta()
        self._write_status(state="completed" if success else "failed", returncode=returncode)
        self.upsert_workbench_meta({"workbench_status": "ready" if success else "failed"})
        return success

    def _bundle_for_inference(self):
        """Internal: Put the inference code and metadata in the model directory.

        SageMaker's entry point is ``training_harness.py``, which runs the model script
        and then bundles the code so the inference container knows what to serve. Local
        runs the model script directly and calls the same bundling step here.

        The harness's other job -- pip installing the script's requirements.txt -- is
        deliberately skipped: local dependencies are the user's environment, not
        something a training run should mutate.
        """
        from workbench.training.training_harness import include_code_and_meta_for_inference

        include_code_and_meta_for_inference(
            model_dir=self.model_dir,
            code_dir=self.script_dir,
            entry_point="generated_model_script.py",
        )

    def _finish_training(self, returncode: int):
        """Internal: Record the outcome of a blocking training run, raising on failure

        Args:
            returncode (int): The model script's exit code

        Raises:
            RuntimeError: If the model script exited non-zero
        """
        if not self._record_outcome(returncode):
            tail = self.training_log(lines=15)
            self.log.error(f"Local training failed for {self.name} (exit {returncode}):\n{tail}")
            raise RuntimeError(f"Local training failed for {self.name} (exit {returncode}), see {self.log_path}")
        self.log.important(f"Local training complete: {self.model_dir}")

    def _write_status(self, state: str, pid: int = None, returncode: int = None):
        """Internal: Write the training status file used to reattach across sessions

        Args:
            state (str): One of training/completed/failed
            pid (int, optional): The training process id
            returncode (int, optional): The training process exit code
        """
        status = self.training_state()
        status["state"] = state
        status["updated"] = datetime.now(timezone.utc).isoformat()
        if pid is not None:
            status["pid"] = pid
            status["started"] = status["updated"]
        if returncode is not None:
            status["returncode"] = returncode
            status["finished"] = status["updated"]
        with open(self.status_path, "w") as fp:
            json.dump(status, fp, indent=4)

    def training_state(self) -> dict:
        """The training status for this model, as recorded on disk.

        A run whose watcher never got to record an outcome -- the session exited, or the
        process died -- is reported as "interrupted" rather than left claiming to be
        training forever. Whether the child finished its work is unknown at that point.

        Returns:
            dict: {state, pid, started, updated, returncode, finished}, empty before any run
        """
        try:
            with open(self.status_path, "r") as fp:
                status = json.load(fp)
        except (FileNotFoundError, json.JSONDecodeError):
            return {}

        if status.get("state") == "training" and not self._pid_alive(status.get("pid")):
            status["state"] = "interrupted"
        return status

    @staticmethod
    def _pid_alive(pid: int) -> bool:
        """Internal: Is a process still running?

        Args:
            pid (int): The process id to check

        Returns:
            bool: True if the process exists
        """
        if not pid:
            return False
        try:
            os.kill(pid, 0)
        except (OSError, ProcessLookupError):
            return False
        return True

    def training_log(self, lines: int = None) -> str:
        """The training log for this model.

        Args:
            lines (int, optional): Return only the last N lines (default: the whole log)

        Returns:
            str: The log contents ("" if the model hasn't trained)
        """
        try:
            with open(self.log_path, "r") as fp:
                content = fp.readlines()
        except OSError:
            return ""
        return "".join(content[-lines:] if lines else content)

    def oof_predictions(self) -> pd.DataFrame:
        """Out-of-fold predictions written by the training run.

        Returns:
            pd.DataFrame: The OOF predictions (empty if the model hasn't trained)
        """
        return self._output_csv("oof_predictions.csv")

    def validation_predictions(self) -> pd.DataFrame:
        """Held-out validation predictions written by the training run.

        Returns:
            pd.DataFrame: The validation predictions (empty if there was no validation set)
        """
        return self._output_csv("val_predictions.csv")

    def _output_csv(self, file_name: str) -> pd.DataFrame:
        """Internal: Read a CSV the training run wrote to the output directory

        Args:
            file_name (str): The file name within the output directory

        Returns:
            pd.DataFrame: The file's contents (empty if missing)
        """
        path = os.path.join(self.output_dir, file_name)
        return pd.read_csv(path) if os.path.isfile(path) else pd.DataFrame()

    def list_inference_runs(self) -> list[str]:
        """List the inference runs for this model.

        Returns:
            list[str]: The cross-fold run from training, then any endpoint captures
        """
        runs = [CROSS_FOLD_RUN] if not self.oof_predictions().empty else []
        for endpoint in self._endpoints():
            runs.extend(name for name in endpoint.list_captures() if name not in runs)
        return runs

    def default_inference_run(self) -> Union[str, None]:
        """Resolve the default inference run for this model.

        Returns:
            Union[str, None]: full_cross_fold -> test_inference -> first run, None if there are none
        """
        from workbench.utils.metrics_utils import default_inference_run

        return default_inference_run(self.list_inference_runs())

    def get_inference_predictions(self, capture_name: str = "default") -> Union[pd.DataFrame, None]:
        """Retrieve the captured predictions for this model.

        Args:
            capture_name (str, optional): A run from list_inference_runs(), or "default"
                to resolve via default_inference_run()

        Returns:
            Union[pd.DataFrame, None]: The predictions, or None if that run doesn't exist
        """
        if capture_name == "default":
            capture_name = self.default_inference_run()
            if capture_name is None:
                self.log.warning(f"No inference runs for {self.name}...")
                return None

        if capture_name == CROSS_FOLD_RUN:
            predictions = self.oof_predictions()
            return predictions if not predictions.empty else None

        for endpoint in self._endpoints():
            if capture_name in endpoint.list_captures():
                return endpoint.get_inference_predictions(capture_name)
        self.log.warning(f"No inference run '{capture_name}' for {self.name}...")
        return None

    def get_inference_metrics(self, capture_name: str = "default") -> Union[pd.DataFrame, None]:
        """Retrieve the inference performance metrics for this model.

        Computed from the run's predictions rather than stored, so there is nothing
        to keep in sync with them.

        Args:
            capture_name (str, optional): A run from list_inference_runs(), or "default"
                to resolve via default_inference_run()

        Returns:
            Union[pd.DataFrame, None]: The metrics, or None if they can't be computed
        """
        from workbench.utils.metrics_utils import compute_metrics_from_predictions, resolve_primary_target

        predictions = self.get_inference_predictions(capture_name)
        if predictions is None:
            return None

        target = resolve_primary_target(self.workbench_meta().get("workbench_model_target"))
        if target is None or target not in predictions.columns:
            self.log.warning(f"No target column in the '{capture_name}' predictions for {self.name}")
            return None

        # Multi-task targets are sparse, so rows the primary target doesn't cover are dropped
        predictions = predictions.dropna(subset=[target])

        class_labels = None
        if self.workbench_meta().get("model_type") == ModelType.CLASSIFIER.value:
            class_labels = sorted(predictions[target].unique().tolist())
        return compute_metrics_from_predictions(predictions, target, class_labels)

    def to_endpoint(self, name: str = None) -> "LocalEndpoint":  # noqa: F821
        """Create a LocalEndpoint that serves this model.

        Args:
            name (str, optional): Endpoint name (defaults to the model name)

        Returns:
            LocalEndpoint: The endpoint serving this model
        """
        from workbench.local.local_endpoint import LocalEndpoint

        return LocalEndpoint.from_model(self, name=name)

    def parent(self):
        """The LocalFeatureSet this model trained on, if it still exists locally"""
        from workbench.local.local_feature_set import LocalFeatureSet

        feature_set = LocalFeatureSet(self.get_input())
        return feature_set if feature_set.exists() else None

    def aws_exists(self) -> bool:
        """Does an AWS Model by this name already exist?

        Returns:
            bool: True if AWS already has this Model
        """
        from workbench.api import Model

        return Model(self.name).exists()

    def _aws_artifact(self):
        """Internal: The AWS Model for this local one"""
        from workbench.api import Model

        return Model(self.name)

    def version_drift(self) -> str:
        """Package versions that differ between this machine and the training image.

        Returns:
            str: A drift report, or "" when everything that matters matches
        """
        from workbench.utils.version_drift import drift_summary

        return drift_summary(self.workbench_meta().get("model_framework", "xgboost"))

    def _publish_self(self, **kwargs):
        """Internal: Train this model in AWS from the published FeatureSet.

        Publishing retrains rather than uploading local artifacts, so the model lands in
        the registry the same way any AWS model does. The row roles recorded at local
        training time are replayed, so AWS trains on the same rows.

        Returns:
            Model: The created AWS Model
        """
        from workbench.api import FeatureSet

        meta = self.workbench_meta()
        feature_set = FeatureSet(self.get_input())
        return feature_set.to_model(
            name=self.name,
            model_type=ModelType(meta["model_type"]),
            model_framework=ModelFramework(meta["model_framework"]),
            target_column=meta.get("workbench_model_target"),
            feature_list=meta.get("workbench_model_features"),
            hyperparameters=meta.get("hyperparameters") or {},
            sample_weights=dict(meta["sample_weights"]) if meta.get("sample_weights") else None,
            validation_ids=meta.get("validation_ids"),
            exclude_ids=meta.get("exclude_ids"),
            **kwargs,
        )

    def publish(self, endpoint: bool = True, **kwargs: Any) -> "Model":  # noqa: F821
        """Publish this model and its lineage to AWS, then deploy an endpoint.

        Args:
            endpoint (bool): Also deploy a serverless endpoint (default True)
            **kwargs: Passed to the AWS training job

        Returns:
            Model: The published AWS Model
        """
        from workbench.api import Endpoint

        aws_model = super().publish(**kwargs)
        if not endpoint:
            return aws_model

        endpoint_name = self._endpoint_name()
        if Endpoint(endpoint_name).exists():
            self.log.important(f"AWS endpoint '{endpoint_name}' already exists, skipping...")
        else:
            self.log.important(f"Deploying endpoint '{endpoint_name}' to AWS...")
            aws_model.to_endpoint(name=endpoint_name)
        return aws_model

    def _endpoints(self) -> list:
        """Internal: The local endpoints serving this model"""
        from workbench.local import storage
        from workbench.local.local_endpoint import LocalEndpoint

        serving = [LocalEndpoint(name) for name in storage.list_artifacts("endpoint")]
        return [endpoint for endpoint in serving if endpoint.model_name == self.name]

    def _endpoint_name(self) -> str:
        """Internal: The name to deploy this model's AWS endpoint under.

        A local endpoint may carry a custom name, so publishing reuses it. Otherwise
        this is the same default AWS uses: the model name.

        Returns:
            str: The endpoint name
        """
        serving = self._endpoints()
        return serving[0].name if serving else self.name

    def delete(self):
        """Delete this model and the endpoints serving it.

        An endpoint is not an independent artifact -- it loads from the model's
        directory and is meaningless once that is gone, so it comes down too. This is
        the only cascade: deleting a FeatureSet leaves its models alone.
        """
        for endpoint in self._endpoints():
            endpoint.delete()
        super().delete()

    def details(self, **kwargs) -> dict:
        """LocalModel Details

        Returns:
            dict: A dictionary of details about the LocalModel
        """
        return {**super().details(), "training": self.training_state()}

__init__(name, **kwargs)

Initialize a LocalModel

Parameters:

Name	Type	Description	Default
name	str	The name of the model	required

Source code in src/workbench/local/local_model.py

def __init__(self, name: str, **kwargs):
    """Initialize a LocalModel

    Args:
        name (str): The name of the model
    """
    Artifact.is_name_valid(name, delimiter="-", lower_case=False)
    super().__init__(name, **kwargs)
    self.script_dir = os.path.join(self.path, "script")
    self.model_dir = os.path.join(self.path, "model_artifacts")
    self.output_dir = os.path.join(self.path, "output")
    self.train_dir = os.path.join(self.path, "input", "train")
    self.log_path = os.path.join(self.output_dir, "training.log")
    self.status_path = os.path.join(self.path, "status.json")

aws_exists()

Does an AWS Model by this name already exist?

Returns:

Name	Type	Description
bool	bool	True if AWS already has this Model

Source code in src/workbench/local/local_model.py

def aws_exists(self) -> bool:
    """Does an AWS Model by this name already exist?

    Returns:
        bool: True if AWS already has this Model
    """
    from workbench.api import Model

    return Model(self.name).exists()

default_inference_run()

Resolve the default inference run for this model.

Returns:

Type	Description
Union[str, None]	Union[str, None]: full_cross_fold -> test_inference -> first run, None if there are none

Source code in src/workbench/local/local_model.py

def default_inference_run(self) -> Union[str, None]:
    """Resolve the default inference run for this model.

    Returns:
        Union[str, None]: full_cross_fold -> test_inference -> first run, None if there are none
    """
    from workbench.utils.metrics_utils import default_inference_run

    return default_inference_run(self.list_inference_runs())

delete()

Delete this model and the endpoints serving it.

An endpoint is not an independent artifact -- it loads from the model's directory and is meaningless once that is gone, so it comes down too. This is the only cascade: deleting a FeatureSet leaves its models alone.

Source code in src/workbench/local/local_model.py

def delete(self):
    """Delete this model and the endpoints serving it.

    An endpoint is not an independent artifact -- it loads from the model's
    directory and is meaningless once that is gone, so it comes down too. This is
    the only cascade: deleting a FeatureSet leaves its models alone.
    """
    for endpoint in self._endpoints():
        endpoint.delete()
    super().delete()

details(**kwargs)

LocalModel Details

Returns:

Name	Type	Description
dict	dict	A dictionary of details about the LocalModel

Source code in src/workbench/local/local_model.py

def details(self, **kwargs) -> dict:
    """LocalModel Details

    Returns:
        dict: A dictionary of details about the LocalModel
    """
    return {**super().details(), "training": self.training_state()}

from_feature_set(feature_set, name, model_type, model_framework, target_column=None, feature_list=None, model_class=None, model_import_str=None, custom_script=None, hyperparameters=None, sample_weights=None, validation_ids=None, exclude_ids=None, wait=True) classmethod

Train a LocalModel from a LocalFeatureSet.

Parameters:

Name	Type	Description	Default
feature_set	LocalFeatureSet	The feature set to train on	required
name	str	The name of the model to create	required
model_type	ModelType	The type of model to create	required
model_framework	ModelFramework	The framework to use	required
target_column	str or list[str]	Target column(s), None for unsupervised	None
feature_list	list	Feature columns; derived from the FeatureSet if omitted	None
model_class	str	Model class for scikit-learn models (e.g. "KMeans")	None
model_import_str	str	Import line for the model class	None
custom_script	str	Path to a custom model script or template	None
hyperparameters	dict	Hyperparameters for the model	None
sample_weights	Union[dict, DataFrame]	id -> framework weight	None
validation_ids	list	ids held out and scored as a validation set	None
exclude_ids	list	ids dropped from training entirely	None
wait	bool	Block until training finishes (default: True)	True

Returns:

Name	Type	Description
LocalModel	LocalModel	The model (trained if wait=True, still training otherwise)

Source code in src/workbench/local/local_model.py

@classmethod
def from_feature_set(
    cls,
    feature_set,
    name: str,
    model_type: ModelType,
    model_framework: ModelFramework,
    target_column: Union[str, list[str]] = None,
    feature_list: list = None,
    model_class: str = None,
    model_import_str: str = None,
    custom_script: str = None,
    hyperparameters: dict = None,
    sample_weights: Union[dict, pd.DataFrame] = None,
    validation_ids: list = None,
    exclude_ids: list = None,
    wait: bool = True,
) -> "LocalModel":
    """Train a LocalModel from a LocalFeatureSet.

    Args:
        feature_set (LocalFeatureSet): The feature set to train on
        name (str): The name of the model to create
        model_type (ModelType): The type of model to create
        model_framework (ModelFramework): The framework to use
        target_column (str or list[str], optional): Target column(s), None for unsupervised
        feature_list (list, optional): Feature columns; derived from the FeatureSet if omitted
        model_class (str, optional): Model class for scikit-learn models (e.g. "KMeans")
        model_import_str (str, optional): Import line for the model class
        custom_script (str, optional): Path to a custom model script or template
        hyperparameters (dict, optional): Hyperparameters for the model
        sample_weights (Union[dict, pd.DataFrame], optional): id -> framework weight
        validation_ids (list, optional): ids held out and scored as a validation set
        exclude_ids (list, optional): ids dropped from training entirely
        wait (bool): Block until training finishes (default: True)

    Returns:
        LocalModel: The model (trained if wait=True, still training otherwise)
    """
    supervised = model_type in (
        ModelType.CLASSIFIER,
        ModelType.REGRESSOR,
        ModelType.UQ_REGRESSOR,
        ModelType.ENSEMBLE_REGRESSOR,
    )
    if target_column is None and supervised:
        raise ValueError("target_column is required for supervised models (pass target_column=...)")

    model = cls(name)
    model._init_dirs(input_name=feature_set.name)

    # Stage the training data: features plus the sample_weight/validation/exclude roles
    train_df = feature_set.training_view(
        sample_weights=sample_weights, validation_ids=validation_ids, exclude_ids=exclude_ids
    )
    train_df.to_csv(os.path.join(model.train_dir, "train.csv"), index=False)
    model.log.important(f"Staged {len(train_df)} training rows for {name}...")

    # Derive the feature list the same way the AWS path does when it isn't given
    target_list = [target_column] if isinstance(target_column, str) else (target_column or [])
    if feature_list is None:
        feature_list = model._derive_feature_list(feature_set, target_list)
    model.log.important(f"Feature List for Modeling: {feature_list}")

    # Generate the script, then keep it with the model so a training run is reproducible
    target_for_template = target_list if model_framework == ModelFramework.CHEMPROP else target_column
    template_params = {
        "model_imports": model_import_str,
        "model_type": model_type,
        "model_framework": model_framework,
        "model_class": model_class,
        "target_column": target_for_template,
        "feature_list": feature_list,
        "compressed_features": [],
        "model_metrics_path": model.output_dir,
        "id_column": feature_set.id_column,
        "hyperparameters": hyperparameters or {},
    }
    script_path = model._build_script(template_params, custom_script)
    shutil.copytree(os.path.dirname(script_path), model.script_dir, dirs_exist_ok=True)

    model.upsert_workbench_meta(
        {
            "workbench_model_features": feature_list,
            "workbench_model_target": target_column,
            "model_type": model_type.value,
            "model_framework": model_framework.value,
            "id_column": feature_set.id_column,
            "hyperparameters": hyperparameters or {},
            # All three row roles are kept, so publish() trains in AWS on the same
            # rows with the same weights. Dropping any of them here would make the
            # published model quietly differ from the local one.
            "sample_weights": cls._weights_as_pairs(sample_weights),
            "validation_ids": list(validation_ids) if validation_ids else None,
            "exclude_ids": list(exclude_ids) if exclude_ids else None,
        }
    )

    model._launch_training(wait=wait)
    return model

get_inference_metrics(capture_name='default')

Retrieve the inference performance metrics for this model.

Computed from the run's predictions rather than stored, so there is nothing to keep in sync with them.

Parameters:

Name	Type	Description	Default
capture_name	str	A run from list_inference_runs(), or "default" to resolve via default_inference_run()	'default'

Returns:

Type	Description
Union[DataFrame, None]	Union[pd.DataFrame, None]: The metrics, or None if they can't be computed

Source code in src/workbench/local/local_model.py

def get_inference_metrics(self, capture_name: str = "default") -> Union[pd.DataFrame, None]:
    """Retrieve the inference performance metrics for this model.

    Computed from the run's predictions rather than stored, so there is nothing
    to keep in sync with them.

    Args:
        capture_name (str, optional): A run from list_inference_runs(), or "default"
            to resolve via default_inference_run()

    Returns:
        Union[pd.DataFrame, None]: The metrics, or None if they can't be computed
    """
    from workbench.utils.metrics_utils import compute_metrics_from_predictions, resolve_primary_target

    predictions = self.get_inference_predictions(capture_name)
    if predictions is None:
        return None

    target = resolve_primary_target(self.workbench_meta().get("workbench_model_target"))
    if target is None or target not in predictions.columns:
        self.log.warning(f"No target column in the '{capture_name}' predictions for {self.name}")
        return None

    # Multi-task targets are sparse, so rows the primary target doesn't cover are dropped
    predictions = predictions.dropna(subset=[target])

    class_labels = None
    if self.workbench_meta().get("model_type") == ModelType.CLASSIFIER.value:
        class_labels = sorted(predictions[target].unique().tolist())
    return compute_metrics_from_predictions(predictions, target, class_labels)

get_inference_predictions(capture_name='default')

Retrieve the captured predictions for this model.

Parameters:

Name	Type	Description	Default
capture_name	str	A run from list_inference_runs(), or "default" to resolve via default_inference_run()	'default'

Returns:

Type	Description
Union[DataFrame, None]	Union[pd.DataFrame, None]: The predictions, or None if that run doesn't exist

Source code in src/workbench/local/local_model.py

def get_inference_predictions(self, capture_name: str = "default") -> Union[pd.DataFrame, None]:
    """Retrieve the captured predictions for this model.

    Args:
        capture_name (str, optional): A run from list_inference_runs(), or "default"
            to resolve via default_inference_run()

    Returns:
        Union[pd.DataFrame, None]: The predictions, or None if that run doesn't exist
    """
    if capture_name == "default":
        capture_name = self.default_inference_run()
        if capture_name is None:
            self.log.warning(f"No inference runs for {self.name}...")
            return None

    if capture_name == CROSS_FOLD_RUN:
        predictions = self.oof_predictions()
        return predictions if not predictions.empty else None

    for endpoint in self._endpoints():
        if capture_name in endpoint.list_captures():
            return endpoint.get_inference_predictions(capture_name)
    self.log.warning(f"No inference run '{capture_name}' for {self.name}...")
    return None

list_inference_runs()

List the inference runs for this model.

Returns:

Type	Description
list[str]	list[str]: The cross-fold run from training, then any endpoint captures

Source code in src/workbench/local/local_model.py

def list_inference_runs(self) -> list[str]:
    """List the inference runs for this model.

    Returns:
        list[str]: The cross-fold run from training, then any endpoint captures
    """
    runs = [CROSS_FOLD_RUN] if not self.oof_predictions().empty else []
    for endpoint in self._endpoints():
        runs.extend(name for name in endpoint.list_captures() if name not in runs)
    return runs

oof_predictions()

Out-of-fold predictions written by the training run.

Returns:

Type	Description
DataFrame	pd.DataFrame: The OOF predictions (empty if the model hasn't trained)

Source code in src/workbench/local/local_model.py

def oof_predictions(self) -> pd.DataFrame:
    """Out-of-fold predictions written by the training run.

    Returns:
        pd.DataFrame: The OOF predictions (empty if the model hasn't trained)
    """
    return self._output_csv("oof_predictions.csv")

parent()

The LocalFeatureSet this model trained on, if it still exists locally

Source code in src/workbench/local/local_model.py

def parent(self):
    """The LocalFeatureSet this model trained on, if it still exists locally"""
    from workbench.local.local_feature_set import LocalFeatureSet

    feature_set = LocalFeatureSet(self.get_input())
    return feature_set if feature_set.exists() else None

publish(endpoint=True, **kwargs)

Publish this model and its lineage to AWS, then deploy an endpoint.

Parameters:

Name	Type	Description	Default
endpoint	bool	Also deploy a serverless endpoint (default True)	True
**kwargs	Any	Passed to the AWS training job	{}

Returns:

Name	Type	Description
Model	Model	The published AWS Model

Source code in src/workbench/local/local_model.py

def publish(self, endpoint: bool = True, **kwargs: Any) -> "Model":  # noqa: F821
    """Publish this model and its lineage to AWS, then deploy an endpoint.

    Args:
        endpoint (bool): Also deploy a serverless endpoint (default True)
        **kwargs: Passed to the AWS training job

    Returns:
        Model: The published AWS Model
    """
    from workbench.api import Endpoint

    aws_model = super().publish(**kwargs)
    if not endpoint:
        return aws_model

    endpoint_name = self._endpoint_name()
    if Endpoint(endpoint_name).exists():
        self.log.important(f"AWS endpoint '{endpoint_name}' already exists, skipping...")
    else:
        self.log.important(f"Deploying endpoint '{endpoint_name}' to AWS...")
        aws_model.to_endpoint(name=endpoint_name)
    return aws_model

to_endpoint(name=None)

Create a LocalEndpoint that serves this model.

Parameters:

Name	Type	Description	Default
name	str	Endpoint name (defaults to the model name)	None

Returns:

Name	Type	Description
LocalEndpoint	LocalEndpoint	The endpoint serving this model

Source code in src/workbench/local/local_model.py

def to_endpoint(self, name: str = None) -> "LocalEndpoint":  # noqa: F821
    """Create a LocalEndpoint that serves this model.

    Args:
        name (str, optional): Endpoint name (defaults to the model name)

    Returns:
        LocalEndpoint: The endpoint serving this model
    """
    from workbench.local.local_endpoint import LocalEndpoint

    return LocalEndpoint.from_model(self, name=name)

training_log(lines=None)

The training log for this model.

Parameters:

Name	Type	Description	Default
lines	int	Return only the last N lines (default: the whole log)	None

Returns:

Name	Type	Description
str	str	The log contents ("" if the model hasn't trained)

Source code in src/workbench/local/local_model.py

def training_log(self, lines: int = None) -> str:
    """The training log for this model.

    Args:
        lines (int, optional): Return only the last N lines (default: the whole log)

    Returns:
        str: The log contents ("" if the model hasn't trained)
    """
    try:
        with open(self.log_path, "r") as fp:
            content = fp.readlines()
    except OSError:
        return ""
    return "".join(content[-lines:] if lines else content)

training_state()

The training status for this model, as recorded on disk.

A run whose watcher never got to record an outcome -- the session exited, or the process died -- is reported as "interrupted" rather than left claiming to be training forever. Whether the child finished its work is unknown at that point.

Returns:

Name	Type	Description
dict	dict	{state, pid, started, updated, returncode, finished}, empty before any run

Source code in src/workbench/local/local_model.py

def training_state(self) -> dict:
    """The training status for this model, as recorded on disk.

    A run whose watcher never got to record an outcome -- the session exited, or the
    process died -- is reported as "interrupted" rather than left claiming to be
    training forever. Whether the child finished its work is unknown at that point.

    Returns:
        dict: {state, pid, started, updated, returncode, finished}, empty before any run
    """
    try:
        with open(self.status_path, "r") as fp:
            status = json.load(fp)
    except (FileNotFoundError, json.JSONDecodeError):
        return {}

    if status.get("state") == "training" and not self._pid_alive(status.get("pid")):
        status["state"] = "interrupted"
    return status

validation_predictions()

Held-out validation predictions written by the training run.

Returns:

Type	Description
DataFrame	pd.DataFrame: The validation predictions (empty if there was no validation set)

Source code in src/workbench/local/local_model.py

def validation_predictions(self) -> pd.DataFrame:
    """Held-out validation predictions written by the training run.

    Returns:
        pd.DataFrame: The validation predictions (empty if there was no validation set)
    """
    return self._output_csv("val_predictions.csv")

version_drift()

Package versions that differ between this machine and the training image.

Returns:

Name	Type	Description
str	str	A drift report, or "" when everything that matters matches

Source code in src/workbench/local/local_model.py

def version_drift(self) -> str:
    """Package versions that differ between this machine and the training image.

    Returns:
        str: A drift report, or "" when everything that matters matches
    """
    from workbench.utils.version_drift import drift_summary

    return drift_summary(self.workbench_meta().get("model_framework", "xgboost"))

LocalEndpoint: In-process inference against a locally trained model.

LocalEndpoint

Bases: LocalArtifact

LocalEndpoint: Workbench Local Endpoint Class

Loads the model bundle once with model_fn and calls predict_fn per inference, which are the same functions the serving container calls.

Common Usage

my_endpoint = LocalEndpoint("my-model")
results = my_endpoint.inference(eval_df)

Source code in src/workbench/local/local_endpoint.py

class LocalEndpoint(LocalArtifact):
    """LocalEndpoint: Workbench Local Endpoint Class

    Loads the model bundle once with `model_fn` and calls `predict_fn` per
    inference, which are the same functions the serving container calls.

    Common Usage:
        ```python
        my_endpoint = LocalEndpoint("my-model")
        results = my_endpoint.inference(eval_df)
        ```
    """

    artifact_type = "endpoint"

    def __init__(self, name: str, **kwargs):
        """Initialize a LocalEndpoint

        Args:
            name (str): The name of the endpoint
        """
        Artifact.is_name_valid(name, delimiter="-", lower_case=False)
        super().__init__(name, **kwargs)
        self.inference_dir = os.path.join(self.path, "inference")
        self._model_bundle = None
        self._inference_module = None

    @classmethod
    def from_model(cls, model: LocalModel, name: str = None) -> "LocalEndpoint":
        """Create a LocalEndpoint that serves a LocalModel.

        Args:
            model (LocalModel): The trained model to serve
            name (str, optional): Endpoint name (defaults to the model name)

        Returns:
            LocalEndpoint: The created endpoint

        Raises:
            ValueError: If the model hasn't trained successfully
        """
        if model.training_state().get("state") != "completed":
            raise ValueError(f"Model {model.name} has not trained successfully, cannot serve it")

        endpoint = cls(name or model.name)
        storage.local_root(create=True)
        os.makedirs(endpoint.inference_dir, exist_ok=True)
        endpoint._init_storage(input_name=model.name)
        endpoint.upsert_workbench_meta({"model_name": model.name})
        endpoint.log.important(f"Local endpoint {endpoint.name} serving {model.name}")
        return endpoint

    @property
    def model_name(self) -> str:
        """The model this endpoint serves"""
        return self.workbench_meta().get("model_name")

    @property
    def model_dir(self) -> str:
        """The model artifacts directory this endpoint loads from.

        Resolved from the model name on every access rather than stored, so the
        endpoint keeps working when the storage root moves or the config changes.
        """
        return LocalModel(self.model_name).model_dir

    def inference(self, eval_df: pd.DataFrame, capture_name: str = None) -> pd.DataFrame:
        """Run inference on a DataFrame.

        Args:
            eval_df (pd.DataFrame): The data to run inference on
            capture_name (str, optional): Store the predictions under this name, which
                makes them available from the model as an inference run

        Returns:
            pd.DataFrame: The predictions
        """
        module = self._load_inference_module()
        bundle = self._load_model_bundle()

        # predict_fn resolves its own supporting files through SM_MODEL_DIR
        previous = os.environ.get("SM_MODEL_DIR")
        os.environ["SM_MODEL_DIR"] = self.model_dir
        try:
            predictions = module.predict_fn(eval_df.copy(), bundle)
        finally:
            if previous is None:
                os.environ.pop("SM_MODEL_DIR", None)
            else:
                os.environ["SM_MODEL_DIR"] = previous

        if capture_name:
            self._capture(predictions, capture_name)
        return predictions

    def _capture(self, predictions: pd.DataFrame, capture_name: str):
        """Internal: Store an inference run under a capture name

        Args:
            predictions (pd.DataFrame): The predictions to store
            capture_name (str): The name to store them under
        """
        capture_dir = os.path.join(self.inference_dir, capture_name)
        os.makedirs(capture_dir, exist_ok=True)
        predictions.to_parquet(os.path.join(capture_dir, "predictions.parquet"), index=False)
        self.log.important(f"Captured {len(predictions)} predictions as '{capture_name}'")

    def get_inference_predictions(self, capture_name: str = "auto_inference") -> Union[pd.DataFrame, None]:
        """Retrieve a captured inference run.

        Args:
            capture_name (str): The capture to retrieve (default: "auto_inference")

        Returns:
            Union[pd.DataFrame, None]: The predictions, or None if that capture doesn't exist
        """
        path = os.path.join(self.inference_dir, capture_name, "predictions.parquet")
        if not os.path.isfile(path):
            self.log.warning(f"No inference capture named '{capture_name}' for {self.name}")
            return None
        return pd.read_parquet(path)

    def list_captures(self) -> list[str]:
        """The inference captures stored on this endpoint.

        Returns:
            list[str]: Sorted capture names
        """
        if not os.path.isdir(self.inference_dir):
            return []
        return sorted(d.name for d in os.scandir(self.inference_dir) if d.is_dir())

    def _load_inference_module(self):
        """Internal: Import the model's inference script, the way the container does.

        Returns:
            module: The imported inference module
        """
        if self._inference_module is not None:
            return self._inference_module

        metadata_path = os.path.join(self.model_dir, "inference-metadata.json")
        if not os.path.isfile(metadata_path):
            raise FileNotFoundError(
                f"No model artifacts for '{self.model_name}' at {self.model_dir}. "
                "The model was deleted or never finished training."
            )
        with open(metadata_path) as fp:
            script = json.load(fp)["inference_script"]

        spec = importlib.util.spec_from_file_location("workbench_local_inference", os.path.join(self.model_dir, script))
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        self._inference_module = module
        return module

    def _load_model_bundle(self) -> dict:
        """Internal: Load the model once via model_fn, then keep it.

        Returns:
            dict: Whatever model_fn returns for this framework
        """
        if self._model_bundle is None:
            self._check_openmp_conflict()
            self.log.info(f"Loading model bundle from {self.model_dir}...")
            self._model_bundle = self._load_inference_module().model_fn(self.model_dir)
        return self._model_bundle

    def _check_openmp_conflict(self):
        """Internal: Refuse to load an XGBoost model into a process that has torch.

        Torch and XGBoost each bring their own OpenMP runtime. Two in one process is
        undefined behavior, and unpickling an XGBoost Booster there segfaults the
        interpreter outright -- no exception, no traceback. Serving happens in separate
        containers in AWS, so the two never meet; in-process local inference is the
        first place they can.

        Raises:
            RuntimeError: If the conflict is present
        """
        if "torch" not in sys.modules:
            return
        if not glob.glob(os.path.join(self.model_dir, "xgb*.joblib")):
            return
        raise RuntimeError(
            f"Cannot serve XGBoost model '{self.model_name}' in this process: torch is already "
            "imported, and loading an XGBoost model alongside it segfaults the interpreter "
            "(duplicate OpenMP runtimes). Run this inference in a fresh process."
        )

    def parent(self):
        """The LocalModel this endpoint serves, if it still exists locally"""
        model = LocalModel(self.get_input())
        return model if model.exists() else None

    def aws_exists(self) -> bool:
        """Does an AWS Endpoint by this name already exist?

        Returns:
            bool: True if AWS already has this Endpoint
        """
        from workbench.api import Endpoint

        return Endpoint(self.name).exists()

    def _aws_artifact(self):
        """Internal: The AWS Endpoint for this local one"""
        from workbench.api import Endpoint

        return Endpoint(self.name)

    def _publish_self(self, **kwargs):
        """Internal: Deploy the published Model as an AWS Endpoint

        Returns:
            Endpoint: The created AWS Endpoint
        """
        from workbench.api import Model

        return Model(self.model_name).to_endpoint(name=self.name, **kwargs)

    def details(self, **kwargs) -> dict:
        """LocalEndpoint Details

        Returns:
            dict: A dictionary of details about the LocalEndpoint
        """
        return {**super().details(), "model_name": self.model_name, "captures": self.list_captures()}

model_dir property

The model artifacts directory this endpoint loads from.

Resolved from the model name on every access rather than stored, so the endpoint keeps working when the storage root moves or the config changes.

model_name property

The model this endpoint serves

__init__(name, **kwargs)

Initialize a LocalEndpoint

Parameters:

Name	Type	Description	Default
name	str	The name of the endpoint	required

Source code in src/workbench/local/local_endpoint.py

def __init__(self, name: str, **kwargs):
    """Initialize a LocalEndpoint

    Args:
        name (str): The name of the endpoint
    """
    Artifact.is_name_valid(name, delimiter="-", lower_case=False)
    super().__init__(name, **kwargs)
    self.inference_dir = os.path.join(self.path, "inference")
    self._model_bundle = None
    self._inference_module = None

aws_exists()

Does an AWS Endpoint by this name already exist?

Returns:

Name	Type	Description
bool	bool	True if AWS already has this Endpoint

Source code in src/workbench/local/local_endpoint.py

def aws_exists(self) -> bool:
    """Does an AWS Endpoint by this name already exist?

    Returns:
        bool: True if AWS already has this Endpoint
    """
    from workbench.api import Endpoint

    return Endpoint(self.name).exists()

details(**kwargs)

LocalEndpoint Details

Returns:

Name	Type	Description
dict	dict	A dictionary of details about the LocalEndpoint

Source code in src/workbench/local/local_endpoint.py

def details(self, **kwargs) -> dict:
    """LocalEndpoint Details

    Returns:
        dict: A dictionary of details about the LocalEndpoint
    """
    return {**super().details(), "model_name": self.model_name, "captures": self.list_captures()}

from_model(model, name=None) classmethod

Create a LocalEndpoint that serves a LocalModel.

Parameters:

Name	Type	Description	Default
model	LocalModel	The trained model to serve	required
name	str	Endpoint name (defaults to the model name)	None

Returns:

Name	Type	Description
LocalEndpoint	LocalEndpoint	The created endpoint

Raises:

Type	Description
ValueError	If the model hasn't trained successfully

Source code in src/workbench/local/local_endpoint.py

@classmethod
def from_model(cls, model: LocalModel, name: str = None) -> "LocalEndpoint":
    """Create a LocalEndpoint that serves a LocalModel.

    Args:
        model (LocalModel): The trained model to serve
        name (str, optional): Endpoint name (defaults to the model name)

    Returns:
        LocalEndpoint: The created endpoint

    Raises:
        ValueError: If the model hasn't trained successfully
    """
    if model.training_state().get("state") != "completed":
        raise ValueError(f"Model {model.name} has not trained successfully, cannot serve it")

    endpoint = cls(name or model.name)
    storage.local_root(create=True)
    os.makedirs(endpoint.inference_dir, exist_ok=True)
    endpoint._init_storage(input_name=model.name)
    endpoint.upsert_workbench_meta({"model_name": model.name})
    endpoint.log.important(f"Local endpoint {endpoint.name} serving {model.name}")
    return endpoint

get_inference_predictions(capture_name='auto_inference')

Retrieve a captured inference run.

Parameters:

Name	Type	Description	Default
capture_name	str	The capture to retrieve (default: "auto_inference")	'auto_inference'

Returns:

Type	Description
Union[DataFrame, None]	Union[pd.DataFrame, None]: The predictions, or None if that capture doesn't exist

Source code in src/workbench/local/local_endpoint.py

def get_inference_predictions(self, capture_name: str = "auto_inference") -> Union[pd.DataFrame, None]:
    """Retrieve a captured inference run.

    Args:
        capture_name (str): The capture to retrieve (default: "auto_inference")

    Returns:
        Union[pd.DataFrame, None]: The predictions, or None if that capture doesn't exist
    """
    path = os.path.join(self.inference_dir, capture_name, "predictions.parquet")
    if not os.path.isfile(path):
        self.log.warning(f"No inference capture named '{capture_name}' for {self.name}")
        return None
    return pd.read_parquet(path)

inference(eval_df, capture_name=None)

Run inference on a DataFrame.

Parameters:

Name	Type	Description	Default
eval_df	DataFrame	The data to run inference on	required
capture_name	str	Store the predictions under this name, which makes them available from the model as an inference run	None

Returns:

Type	Description
DataFrame	pd.DataFrame: The predictions

Source code in src/workbench/local/local_endpoint.py

def inference(self, eval_df: pd.DataFrame, capture_name: str = None) -> pd.DataFrame:
    """Run inference on a DataFrame.

    Args:
        eval_df (pd.DataFrame): The data to run inference on
        capture_name (str, optional): Store the predictions under this name, which
            makes them available from the model as an inference run

    Returns:
        pd.DataFrame: The predictions
    """
    module = self._load_inference_module()
    bundle = self._load_model_bundle()

    # predict_fn resolves its own supporting files through SM_MODEL_DIR
    previous = os.environ.get("SM_MODEL_DIR")
    os.environ["SM_MODEL_DIR"] = self.model_dir
    try:
        predictions = module.predict_fn(eval_df.copy(), bundle)
    finally:
        if previous is None:
            os.environ.pop("SM_MODEL_DIR", None)
        else:
            os.environ["SM_MODEL_DIR"] = previous

    if capture_name:
        self._capture(predictions, capture_name)
    return predictions

list_captures()

The inference captures stored on this endpoint.

Returns:

Type	Description
list[str]	list[str]: Sorted capture names

Source code in src/workbench/local/local_endpoint.py

def list_captures(self) -> list[str]:
    """The inference captures stored on this endpoint.

    Returns:
        list[str]: Sorted capture names
    """
    if not os.path.isdir(self.inference_dir):
        return []
    return sorted(d.name for d in os.scandir(self.inference_dir) if d.is_dir())

parent()

The LocalModel this endpoint serves, if it still exists locally

Source code in src/workbench/local/local_endpoint.py

def parent(self):
    """The LocalModel this endpoint serves, if it still exists locally"""
    model = LocalModel(self.get_input())
    return model if model.exists() else None

LocalMeta: Listings for the artifacts in local storage.

A directory glob plus a meta.json read per artifact. No caching tier and no artifact objects constructed, so listing stays cheap.

LocalMeta

LocalMeta: Workbench Local Metadata Class

Common Usage

meta = LocalMeta()
meta.data_sources()
meta.feature_sets()
meta.models()
meta.endpoints()

Source code in src/workbench/local/local_meta.py

class LocalMeta:
    """LocalMeta: Workbench Local Metadata Class

    Common Usage:
        ```python
        meta = LocalMeta()
        meta.data_sources()
        meta.feature_sets()
        meta.models()
        meta.endpoints()
        ```
    """

    def __init__(self):
        """Initialize the LocalMeta class"""
        self.log = Artifact.log

    def data_sources(self) -> pd.DataFrame:
        """Get a summary of the local Data Sources

        Returns:
            pd.DataFrame: A summary of the local Data Sources
        """
        return self._summary("data_source", extra={"Rows": "num_rows", "Columns": "num_columns"})

    def feature_sets(self) -> pd.DataFrame:
        """Get a summary of the local Feature Sets

        Returns:
            pd.DataFrame: A summary of the local Feature Sets
        """
        return self._summary("feature_set", extra={"Rows": "num_rows", "Columns": "num_columns", "Id": "id_column"})

    def models(self) -> pd.DataFrame:
        """Get a summary of the local Models

        Returns:
            pd.DataFrame: A summary of the local Models
        """
        return self._summary("model", extra={"Type": "model_type", "Framework": "model_framework"})

    def endpoints(self) -> pd.DataFrame:
        """Get a summary of the local Endpoints

        Returns:
            pd.DataFrame: A summary of the local Endpoints
        """
        return self._summary("endpoint", extra={"Model": "model_name"})

    def _summary(self, artifact_type: str, extra: dict = None) -> pd.DataFrame:
        """Internal: Build a summary DataFrame for one artifact type.

        Args:
            artifact_type (str): One of the keys in storage.SUBDIRS
            extra (dict, optional): Column name -> metadata key, appended per type

        Returns:
            pd.DataFrame: One row per artifact
        """
        extra = extra or {}
        columns = ["Name", "Health", "Owner", "Created", "Modified", "Status", "Input", "Tags"] + list(extra)

        rows = []
        for name in storage.list_artifacts(artifact_type):
            path = storage.artifact_path(artifact_type, name)
            meta = self._read_meta(path)
            if meta is None:
                continue

            row = {
                "Name": name,
                "Health": ", ".join(self._split(meta.get("workbench_health_tags"))) or "healthy",
                "Owner": meta.get("workbench_owner", "unknown"),
                "Created": self._created(meta, path),
                "Modified": self._modified(path),
                "Status": meta.get("workbench_status", "unknown"),
                "Input": meta.get("workbench_input", "unknown"),
                "Tags": ", ".join(self._split(meta.get("workbench_tags"))),
            }
            # num_columns isn't stored; it's the length of the stored column list
            for column, key in extra.items():
                row[column] = len(meta.get("columns", [])) if key == "num_columns" else meta.get(key)
            rows.append(row)

        return pd.DataFrame(rows, columns=columns)

    @staticmethod
    def _read_meta(path: str) -> dict:
        """Internal: Read an artifact's meta.json

        Args:
            path (str): The artifact directory

        Returns:
            dict: The metadata, or None if unreadable
        """
        try:
            with open(os.path.join(path, "meta.json"), "r") as fp:
                return json.load(fp)
        except (FileNotFoundError, json.JSONDecodeError):
            return None

    @staticmethod
    def _split(tags: str) -> list[str]:
        """Internal: Split a delimited tag string

        Args:
            tags (str): The delimited tag string

        Returns:
            list[str]: The individual tags
        """
        return tags.split(Artifact.tag_delimiter) if tags else []

    @staticmethod
    def _created(meta: dict, path: str) -> datetime:
        """Internal: Creation time from metadata, falling back to the meta.json mtime

        Args:
            meta (dict): The artifact metadata
            path (str): The artifact directory

        Returns:
            datetime: The creation time
        """
        created = meta.get("workbench_created")
        if created:
            return datetime.fromisoformat(created)
        return datetime.fromtimestamp(os.path.getmtime(os.path.join(path, "meta.json")), tz=timezone.utc)

    @staticmethod
    def _modified(path: str) -> datetime:
        """Internal: Most recent modification time in the artifact directory

        Args:
            path (str): The artifact directory

        Returns:
            datetime: The modification time
        """
        newest = storage.newest_mtime(path)
        return datetime.fromtimestamp(newest, tz=timezone.utc) if newest else None

__init__()

Initialize the LocalMeta class

Source code in src/workbench/local/local_meta.py

def __init__(self):
    """Initialize the LocalMeta class"""
    self.log = Artifact.log

data_sources()

Get a summary of the local Data Sources

Returns:

Type	Description
DataFrame	pd.DataFrame: A summary of the local Data Sources

Source code in src/workbench/local/local_meta.py

def data_sources(self) -> pd.DataFrame:
    """Get a summary of the local Data Sources

    Returns:
        pd.DataFrame: A summary of the local Data Sources
    """
    return self._summary("data_source", extra={"Rows": "num_rows", "Columns": "num_columns"})

endpoints()

Get a summary of the local Endpoints

Returns:

Type	Description
DataFrame	pd.DataFrame: A summary of the local Endpoints

Source code in src/workbench/local/local_meta.py

def endpoints(self) -> pd.DataFrame:
    """Get a summary of the local Endpoints

    Returns:
        pd.DataFrame: A summary of the local Endpoints
    """
    return self._summary("endpoint", extra={"Model": "model_name"})

feature_sets()

Get a summary of the local Feature Sets

Returns:

Type	Description
DataFrame	pd.DataFrame: A summary of the local Feature Sets

Source code in src/workbench/local/local_meta.py

def feature_sets(self) -> pd.DataFrame:
    """Get a summary of the local Feature Sets

    Returns:
        pd.DataFrame: A summary of the local Feature Sets
    """
    return self._summary("feature_set", extra={"Rows": "num_rows", "Columns": "num_columns", "Id": "id_column"})

models()

Get a summary of the local Models

Returns:

Type	Description
DataFrame	pd.DataFrame: A summary of the local Models

Source code in src/workbench/local/local_meta.py

def models(self) -> pd.DataFrame:
    """Get a summary of the local Models

    Returns:
        pd.DataFrame: A summary of the local Models
    """
    return self._summary("model", extra={"Type": "model_type", "Framework": "model_framework"})