Tracking provenance in data science scripts

Abstract

Embodiments described herein enable tracking machine learning (“ML”) model data provenance. In particular, a computing device is configured to accept ML model code that, when executed, instantiates and trains an ML model, to parse the ML model code into a workflow intermediate representation (WIR), to semantically annotate the WIR to provide an annotated WIR, and to identify, based on the annotated WIR and ML API corresponding to the ML model code, data from at least one data source that is relied upon by the ML model code when training the ML model. A WIR may be generated from an abstract syntax tree (AST) based on the ML model code, generating provenance relationships (PRs) based at least in part on relationships between nodes of the AST, wherein a PR comprises one or more input variables, an operation, a caller, and one or more output variables.

What is claimed is: 
     
         1 . A machine learning data provenance tracking system, comprising:
 a memory device configured to store a queryable knowledge base (KB), the queryable KB comprising:
 a first machine learning (ML) application programming interface (API) annotation corresponding to a semantic identity of, and a relationship between, an input and an output of a first operation provided by a first ML API, the first ML API when executed trains a ML API model and 
 a table of operations corresponding to ML code operations that perform data selection or exclusion operations; and 
   a data source identifier executed by a processing circuit, the data source identifier comprising:
 an ML model analyzer configured to:
 obtain a workflow intermediate representation (WIR) of ML model code, the ML code using the first ML API, and 
 semantically annotate the WIR based on the first ML API by querying the queryable KB, and 
 
 a provenance tracker configured to identify, based on the annotated WIR and the first ML API, data from a data source that is relied upon by the ML model code when training the ML model. 
   
     
     
         2 . The machine learning data provenance tracking system of  claim 1 , wherein the WIR comprises a directed graph comprising a plurality of provenance relationships (PRs), each PR of the plurality of PRs corresponding to a relationship between nodes of an abstract syntax tree (AST) corresponding to the ML model code. 
     
     
         3 . The machine learning data provenance tracking system of  claim 2 , wherein to semantically annotate the WIR, the ML model analyzer is further configured to:
 forward traverse the directed graph to find a first PR of the plurality of PRs;   query the queryable KB for a first input annotation corresponding to an input variable of the first PR and a first output annotation corresponding to an output variable of the first PR; and   add the first input annotation and the first output annotation to the first PR.   
     
     
         4 . The machine learning data provenance tracking system of  claim 3  wherein the ML model analyzer is configured to query the queryable KB for the first input annotation corresponding to the input variable of the first PR and the first output annotation corresponding to the output variable of the first PR based on a caller and an operation of the first PR. 
     
     
         5 . The machine learning data provenance tracking system of  claim 3  wherein to semantically annotate the WIR, the ML model analyzer is further configured to:
 forward traverse the directed graph to find a second PR of the plurality of PRs; 
 query the queryable KB for a second input annotation corresponding to an input variable of the second PR and a second output annotation corresponding to an output variable of the second PR; and 
 add the second input annotation and the second output annotation to the second PR. 
 
     
     
         6 . The machine learning data provenance tracking system of  claim 5  wherein to semantically annotate the WIR, the ML model analyzer is further configured to:
 starting with the second PR, backward traverse the WIR to find a third PR of the plurality of PRs, wherein the input variable of the second PR is an output variable of the third PR; 
 query, based on the second input annotation, the queryable KB for a third input annotation corresponding to an input variable of the second PR; and 
 add the third input annotation to the second PR. 
 
     
     
         7 . The machine learning data provenance tracking system of  claim 1 , wherein to identify the data, the provenance tracker is further configured to:
 determine a feature variable of a first PR, the feature variable annotated as a feature and corresponding to a data selection or exclusion operation; and   starting with the first PR, recursively backwards traverse the directed graph to identify the data from the data source that corresponds to the feature variable.   
     
     
         8 . A method for tracking machine learning data provenance, comprising:
 obtaining a workflow intermediate representation (WIR) of machine learning (ML) model code, the ML code using a first ML application programming interface (API) that when executed trains a ML model;   semantically annotating the WIR based on the first ML API by querying a queryable knowledge base (KB), the queryable KB comprising:
 a first ML API annotation corresponding to a semantic identity of, and a relationship between, an input and an output of a first operation provided by the first ML API, and 
 a table of operations corresponding to ML code operations that perform data selection or exclusion operations; and 
   identifying, based on the annotated WIR and the first ML API, data from a data source that is relied upon by the ML model code when training the ML model.   
     
     
         9 . The method of  claim 8 , wherein the WIR comprises a directed graph comprising a plurality of provenance relationships (PRs), each PR of the plurality of PRs corresponding to a relationship between nodes of an abstract syntax tree (AST) corresponding to the ML model code. 
     
     
         10 . The method of  claim 9 , wherein said semantically annotating the WIR comprises:
 forward traversing the directed graph to find a first PR of the plurality of PRs;   querying the queryable KB for a first input annotation corresponding to an input variable of the first PR and a first output annotation corresponding to an output variable of the first PR; and   adding the first input annotation and the first output annotation to the first PR.   
     
     
         11 . The method of  claim 10 , wherein said querying the queryable KB for a first input annotation corresponding to an input variable of the first PR and a first output annotation corresponding to an output variable of the first PR is based on a caller and an operation of the first PR. 
     
     
         12 . The method of  claim 10 , wherein said semantically annotating the WIR further comprises:
 forward traversing the directed graph to find a second PR of the plurality of PRs;   querying the queryable KB for a second input annotation corresponding to an input variable of the second PR and a second output annotation corresponding to an output variable of the second PR; and   adding the second input annotation and the second output annotation to the second PR.   
     
     
         13 . The method of  claim 12 , wherein said semantically annotating the WIR further comprises:
 starting with the second PR, backward traversing the WIR to find a third PR of the plurality of PRs, wherein the input variable of the second PR is an output variable of the third PR;   querying, based on the second input annotation, the queryable KB for a third input annotation corresponding to an input variable of the second PR; and   adding the third input annotation to the second PR.   
     
     
         14 . The method of  claim 8 , wherein said identifying the data comprises:
 determining a feature variable of a first PR, the feature variable annotated as a feature and corresponding to a data selection or exclusion operation; and   starting with the first PR, recursively backwards traversing the directed graph to identify the data from the data source that corresponds to the feature variable.   
     
     
         15 . A computer-readable storage medium having program instructions recorded thereon that, when executed a processor circuit, perform a method, the method comprising:
 obtaining a workflow intermediate representation (WIR) of machine learning (ML) model code, the ML code using a first ML application programming interface (API) that when executed trains a ML model;   semantically annotating the WIR based on the first ML API by querying a queryable knowledge base (KB), the queryable KB comprising:
 a first ML API annotation corresponding to a semantic identity of, and a relationship between, an input and an output of a first operation provided by the first ML API, and 
 a table of operations corresponding to ML code operations that perform data selection or exclusion operations; and 
   identifying, based on the annotated WIR and the first ML API, data from a data source that is relied upon by the ML model code when training the ML model.   
     
     
         16 . The computer-readable storage medium of  claim 15 , wherein the WIR comprises a directed graph comprising a plurality of provenance relationships (PRs), each PR of the plurality of PRs corresponding to a relationship between nodes of an abstract syntax tree (AST) corresponding to the ML model code. 
     
     
         17 . The computer-readable storage medium of  claim 16 , wherein said semantically annotating the WIR comprises:
 forward traversing the directed graph to find a first PR of the plurality of PRs;   querying the queryable KB for a first input annotation corresponding to an input variable of the first PR and a first output annotation corresponding to an output variable of the first PR; and   adding the first input annotation and the first output annotation to the first PR.   
     
     
         18 . The computer-readable storage medium of  claim 17 , wherein said querying the queryable KB for a first input annotation corresponding to an input variable of the first PR and a first output annotation corresponding to an output variable of the first PR is based on a caller and an operation of the first PR. 
     
     
         19 . The computer-readable storage medium of  claim 17 , wherein said semantically annotating the WIR further comprises:
 forward traversing the directed graph to find a second PR of the plurality of PRs;   querying the queryable KB for a second input annotation corresponding to an input variable of the second PR and a second output annotation corresponding to an output variable of the second PR; and   adding the second input annotation and the second output annotation to the second PR.   
     
     
         20 . The computer-readable storage medium of  claim 19 , wherein said semantically annotating the WIR further comprises:
 starting with the second PR, backward traversing the WIR to find a third PR of the plurality of PRs, wherein the input variable of the second PR is an output variable of the third PR;   querying, based on the second input annotation, the queryable KB for a third input annotation corresponding to an input variable of the second PR; and   adding the third input annotation to the second PR.