orient_express 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,175 @@
1
+ Metadata-Version: 2.1
2
+ Name: orient_express
3
+ Version: 0.2.2
4
+ Summary: A library to simplify model deployment to Vertex AI
5
+ Author: Alexey Zankevich
6
+ Author-email: alex.zankevich@shiftsmart.com
7
+ Requires-Python: >=3.9
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: Programming Language :: Python :: 3.9
10
+ Classifier: Programming Language :: Python :: 3.10
11
+ Classifier: Programming Language :: Python :: 3.11
12
+ Classifier: Programming Language :: Python :: 3.12
13
+ Requires-Dist: google-cloud-aiplatform
14
+ Requires-Dist: google-cloud-storage
15
+ Requires-Dist: pandas
16
+ Requires-Dist: scikit-learn (>=1.5.2,<2.0.0)
17
+ Requires-Dist: xgboost (>=2.1.2,<3.0.0)
18
+ Description-Content-Type: text/markdown
19
+
20
+ # Orient Express
21
+ A library to accelerate model deployments to Vertex AI directly from colab notebooks
22
+
23
+ ![train-resized](https://github.com/user-attachments/assets/f1ed32ec-07d9-4d48-8b96-3323db6b5091)
24
+
25
+ ## Installation
26
+
27
+ ```
28
+ pip install orient_express
29
+ ```
30
+
31
+ ## Example
32
+
33
+ ### Train Model
34
+
35
+ Train a regular model. In the example below, it's xgboost model, trained on the Titanic dataset.
36
+
37
+ ```python
38
+
39
+ # Import necessary libraries
40
+ import xgboost as xgb
41
+ from sklearn.model_selection import train_test_split
42
+ from sklearn.metrics import accuracy_score, confusion_matrix
43
+ from sklearn.pipeline import Pipeline
44
+ from sklearn.compose import ColumnTransformer
45
+ from sklearn.preprocessing import StandardScaler, OneHotEncoder
46
+ from sklearn.impute import SimpleImputer
47
+
48
+ # Load the Titanic dataset
49
+ data = sns.load_dataset('titanic').dropna(subset=['survived']) # Dropping rows with missing target labels
50
+
51
+ # Select features and target
52
+ X = data[['pclass', 'sex', 'age', 'sibsp', 'parch', 'fare', 'embarked']]
53
+ y = data['survived']
54
+
55
+ # Define preprocessing for numeric columns (impute missing values and scale features)
56
+ numeric_features = ['age', 'fare', 'sibsp', 'parch']
57
+ numeric_transformer = Pipeline(steps=[
58
+ ('imputer', SimpleImputer(strategy='median')),
59
+ ('scaler', StandardScaler())
60
+ ])
61
+
62
+ # Define preprocessing for categorical columns (impute missing values and one-hot encode)
63
+ categorical_features = ['pclass', 'sex', 'embarked']
64
+ categorical_transformer = Pipeline(steps=[
65
+ ('imputer', SimpleImputer(strategy='most_frequent')),
66
+ ('onehot', OneHotEncoder(handle_unknown='ignore'))
67
+ ])
68
+
69
+ # Combine preprocessing steps
70
+ preprocessor = ColumnTransformer(
71
+ transformers=[
72
+ ('num', numeric_transformer, numeric_features),
73
+ ('cat', categorical_transformer, categorical_features)
74
+ ])
75
+
76
+ # Create a pipeline that first transforms the data, then trains an XGBoost model
77
+ model = Pipeline(steps=[
78
+ ('preprocessor', preprocessor),
79
+ ('classifier', xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss'))
80
+ ])
81
+
82
+ # Split the dataset into training and test sets
83
+ X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
84
+
85
+ # Train the model
86
+ model.fit(X_train, y_train)
87
+ ```
88
+
89
+ ## Upload Model To Model Registry
90
+
91
+ ```python
92
+
93
+ model_wrapper = ModelExpress(model=model,
94
+ project_name='my-project-name',
95
+ region='us-central1',
96
+ bucket_name='my-artifacts-bucket',
97
+ model_name='titanic')
98
+ model_wrapper.upload()
99
+ ```
100
+
101
+ ## Local Inference (Without Online Prediction Endpoint)
102
+
103
+ The following code will download the last model from the model registry and run the inference locally.
104
+
105
+ ```python
106
+
107
+ # create input dataframe
108
+ titanic_data = {
109
+ "pclass": [1], # Passenger class (1st, 2nd, 3rd)
110
+ "sex": ["female"], # Gender
111
+ "age": [29], # Age
112
+ "sibsp": [0], # Number of siblings/spouses aboard
113
+ "parch": [0], # Number of parents/children aboard
114
+ "fare": [100.0], # Ticket fare
115
+ "embarked": ["S"] # Port of Embarkation (C = Cherbourg, Q = Queenstown, S = Southampton)
116
+ }
117
+ input_df = pd.DataFrame(titanic_data)
118
+
119
+ # init the model wrapper
120
+ model_wrapper = ModelExpress(project_name='my-project-name',
121
+ region='us-central1',
122
+ model_name='titanic')
123
+
124
+ # Run inference locally
125
+ # It will download the most recent version from the model registry automatically
126
+ model_wrapper.local_predict(input_df)
127
+ ```
128
+
129
+ ## Pin Model Version
130
+
131
+ In many cases, the pipeline should be pinned to a specific model version so the model can only
132
+ be updated explicitly. Just pass a `model_version` parameter when instantiating the ModelExpress wrapper.
133
+
134
+ ```python
135
+
136
+ # init the model wrapper
137
+ model_wrapper = ModelExpress(project_name='my-project-name',
138
+ region='us-central1',
139
+ model_name='titanic',
140
+ model_version=11)
141
+ ```
142
+
143
+ ## Remote Inference (With Online Prediction Endpoint)
144
+
145
+ Make sure the model is deployed:
146
+ ```python
147
+
148
+ model_wrapper = ModelExpress(model=model,
149
+ project_name='my-project-name',
150
+ region='us-central1',
151
+ bucket_name='my-artifacts-bucket',
152
+ model_name='titanic')
153
+
154
+ # upload the version to the registry and deploy it to the endpoint
155
+ model_wrapper.deploy()
156
+ ```
157
+
158
+ Run inference with `remote_predict` method. It will make a remote call to the endpoint without fetching the model locally.
159
+
160
+ ```python
161
+
162
+ titanic_data = {
163
+ "pclass": [1], # Passenger class (1st, 2nd, 3rd)
164
+ "sex": ["female"], # Gender
165
+ "age": [29], # Age
166
+ "sibsp": [0], # Number of siblings/spouses aboard
167
+ "parch": [0], # Number of parents/children aboard
168
+ "fare": [100.0], # Ticket fare
169
+ "embarked": ["S"] # Port of Embarkation (C = Cherbourg, Q = Queenstown, S = Southampton)
170
+ }
171
+ df = pd.DataFrame(titanic_data)
172
+
173
+ model_wrapper.remote_predict(df)
174
+ ```
175
+
@@ -0,0 +1,155 @@
1
+ # Orient Express
2
+ A library to accelerate model deployments to Vertex AI directly from colab notebooks
3
+
4
+ ![train-resized](https://github.com/user-attachments/assets/f1ed32ec-07d9-4d48-8b96-3323db6b5091)
5
+
6
+ ## Installation
7
+
8
+ ```
9
+ pip install orient_express
10
+ ```
11
+
12
+ ## Example
13
+
14
+ ### Train Model
15
+
16
+ Train a regular model. In the example below, it's xgboost model, trained on the Titanic dataset.
17
+
18
+ ```python
19
+
20
+ # Import necessary libraries
21
+ import xgboost as xgb
22
+ from sklearn.model_selection import train_test_split
23
+ from sklearn.metrics import accuracy_score, confusion_matrix
24
+ from sklearn.pipeline import Pipeline
25
+ from sklearn.compose import ColumnTransformer
26
+ from sklearn.preprocessing import StandardScaler, OneHotEncoder
27
+ from sklearn.impute import SimpleImputer
28
+
29
+ # Load the Titanic dataset
30
+ data = sns.load_dataset('titanic').dropna(subset=['survived']) # Dropping rows with missing target labels
31
+
32
+ # Select features and target
33
+ X = data[['pclass', 'sex', 'age', 'sibsp', 'parch', 'fare', 'embarked']]
34
+ y = data['survived']
35
+
36
+ # Define preprocessing for numeric columns (impute missing values and scale features)
37
+ numeric_features = ['age', 'fare', 'sibsp', 'parch']
38
+ numeric_transformer = Pipeline(steps=[
39
+ ('imputer', SimpleImputer(strategy='median')),
40
+ ('scaler', StandardScaler())
41
+ ])
42
+
43
+ # Define preprocessing for categorical columns (impute missing values and one-hot encode)
44
+ categorical_features = ['pclass', 'sex', 'embarked']
45
+ categorical_transformer = Pipeline(steps=[
46
+ ('imputer', SimpleImputer(strategy='most_frequent')),
47
+ ('onehot', OneHotEncoder(handle_unknown='ignore'))
48
+ ])
49
+
50
+ # Combine preprocessing steps
51
+ preprocessor = ColumnTransformer(
52
+ transformers=[
53
+ ('num', numeric_transformer, numeric_features),
54
+ ('cat', categorical_transformer, categorical_features)
55
+ ])
56
+
57
+ # Create a pipeline that first transforms the data, then trains an XGBoost model
58
+ model = Pipeline(steps=[
59
+ ('preprocessor', preprocessor),
60
+ ('classifier', xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss'))
61
+ ])
62
+
63
+ # Split the dataset into training and test sets
64
+ X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
65
+
66
+ # Train the model
67
+ model.fit(X_train, y_train)
68
+ ```
69
+
70
+ ## Upload Model To Model Registry
71
+
72
+ ```python
73
+
74
+ model_wrapper = ModelExpress(model=model,
75
+ project_name='my-project-name',
76
+ region='us-central1',
77
+ bucket_name='my-artifacts-bucket',
78
+ model_name='titanic')
79
+ model_wrapper.upload()
80
+ ```
81
+
82
+ ## Local Inference (Without Online Prediction Endpoint)
83
+
84
+ The following code will download the last model from the model registry and run the inference locally.
85
+
86
+ ```python
87
+
88
+ # create input dataframe
89
+ titanic_data = {
90
+ "pclass": [1], # Passenger class (1st, 2nd, 3rd)
91
+ "sex": ["female"], # Gender
92
+ "age": [29], # Age
93
+ "sibsp": [0], # Number of siblings/spouses aboard
94
+ "parch": [0], # Number of parents/children aboard
95
+ "fare": [100.0], # Ticket fare
96
+ "embarked": ["S"] # Port of Embarkation (C = Cherbourg, Q = Queenstown, S = Southampton)
97
+ }
98
+ input_df = pd.DataFrame(titanic_data)
99
+
100
+ # init the model wrapper
101
+ model_wrapper = ModelExpress(project_name='my-project-name',
102
+ region='us-central1',
103
+ model_name='titanic')
104
+
105
+ # Run inference locally
106
+ # It will download the most recent version from the model registry automatically
107
+ model_wrapper.local_predict(input_df)
108
+ ```
109
+
110
+ ## Pin Model Version
111
+
112
+ In many cases, the pipeline should be pinned to a specific model version so the model can only
113
+ be updated explicitly. Just pass a `model_version` parameter when instantiating the ModelExpress wrapper.
114
+
115
+ ```python
116
+
117
+ # init the model wrapper
118
+ model_wrapper = ModelExpress(project_name='my-project-name',
119
+ region='us-central1',
120
+ model_name='titanic',
121
+ model_version=11)
122
+ ```
123
+
124
+ ## Remote Inference (With Online Prediction Endpoint)
125
+
126
+ Make sure the model is deployed:
127
+ ```python
128
+
129
+ model_wrapper = ModelExpress(model=model,
130
+ project_name='my-project-name',
131
+ region='us-central1',
132
+ bucket_name='my-artifacts-bucket',
133
+ model_name='titanic')
134
+
135
+ # upload the version to the registry and deploy it to the endpoint
136
+ model_wrapper.deploy()
137
+ ```
138
+
139
+ Run inference with `remote_predict` method. It will make a remote call to the endpoint without fetching the model locally.
140
+
141
+ ```python
142
+
143
+ titanic_data = {
144
+ "pclass": [1], # Passenger class (1st, 2nd, 3rd)
145
+ "sex": ["female"], # Gender
146
+ "age": [29], # Age
147
+ "sibsp": [0], # Number of siblings/spouses aboard
148
+ "parch": [0], # Number of parents/children aboard
149
+ "fare": [100.0], # Ticket fare
150
+ "embarked": ["S"] # Port of Embarkation (C = Cherbourg, Q = Queenstown, S = Southampton)
151
+ }
152
+ df = pd.DataFrame(titanic_data)
153
+
154
+ model_wrapper.remote_predict(df)
155
+ ```
@@ -0,0 +1 @@
1
+ from .model_wrapper import ModelExpress
@@ -0,0 +1,197 @@
1
+ import os
2
+ import joblib
3
+ import logging
4
+ import pandas as pd
5
+ from google.cloud import storage
6
+ from google.cloud import aiplatform
7
+
8
+
9
+ class ModelExpress:
10
+ def __init__(
11
+ self,
12
+ model_name,
13
+ project_name,
14
+ bucket_name,
15
+ model_version=None,
16
+ model=None,
17
+ region="us-central1",
18
+ serialized_model_path="model.joblib",
19
+ docker_image_uri="us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-7:latest",
20
+ endpoint_name=None,
21
+ machine_type="n1-standard-4",
22
+ min_replica_count=1,
23
+ max_replica_count=1,
24
+ ):
25
+ self.model = model
26
+ self.model_name = model_name
27
+ self.model_version = model_version
28
+ self.region = region
29
+ self.project_name = project_name
30
+ self.bucket_name = bucket_name
31
+ self.serialized_model_path = serialized_model_path
32
+ self.docker_image_uri = docker_image_uri
33
+ self.machine_type = machine_type
34
+ self.min_replica_count = min_replica_count
35
+ self.max_replica_count = max_replica_count
36
+ self.endpoint = None
37
+
38
+ if not endpoint_name:
39
+ self.endpoint_name = f"model-xpress-{model_name}"
40
+ else:
41
+ self.endpoint_name = endpoint_name
42
+
43
+ def colab_auth(self):
44
+ from google.colab import auth
45
+
46
+ auth.authenticate_user()
47
+
48
+ def _vertex_init(self):
49
+ aiplatform.init(project=self.project_name, location=self.region)
50
+
51
+ def get_latest_vertex_model(self, model_name):
52
+ self._vertex_init()
53
+
54
+ # Search for models with the specified display name
55
+ models = aiplatform.Model.list(filter=f"display_name={model_name}")
56
+
57
+ if not models:
58
+ return None # Return None if no model with the given name exists
59
+
60
+ # Sort models by update time in descending order to get the latest version
61
+ latest_model = sorted(models, key=lambda x: x.update_time, reverse=True)[0]
62
+ return latest_model
63
+
64
+ # upload model to vertex ai model registry
65
+ def upload(self):
66
+ joblib.dump(self.model, self.serialized_model_path)
67
+ logging.info(f"Model saved to {self.serialized_model_path}")
68
+
69
+ # Initialize the GCS client
70
+ client = storage.Client()
71
+ bucket = client.bucket(self.bucket_name)
72
+
73
+ last_model = self.get_latest_vertex_model(self.model_name)
74
+ if last_model:
75
+ last_version = last_model.version_id
76
+ new_version = int(last_version) + 1
77
+ else:
78
+ new_version = 1
79
+
80
+ # Upload the model file
81
+ blob = bucket.blob(
82
+ self.get_artifacts_path(new_version, self.serialized_model_path)
83
+ )
84
+ blob.upload_from_filename(self.serialized_model_path)
85
+
86
+ return self.create_model_version(new_version, last_model)
87
+
88
+ def get_artifacts_path(self, version, file_name=None):
89
+ dir_name = f"models/{self.model_name}/{version}"
90
+ if file_name:
91
+ return f"{dir_name}/{file_name}"
92
+
93
+ return f"{dir_name}/"
94
+
95
+ def create_model_version(self, version_number, last_version):
96
+ artifact_uri = (
97
+ f"gs://{self.bucket_name}/{self.get_artifacts_path(version_number)}"
98
+ )
99
+
100
+ if last_version:
101
+ parent_model = f"projects/{self.project_name}/locations/{self.region}/models/{last_version.name}"
102
+ else:
103
+ parent_model = None
104
+
105
+ release = aiplatform.Model.upload(
106
+ display_name=self.model_name,
107
+ artifact_uri=artifact_uri,
108
+ parent_model=parent_model,
109
+ serving_container_image_uri=self.docker_image_uri,
110
+ sync=True,
111
+ )
112
+
113
+ return release
114
+
115
+ def get_or_create_endpoint(self):
116
+ endpoint = self.get_endpoint()
117
+ if endpoint:
118
+ return endpoint
119
+ else:
120
+ return self.create_endpoint()
121
+
122
+ def get_endpoint(self):
123
+ endpoints = aiplatform.Endpoint.list(
124
+ filter=f"display_name={self.endpoint_name}", order_by="create_time"
125
+ )
126
+ if endpoints:
127
+ return endpoints[0]
128
+
129
+ def create_endpoint(self):
130
+ return aiplatform.Endpoint.create(
131
+ display_name=self.endpoint_name,
132
+ project=self.project_name,
133
+ location=self.region,
134
+ )
135
+
136
+ def deploy(self):
137
+ self._vertex_init()
138
+
139
+ endpoint = self.get_or_create_endpoint()
140
+ model_version = self.upload()
141
+ model_version.deploy(
142
+ endpoint=endpoint,
143
+ machine_type=self.machine_type,
144
+ min_replica_count=self.min_replica_count,
145
+ max_replica_count=self.max_replica_count,
146
+ traffic_percentage=100,
147
+ )
148
+
149
+ def remote_predict(self, input_df):
150
+ if not self.endpoint:
151
+ endpoint = self.get_endpoint()
152
+ if not endpoint:
153
+ raise Exception(
154
+ f"Endpoint '{self.endpoint_name}' not found. Please deploy the model first."
155
+ )
156
+ self.endpoint = endpoint
157
+
158
+ instances = self.df_to_features(input_df)
159
+ predictions = self.endpoint.predict(instances=instances)
160
+ return predictions.predictions
161
+
162
+ def local_predict(self, input_df):
163
+ if not self.model:
164
+ self.load_model_from_registry()
165
+
166
+ return self.model.predict(input_df)
167
+
168
+ def load_model_from_registry(self):
169
+ if self.model_version:
170
+ vertex_model = aiplatform.Model(
171
+ model_name=self.model_name, version=self.model_version
172
+ )
173
+
174
+ else:
175
+ vertex_model = self.get_latest_vertex_model(self.model_name)
176
+
177
+ if not vertex_model:
178
+ raise Exception(f"Model '{self.model_name}' not found in the registry.")
179
+
180
+ artifact_uri = vertex_model.gca_resource.artifact_uri
181
+ self.download_artifacts(artifact_uri)
182
+
183
+ self.model = joblib.load(self.serialized_model_path)
184
+
185
+ def download_artifacts(self, artifact_uri):
186
+ storage_client = storage.Client()
187
+ bucket_name, artifact_path = artifact_uri.replace("gs://", "").split("/", 1)
188
+ bucket = storage_client.bucket(bucket_name)
189
+
190
+ blobs = bucket.list_blobs(prefix=artifact_path)
191
+ for blob in blobs:
192
+ artifact_path = blob.name.split("/")[-1]
193
+ blob.download_to_filename(artifact_path)
194
+
195
+ def df_to_features(self, df: pd.DataFrame):
196
+ #
197
+ return df.to_dict(orient="records")
@@ -0,0 +1,22 @@
1
+ [tool.poetry]
2
+ name = "orient_express"
3
+ version = "0.2.2"
4
+ description = "A library to simplify model deployment to Vertex AI"
5
+ authors = ["Alexey Zankevich <alex.zankevich@shiftsmart.com>"]
6
+ readme = "README.md"
7
+
8
+ [tool.poetry.dependencies]
9
+ python = ">=3.9"
10
+ google-cloud-aiplatform = "*"
11
+ google-cloud-storage = "*"
12
+ pandas = "*"
13
+ scikit-learn = "^1.5.2"
14
+ xgboost = "^2.1.2"
15
+
16
+ [tool.poetry.dev-dependencies]
17
+ black = "24.10.0"
18
+ pytest = "8.3.3"
19
+
20
+ [build-system]
21
+ requires = ["poetry-core"]
22
+ build-backend = "poetry.core.masonry.api"