azureml_feature_store_entity_and_featureset_registration.py

1import os
2from azure.ai.ml import MLClient
3from azure.ai.ml.entities import (
4    FeatureStore,
5    FeatureSet,
6    FeatureSetSpecification,
7    Entity,
8)
9from azure.identity import DefaultAzureCredential
10
11# 1. Connect to the Azure ML Feature Store
12subscription_id = "<SUBSCRIPTION_ID>"
13resource_group = "<RESOURCE_GROUP>"
14featurestore_name = "<FEATURESTORE_NAME>"
15
16ml_client = MLClient(
17    DefaultAzureCredential(),
18    subscription_id,
19    resource_group,
20    featurestore_name
21)
22
23# 2. Define an Entity (The logical grouping for features, e.g., 'account')
24account_entity_name = "account"
25account_entity = Entity(
26    name=account_entity_name,
27    version="1",
28    index_columns=[{"name": "accountID", "type": "String"}],
29    description="Account entity for banking features",
30)
31
32ml_client.entities.begin_create_or_update(account_entity)
33
34# 3. Define the Feature Set Specification
35# This assumes you have a 'feature_set_src' folder with 'spec.yaml' and transformation code
36feature_set_spec = FeatureSetSpecification(
37    path="./feature_set_src",
38    read_config=None # Defined in spec.yaml
39)
40
41# 4. Register the Feature Set
42feature_set_name = "transactions_featureset"
43feature_set = FeatureSet(
44    name=feature_set_name,
45    version="1",
46    specification=feature_set_spec,
47    entities=[f"azureml:{account_entity_name}:1"],
48)
49
50poller = ml_client.feature_sets.begin_create_or_update(feature_set)
51print(f"Feature set registration started: {poller.result().name}")
52
53# 5. List Feature Sets to verify
54feature_sets = ml_client.feature_sets.list(name=feature_set_name)
55for fs in feature_sets:
56    print(f"Registered Feature Set: {fs.name}, Version: {fs.version}")