How Data Capture Transforms ML Observability

Key Takeaways

[Music] welcome push [Music] car hey everyone how's it going good all right let's get started let me share my screen all right let's get started so hi everyone uh my name is pushkar gar um I work as a staff uh machine learning platform engineer at Clary and um excited to be presenting this topic which is a slightly Niche area but um like more focused in terms of like if you talk about data quality but I'm going to talk about uh data capture on model endpoints and how that helps in um running data observability on uh like the models deployed in uh in production so let's get started um so the agenda for today uh in in my talk is going to be so we'll we'll go over the different data components that you may see in um in a data platform which Powers all of the model building capabilities in the M AI nml platform um then we'll go over the importance of monitoring uh data quality and then in general model monitoring as well um uh uh after that we'll we'll move on to like how do we measure actually measure data quality with with metrics uh what are the metrics that we need um how do we build that infrastructure to um identify and capture those metrics um so yeah uh the next topic we're going to cover is on the managed data capture so there are uh certain uh tools which we which if we use uh to uh deploy our models they provide managed data capture meaning like automatic data capture uh and what are some of the limitations Rel relateded to that and then uh what is the need for a custom what may be the need for a custom data capture architecture uh and then how do we go about Implement uh implementing that and then it will all we try to uh bring all of that together into like how does that fit it into the downstream systems and then uh some of the best practices and some some lessons learned uh around like implementing custom data capture so let's start with uh the different data components right so um uh as any company uh would already have in production their uh production databases right and then with with the Advent of generative AI nowadays we are seeing a lot of the uh use cases for utilizing the data which is stored in silos in um in like S3 buckets or like object storage which can be uh like files um call transcripts email bodies anything like that right so how do we bring all of that data uh together into into one component which can be used to u to run Transformations uh and then ultimately build features uh so that we can uh use those uh features to ultimately build models so in in in this diagram I've tried to like separate out the data world and then the ml world so that it it becomes easier to uh you know specify responsibilities and uh also Implement like data quality check so in in this case it's almost possible to implement data quality checks at all of the uh checkpoints wherever data is flowing and ultimately getting persisted so we start with the production uh databases uh and then uh the object storage to bring together the structured and unstructured data into the data Lake uh for example and then uh following The Medallion architecture that has been popularized by data bricks we can uh transform that data uh o over certain schemas and then move from like a bronze to silver to gold uh schema which uh which gives you like the confidence that uh the that this data will always be ready for like consuming uh into like model building so once that data has been transformed into the modeled layer or the gold layer uh you can sync that into an offline feature store uh and use that to train the model so run training jobs on the offline feature store develop models U and then ultimately deploy those to a serving endpoint which serves prediction requests from Upstream services so like any kind of uh model prediction that can be done uh once the model is deployed to that serving endpoint it can take care of that and there may well also be a need to query real-time features from uh from an online feature store which ingests realtime features from Upstream Services um that may or may not go through transformation so like your real-time Services may be uh creating data uh for a particular user that you want to do a prediction for uh so it's necessary to get the most updated data uh that may go through transformation so ultimately all that comes together in the online feature store and then the serving endpoint can uh use that uh to get the latest data about a user and uh or any product and then uh make a prediction so the the uh focus of the talk today will be on the serving endpoint like on like the data quality on the serving endpoint and how do we enable that by enabling data capture in the first place so uh yeah excited about that so the importance of monitoring like as um all of you already know right uh models production quality is as good as the data quality right so garbage in garbage out whatever is fed to the model is uh like uh comparable to the prediction that is generated by the model so the the the other importance of the uh of monitoring in general and data quality in more to be more specific is to identify training serving skew when we train the model uh on a on a certain data set it identifies the relationship between the independent variables and the dependent variable um and then it attaches a certain weight to all of those variables to come up with a certain equation uh and then ultimately give you for a certain set of features using those weights give you uh U the the output which is which is the dependent variable now if the if the uh if the data in production is much different right if the if the distribution of the data in production is much different than it was in training you can expect there needs to be some adjustment to the weights that needs to be done for an accurate prediction so to be able to identify that that has happened in the produ ction data meaning that there has been a distribution drift uh in the train in the inference data uh we need to be able to uh capture data and then uh ultimately identify that and then retrain the models on on the latest updated data to capture those changes the next uh type of identification that we need to do is concept drift so the data in itself may not have a large distribution change but the relationship ship between the data and the dependent like the features and the dependent variable may change over time due to um like for example covid right so um uh when covid happened a lot of the models which were in production would not have performed uh uh up to the mark because the the overall situation has changed and then the the the um the relationship between the features and the uh prediction has changed so it's it's it's necessary uh monitoring is necessary for identifying that use case as well and then finally uh like I said the the data quality checks can be implemented at all um um checkpoints of of a data pipeline but um but on the serving endpoint is sort of the Last Frontier where you can ident still identify issues in the data pipeline so anything which may have been ignored or which may not have been captured uh in terms of failures or in terms of uh having um breakages in the data pipeline can ultimately be identified in the serving endpoint if if you decide to have monitoring there so so we've established the importance of data quality now let's see how do we measure data quality right so there are different types of metrics and these are not the exhaustive list of metrics but this is something that um is this is something that can be easily um uh enabled on a serving endpoint to identify uh the usual data quality metrics so uh let's start with the data type check right so data type so in here what we're trying to do is we're trying to take the training data and compare it with the inference data to identify um these checks so what what we are doing is we're looking at every column in the training data comparing it with the corresponding column in the inference data and then identify the so data type check is is has each column uh if each column has the same data type as it was present in the training data completeness check how many nulls are there as compared to um uh the training data right so if if you're getting more nulls uh as part of your inference data set that it it means that there is some kind of broken uh either like there is a broken pipeline or there has been some change which which has not been accounted for uh and then this this issue needs to be fixed so to be able to identify that that is happening is the is the first Frontier and then move on from there uh to to rectify it is the next uh and then the I think the most important one is like the Baseline drift check so how do we how do we measure that there has been a distribution change right so to be able to measure that we we look at um each column and then compare the distribution of the data in the training data set uh to the inference data set to see if there has been been a distribution uh change in the data uh and um and Trigger some alerts if there has been like a significant change now there are um uh buil-in algorithms that that can be used it's it's actually non-trivial to implement it um uh uh if you if you go the native way right if you try to implement it uh using an algorithm but there are like libraries like DQ which support uh identifying uh these uh uh the the the drift in the data by just using the library directly uh and supplying it with like different uh the two columns to compare so this is something that can be uh used and it's it's a it's a good metric to identify uh and then these last two are U uh straightforward like if you if there is any missing column in itself or if there is like any extra column that you did not expect in the inference data um then that needs to be identified and rectified as well okay so once once we have established how do we uh measure data quality by generating the metrics right now how do we generate those metrics so let's start with the um training data here so so you have a training data set you run a training job build a model deploy the model to an endpoint and then that model endpoint serves prediction requests from the prediction service um to to be able to generate those metrics we need the training data and we need the inference data so that can the inference so we already have the training data uh and we can use data capture to capture the requests that are coming into the model endpoint uh to create that inference data set uh once both of these are available uh we can pass them on to a monitoring job and uh that monitoring job can uh look at each column in the training data and inference data and then generate the metrics that we just discussed and then uh assist them to your uh choice of uh observability tool um yeah it's important to version training data um because a model goes through like multiple experiments multiple U AB testing so uh it's necessary to uh capture the the current version of the model was deployed with which version of the training data in a model registry and then the monitoring job can pull that data from the training uh that model registry uh uh and then load that from from the uh from S3 and then it it is also possible to generate like firsthand uh uh Baseline and constraints using that trending data and then only use those constraints uh to look at the inference data and then for for each column right let's say um what is the range bound of of the variables of of the uh features that you see for a particular column what was that compared to that training data set so it's possible to optimize it even further but this is like an abstraction on on on that front uh okay so now let's talk about data capture right so we identified that we can enable data capture on a Model endpoint uh and then that can help us in capturing the inference data so there are automatic ways of capturing data so if you for example if you deploy a sagemaker endpoint your model code gets deployed uh on an ec2 instance and it would automatically provide you a server uh a flask server uh which is uh which is the uh which which uh gets all of the prediction requests uh in HTTP or grpc whatever you have it and it will um it will forward the requests to the model model will make what model will uh use those features to make the prediction and it may also query the online feature store like we discussed and then um it will make the prediction return the request to the server and then the server will persist the combination of the features and the prediction so uh as it received from the prediction service and the model so so imagine if you're if you're using uh protuff as the object uh in the request to the server so the server will persist that uh and then the model will ALS so model will will deserialize those requests uh to make the features in in plain text uh it will generate the prediction uh serialize the request again uh and then send it back to the to the server to send it back to the prediction service so the in in here the server will um capture the request um in in prot format if if that is the case so that's how the managed data capture works now what is the need for coding custom data capture right so so first of all it it doesn't work for realtime features in in this particular case for example with sagemaker so whatever is uh sent to the server that's what is captured in in in like the inference data set uh for data capture if we go like the managed way so imagine uh the prediction service sending just a user ID to the server and then the model querying the online feature store for the latest uh values or the latest features Ed to that particular user um it would not capture those requests uh or those features that the model um got from the online feature store uh and then it will just persist that user ID and then the corresponding prediction so there are ways to uh to still um like recreate that data set by doing time travel queries so ultimately what happens is whatever is the data that comes in the online features store gets synced in the offline feature store at a specified Cadence so what you can do is you can run uh another job which uh which parses all of the data in in the data capture logs uh and takes the uh request time and then regenerates the set of features from the offline feature store by doing time travel queries that whatever is the feature set that the model would have used at that time to make the prediction you can recreate that data set and then run the monitoring job to have like the full set of features to compare against the training data and uh one thing that I identified was that it does not work for uh multi-production models so like I um what I mean by multi-production model is like for example a ranking model where you would have to predict a set of requests it's not a multimodel endpoint uh per se but it's a uh it's one model but it is serving uh multiple requests so like ranking for example you would have um multiple items that you would want to rank uh and then generate generate a score for all of them and then rank them um uh in in order of user interactivity or whatever so uh so it it uh in that case uh the the the persistence to S3 for the data capture would involve like a whole array of objects and then um it at that time Sage maker was not able to handle that so there was this another uh need there and then like I discussed earlier as well if you have if you're using phobu as your request object you need to keep track of the uh the the serialization and der serialization class to be able to uh deserialize the requests that are captured in S3 from uh whenever you run the monitoring job to get uh the data in plane text so these are some of the limitations which um uh because of which uh we implemented Uh custom data capture excuse me now let's uh look at custom data capture right uh so how do we how did we go about implementing Uh custom data capture so the uh you can Implement a data capture buffer on the model itself so uh it will use the uh Hardware of the model endpoint that it is running on like in terms of memory uh and compute uh but you can still code that within the model code so so what what we did was we created a package and uh that package has to be imported by any data scientist who is looking to implement this data capture and what can be done is you instantiate that buffer and then you uh after the after the model has made the prediction right so it let's say it queried something from the online feature store it made that prediction and then ultimately it's able to um uh get everything in plain text and then just before serializing the request to send back to the server it can capture all of that data um it can uh yeah it can add append that data to the to the buffer that it instantiated so and then the buffer from time to time what it will do is it will uh flush it out to to S3 uh based on uh different knobs that that you want would want to enable right so what we did was like buffer memory size message count and then the active buffer time so uh let's say from after every five minutes we would purist the data to uh to S3 so this is something that we did for uh implementing custom data capture um and then let's let's go to like how do how do Downstream systems use this so like we discussed training data uh and then the inference data so the inference data is generated by the data capture monitoring job uh takes both of these as input generates the metrics pushes the metrics to your observability tool and then you can set thresholds for for um these metrics to go if that if they go over a certain value you would send alerts uh and then those alerts go to a notification queue and then that notification queue uh can have subscribers Downstream which on notification send um uh send an alert to page of Duty or slack uh and then retraining can be done basically too so that that sort of completes the loop on uh observability so you started with data capture identified uh created the metrics identified if there is anything wrong in the model and then completed the loop with retraining uh and then this data capture can also be used for uh running model quality right so you can you can imagine that the inference data can be augmented with the ground Ruth and then ultimately that data can be joined with the uh data captured in production so so the monitoring job will have the input of for this set of features what was the prediction that the model did and then what was the actual ground truth and use those to compare uh and then generate model quality monitoring metrics as well um so this is how data capture overall comes together and then creates the foundation of running uh observability what ml observability in particular so some of the Lessons Learned around uh implementing Uh custom data capture is uh like batching requests is obviously a very good idea uh you would not want to persist every request that you get uh uh to to S3 because that hampers your iio performance uh experimenting with different knobs to find the right threshold is a good idea like like discussed um memory size uh the total time after which uh uh the the buffer would write data S3 or the message count whatever may work for you but integrating those knobs into the code is is a good idea sampling of request is is something that you can save money on maybe not a lot but like in terms of compute uh it it there can be some uh savings right so if you there is not a need to capture everything uh so in in our case like 75% of the requests are also able to uh uh identify if if there is any kind of data drift and then uh performance benchmarking is also important because this is something that is running on the uh on the endpoint itself so it's important to be able to set some SLA around how much um how much latency uh this data capture buffer may add so it's important to do benchmarking to identify like the average Max 99% High latency that that may be added as as part of the DAT capture buer uh yeah that's that's my time I think uh thank you very much for listening in happy to answer any questions if there are awesome thank you so much