azureml_featurestore_init_and_featureset_registration_quickstart.py

1from azure.ai.ml import MLClient
2from azure.ai.ml.entities import (
3    FeatureStore,
4    FeatureSet,
5    FeatureSetSpecification,
6)
7from azure.identity import DefaultAzureCredential
8
9# Basic configuration
10subscription_id = "<SUBSCRIPTION_ID>"
11resource_group = "<RESOURCE_GROUP>"
12location = "<LOCATION>"
13featurestore_name = "<FEATURESTORE_NAME>"
14
15# 1. Initialize the MLClient for Feature Store operations
16ml_client = MLClient(
17    DefaultAzureCredential(), 
18    subscription_id, 
19    resource_group
20)
21
22# 2. Define and create a Feature Store (if it doesn't exist)
23fs_entity = FeatureStore(
24    name=featurestore_name, 
25    location=location,
26    description="Quickstart feature store"
27)
28
29# Note: In a real scenario, this is often pre-provisioned via CLI or UI
30# fs_poller = ml_client.feature_stores.begin_create(fs_entity)
31# fs_poller.result()
32
33# 3. Reference an existing Feature Store client
34featurestore_client = MLClient(
35    DefaultAzureCredential(),
36    subscription_id,
37    resource_group,
38    featurestore_name,
39)
40
41# 4. Define a Feature Set Specification (pointing to your local source code/transformation)
42# This assumes you have a directory './feature_set_src' containing 'spec.yaml' and transformation code
43feature_set_spec = FeatureSetSpecification(path="./feature_set_src")
44
45# 5. Define the Feature Set metadata
46feature_set = FeatureSet(
47    name="transactions_featureset",
48    version="1.0.0",
49    description="Feature set for transaction data",
50    specification=feature_set_spec
51)
52
53# 6. Register the Feature Set to the Store
54feature_set_poller = featurestore_client.feature_sets.begin_create_or_update(feature_set)
55print(f"Feature set registration started: {feature_set_poller.result().name}")
56
57# 7. List features in the registered feature set
58features = featurestore_client.features.list(feature_set_name="transactions_featureset", version="1.0.0")
59for feature in features:
60    print(f"Feature name: {feature.name}")