orient_express 0.2.2__py3-none-any.whl

orient_express/__init__.py

@@ -0,0 +1 @@
+from .model_wrapper import ModelExpress

orient_express/model_wrapper.py

@@ -0,0 +1,197 @@
+import os
+import joblib
+import logging
+import pandas as pd
+from google.cloud import storage
+from google.cloud import aiplatform
+
+
+class ModelExpress:
+    def __init__(
+        self,
+        model_name,
+        project_name,
+        bucket_name,
+        model_version=None,
+        model=None,
+        region="us-central1",
+        serialized_model_path="model.joblib",
+        docker_image_uri="us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-7:latest",
+        endpoint_name=None,
+        machine_type="n1-standard-4",
+        min_replica_count=1,
+        max_replica_count=1,
+    ):
+        self.model = model
+        self.model_name = model_name
+        self.model_version = model_version
+        self.region = region
+        self.project_name = project_name
+        self.bucket_name = bucket_name
+        self.serialized_model_path = serialized_model_path
+        self.docker_image_uri = docker_image_uri
+        self.machine_type = machine_type
+        self.min_replica_count = min_replica_count
+        self.max_replica_count = max_replica_count
+        self.endpoint = None
+
+        if not endpoint_name:
+            self.endpoint_name = f"model-xpress-{model_name}"
+        else:
+            self.endpoint_name = endpoint_name
+
+    def colab_auth(self):
+        from google.colab import auth
+
+        auth.authenticate_user()
+
+    def _vertex_init(self):
+        aiplatform.init(project=self.project_name, location=self.region)
+
+    def get_latest_vertex_model(self, model_name):
+        self._vertex_init()
+
+        # Search for models with the specified display name
+        models = aiplatform.Model.list(filter=f"display_name={model_name}")
+
+        if not models:
+            return None  # Return None if no model with the given name exists
+
+        # Sort models by update time in descending order to get the latest version
+        latest_model = sorted(models, key=lambda x: x.update_time, reverse=True)[0]
+        return latest_model
+
+    # upload model to vertex ai model registry
+    def upload(self):
+        joblib.dump(self.model, self.serialized_model_path)
+        logging.info(f"Model saved to {self.serialized_model_path}")
+
+        # Initialize the GCS client
+        client = storage.Client()
+        bucket = client.bucket(self.bucket_name)
+
+        last_model = self.get_latest_vertex_model(self.model_name)
+        if last_model:
+            last_version = last_model.version_id
+            new_version = int(last_version) + 1
+        else:
+            new_version = 1
+
+        # Upload the model file
+        blob = bucket.blob(
+            self.get_artifacts_path(new_version, self.serialized_model_path)
+        )
+        blob.upload_from_filename(self.serialized_model_path)
+
+        return self.create_model_version(new_version, last_model)
+
+    def get_artifacts_path(self, version, file_name=None):
+        dir_name = f"models/{self.model_name}/{version}"
+        if file_name:
+            return f"{dir_name}/{file_name}"
+
+        return f"{dir_name}/"
+
+    def create_model_version(self, version_number, last_version):
+        artifact_uri = (
+            f"gs://{self.bucket_name}/{self.get_artifacts_path(version_number)}"
+        )
+
+        if last_version:
+            parent_model = f"projects/{self.project_name}/locations/{self.region}/models/{last_version.name}"
+        else:
+            parent_model = None
+
+        release = aiplatform.Model.upload(
+            display_name=self.model_name,
+            artifact_uri=artifact_uri,
+            parent_model=parent_model,
+            serving_container_image_uri=self.docker_image_uri,
+            sync=True,
+        )
+
+        return release
+
+    def get_or_create_endpoint(self):
+        endpoint = self.get_endpoint()
+        if endpoint:
+            return endpoint
+        else:
+            return self.create_endpoint()
+
+    def get_endpoint(self):
+        endpoints = aiplatform.Endpoint.list(
+            filter=f"display_name={self.endpoint_name}", order_by="create_time"
+        )
+        if endpoints:
+            return endpoints[0]
+
+    def create_endpoint(self):
+        return aiplatform.Endpoint.create(
+            display_name=self.endpoint_name,
+            project=self.project_name,
+            location=self.region,
+        )
+
+    def deploy(self):
+        self._vertex_init()
+
+        endpoint = self.get_or_create_endpoint()
+        model_version = self.upload()
+        model_version.deploy(
+            endpoint=endpoint,
+            machine_type=self.machine_type,
+            min_replica_count=self.min_replica_count,
+            max_replica_count=self.max_replica_count,
+            traffic_percentage=100,
+        )
+
+    def remote_predict(self, input_df):
+        if not self.endpoint:
+            endpoint = self.get_endpoint()
+            if not endpoint:
+                raise Exception(
+                    f"Endpoint '{self.endpoint_name}' not found. Please deploy the model first."
+                )
+            self.endpoint = endpoint
+
+        instances = self.df_to_features(input_df)
+        predictions = self.endpoint.predict(instances=instances)
+        return predictions.predictions
+
+    def local_predict(self, input_df):
+        if not self.model:
+            self.load_model_from_registry()
+
+        return self.model.predict(input_df)
+
+    def load_model_from_registry(self):
+        if self.model_version:
+            vertex_model = aiplatform.Model(
+                model_name=self.model_name, version=self.model_version
+            )
+
+        else:
+            vertex_model = self.get_latest_vertex_model(self.model_name)
+
+        if not vertex_model:
+            raise Exception(f"Model '{self.model_name}' not found in the registry.")
+
+        artifact_uri = vertex_model.gca_resource.artifact_uri
+        self.download_artifacts(artifact_uri)
+
+        self.model = joblib.load(self.serialized_model_path)
+
+    def download_artifacts(self, artifact_uri):
+        storage_client = storage.Client()
+        bucket_name, artifact_path = artifact_uri.replace("gs://", "").split("/", 1)
+        bucket = storage_client.bucket(bucket_name)
+
+        blobs = bucket.list_blobs(prefix=artifact_path)
+        for blob in blobs:
+            artifact_path = blob.name.split("/")[-1]
+            blob.download_to_filename(artifact_path)
+
+    def df_to_features(self, df: pd.DataFrame):
+        #
+        return df.to_dict(orient="records")

orient_express-0.2.2.dist-info/METADATA

@@ -0,0 +1,175 @@
+Metadata-Version: 2.1
+Name: orient_express
+Version: 0.2.2
+Summary: A library to simplify model deployment to Vertex AI
+Author: Alexey Zankevich
+Author-email: alex.zankevich@shiftsmart.com
+Requires-Python: >=3.9
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: google-cloud-aiplatform
+Requires-Dist: google-cloud-storage
+Requires-Dist: pandas
+Requires-Dist: scikit-learn (>=1.5.2,<2.0.0)
+Requires-Dist: xgboost (>=2.1.2,<3.0.0)
+Description-Content-Type: text/markdown
+
+# Orient Express
+A library to accelerate model deployments to Vertex AI directly from colab notebooks
+
+![train-resized](https://github.com/user-attachments/assets/f1ed32ec-07d9-4d48-8b96-3323db6b5091)
+
+## Installation
+
+```
+pip install orient_express
+```
+
+## Example
+
+### Train Model
+
+Train a regular model. In the example below, it's xgboost model, trained on the Titanic dataset.
+
+```python
+
+# Import necessary libraries
+import xgboost as xgb
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, confusion_matrix
+from sklearn.pipeline import Pipeline
+from sklearn.compose import ColumnTransformer
+from sklearn.preprocessing import StandardScaler, OneHotEncoder
+from sklearn.impute import SimpleImputer
+
+# Load the Titanic dataset
+data = sns.load_dataset('titanic').dropna(subset=['survived'])  # Dropping rows with missing target labels
+
+# Select features and target
+X = data[['pclass', 'sex', 'age', 'sibsp', 'parch', 'fare', 'embarked']]
+y = data['survived']
+
+# Define preprocessing for numeric columns (impute missing values and scale features)
+numeric_features = ['age', 'fare', 'sibsp', 'parch']
+numeric_transformer = Pipeline(steps=[
+    ('imputer', SimpleImputer(strategy='median')),
+    ('scaler', StandardScaler())
+])
+
+# Define preprocessing for categorical columns (impute missing values and one-hot encode)
+categorical_features = ['pclass', 'sex', 'embarked']
+categorical_transformer = Pipeline(steps=[
+    ('imputer', SimpleImputer(strategy='most_frequent')),
+    ('onehot', OneHotEncoder(handle_unknown='ignore'))
+])
+
+# Combine preprocessing steps
+preprocessor = ColumnTransformer(
+    transformers=[
+        ('num', numeric_transformer, numeric_features),
+        ('cat', categorical_transformer, categorical_features)
+    ])
+
+# Create a pipeline that first transforms the data, then trains an XGBoost model
+model = Pipeline(steps=[
+    ('preprocessor', preprocessor),
+    ('classifier', xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss'))
+])
+
+# Split the dataset into training and test sets
+X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+# Train the model
+model.fit(X_train, y_train)
+```
+
+## Upload Model To Model Registry
+
+```python
+
+model_wrapper  = ModelExpress(model=model,
+                             project_name='my-project-name',
+                             region='us-central1',
+                             bucket_name='my-artifacts-bucket',
+                             model_name='titanic')
+model_wrapper.upload()
+```
+
+## Local Inference (Without Online Prediction Endpoint)
+
+The following code will download the last model from the model registry and run the inference locally.
+
+```python
+
+# create input dataframe
+titanic_data = {
+    "pclass": [1],          # Passenger class (1st, 2nd, 3rd)
+    "sex": ["female"],      # Gender
+    "age": [29],            # Age
+    "sibsp": [0],           # Number of siblings/spouses aboard
+    "parch": [0],           # Number of parents/children aboard
+    "fare": [100.0],        # Ticket fare
+    "embarked": ["S"]       # Port of Embarkation (C = Cherbourg, Q = Queenstown, S = Southampton)
+}
+input_df = pd.DataFrame(titanic_data)
+
+# init the model wrapper
+model_wrapper  = ModelExpress(project_name='my-project-name',
+                             region='us-central1',
+                             model_name='titanic')
+
+# Run inference locally
+# It will download the most recent version from the model registry automatically
+model_wrapper.local_predict(input_df)
+```
+
+## Pin Model Version
+
+In many cases, the pipeline should be pinned to a specific model version so the model can only
+be updated explicitly. Just pass a `model_version` parameter when instantiating the ModelExpress wrapper.
+
+```python
+
+# init the model wrapper
+model_wrapper  = ModelExpress(project_name='my-project-name',
+                             region='us-central1',
+                             model_name='titanic',
+                             model_version=11)
+```
+
+## Remote Inference (With Online Prediction Endpoint)
+
+Make sure the model is deployed:
+```python
+
+model_wrapper  = ModelExpress(model=model,
+                             project_name='my-project-name',
+                             region='us-central1',
+                             bucket_name='my-artifacts-bucket',
+                             model_name='titanic')
+
+# upload the version to the registry and deploy it to the endpoint
+model_wrapper.deploy()
+```
+
+Run inference with `remote_predict` method. It will make a remote call to the endpoint without fetching the model locally.
+
+```python
+
+titanic_data = {
+    "pclass": [1],             # Passenger class (1st, 2nd, 3rd)
+    "sex": ["female"],         # Gender
+    "age": [29],               # Age
+    "sibsp": [0],              # Number of siblings/spouses aboard
+    "parch": [0],              # Number of parents/children aboard
+    "fare": [100.0],           # Ticket fare
+    "embarked": ["S"]          # Port of Embarkation (C = Cherbourg, Q = Queenstown, S = Southampton)
+}
+df = pd.DataFrame(titanic_data)
+
+model_wrapper.remote_predict(df)
+```
+

orient_express-0.2.2.dist-info/RECORD

@@ -0,0 +1,5 @@
+orient_express/__init__.py,sha256=I7VRbcJKqePRPtOiZl3ro9gwwGWiBcj-DvaWPHIEUqQ,40
+orient_express/model_wrapper.py,sha256=FnIq4ExP-rIpik0LrhlMfOyszzbeNQLRPni8dsD640M,6385
+orient_express-0.2.2.dist-info/METADATA,sha256=e1bbcbfIprg4xA-8fXAmDMpmkrr3rVdpbeS98aQNqFY,5821
+orient_express-0.2.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+orient_express-0.2.2.dist-info/RECORD,,

orient_express-0.2.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.9.0
+Root-Is-Purelib: true
+Tag: py3-none-any