sagemaker

@staticmethod
def fetch_sm_hosting_od_price(instance_type: str, region: str) -> float:
    """Look up USD hourly rates for on-demand SM hosting instances from the AWS Price List API

    This function assumes:

    1. You're using standard "on-demand" pricing - no savings plans or private pricing
    2. No free tier or volume discounts are applicable to this usage
    3. Your pricing is provided in USD

    Args:
        instance_type: Amazon SageMaker instance type, e.g. 'ml.g5.4xlarge'
        region: AWS region where the endpoint is running, e.g. 'us-east-1'

    Returns:
        price_per_hour: The standard, on-demand hourly price for the given instance type in USD
    """
    unit_price, _ = _fetch_single_product_ondemand_unit_price(
        service_code="AmazonSageMaker",
        attribute_values={"component": "Hosting", "instanceName": instance_type},
        region=region,
        dim_unit="Hrs",
    )
    return unit_price

@classmethod
def from_endpoint(
    cls, endpoint_name: str, region: str | None = None
) -> dict[str, SageMakerRTEndpointCompute]:
    """Configure SageMakerRTEndpointCompute dimension(s) from an existing SageMaker endpoint

    NOTE: You'll need IAM permissions to `sagemaker:DescribeEndpoint` and
    `sagemaker:DescribeEndpointConfig` to use this method.

    This function returns a dictionary rather than a single `SageMakerRTEndpointCompute`,
    because different "variants" deployed behind an endpoint may be backed by clusters of
    different instance types, and therefore need separate dimensions.

    Instance counts will be retrieved at the point in time this method is called, so watch out
    if you have auto-scaling enabled on your endpoint.

    Args:
        endpoint_name: Name of the SageMaker endpoint
        region: AWS region where the endpoint is running (default: current region)

    Returns:
        run_dims: A dictionary containing one or more dimensions, to pass to your
            `CostModel(run_dims=...)`. If the endpoint has only one "Variant", the returned
            dict will have a single key (cost dimension name) `SageMakerRTEndpointCompute`.
            Otherwise, keys will be generated as `{variant_name}SageMakerRTEndpointCompute`
    """
    sagemaker = boto3.client("sagemaker", region_name=region)
    endpoint_desc = sagemaker.describe_endpoint(EndpointName=endpoint_name)
    ep_vars: dict[str, dict] = {
        e["VariantName"]: e for e in endpoint_desc["ProductionVariants"]
    }
    endpoint_config = sagemaker.describe_endpoint_config(
        EndpointConfigName=endpoint_desc["EndpointConfigName"]
    )
    cfg_vars: dict[str, dict] = {
        e["VariantName"]: e for e in endpoint_config["ProductionVariants"]
    }

    if len(ep_vars) == 1:
        variant = next(iter(ep_vars.values()))
        variant_config = cfg_vars[variant["VariantName"]]
        return {
            cls.__name__: cls(
                instance_type=variant_config["InstanceType"],
                instance_count=variant["CurrentInstanceCount"],
                region=sagemaker.meta.region_name,
            )
        }
    else:
        return {
            f"{var_name}_{cls.__name__}": cls(
                instance_type=cfg_vars[var_name]["InstanceType"],
                instance_count=ep_vars[var_name]["CurrentInstanceCount"],
                region=sagemaker.meta.region_name,
            )
            for var_name in ep_vars.keys()
        }

@staticmethod
def fetch_sm_hosting_ebs_price(region: str) -> float:
    """Look up hourly USD rate for SageMaker hosting EBS storage from the AWS Price List API

    The API actually lists this rate as monthly, so we take an assumption of 30days * 24hrs to
    convert.

    Args:
        region: AWS region where the endpoint is running, e.g. 'us-east-1'

    Returns:
        price_per_gb_hour: The standard, on-demand price per GB-hour in USD
    """
    unit_price, _ = _fetch_single_product_ondemand_unit_price(
        service_code="AmazonSageMaker",
        attribute_values={"volumeType": "General Purpose-Hosting"},
        region=region,
        dim_unit="GB-Mo",
    )

    return unit_price / (30 * 24)  # Convert GB-mo to GB-hr

@classmethod
def from_endpoint(
    cls,
    endpoint_name: str,
    region: str | None = None,
    merge_variants: bool = False,
) -> dict[str, SageMakerRTEndpointStorage]:
    """Configure SageMakerRTEndpointStorage dimension(s) from an existing SageMaker endpoint

    NOTE: You'll need IAM permissions to `sagemaker:DescribeEndpoint` and
    `sagemaker:DescribeEndpointConfig` to use this method.

    This function returns a dictionary rather than a single `SageMakerRTEndpointStorage`,
    because different "variants" deployed behind an endpoint may be reported separately, if
    more than one are present.

    Instance counts will be retrieved at the point in time this method is called, so watch out
    if you have auto-scaling enabled on your endpoint.

    Args:
        endpoint_name: Name of the SageMaker endpoint
        region: AWS region where the endpoint is running (default: current region)
        merge_variants: Set `True` to merge multiple "variants" into a single storage dimension
            (since the hourly rate is constant across instance types). By default, if multiple
            variants are configured these will be reported as separate cost dimensions.

    Returns:
        run_dims: A dictionary containing zero or more dimensions, to pass to your
            `CostModel(run_dims=...)`. If the endpoint has only one "Variant", the returned
            dict will have a single key (cost dimension name) `SageMakerRTEndpointCompute`.
            Otherwise, keys will be generated as `{variant_name}SageMakerRTEndpointCompute`.
            Any variants with 0 EBS storage configured will be omitted from the result.
    """
    sagemaker = boto3.client("sagemaker", region_name=region)
    region_final = sagemaker.meta.region_name
    endpoint_desc = sagemaker.describe_endpoint(EndpointName=endpoint_name)
    ep_vars: dict[str, dict] = {
        e["VariantName"]: e for e in endpoint_desc["ProductionVariants"]
    }
    endpoint_config = sagemaker.describe_endpoint_config(
        EndpointConfigName=endpoint_desc["EndpointConfigName"]
    )
    cfg_vars: dict[str, dict] = {
        e["VariantName"]: e for e in endpoint_config["ProductionVariants"]
    }

    # variant.VolumeSizeInGB seems to be set automatically for EBS-enabled instance types, but
    # missing for non-EBS (storage-included) instance types.
    vars_with_ebs = [
        name for name, cfg in cfg_vars.items() if cfg.get("VolumeSizeInGB")
    ]

    if len(vars_with_ebs) == 0:
        return {}
    elif len(ep_vars) == 1:
        variant = next(iter(ep_vars.values()))
        variant_config = cfg_vars[variant["VariantName"]]
        return {
            cls.__name__: cls(
                gbs_provisioned=variant_config["VolumeSizeInGB"]
                * variant["CurrentInstanceCount"],
                region=region_final,
            )
        }
    elif merge_variants:
        total_gbs = sum(
            [
                cfg_vars[var_name].get("VolumeSizeInGB", 0)
                * ep_vars[var_name]["CurrentInstanceCount"]
                for var_name in ep_vars.keys()
            ]
        )
        return {cls.__name__: cls(gbs_provisioned=total_gbs, region=region_final)}
    else:
        return {
            f"{var_name}_{cls.__name__}": cls(
                gbs_provisioned=cfg_vars[var_name]["VolumeSizeInGB"]
                * ep_vars[var_name]["CurrentInstanceCount"],
                region=region_final,
            )
            for var_name in ep_vars.keys()
            if var_name in vars_with_ebs
        }