py-sdk

BharatML Stack Python SDKs

A collection of independent Python packages for interacting with BharatMLStack components. 🚀

📦 Separated Packages

The BharatML Stack Python SDK has been separated into 3 independent packages for better modularity and focused dependencies:

Package	Purpose	PyPI
bharatml_commons	Common utilities and protobuf definitions
spark_feature_push_client	Apache Spark-based data pipeline client
grpc_feature_client	High-performance gRPC client for real-time operations

Key Features

• Feature metadata retrieval
• Protobuf serialization of feature values and produce to Apache Kafka
• Support for features of different various data types:
  • Scalar types (FP32, FP64, Int32, Int64, UInt32, UInt64, String, Bool)
  • Vector types (Vectors of each of the above Scalar Types)
• Kafka integration with configurable settings

🚀 Quick Installation

Install all packages:

pip install bharatml_commons spark_feature_push_client grpc_feature_client

Install only what you need:

# For Spark-based data pipelines
pip install bharatml_commons spark_feature_push_client
# For real-time gRPC operations
pip install bharatml_commons grpc_feature_client
# For common utilities only
pip install bharatml_commons

📖 Documentation

• bharatml_commons README: Common utilities and protobuf definitions
• spark_feature_push_client README: Spark-based data pipeline client
• grpc_feature_client README: High-performance gRPC client

Prerequisites

• Python 3.7+ (tested on Python 3.7, 3.8, 3.9, 3.10, 3.11, 3.12)
• (Optional) Apache Spark 3.0+ & spark-sql-kafka for Spark-based functionality

Quick Start

1. bharatml_commons

Foundation package with shared utilities, protobuf definitions, and base classes.

frombharatml_commonsimportFeatureMetadataClient, clean_column_namefrombharatml_commons.proto.persist.persist_pb2importQuery# HTTP client for metadata operationsclient=FeatureMetadataClient(url, job_id, token)
metadata=client.get_feature_metadata(["user_features"])
# Utility functionsclean_name=clean_column_name("feature@name#1")

2. spark_feature_push_client

Apache Spark client for batch data pipelines - reading from data sources and pushing to Kafka.

fromspark_feature_push_clientimportOnlineFeatureStorePyClient# Initialize client for data pipelineclient=OnlineFeatureStorePyClient(metadata_url, job_id, job_token)
# Process Spark DataFrame → Protobuf → Kafkaproto_df=client.generate_df_with_protobuf_messages(spark_df)
client.write_protobuf_df_to_kafka(proto_df, kafka_servers, topic)

3. grpc_feature_client

High-performance gRPC client for real-time feature operations with direct API access.

fromgrpc_feature_clientimportGRPCFeatureClient, GRPCClientConfig# Configure for real-time operationsconfig=GRPCClientConfig(server_address, job_id, job_token)
client=GRPCFeatureClient(config)
# Direct API operationsresult=client.persist_features(entity_label, keys_schema, feature_groups, data)
features=client.retrieve_decoded_features(entity_label, feature_groups, keys, entity_keys)

Usage

Spark Feature Push Client Usage

fromspark_feature_push_clientimportOnlineFeatureStorePyClient# Initialize the clientclient=OnlineFeatureStorePyClient(
features_metadata_source_url="your_features_metadata_source_url",
job_id="your_job_id",
job_token="your_job_token"
)
# Get feature details
(
offline_src_type_columns,
offline_col_to_default_values_map,
entity_column_names
) =client.get_features_details()

Push Feature Values from Offline sources to online-feature-store via Spark -> Kafka

Supported Offline Sources

1. Table (Hive/Delta)
2. Parquet folder stored in Cloud Storage (AWS/GCS/ADLS)
3. Delta folder stored in Cloud Storage (AWS/GCS/ADLS)

Refer to the examples for detailed examples of how to configure a job and push the feature values

Following is a simple flow / outline of the steps involved in above example:

# create a new onlineFeatureStore clientopy_client=OnlineFeatureStorePyClient(features_metadata_source_url, job_id, job_token) # get the features detailsfeature_mapping, offline_col_to_default_values_map, onfs_fg_to_onfs_feat_map, onfs_fg_to_ofs_feat_map, fg_to_datatype_map, entity_label, entity_column_names=opy_client.get_features_details(fgs_to_consider)
# read the data from different sourcesdf=get_features_from_all_sources(spark, entity_column_names, feature_mapping, offline_col_to_default_values_map)
# serialize of protobuf binaryproto_df=opy_client.generate_df_with_protobuf_messages(df, intra_batch_size=20) # Produce data to kafka so that consumers write features to Online Feature Storeopy_client.write_protobuf_df_to_kafka(proto_df, kafka_bootstrap_servers, kafka_topic, additional_options)

🏗️ SDK Architecture

The multi-SDK architecture provides:

py-sdk/
├── src/
│ ├── spark_feature_push_client/ # Spark-based data pipeline
│ │ ├── utils/helpers.py # Spark-specific utilities
│ │ ├── __init__.py
│ │ └── client.py # Batch ETL operations
│ ├── grpc_feature_client/ # gRPC real-time operations
│ │ ├── config.py # gRPC configuration
│ │ ├── client.py # Real-time API operations
│ │ ├── README.md # gRPC documentation
│ │ └── __init__.py
│ ├── bharatml_common/ # Shared utilities & protobuf
│ │ ├── proto/ # ✅ Protobuf definitions
│ │ │ ├── persist.proto # Persist operation schema
│ │ │ ├── retrieve.proto # Retrieve operation schema
│ │ │ ├── persist/persist_pb2.py # Generated Python files
│ │ │ ├── retrieve/retrieve_pb2.py
│ │ │ └── generate_proto.py # Code generation script
│ │ ├── http_client.py # HTTP client utilities
│ │ ├── feature_metadata_client.py # ✅ Feature metadata client
│ │ ├── column_utils.py # Column processing
│ │ ├── feature_utils.py # Feature processing
│ │ ├── sdk_template.py # Template for new SDKs
│ │ └── __init__.py
├── README.md
└── pyproject.toml

🛠️ Creating New SDKs

To create a new SDK in this project:

1. Create the SDK directory structure:

   src/your_new_sdk/
   ├── __init__.py # Main exports
   ├── client.py # Main client class
   ├── config.py # Configuration classes
   └── utils/ # SDK-specific utilities
   

2. Use shared utilities:

   frombharatml_common.http_clientimportBharatMLHTTPClientfrombharatml_common.sdk_templateimportBaseSDKClient

3. Update pyproject.toml:

   [tool.hatch.build.targets.wheel]
   packages = [
   "src/spark_feature_push_client",
   "src/bharatml_common",
   "src/your_new_sdk"# Add your new SDK
   ]

Development

Setting up Development Environment

# Clone the repository
git clone https://github.com/Meesho/BharatMLStack.git
cd BharatMLStack/py-sdk
# Install in development mode
pip install -e .# Install development dependencies
pip install build pytest flake8 black isort mypy

Running Tests

# Run all tests
pytest tests/ -v
# Run tests with coverage
pytest tests/ --cov=src --cov-report=html
# Run specific test file
pytest tests/test_client.py -v

Code Quality

# Format code with black
black src/
# Sort imports with isort
isort src/
# Lint with flake8
flake8 src/
# Type checking with mypy
mypy src/ --ignore-missing-imports

Building the Package

# Build the package
python -m build
# Check package metadata
pip install twine
twine check dist/*

Code Style Guidelines

• Follow PEP 8 style guidelines
• Use Black for code formatting
• Use isort for import sorting
• Add type hints where possible
• Write docstrings for public functions and classes
• Keep line length to 88 characters (Black default)

Complete Architecture Example

Here's how the Spark and gRPC clients work together in a complete ML feature pipeline:

# 1. BATCH PIPELINE (Daily ETL Job)fromspark_feature_push_clientimportOnlineFeatureStorePyClient# Process batch data with Sparkspark_client=OnlineFeatureStorePyClient(
features_metadata_source_url="https://metadata.example.com",
job_id="daily-batch-etl",
job_token="pipeline-token"
)
# Read from data warehouse, transform, and push to Kafkafeature_details=spark_client.get_features_details()
historical_df=spark.sql("SELECT * FROM feature_warehouse.user_features")
proto_df=spark_client.generate_df_with_protobuf_messages(historical_df)
spark_client.write_protobuf_df_to_kafka(proto_df, kafka_brokers, "features.batch")
# 2. REAL-TIME SERVICE (Model Inference API)fromgrpc_feature_clientimportGRPCFeatureClient, GRPCClientConfig# Configure gRPC client for real-time operationsgrpc_config=GRPCClientConfig(
server_address="feature-store.example.com:50051",
job_id="predator",
job_token="api-token"
)
grpc_client=GRPCFeatureClient(grpc_config)
# Persist real-time features from user interactionsgrpc_client.persist_features(
entity_label="user_interaction",
keys_schema=["user_id", "session_id"],
feature_group_schemas=[{"label": "realtime_features", "feature_labels": ["click_count", "page_views"]}],
data_rows=[{"user_id": "u123", "session_id": "s456", "click_count": 5, "page_views": 3}]
)
# Retrieve features for ML model inferencefeatures=grpc_client.retrieve_decoded_features(
entity_label="user_interaction", feature_groups=[{"label": "user_features", "feature_labels": ["age", "location"]}],
keys_schema=["user_id"],
entity_keys=[["u123"], ["u456"]]
)
# Use features in ML modelmodel_input=prepare_features(features)
prediction=ml_model.predict(model_input)
# 3. FEATURE METADATA CLIENT (For REST API access)frombharatml_commonimportFeatureMetadataClient# Use metadata client for feature metadata operationsmetadata_client=FeatureMetadataClient("https://api.example.com", "http-job", "http-token")
metadata=metadata_client.get_feature_metadata(["user_features"])
health=metadata_client.health_check()

🎯 When to Use Which Package

Use Case	Package	Why
Daily ETL Jobs	spark_feature_push_client	Distributed processing, handles large datasets efficiently
Historical Backfill	spark_feature_push_client	Batch processing from data warehouses/lakes
Real-time Inference	grpc_feature_client	Low latency, direct API access
Feature Store Updates	grpc_feature_client	Direct persist/retrieve operations
Model Training	spark_feature_push_client	Process training datasets at scale
Model Serving	grpc_feature_client	Real-time feature retrieval for predictions
Metadata Operations	bharatml_commons	HTTP-based metadata queries
Utility Functions	bharatml_commons	Column cleaning, feature processing

📖 Documentation

• bharatml_commons README: Common utilities and protobuf definitions
• spark_feature_push_client README: Spark-based data pipeline client
• grpc_feature_client README: High-performance gRPC client
• Migration Guide: Migrating from the old unified package

🤝 Contributing

We welcome contributions! Please see our Contributing Guide for details.

📄 License

Licensed under the BharatMLStack Business Source License 1.1. See LICENSE for details.

🔗 Links

• BharatML Stack: Main repository
• PyPI Packages: All published packages
• Issues: Bug reports and feature requests
• Discussions: Community discussions

---

Built with ❤️ for the ML community from Meesho

If you find this useful, ⭐️ the repo — your support means the world to us!