Skip to content

feature/feature-store-telemetry #5557

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nargokul merged 1 commit into from
Feb 18, 2026
Merged

feature/feature-store-telemetry #5557

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions src/sagemaker/feature_store/feature_group.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,8 @@
TargetStoreEnum,
)
from sagemaker.utils import resolve_value_from_config, format_tags, Tags
from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
from sagemaker.telemetry.constants import Feature

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -672,6 +674,7 @@ class FeatureGroup:
FeatureTypeEnum.STRING.value: "STRING",
}

@_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.create")
def create(
self,
s3_uri: Union[str, bool],
Expand Down Expand Up @@ -788,10 +791,12 @@ def create(

return self.sagemaker_session.create_feature_group(**create_feature_store_args)

@_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.delete")
def delete(self):
"""Delete a FeatureGroup."""
self.sagemaker_session.delete_feature_group(feature_group_name=self.name)

@_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.describe")
def describe(self, next_token: str = None) -> Dict[str, Any]:
"""Describe a FeatureGroup.

Expand All @@ -805,6 +810,7 @@ def describe(self, next_token: str = None) -> Dict[str, Any]:
feature_group_name=self.name, next_token=next_token
)

@_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.update")
def update(
self,
feature_additions: Sequence[FeatureDefinition] = None,
Expand Down Expand Up @@ -843,6 +849,9 @@ def update(
throughput_config=throughput_config_parameter,
)

@_telemetry_emitter(
feature=Feature.FEATURE_STORE_V2, func_name="feature_group.update_feature_metadata"
)
def update_feature_metadata(
self,
feature_name: str,
Expand Down Expand Up @@ -871,6 +880,9 @@ def update_feature_metadata(
parameter_removals=(parameter_removals or []),
)

@_telemetry_emitter(
feature=Feature.FEATURE_STORE_V2, func_name="feature_group.describe_feature_metadata"
)
def describe_feature_metadata(self, feature_name: str) -> Dict[str, Any]:
"""Describe feature metadata by feature name.

Expand Down Expand Up @@ -1038,6 +1050,7 @@ def load_feature_definitions(
self.feature_definitions = feature_definitions
return self.feature_definitions

@_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.get_record")
def get_record(
self,
record_identifier_value_as_string: str,
Expand All @@ -1057,6 +1070,7 @@ def get_record(
feature_names=feature_names,
).get("Record")

@_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.put_record")
def put_record(
self,
record: Sequence[FeatureValue],
Expand All @@ -1080,6 +1094,7 @@ def put_record(
ttl_duration=ttl_duration.to_dict() if ttl_duration is not None else None,
)

@_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.delete_record")
def delete_record(
self,
record_identifier_value_as_string: str,
Expand Down
9 changes: 9 additions & 0 deletions src/sagemaker/feature_store/feature_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,8 @@
SortOrderEnum,
Identifier,
)
from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
from sagemaker.telemetry.constants import Feature


@attr.s
Expand All @@ -47,6 +49,7 @@ class FeatureStore:

sagemaker_session: Session = attr.ib(default=Session)

@_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_store.create_dataset")
def create_dataset(
self,
base: Union[FeatureGroup, pd.DataFrame],
Expand Down Expand Up @@ -92,6 +95,9 @@ def create_dataset(
kms_key_id,
)

@_telemetry_emitter(
feature=Feature.FEATURE_STORE_V2, func_name="feature_store.list_feature_groups"
)
def list_feature_groups(
self,
name_contains: str = None,
Expand Down Expand Up @@ -137,6 +143,9 @@ def list_feature_groups(
next_token=next_token,
)

@_telemetry_emitter(
feature=Feature.FEATURE_STORE_V2, func_name="feature_store.batch_get_record"
)
def batch_get_record(
self,
identifiers: Sequence[Identifier],
Expand Down
1 change: 1 addition & 0 deletions src/sagemaker/telemetry/constants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ class Feature(Enum):
# Note: HyperPod CLI uses codes 6 and 7
JUMPSTART_V2 = 8 # Added to support JumpStart telemetry
MLOPS_V2 = 9 # Added to support MLOps telemetry
FEATURE_STORE_V2 = 10

def __str__(self): # pylint: disable=E0307
"""Return the feature name."""
Expand Down
1 change: 1 addition & 0 deletions src/sagemaker/telemetry/telemetry_logging.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
# Note: HyperPod CLI uses codes 6 and 7
str(Feature.JUMPSTART_V2): 8,
str(Feature.MLOPS_V2): 9,
str(Feature.FEATURE_STORE_V2): 10,
}

STATUS_TO_CODE = {
Expand Down
Loading