Features To Model

API Classes

For most users the API Classes will provide all the general functionality to create a full AWS ML Pipeline

FeaturesToModel: Train/Create a Model from a Feature Set

CapacityTimeout

Bases: RuntimeError

A training job was stopped without ever getting the instance it asked for.

Source code in src/workbench/core/transforms/features_to_model/features_to_model.py

class CapacityTimeout(RuntimeError):
    """A training job was stopped without ever getting the instance it asked for."""

FeaturesToModel

Bases: Transform

FeaturesToModel: Train/Create a Model from a FeatureSet

Common Usage

from workbench.core.transforms.features_to_model.features_to_model import FeaturesToModel
to_model = FeaturesToModel(feature_name, model_name, model_type=ModelType)
to_model.set_output_tags(["abalone", "public", "whatever"])
to_model.transform(target_column="class_number_of_rings",
                   feature_list=["my", "best", "features"])

Source code in src/workbench/core/transforms/features_to_model/features_to_model.py

class FeaturesToModel(Transform):
    """FeaturesToModel: Train/Create a Model from a FeatureSet

    Common Usage:
        ```python
        from workbench.core.transforms.features_to_model.features_to_model import FeaturesToModel
        to_model = FeaturesToModel(feature_name, model_name, model_type=ModelType)
        to_model.set_output_tags(["abalone", "public", "whatever"])
        to_model.transform(target_column="class_number_of_rings",
                           feature_list=["my", "best", "features"])
        ```
    """

    def __init__(
        self,
        feature_name: str,
        model_name: str,
        model_type: ModelType,
        model_framework: ModelFramework,
        model_class=None,
        model_import_str=None,
        custom_script=None,
        custom_args=None,
        training_image="base_training",
        inference_image="base_inference",
        inference_arch="x86_64",
    ):
        """FeaturesToModel Initialization
        Args:
            feature_name (str): Name of the FeatureSet to use as input
            model_name (str): Name of the Model to create as output
            model_type (ModelType): ModelType.REGRESSOR or ModelType.CLASSIFIER, etc.
            model_framework (ModelFramework): SKLEARN, XGBOOST, PYTORCH, CHEMPROP, TRANSFORMER
            model_class (str, optional): The scikit model (e.g. KNeighborsRegressor) (default None)
            model_import_str (str, optional): The import string for the model (default None)
            custom_script (str, optional): Custom script to use for the model (default None)
            custom_args (dict, optional): Custom arguments to pass to custom model scripts (default None)
            training_image (str, optional): Training image (default "training")
            inference_image (str, optional): Inference image (default "inference")
            inference_arch (str, optional): Inference architecture (default "x86_64")
        """

        # Make sure the model_name is a valid name
        Artifact.is_name_valid(model_name, delimiter="-", lower_case=False)

        # Call superclass init
        super().__init__(feature_name, model_name)

        # Set up all my instance attributes
        self.input_type = TransformInput.FEATURE_SET
        self.output_type = TransformOutput.MODEL
        self.model_type = model_type
        self.model_framework = model_framework
        self.model_class = model_class
        self.model_import_str = model_import_str
        self.custom_script = str(custom_script) if custom_script else None
        self.custom_args = custom_args if custom_args else {}
        self.model_trainer = None
        self.training_job_name = None
        self.model_description = None
        self.model_training_root = f"{self.models_s3_path}/{self.output_name}/training"
        self.model_feature_list = None
        self.target_column = None
        self.class_labels = None
        self.training_image = training_image
        self.inference_image = inference_image
        self.inference_arch = inference_arch

    def transform_impl(
        self,
        target_column: Union[str, list[str]] = None,
        description: str = None,
        feature_list: list = None,
        sample_weights: Union[dict, pd.DataFrame] = None,
        validation_ids: list = None,
        exclude_ids: list = None,
        **kwargs,
    ):
        """Generic Features to Model: Note you should create a new class and inherit from
        this one to include specific logic for your Feature Set/Model
        Args:
            target_column (str or list[str]): Column name(s) of the target variable(s)
            description (str): Description of the model (optional)
            feature_list (list[str]): A list of columns for the features (default None, will try to guess)
            sample_weights (dict | pd.DataFrame): Sparse per-id sample weights as a
                ``{id: weight}`` dict or a ``[id_column, "sample_weight"]`` DataFrame. A pure
                framework weight forwarded as-is to the model script; ids not listed default to
                1.0. It no longer encodes any role (default None: all rows weight 1.0).
            validation_ids (list): Ids designated as a held-out validation set — kept in the
                training view and marked, but routed out of training and scored as an honest
                held-out set by the model script (default None: no validation rows).
            exclude_ids (list): Ids to drop from the training view entirely (outliers/anomalies);
                no model ever sees them. Takes precedence over ``validation_ids`` on overlap
                (default None: nothing excluded).
        """
        supervised_types = (
            ModelType.CLASSIFIER,
            ModelType.REGRESSOR,
            ModelType.UQ_REGRESSOR,
            ModelType.ENSEMBLE_REGRESSOR,
        )
        if target_column is None and self.model_type in supervised_types:
            raise ValueError("target_column is required for supervised models (pass target_column=...)")

        # Set our model description
        self.model_description = description if description is not None else f"Model created from {self.input_name}"

        # Get our Feature Set and build the model's own training view
        feature_set = FeatureSetCore(self.input_name)
        short_id = uuid.uuid4().hex[:6]
        self.model_training_view_name = f"{self.output_name.replace('-', '_')}_training_{short_id}".lower()
        self.log.important(f"Creating Model Training View: {self.model_training_view_name}...")
        self._create_model_training_view(
            feature_set, self.model_training_view_name, sample_weights, validation_ids, exclude_ids
        )

        # Delete the existing model (if it exists)
        self.log.important(f"Trying to delete existing model {self.output_name}...")
        ModelCore.managed_delete(self.output_name)

        # Create S3 training data from the model-owned snapshot (not the shared training view)
        base_table = feature_set.data_source.table
        model_view_table = f"{base_table}___{self.model_training_view_name}"
        s3_training_path = feature_set.create_s3_training_data(source_table=model_view_table)
        self.log.info(f"Created new training data {s3_training_path}...")

        # Report the target column(s)
        self.target_column = target_column
        # Normalize target_column to a list for internal use
        target_list = [target_column] if isinstance(target_column, str) else (target_column or [])
        self.log.info(f"Target column(s): {self.target_column}")

        # Did they specify a feature list?
        if feature_list:
            # AWS Feature Groups will also add these implicit columns, so remove them
            aws_cols = ["write_time", "api_invocation_time", "is_deleted", "event_time", "training"]
            feature_list = [c for c in feature_list if c not in aws_cols]

        # If they didn't specify a feature list, try to guess it
        else:
            # Try to figure out features with this logic
            # - Don't include id, event_time, __index_level_0__, or training columns
            # - Don't include AWS generated columns (e.g. write_time, api_invocation_time, is_deleted)
            # - Don't include the target columns
            # - Don't include any columns that are of type string or timestamp
            # - The rest of the columns are assumed to be features
            self.log.warning("Guessing at the feature list, HIGHLY RECOMMENDED to specify an explicit feature list!")
            all_columns = feature_set.columns
            filter_list = [
                "id",
                "auto_id",
                "__index_level_0__",
                "write_time",
                "api_invocation_time",
                "is_deleted",
                "event_time",
                "training",
            ] + target_list
            feature_list = [c for c in all_columns if c not in filter_list]

            # AWS Feature Store has 3 user column types (String, Integral, Fractional)
            # and two internal types (Timestamp and Boolean). A Feature List for
            # modeling can only contain Integral and Fractional types.
            remove_columns = []
            column_details = feature_set.column_details()
            for column_name in feature_list:
                if column_details[column_name] not in ["Integral", "Fractional"]:
                    self.log.warning(
                        f"Removing {column_name} from feature list, improper type {column_details[column_name]}"
                    )
                    remove_columns.append(column_name)

            # Remove the columns that are not Integral or Fractional
            feature_list = [c for c in feature_list if c not in remove_columns]

        # Set the final feature list
        self.model_feature_list = feature_list
        self.log.important(f"Feature List for Modeling: {self.model_feature_list}")

        # Set up our parameters for the model script
        # ChemProp expects target_column as a list; other templates expect a string
        target_for_template = target_list if self.model_framework == ModelFramework.CHEMPROP else self.target_column
        template_params = {
            "model_imports": self.model_import_str,
            "model_type": self.model_type,
            "model_framework": self.model_framework,
            "model_class": self.model_class,
            "target_column": target_for_template,
            "feature_list": self.model_feature_list,
            "compressed_features": feature_set.get_compressed_features(),
            "model_metrics_s3_path": self.model_training_root,
            "id_column": feature_set.id_column,
            "hyperparameters": kwargs.get("hyperparameters", {}),
        }

        # Caller-supplied template params (used by MetaEndpoint to inject DAG
        # config, region, s3_bucket; available for any framework that needs
        # extra placeholders beyond the standard set above).
        if self.custom_args:
            template_params.update(self.custom_args)

        # Custom Script
        if self.custom_script:
            script_path = self.custom_script
            if self.custom_script.endswith(".template"):
                # Model Type is an enumerated type, so we need to convert it to a string
                template_params["model_type"] = template_params["model_type"].value
                script_path = fill_template(self.custom_script, template_params, "generated_model_script.py")

            # Ensure training_harness.py is in the custom script's source directory
            source_dir = Path(script_path).parent
            harness_dst = source_dir / "training_harness.py"
            if not harness_dst.exists():
                harness_src = pkg_resources.files("workbench") / "training" / "training_harness.py"
                shutil.copy(str(harness_src), harness_dst)
                self.log.info(f"Copied training_harness.py into {source_dir}")

            self.log.info(f"Custom script path: {script_path}")

        # We're using one of the built-in model script templates
        else:
            # Generate our model script
            script_path = generate_model_script(template_params)

        # Metric Definitions for Regression (matches model script output format)
        if self.model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]:
            metric_definitions = [
                {"Name": "rmse", "Regex": r"rmse: ([0-9.]+)"},
                {"Name": "mae", "Regex": r"mae: ([0-9.]+)"},
                {"Name": "medae", "Regex": r"medae: ([0-9.]+)"},
                {"Name": "r2", "Regex": r"r2: ([0-9.-]+)"},
                {"Name": "spearmanr", "Regex": r"spearmanr: ([0-9.-]+)"},
                {"Name": "support", "Regex": r"support: ([0-9]+)"},
            ]

        # Metric Definitions for Classification
        elif self.model_type == ModelType.CLASSIFIER:
            # We need to get creative with the Classification Metrics
            # Note: Classification only supports single target
            class_target = target_list[0] if target_list else self.target_column

            # Grab all the target column class values (class labels)
            table = feature_set.data_source.table
            self.class_labels = feature_set.query(f'select DISTINCT {class_target} FROM "{table}"')[
                class_target
            ].to_list()

            # Sanity check on the targets
            if len(self.class_labels) > 10:
                msg = f"Too many target classes ({len(self.class_labels)}) for classification, aborting!"
                self.log.critical(msg)
                raise ValueError(msg)

            # Dynamically create the metric definitions (per-class precision/recall/f1/support)
            # Note: Confusion matrix metrics are skipped to stay under SageMaker's 40 metric limit
            metrics = ["precision", "recall", "f1", "support"]
            metric_definitions = []
            for t in self.class_labels:
                for m in metrics:
                    metric_definitions.append({"Name": f"Metrics:{t}:{m}", "Regex": f"Metrics:{t}:{m} ([0-9.]+)"})

        # If the model type is UNKNOWN, our metric_definitions will be empty
        else:
            self.log.important(f"ModelType is {self.model_type}, skipping metric_definitions...")
            metric_definitions = []

        # Take the full script path and extract the entry point and source directory
        entry_point = str(Path(script_path).name)
        source_dir = str(Path(script_path).parent)

        # Create a Sagemaker Model with our script
        image = ModelImages.get_image_uri(self.sm_session.boto_region_name, self.training_image)

        # Use user-specified instance or pick a ladder from the workload.
        hpo_requested = (kwargs.get("hyperparameters") or {}).get("hpo") is not None
        train_instance_type = kwargs.get("training_instance")
        if train_instance_type:
            instance_ladder = [train_instance_type]
            self.log.important(f"Using user-specified instance {train_instance_type}")
        else:
            gpu_framework = self.model_framework in [ModelFramework.CHEMPROP, ModelFramework.PYTORCH]
            workload = training_workload(kwargs.get("hyperparameters"), gpu_framework=gpu_framework)
            instance_ladder = INSTANCE_LADDERS[workload]
            self.log.important(f"Using {workload} instances {instance_ladder} for {self.model_framework.value}")

        # Convert metric definitions to V3 MetricDefinition objects
        v3_metric_definitions = [MetricDefinition(name=m["Name"], regex=m["Regex"]) for m in metric_definitions]

        # PRM attribution tag: connects the training-job compute to our Marketplace listing
        prm_tag = Tag(key="aws-apn-id", value=f"pc:{AWS_MARKETPLACE_PRODUCT_CODE}")

        self.log.important(f"Training the Model {self.output_name} with Training Image {image}...")
        _suppress_sagemaker_logging()

        for rung, train_instance_type in enumerate(instance_ladder):
            last_rung = rung == len(instance_ladder) - 1

            # The final rung queues for capacity as long as it takes; the rungs above it
            # give up so the ladder can move on.
            stopping_condition = StoppingCondition(max_runtime_in_seconds=(24 if hpo_requested else 6) * 3600)
            if not last_rung:
                stopping_condition.max_pending_time_in_seconds = MAX_PENDING_SECONDS

            # Create ModelTrainer (V3 replacement for Estimator)
            # Use command= to run our entrypoint wrapper, which executes the model script
            # and then bundles inference code/metadata into the model artifacts
            self.model_trainer = ModelTrainer(
                training_image=image,
                # The image carries no site config, so hand the container the bucket it needs
                # (chemprop foundation weights, df_store, ...) explicitly.
                environment={"WORKBENCH_BUCKET": self.workbench_bucket},
                source_code=SourceCode(
                    source_dir=source_dir,
                    command=f"python training_harness.py {entry_point}",
                ),
                compute=Compute(instance_type=train_instance_type, instance_count=1),
                output_data_config=OutputDataConfig(s3_output_path=self.model_training_root, compression_type="GZIP"),
                stopping_condition=stopping_condition,
                base_job_name=self.output_name,
                role=self.workbench_role_arn,
                sagemaker_session=self.sm_session,
                tags=[prm_tag],
            )
            self.model_trainer.with_metric_definitions(v3_metric_definitions)
            input_data = self.model_trainer.create_input_data_channel("train", s3_training_path)

            # Submit the job and own the poll loop ourselves. The SDK's wait=True polls via a
            # global client singleton that bypasses our hardened session and never retries a
            # transient HttpTimeoutException on DescribeTrainingJob (see _wait_for_training_job).
            # Silence the SDK's stray "not displaying logs" warning that wait=False emits.
            logging.getLogger("sagemaker.train.model_trainer").setLevel(logging.ERROR)
            self.log.important(f"Submitting training job on {train_instance_type}...")
            try:
                self.model_trainer.train(input_data_config=[input_data], wait=False)
            except ClientError as e:
                # The account's quota for this type is already fully in use by another job
                if e.response.get("Error", {}).get("Code") == "ResourceLimitExceeded" and not last_rung:
                    self.log.warning(f"{train_instance_type} quota exhausted, dropping to next instance...")
                    continue
                raise

            # Capture the actual training job name (ModelTrainer appends a timestamp to base_job_name)
            self.training_job_name = self.model_trainer._latest_training_job.training_job_name
            try:
                self._wait_for_training_job(capacity_timeout=None if last_rung else CAPACITY_WAIT_SECONDS)
                break
            except CapacityTimeout:
                if last_rung:
                    raise
                self.log.warning(f"Dropping from {train_instance_type} to {instance_ladder[rung + 1]}...")

        # Now delete the training data
        self.log.info(f"Deleting training data {s3_training_path}...")
        wr.s3.delete_objects(
            [s3_training_path, s3_training_path.replace(".csv", ".csv.metadata")],
            boto3_session=self.boto3_session,
        )

        # Create Model and officially Register
        self.log.important(f"Creating new model {self.output_name}...")
        self.create_and_register_model(**kwargs)

    def _wait_for_training_job(self, poll_interval: int = 30, capacity_timeout: int = None):
        """Poll the training job to completion using our hardened SageMaker client.

        Retries transient DescribeTrainingJob control-plane errors (HttpTimeoutException,
        throttling) that botocore won't retry on its own, so a blip doesn't fail the job.

        Args:
            poll_interval (int): Seconds between DescribeTrainingJob polls (default 30)
            capacity_timeout (int): Stop the job after this many seconds queued for an
                instance. None (default) queues for as long as SageMaker will.

        Raises:
            CapacityTimeout: If the job was stopped while still waiting for an instance
            RuntimeError: If the training job reaches a terminal status other than Completed
        """
        terminal = {"Completed", "Failed", "Stopped"}
        # Any secondary status past the queue means SageMaker handed us the instance
        got_capacity = {"LaunchingMLInstances", "PreparingTrainingStack", "Downloading", "Training"}
        last_status = None
        while True:
            try:
                desc = self.sm_client.describe_training_job(TrainingJobName=self.training_job_name)
            except ClientError as e:
                code = e.response.get("Error", {}).get("Code", "")
                if code in TRANSIENT_DESCRIBE_ERRORS:
                    self.log.warning(f"Transient {code} polling {self.training_job_name}; retrying...")
                    time.sleep(poll_interval)
                    continue
                raise

            status = desc["TrainingJobStatus"]
            if status != last_status:
                self.log.info(f"Training job {self.training_job_name} status: {status}")
                last_status = status
            if status in terminal:
                if status == "Completed":
                    return
                # max_pending_time_in_seconds expiring stops the job without a FailureReason,
                # so a stop with no transition past the queue is a capacity timeout.
                reached = {t["Status"] for t in desc.get("SecondaryStatusTransitions", [])}
                if status == "Stopped" and not (reached & got_capacity):
                    raise CapacityTimeout(f"Training job {self.training_job_name} never got an instance")
                reason = desc.get("FailureReason", "(no reason provided)")
                raise RuntimeError(f"Training job {self.training_job_name} {status}: {reason}")

            # Queued for an instance the account may not get: stop it so the caller can try
            # elsewhere. SageMaker's own floor for this is 2 hours, hence our own clock.
            if capacity_timeout and desc["SecondaryStatus"] == "Pending":
                queued_since = desc["SecondaryStatusTransitions"][-1]["StartTime"]
                queued_for = (datetime.now(timezone.utc) - queued_since).total_seconds()
                if queued_for > capacity_timeout:
                    self.log.warning(
                        f"No {desc['ResourceConfig']['InstanceType']} capacity after "
                        f"{queued_for / 60:.0f} minutes, stopping {self.training_job_name}..."
                    )
                    try:
                        self.sm_client.stop_training_job(TrainingJobName=self.training_job_name)
                    except ClientError as e:
                        # Without StopTrainingJob permission the job's own max_pending_time
                        # is the only way out, so keep polling and let it expire.
                        self.log.error(f"Cannot stop {self.training_job_name} ({e}); waiting on max_pending_time")
                        capacity_timeout = None
                    else:
                        raise CapacityTimeout(f"Training job {self.training_job_name} never got an instance")
            time.sleep(poll_interval)

    def _create_model_training_view(
        self,
        feature_set: FeatureSetCore,
        model_view_name: str,
        sample_weights: Union[dict, pd.DataFrame] = None,
        validation_ids: list = None,
        exclude_ids: list = None,
    ):
        """Create a model-owned training view of the FeatureSet's features plus per-row role columns.

        The view always exposes ``sample_weight`` (framework weight), ``validation`` (held-out
        marker), and ``exclude`` (drop marker). When any role input is provided, the non-default
        rows are stored in a sparse model-scoped roles table that the view left-joins; excluded
        rows are filtered out. See ``view_utils.create_model_training_view`` for the SQL.

        Args:
            feature_set (FeatureSetCore): The source FeatureSet
            model_view_name (str): The model training view name (e.g. "my_model_training")
            sample_weights (dict | pd.DataFrame): Sparse per-id weights (``{id: weight}`` or a
                ``[id_column, "sample_weight"]`` DataFrame). None means all rows weight 1.0.
            validation_ids (list): Ids marked as held-out validation rows (kept, not trained).
            exclude_ids (list): Ids dropped from the view entirely (precedence over validation).
        """
        from workbench.core.views.view_utils import create_model_training_view

        id_column = feature_set.id_column

        # Feature columns (exclude AWS-generated columns)
        aws_cols = ["write_time", "api_invocation_time", "is_deleted", "event_time"]
        feature_columns = [c for c in feature_set.columns if c not in aws_cols]

        # Normalize the three role inputs into one sparse roles DataFrame (or None), then build the view
        roles_df = self._normalize_row_roles(sample_weights, validation_ids, exclude_ids, id_column)
        create_model_training_view(feature_set.data_source, model_view_name, feature_columns, id_column, roles_df)

    @staticmethod
    def _normalize_row_roles(
        sample_weights: Union[dict, pd.DataFrame, None],
        validation_ids: Union[list, None],
        exclude_ids: Union[list, None],
        id_column: str,
    ) -> Union[pd.DataFrame, None]:
        """Normalize the three role inputs into a sparse roles DataFrame.

        Produces columns ``[id_column, "sample_weight", "validation", "exclude"]`` containing only
        ids that differ from the defaults (weight 1.0, not validation, not excluded).

        Args:
            sample_weights (dict | pd.DataFrame | None): A ``{id: weight}`` dict or a
                ``[id_column, "sample_weight"]`` DataFrame.
            validation_ids (list | None): Ids to mark as held-out validation.
            exclude_ids (list | None): Ids to drop from the view.
            id_column (str): The FeatureSet id column name.

        Returns:
            pd.DataFrame | None: The sparse roles frame, or None when everything is at defaults.
        """
        # Normalize weights into a plain {id: weight} mapping
        weights: dict = {}
        if isinstance(sample_weights, dict):
            weights = dict(sample_weights)
        elif isinstance(sample_weights, pd.DataFrame) and not sample_weights.empty:
            weights = dict(zip(sample_weights[id_column], sample_weights["sample_weight"]))

        validation_set = set(validation_ids or [])
        exclude_set = set(exclude_ids or [])

        all_ids = set(weights) | validation_set | exclude_set
        if not all_ids:
            return None

        rows = [
            {
                id_column: i,
                "sample_weight": float(weights.get(i, 1.0)),
                "validation": i in validation_set,
                "exclude": i in exclude_set,
            }
            for i in all_ids
        ]
        return pd.DataFrame(rows, columns=[id_column, "sample_weight", "validation", "exclude"])

    def post_transform(self, **kwargs):
        """Post-Transform: Calling onboard() on the Model"""
        self.log.info("Post-Transform: Calling onboard() on the Model...")
        time.sleep(3)  # Give AWS time to complete Model register

        # Store the model metadata information
        output_model = ModelCore(self.output_name)
        output_model._set_model_type(self.model_type)
        output_model._set_model_framework(self.model_framework)
        output_model.upsert_workbench_meta({"workbench_model_features": self.model_feature_list})
        output_model.upsert_workbench_meta({"workbench_model_target": self.target_column})
        output_model.upsert_workbench_meta({"workbench_training_view": self.model_training_view_name})

        # Persist hyperparameters to meta. Template frameworks write the resolved set
        # (defaults + overrides) to S3 during training; custom scripts don't, so fall
        # back to the passed kwarg. Prefer S3 so the stored set is complete.
        training_hp = read_s3_json(
            f"{output_model.model_training_path}/hyperparameters.json", output_model.boto3_session
        )
        hyperparameters = training_hp if training_hp is not None else kwargs.get("hyperparameters")
        if hyperparameters:
            output_model.upsert_workbench_meta({"workbench_hyperparameters": hyperparameters})

        # Store the class labels (if they exist)
        if self.class_labels:
            output_model.set_class_labels(self.class_labels)

        # Call the Model onboard method
        output_model.onboard_with_args(self.model_type, self.target_column, self.model_feature_list)

    def create_and_register_model(self, aws_region=None, **kwargs):
        """Create and Register the Model

        Args:
            aws_region (str, optional): AWS Region to use (default None)
            **kwargs (dict): Additional keyword arguments to pass to the model registration
        """
        from sagemaker.core.resources import TrainingJob

        # Get the metadata/tags to push into AWS
        aws_tags = self.get_aws_tags()

        # Create model group (if it doesn't already exist)
        try:
            ModelPackageGroup.create(
                model_package_group_name=self.output_name,
                model_package_group_description=self.model_description,
                tags=aws_tags,
                session=self.boto3_session,
            )
        except Exception:
            self.log.info(f"Model Package Group {self.output_name} may already exist, continuing...")

        # Get the model artifacts URL from the completed training job
        training_job = TrainingJob.get(self.training_job_name, session=self.boto3_session)
        model_data_url = training_job.model_artifacts.s3_model_artifacts

        # Get the inference image URI
        image = ModelImages.get_image_uri(
            self.sm_session.boto_region_name, self.inference_image, architecture=self.inference_arch
        )
        self.log.important(f"Registering model {self.output_name} with Inference Image {image}...")

        # FIXME: V3 SDK Bug — ModelPackage.create() does response["ModelPackageName"] at line 25047
        # of resources.py, but the AWS CreateModelPackage API only returns "ModelPackageArn".
        # This causes a KeyError on every versioned model package creation. Using boto3 as workaround.
        # No existing GitHub issue — consider filing on https://github.com/aws/sagemaker-core/issues
        self.sm_client.create_model_package(
            ModelPackageGroupName=self.output_name,
            ModelPackageDescription=self.model_description,
            InferenceSpecification={
                "Containers": [{"Image": image, "ModelDataUrl": model_data_url}],
                "SupportedContentTypes": ["text/csv"],
                "SupportedResponseMIMETypes": ["text/csv"],
            },
            ModelApprovalStatus="Approved",
        )

__init__(feature_name, model_name, model_type, model_framework, model_class=None, model_import_str=None, custom_script=None, custom_args=None, training_image='base_training', inference_image='base_inference', inference_arch='x86_64')

FeaturesToModel Initialization Args: feature_name (str): Name of the FeatureSet to use as input model_name (str): Name of the Model to create as output model_type (ModelType): ModelType.REGRESSOR or ModelType.CLASSIFIER, etc. model_framework (ModelFramework): SKLEARN, XGBOOST, PYTORCH, CHEMPROP, TRANSFORMER model_class (str, optional): The scikit model (e.g. KNeighborsRegressor) (default None) model_import_str (str, optional): The import string for the model (default None) custom_script (str, optional): Custom script to use for the model (default None) custom_args (dict, optional): Custom arguments to pass to custom model scripts (default None) training_image (str, optional): Training image (default "training") inference_image (str, optional): Inference image (default "inference") inference_arch (str, optional): Inference architecture (default "x86_64")

Source code in src/workbench/core/transforms/features_to_model/features_to_model.py

def __init__(
    self,
    feature_name: str,
    model_name: str,
    model_type: ModelType,
    model_framework: ModelFramework,
    model_class=None,
    model_import_str=None,
    custom_script=None,
    custom_args=None,
    training_image="base_training",
    inference_image="base_inference",
    inference_arch="x86_64",
):
    """FeaturesToModel Initialization
    Args:
        feature_name (str): Name of the FeatureSet to use as input
        model_name (str): Name of the Model to create as output
        model_type (ModelType): ModelType.REGRESSOR or ModelType.CLASSIFIER, etc.
        model_framework (ModelFramework): SKLEARN, XGBOOST, PYTORCH, CHEMPROP, TRANSFORMER
        model_class (str, optional): The scikit model (e.g. KNeighborsRegressor) (default None)
        model_import_str (str, optional): The import string for the model (default None)
        custom_script (str, optional): Custom script to use for the model (default None)
        custom_args (dict, optional): Custom arguments to pass to custom model scripts (default None)
        training_image (str, optional): Training image (default "training")
        inference_image (str, optional): Inference image (default "inference")
        inference_arch (str, optional): Inference architecture (default "x86_64")
    """

    # Make sure the model_name is a valid name
    Artifact.is_name_valid(model_name, delimiter="-", lower_case=False)

    # Call superclass init
    super().__init__(feature_name, model_name)

    # Set up all my instance attributes
    self.input_type = TransformInput.FEATURE_SET
    self.output_type = TransformOutput.MODEL
    self.model_type = model_type
    self.model_framework = model_framework
    self.model_class = model_class
    self.model_import_str = model_import_str
    self.custom_script = str(custom_script) if custom_script else None
    self.custom_args = custom_args if custom_args else {}
    self.model_trainer = None
    self.training_job_name = None
    self.model_description = None
    self.model_training_root = f"{self.models_s3_path}/{self.output_name}/training"
    self.model_feature_list = None
    self.target_column = None
    self.class_labels = None
    self.training_image = training_image
    self.inference_image = inference_image
    self.inference_arch = inference_arch

create_and_register_model(aws_region=None, **kwargs)

Create and Register the Model

Parameters:

Name	Type	Description	Default
aws_region	str	AWS Region to use (default None)	None
**kwargs	dict	Additional keyword arguments to pass to the model registration	{}

Source code in src/workbench/core/transforms/features_to_model/features_to_model.py

def create_and_register_model(self, aws_region=None, **kwargs):
    """Create and Register the Model

    Args:
        aws_region (str, optional): AWS Region to use (default None)
        **kwargs (dict): Additional keyword arguments to pass to the model registration
    """
    from sagemaker.core.resources import TrainingJob

    # Get the metadata/tags to push into AWS
    aws_tags = self.get_aws_tags()

    # Create model group (if it doesn't already exist)
    try:
        ModelPackageGroup.create(
            model_package_group_name=self.output_name,
            model_package_group_description=self.model_description,
            tags=aws_tags,
            session=self.boto3_session,
        )
    except Exception:
        self.log.info(f"Model Package Group {self.output_name} may already exist, continuing...")

    # Get the model artifacts URL from the completed training job
    training_job = TrainingJob.get(self.training_job_name, session=self.boto3_session)
    model_data_url = training_job.model_artifacts.s3_model_artifacts

    # Get the inference image URI
    image = ModelImages.get_image_uri(
        self.sm_session.boto_region_name, self.inference_image, architecture=self.inference_arch
    )
    self.log.important(f"Registering model {self.output_name} with Inference Image {image}...")

    # FIXME: V3 SDK Bug — ModelPackage.create() does response["ModelPackageName"] at line 25047
    # of resources.py, but the AWS CreateModelPackage API only returns "ModelPackageArn".
    # This causes a KeyError on every versioned model package creation. Using boto3 as workaround.
    # No existing GitHub issue — consider filing on https://github.com/aws/sagemaker-core/issues
    self.sm_client.create_model_package(
        ModelPackageGroupName=self.output_name,
        ModelPackageDescription=self.model_description,
        InferenceSpecification={
            "Containers": [{"Image": image, "ModelDataUrl": model_data_url}],
            "SupportedContentTypes": ["text/csv"],
            "SupportedResponseMIMETypes": ["text/csv"],
        },
        ModelApprovalStatus="Approved",
    )

post_transform(**kwargs)

Post-Transform: Calling onboard() on the Model

Source code in src/workbench/core/transforms/features_to_model/features_to_model.py

def post_transform(self, **kwargs):
    """Post-Transform: Calling onboard() on the Model"""
    self.log.info("Post-Transform: Calling onboard() on the Model...")
    time.sleep(3)  # Give AWS time to complete Model register

    # Store the model metadata information
    output_model = ModelCore(self.output_name)
    output_model._set_model_type(self.model_type)
    output_model._set_model_framework(self.model_framework)
    output_model.upsert_workbench_meta({"workbench_model_features": self.model_feature_list})
    output_model.upsert_workbench_meta({"workbench_model_target": self.target_column})
    output_model.upsert_workbench_meta({"workbench_training_view": self.model_training_view_name})

    # Persist hyperparameters to meta. Template frameworks write the resolved set
    # (defaults + overrides) to S3 during training; custom scripts don't, so fall
    # back to the passed kwarg. Prefer S3 so the stored set is complete.
    training_hp = read_s3_json(
        f"{output_model.model_training_path}/hyperparameters.json", output_model.boto3_session
    )
    hyperparameters = training_hp if training_hp is not None else kwargs.get("hyperparameters")
    if hyperparameters:
        output_model.upsert_workbench_meta({"workbench_hyperparameters": hyperparameters})

    # Store the class labels (if they exist)
    if self.class_labels:
        output_model.set_class_labels(self.class_labels)

    # Call the Model onboard method
    output_model.onboard_with_args(self.model_type, self.target_column, self.model_feature_list)

transform_impl(target_column=None, description=None, feature_list=None, sample_weights=None, validation_ids=None, exclude_ids=None, **kwargs)

Generic Features to Model: Note you should create a new class and inherit from this one to include specific logic for your Feature Set/Model Args: target_column (str or list[str]): Column name(s) of the target variable(s) description (str): Description of the model (optional) feature_list (list[str]): A list of columns for the features (default None, will try to guess) sample_weights (dict | pd.DataFrame): Sparse per-id sample weights as a {id: weight} dict or a [id_column, "sample_weight"] DataFrame. A pure framework weight forwarded as-is to the model script; ids not listed default to 1.0. It no longer encodes any role (default None: all rows weight 1.0). validation_ids (list): Ids designated as a held-out validation set — kept in the training view and marked, but routed out of training and scored as an honest held-out set by the model script (default None: no validation rows). exclude_ids (list): Ids to drop from the training view entirely (outliers/anomalies); no model ever sees them. Takes precedence over validation_ids on overlap (default None: nothing excluded).

Source code in src/workbench/core/transforms/features_to_model/features_to_model.py

def transform_impl(
    self,
    target_column: Union[str, list[str]] = None,
    description: str = None,
    feature_list: list = None,
    sample_weights: Union[dict, pd.DataFrame] = None,
    validation_ids: list = None,
    exclude_ids: list = None,
    **kwargs,
):
    """Generic Features to Model: Note you should create a new class and inherit from
    this one to include specific logic for your Feature Set/Model
    Args:
        target_column (str or list[str]): Column name(s) of the target variable(s)
        description (str): Description of the model (optional)
        feature_list (list[str]): A list of columns for the features (default None, will try to guess)
        sample_weights (dict | pd.DataFrame): Sparse per-id sample weights as a
            ``{id: weight}`` dict or a ``[id_column, "sample_weight"]`` DataFrame. A pure
            framework weight forwarded as-is to the model script; ids not listed default to
            1.0. It no longer encodes any role (default None: all rows weight 1.0).
        validation_ids (list): Ids designated as a held-out validation set — kept in the
            training view and marked, but routed out of training and scored as an honest
            held-out set by the model script (default None: no validation rows).
        exclude_ids (list): Ids to drop from the training view entirely (outliers/anomalies);
            no model ever sees them. Takes precedence over ``validation_ids`` on overlap
            (default None: nothing excluded).
    """
    supervised_types = (
        ModelType.CLASSIFIER,
        ModelType.REGRESSOR,
        ModelType.UQ_REGRESSOR,
        ModelType.ENSEMBLE_REGRESSOR,
    )
    if target_column is None and self.model_type in supervised_types:
        raise ValueError("target_column is required for supervised models (pass target_column=...)")

    # Set our model description
    self.model_description = description if description is not None else f"Model created from {self.input_name}"

    # Get our Feature Set and build the model's own training view
    feature_set = FeatureSetCore(self.input_name)
    short_id = uuid.uuid4().hex[:6]
    self.model_training_view_name = f"{self.output_name.replace('-', '_')}_training_{short_id}".lower()
    self.log.important(f"Creating Model Training View: {self.model_training_view_name}...")
    self._create_model_training_view(
        feature_set, self.model_training_view_name, sample_weights, validation_ids, exclude_ids
    )

    # Delete the existing model (if it exists)
    self.log.important(f"Trying to delete existing model {self.output_name}...")
    ModelCore.managed_delete(self.output_name)

    # Create S3 training data from the model-owned snapshot (not the shared training view)
    base_table = feature_set.data_source.table
    model_view_table = f"{base_table}___{self.model_training_view_name}"
    s3_training_path = feature_set.create_s3_training_data(source_table=model_view_table)
    self.log.info(f"Created new training data {s3_training_path}...")

    # Report the target column(s)
    self.target_column = target_column
    # Normalize target_column to a list for internal use
    target_list = [target_column] if isinstance(target_column, str) else (target_column or [])
    self.log.info(f"Target column(s): {self.target_column}")

    # Did they specify a feature list?
    if feature_list:
        # AWS Feature Groups will also add these implicit columns, so remove them
        aws_cols = ["write_time", "api_invocation_time", "is_deleted", "event_time", "training"]
        feature_list = [c for c in feature_list if c not in aws_cols]

    # If they didn't specify a feature list, try to guess it
    else:
        # Try to figure out features with this logic
        # - Don't include id, event_time, __index_level_0__, or training columns
        # - Don't include AWS generated columns (e.g. write_time, api_invocation_time, is_deleted)
        # - Don't include the target columns
        # - Don't include any columns that are of type string or timestamp
        # - The rest of the columns are assumed to be features
        self.log.warning("Guessing at the feature list, HIGHLY RECOMMENDED to specify an explicit feature list!")
        all_columns = feature_set.columns
        filter_list = [
            "id",
            "auto_id",
            "__index_level_0__",
            "write_time",
            "api_invocation_time",
            "is_deleted",
            "event_time",
            "training",
        ] + target_list
        feature_list = [c for c in all_columns if c not in filter_list]

        # AWS Feature Store has 3 user column types (String, Integral, Fractional)
        # and two internal types (Timestamp and Boolean). A Feature List for
        # modeling can only contain Integral and Fractional types.
        remove_columns = []
        column_details = feature_set.column_details()
        for column_name in feature_list:
            if column_details[column_name] not in ["Integral", "Fractional"]:
                self.log.warning(
                    f"Removing {column_name} from feature list, improper type {column_details[column_name]}"
                )
                remove_columns.append(column_name)

        # Remove the columns that are not Integral or Fractional
        feature_list = [c for c in feature_list if c not in remove_columns]

    # Set the final feature list
    self.model_feature_list = feature_list
    self.log.important(f"Feature List for Modeling: {self.model_feature_list}")

    # Set up our parameters for the model script
    # ChemProp expects target_column as a list; other templates expect a string
    target_for_template = target_list if self.model_framework == ModelFramework.CHEMPROP else self.target_column
    template_params = {
        "model_imports": self.model_import_str,
        "model_type": self.model_type,
        "model_framework": self.model_framework,
        "model_class": self.model_class,
        "target_column": target_for_template,
        "feature_list": self.model_feature_list,
        "compressed_features": feature_set.get_compressed_features(),
        "model_metrics_s3_path": self.model_training_root,
        "id_column": feature_set.id_column,
        "hyperparameters": kwargs.get("hyperparameters", {}),
    }

    # Caller-supplied template params (used by MetaEndpoint to inject DAG
    # config, region, s3_bucket; available for any framework that needs
    # extra placeholders beyond the standard set above).
    if self.custom_args:
        template_params.update(self.custom_args)

    # Custom Script
    if self.custom_script:
        script_path = self.custom_script
        if self.custom_script.endswith(".template"):
            # Model Type is an enumerated type, so we need to convert it to a string
            template_params["model_type"] = template_params["model_type"].value
            script_path = fill_template(self.custom_script, template_params, "generated_model_script.py")

        # Ensure training_harness.py is in the custom script's source directory
        source_dir = Path(script_path).parent
        harness_dst = source_dir / "training_harness.py"
        if not harness_dst.exists():
            harness_src = pkg_resources.files("workbench") / "training" / "training_harness.py"
            shutil.copy(str(harness_src), harness_dst)
            self.log.info(f"Copied training_harness.py into {source_dir}")

        self.log.info(f"Custom script path: {script_path}")

    # We're using one of the built-in model script templates
    else:
        # Generate our model script
        script_path = generate_model_script(template_params)

    # Metric Definitions for Regression (matches model script output format)
    if self.model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]:
        metric_definitions = [
            {"Name": "rmse", "Regex": r"rmse: ([0-9.]+)"},
            {"Name": "mae", "Regex": r"mae: ([0-9.]+)"},
            {"Name": "medae", "Regex": r"medae: ([0-9.]+)"},
            {"Name": "r2", "Regex": r"r2: ([0-9.-]+)"},
            {"Name": "spearmanr", "Regex": r"spearmanr: ([0-9.-]+)"},
            {"Name": "support", "Regex": r"support: ([0-9]+)"},
        ]

    # Metric Definitions for Classification
    elif self.model_type == ModelType.CLASSIFIER:
        # We need to get creative with the Classification Metrics
        # Note: Classification only supports single target
        class_target = target_list[0] if target_list else self.target_column

        # Grab all the target column class values (class labels)
        table = feature_set.data_source.table
        self.class_labels = feature_set.query(f'select DISTINCT {class_target} FROM "{table}"')[
            class_target
        ].to_list()

        # Sanity check on the targets
        if len(self.class_labels) > 10:
            msg = f"Too many target classes ({len(self.class_labels)}) for classification, aborting!"
            self.log.critical(msg)
            raise ValueError(msg)

        # Dynamically create the metric definitions (per-class precision/recall/f1/support)
        # Note: Confusion matrix metrics are skipped to stay under SageMaker's 40 metric limit
        metrics = ["precision", "recall", "f1", "support"]
        metric_definitions = []
        for t in self.class_labels:
            for m in metrics:
                metric_definitions.append({"Name": f"Metrics:{t}:{m}", "Regex": f"Metrics:{t}:{m} ([0-9.]+)"})

    # If the model type is UNKNOWN, our metric_definitions will be empty
    else:
        self.log.important(f"ModelType is {self.model_type}, skipping metric_definitions...")
        metric_definitions = []

    # Take the full script path and extract the entry point and source directory
    entry_point = str(Path(script_path).name)
    source_dir = str(Path(script_path).parent)

    # Create a Sagemaker Model with our script
    image = ModelImages.get_image_uri(self.sm_session.boto_region_name, self.training_image)

    # Use user-specified instance or pick a ladder from the workload.
    hpo_requested = (kwargs.get("hyperparameters") or {}).get("hpo") is not None
    train_instance_type = kwargs.get("training_instance")
    if train_instance_type:
        instance_ladder = [train_instance_type]
        self.log.important(f"Using user-specified instance {train_instance_type}")
    else:
        gpu_framework = self.model_framework in [ModelFramework.CHEMPROP, ModelFramework.PYTORCH]
        workload = training_workload(kwargs.get("hyperparameters"), gpu_framework=gpu_framework)
        instance_ladder = INSTANCE_LADDERS[workload]
        self.log.important(f"Using {workload} instances {instance_ladder} for {self.model_framework.value}")

    # Convert metric definitions to V3 MetricDefinition objects
    v3_metric_definitions = [MetricDefinition(name=m["Name"], regex=m["Regex"]) for m in metric_definitions]

    # PRM attribution tag: connects the training-job compute to our Marketplace listing
    prm_tag = Tag(key="aws-apn-id", value=f"pc:{AWS_MARKETPLACE_PRODUCT_CODE}")

    self.log.important(f"Training the Model {self.output_name} with Training Image {image}...")
    _suppress_sagemaker_logging()

    for rung, train_instance_type in enumerate(instance_ladder):
        last_rung = rung == len(instance_ladder) - 1

        # The final rung queues for capacity as long as it takes; the rungs above it
        # give up so the ladder can move on.
        stopping_condition = StoppingCondition(max_runtime_in_seconds=(24 if hpo_requested else 6) * 3600)
        if not last_rung:
            stopping_condition.max_pending_time_in_seconds = MAX_PENDING_SECONDS

        # Create ModelTrainer (V3 replacement for Estimator)
        # Use command= to run our entrypoint wrapper, which executes the model script
        # and then bundles inference code/metadata into the model artifacts
        self.model_trainer = ModelTrainer(
            training_image=image,
            # The image carries no site config, so hand the container the bucket it needs
            # (chemprop foundation weights, df_store, ...) explicitly.
            environment={"WORKBENCH_BUCKET": self.workbench_bucket},
            source_code=SourceCode(
                source_dir=source_dir,
                command=f"python training_harness.py {entry_point}",
            ),
            compute=Compute(instance_type=train_instance_type, instance_count=1),
            output_data_config=OutputDataConfig(s3_output_path=self.model_training_root, compression_type="GZIP"),
            stopping_condition=stopping_condition,
            base_job_name=self.output_name,
            role=self.workbench_role_arn,
            sagemaker_session=self.sm_session,
            tags=[prm_tag],
        )
        self.model_trainer.with_metric_definitions(v3_metric_definitions)
        input_data = self.model_trainer.create_input_data_channel("train", s3_training_path)

        # Submit the job and own the poll loop ourselves. The SDK's wait=True polls via a
        # global client singleton that bypasses our hardened session and never retries a
        # transient HttpTimeoutException on DescribeTrainingJob (see _wait_for_training_job).
        # Silence the SDK's stray "not displaying logs" warning that wait=False emits.
        logging.getLogger("sagemaker.train.model_trainer").setLevel(logging.ERROR)
        self.log.important(f"Submitting training job on {train_instance_type}...")
        try:
            self.model_trainer.train(input_data_config=[input_data], wait=False)
        except ClientError as e:
            # The account's quota for this type is already fully in use by another job
            if e.response.get("Error", {}).get("Code") == "ResourceLimitExceeded" and not last_rung:
                self.log.warning(f"{train_instance_type} quota exhausted, dropping to next instance...")
                continue
            raise

        # Capture the actual training job name (ModelTrainer appends a timestamp to base_job_name)
        self.training_job_name = self.model_trainer._latest_training_job.training_job_name
        try:
            self._wait_for_training_job(capacity_timeout=None if last_rung else CAPACITY_WAIT_SECONDS)
            break
        except CapacityTimeout:
            if last_rung:
                raise
            self.log.warning(f"Dropping from {train_instance_type} to {instance_ladder[rung + 1]}...")

    # Now delete the training data
    self.log.info(f"Deleting training data {s3_training_path}...")
    wr.s3.delete_objects(
        [s3_training_path, s3_training_path.replace(".csv", ".csv.metadata")],
        boto3_session=self.boto3_session,
    )

    # Create Model and officially Register
    self.log.important(f"Creating new model {self.output_name}...")
    self.create_and_register_model(**kwargs)

training_workload(hyperparameters, *, gpu_framework)

Which :data:INSTANCE_LADDERS entry a training job should ask for.

Only a parallel search needs a multi-GPU box; a serial one (optuna, or an explicit max_parallel=1) stays on the normal single-GPU instance. The job derives its own concurrency from the cards it lands on, so an unset max_parallel means "pack the box" and has to be given a box worth packing.

Parameters:

Name	Type	Description	Default
hyperparameters	Union[dict, None]	the model's hyperparameters, or None.	required
gpu_framework	bool	whether the framework trains on a GPU.	required

Returns:

Name	Type	Description
str	str	a key into :data:INSTANCE_LADDERS.

Source code in src/workbench/core/transforms/features_to_model/features_to_model.py

def training_workload(hyperparameters: Union[dict, None], *, gpu_framework: bool) -> str:
    """Which :data:`INSTANCE_LADDERS` entry a training job should ask for.

    Only a *parallel* search needs a multi-GPU box; a serial one (optuna, or an explicit
    ``max_parallel=1``) stays on the normal single-GPU instance. The job derives its own
    concurrency from the cards it lands on, so an unset ``max_parallel`` means "pack the
    box" and has to be given a box worth packing.

    Args:
        hyperparameters: the model's hyperparameters, or None.
        gpu_framework: whether the framework trains on a GPU.

    Returns:
        str: a key into :data:`INSTANCE_LADDERS`.
    """
    hpo = (hyperparameters or {}).get("hpo")
    if hpo is None:  # an empty block is a real request: a search on every default
        return "gpu" if gpu_framework else "cpu"
    if not gpu_framework:
        return "cpu_hpo"
    serial = hpo.get("backend", "auto") == "optuna" or hpo.get("max_parallel") == 1
    return "gpu" if serial else "gpu_parallel_hpo"

Supported Models

Currently Workbench supports XGBoost (classifier/regressor), and Scikit Learn models. Those models can be created by just specifying different parameters to the FeaturesToModel class. The main issue with the supported models is they are vanilla versions with default parameters, any customization should be done with Custom Models

XGBoost

from workbench.core.transforms.features_to_model.features_to_model import FeaturesToModel

# XGBoost Regression Model
input_name = "abalone_features"
output_name = "abalone-regression"
to_model = FeaturesToModel(input_name, output_name, model_type=ModelType.REGRESSOR)
to_model.set_output_tags(["abalone", "public"])
to_model.transform(target_column="class_number_of_rings", description="Abalone Regression")

# XGBoost Classification Model
input_name = "wine_features"
output_name = "wine-classification"
to_model = FeaturesToModel(input_name, output_name, ModelType.CLASSIFIER)
to_model.set_output_tags(["wine", "public"])
to_model.transform(target_column="wine_class", description="Wine Classification")

# Quantile Regression Model (Abalone)
input_name = "abalone_features"
output_name = "abalone-quantile-reg"
to_model = FeaturesToModel(input_name, output_name, ModelType.UQ_REGRESSOR)
to_model.set_output_tags(["abalone", "quantiles"])
to_model.transform(target_column="class_number_of_rings", description="Abalone Quantile Regression")

Scikit-Learn

from workbench.core.transforms.features_to_model.features_to_model import FeaturesToModel

# Scikit-Learn Kmeans Clustering Model
input_name = "wine_features"
output_name = "wine-clusters"
to_model = FeaturesToModel(
    input_name,
    output_name,
    model_class="KMeans",  # Clustering algorithm
    model_import_str="from sklearn.cluster import KMeans",  # Import statement for KMeans
    model_type=ModelType.CLUSTERER,
)
to_model.set_output_tags(["wine", "clustering"])
to_model.transform(target_column=None, description="Wine Clustering")

# Scikit-Learn HDBSCAN Clustering Model
input_name = "wine_features"
output_name = "wine-clusters-hdbscan"
to_model = FeaturesToModel(
    input_name,
    output_name,
    model_class="HDBSCAN",  # Density-based clustering algorithm
    model_import_str="from sklearn.cluster import HDBSCAN",
    model_type=ModelType.CLUSTERER,
)
to_model.set_output_tags(["wine", "density-based clustering"])
to_model.transform(target_column=None, description="Wine Clustering with HDBSCAN")

# Scikit-Learn 2D Projection Model using UMAP
input_name = "wine_features"
output_name = "wine-2d-projection"
to_model = FeaturesToModel(
    input_name,
    output_name,
    model_class="UMAP",
    model_import_str="from umap import UMAP",
    model_type=ModelType.PROJECTION,
)
to_model.set_output_tags(["wine", "2d-projection"])
to_model.transform(target_column=None, description="Wine 2D Projection")

Custom Models

For custom models we recommend the following steps:

Experimental

The Workbench Custom Models are currently in experimental mode so have fun but expect issues. Requires workbench >= 0.8.60. Feel free to submit issues to Workbench Github

• Copy the example custom model script into your own directory
  • See: Custom Model Script
• Make a requirements.txt and put into the same directory
• Train/deploy the ^existing^ example
  • This is an important step, don't skip it
  • If the existing model script trains/deploys your in great shape for the next step, if it doesn't then now is a good time to debug AWS account/permissions/etc.
• Now customize the model script
• Train/deploy your custom script

Training/Deploying Custom Models

from workbench.api import ModelType
from workbench.core.transforms.features_to_model.features_to_model import FeaturesToModel

# Note this directory should also have a requirements.txt in it
my_custom_script = "/full/path/to/my/directory/my_custom_script.py"
input_name = "wine_features"    # FeatureSet you want to use
output_name = "my-custom-model" # change to whatever
target_column = "wine-class"    # change to whatever
to_model = FeaturesToModel(input_name, output_name,
                           model_type=ModelType.CLASSIFIER, 
                           custom_script=my_custom_script)
to_model.set_output_tags(["your", "tags"])
to_model.transform(target_column=target_column, description="Custom Model")

Custom Models: Create an Endpoint/Run Inference

from workbench.api import Model, Endpoint

model = Model("my-custom-model")
end = model.to_endpoint()   # Note: This takes a while

# Now run inference on my custom model :)
end.test_inference()

# Run inference with my own dataframe
df = fs.pull_dataframe()  # Or whatever dataframe
end.inference(df)

Questions?

The SuperCowPowers team is happy to answer any questions you may have about AWS and Workbench. Please contact us at workbench@supercowpowers.com or on chat us up on Discord