improvements from external prs

-fix corner case where dr converter failed when initializing

Closes twitter/the-algorithm#550

navi/README.md

@@ -31,6 +31,11 @@ In navi/navi, you can run the following commands:
 - `scripts/run_onnx.sh` for [Onnx](https://onnx.ai/)
 
 Do note that you need to create a models directory and create some versions, preferably using epoch time, e.g., `1679693908377`.
+so the models structure looks like:
+  models/
+       -web_click
+        - 1809000
+        - 1809010
 
 ## Build
 You can adapt the above scripts to build using Cargo.

navi/dr_transform/Cargo.toml

@@ -3,7 +3,6 @@ name = "dr_transform"
 version = "0.1.0"
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [dependencies]
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
@@ -12,7 +11,6 @@ bpr_thrift = { path = "../thrift_bpr_adapter/thrift/"}
 segdense = { path = "../segdense/"}
 thrift = "0.17.0"
 ndarray = "0.15"
-ort = {git ="https://github.com/pykeio/ort.git", tag="v1.14.2"}
 base64 = "0.20.0"
 npyz = "0.7.2"
 log = "0.4.17"
@@ -21,6 +19,11 @@ prometheus = "0.13.1"
 once_cell = "1.17.0"
 rand = "0.8.5"
 itertools = "0.10.5"
+anyhow = "1.0.70"
+[target.'cfg(not(target_os="linux"))'.dependencies]
+ort = {git ="https://github.com/pykeio/ort.git", features=["profiling"], tag="v1.14.6"}
+[target.'cfg(target_os="linux")'.dependencies]
+ort = {git ="https://github.com/pykeio/ort.git", features=["profiling", "tensorrt", "cuda", "copy-dylibs"], tag="v1.14.6"}
 [dev-dependencies]
 criterion = "0.3.0"
 

navi/dr_transform/src/all_config.rs

@@ -44,5 +44,6 @@ pub struct RenamedFeatures {
 }
 
 pub fn parse(json_str: &str) -> Result<AllConfig, Error> {
-    serde_json::from_str(json_str)
+    let all_config: AllConfig = serde_json::from_str(json_str)?;
+    Ok(all_config)
 }

navi/dr_transform/src/converter.rs

@@ -2,6 +2,9 @@ use std::collections::BTreeSet;
 use std::fmt::{self, Debug, Display};
 use std::fs;
 
+use crate::all_config;
+use crate::all_config::AllConfig;
+use anyhow::{bail, Context};
 use bpr_thrift::data::DataRecord;
 use bpr_thrift::prediction_service::BatchPredictionRequest;
 use bpr_thrift::tensor::GeneralTensor;
@@ -16,8 +19,6 @@ use segdense::util;
 use thrift::protocol::{TBinaryInputProtocol, TSerializable};
 use thrift::transport::TBufferChannel;
 
-use crate::{all_config, all_config::AllConfig};
-
 pub fn log_feature_match(
     dr: &DataRecord,
     seg_dense_config: &DensificationTransformSpec,
@@ -28,20 +29,24 @@ pub fn log_feature_match(
 
     for (feature_id, feature_value) in dr.continuous_features.as_ref().unwrap() {
         debug!(
-            "{dr_type} - Continuous Datarecord => Feature ID: {feature_id}, Feature value: {feature_value}"
+            "{} - Continous Datarecord => Feature ID: {}, Feature value: {}",
+            dr_type, feature_id, feature_value
         );
         for input_feature in &seg_dense_config.cont.input_features {
             if input_feature.feature_id == *feature_id {
-                debug!("Matching input feature: {input_feature:?}")
+                debug!("Matching input feature: {:?}", input_feature)
             }
         }
     }
 
     for feature_id in dr.binary_features.as_ref().unwrap() {
-        debug!("{dr_type} - Binary Datarecord => Feature ID: {feature_id}");
+        debug!(
+            "{} - Binary Datarecord => Feature ID: {}",
+            dr_type, feature_id
+        );
         for input_feature in &seg_dense_config.binary.input_features {
             if input_feature.feature_id == *feature_id {
-                debug!("Found input feature: {input_feature:?}")
+                debug!("Found input feature: {:?}", input_feature)
             }
         }
     }
@@ -90,18 +95,19 @@ impl BatchPredictionRequestToTorchTensorConverter {
         model_version: &str,
         reporting_feature_ids: Vec<(i64, &str)>,
         register_metric_fn: Option<impl Fn(&HistogramVec)>,
-    ) -> BatchPredictionRequestToTorchTensorConverter {
-        let all_config_path = format!("{model_dir}/{model_version}/all_config.json");
-        let seg_dense_config_path =
-            format!("{model_dir}/{model_version}/segdense_transform_spec_home_recap_2022.json");
-        let seg_dense_config = util::load_config(&seg_dense_config_path);
+    ) -> anyhow::Result<BatchPredictionRequestToTorchTensorConverter> {
+        let all_config_path = format!("{}/{}/all_config.json", model_dir, model_version);
+        let seg_dense_config_path = format!(
+            "{}/{}/segdense_transform_spec_home_recap_2022.json",
+            model_dir, model_version
+        );
+        let seg_dense_config = util::load_config(&seg_dense_config_path)?;
         let all_config = all_config::parse(
             &fs::read_to_string(&all_config_path)
-                .unwrap_or_else(|error| panic!("error loading all_config.json - {error}")),
-        )
-        .unwrap();
+                .with_context(|| "error loading all_config.json - ")?,
+        )?;
 
-        let feature_mapper = util::load_from_parsed_config_ref(&seg_dense_config);
+        let feature_mapper = util::load_from_parsed_config(seg_dense_config.clone())?;
 
         let user_embedding_feature_id = Self::get_feature_id(
             &all_config
@@ -131,11 +137,11 @@ impl BatchPredictionRequestToTorchTensorConverter {
         let (discrete_feature_metrics, continuous_feature_metrics) = METRICS.get_or_init(|| {
             let discrete = HistogramVec::new(
                 HistogramOpts::new(":navi:feature_id:discrete", "Discrete Feature ID values")
-                    .buckets(Vec::from([
-                        0.0f64, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0, 110.0,
+                    .buckets(Vec::from(&[
+                        0.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0, 110.0,
                         120.0, 130.0, 140.0, 150.0, 160.0, 170.0, 180.0, 190.0, 200.0, 250.0,
                         300.0, 500.0, 1000.0, 10000.0, 100000.0,
-                    ])),
+                    ] as &'static [f64])),
                 &["feature_id"],
             )
             .expect("metric cannot be created");
@@ -144,18 +150,18 @@ impl BatchPredictionRequestToTorchTensorConverter {
                     ":navi:feature_id:continuous",
                     "continuous Feature ID values",
                 )
-                .buckets(Vec::from([
-                    0.0f64, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0, 110.0,
-                    120.0, 130.0, 140.0, 150.0, 160.0, 170.0, 180.0, 190.0, 200.0, 250.0, 300.0,
-                    500.0, 1000.0, 10000.0, 100000.0,
-                ])),
+                .buckets(Vec::from(&[
+                    0.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0, 110.0, 120.0,
+                    130.0, 140.0, 150.0, 160.0, 170.0, 180.0, 190.0, 200.0, 250.0, 300.0, 500.0,
+                    1000.0, 10000.0, 100000.0,
+                ] as &'static [f64])),
                 &["feature_id"],
             )
             .expect("metric cannot be created");
-            if let Some(r) = register_metric_fn {
+            register_metric_fn.map(|r| {
                 r(&discrete);
                 r(&continuous);
-            }
+            });
             (discrete, continuous)
         });
 
@@ -164,13 +170,16 @@ impl BatchPredictionRequestToTorchTensorConverter {
 
         for (feature_id, feature_type) in reporting_feature_ids.iter() {
             match *feature_type {
-                "discrete" => discrete_features_to_report.insert(*feature_id),
-                "continuous" => continuous_features_to_report.insert(*feature_id),
-                _ => panic!("Invalid feature type {feature_type} for reporting metrics!"),
+                "discrete" => discrete_features_to_report.insert(feature_id.clone()),
+                "continuous" => continuous_features_to_report.insert(feature_id.clone()),
+                _ => bail!(
+                    "Invalid feature type {} for reporting metrics!",
+                    feature_type
+                ),
             };
         }
 
-        BatchPredictionRequestToTorchTensorConverter {
+        Ok(BatchPredictionRequestToTorchTensorConverter {
             all_config,
             seg_dense_config,
             all_config_path,
@@ -183,7 +192,7 @@ impl BatchPredictionRequestToTorchTensorConverter {
             continuous_features_to_report,
             discrete_feature_metrics,
             continuous_feature_metrics,
-        }
+        })
     }
 
     fn get_feature_id(feature_name: &str, seg_dense_config: &Root) -> i64 {
@@ -218,43 +227,45 @@ impl BatchPredictionRequestToTorchTensorConverter {
         let mut working_set = vec![0 as f32; total_size];
         let mut bpr_start = 0;
         for (bpr, &bpr_end) in bprs.iter().zip(batch_size) {
-            if bpr.common_features.is_some()
-                && bpr.common_features.as_ref().unwrap().tensors.is_some()
-                && bpr
-                    .common_features
-                    .as_ref()
-                    .unwrap()
-                    .tensors
-                    .as_ref()
-                    .unwrap()
-                    .contains_key(&feature_id)
-            {
-                let source_tensor = bpr
-                    .common_features
-                    .as_ref()
-                    .unwrap()
-                    .tensors
-                    .as_ref()
-                    .unwrap()
-                    .get(&feature_id)
-                    .unwrap();
-                let tensor = match source_tensor {
-                    GeneralTensor::FloatTensor(float_tensor) =>
-                    //Tensor::of_slice(
+            if bpr.common_features.is_some() {
+                if bpr.common_features.as_ref().unwrap().tensors.is_some() {
+                    if bpr
+                        .common_features
+                        .as_ref()
+                        .unwrap()
+                        .tensors
+                        .as_ref()
+                        .unwrap()
+                        .contains_key(&feature_id)
                     {
-                        float_tensor
-                            .floats
-                            .iter()
-                            .map(|x| x.into_inner() as f32)
-                            .collect::<Vec<_>>()
-                    }
-                    _ => vec![0 as f32; cols],
-                };
+                        let source_tensor = bpr
+                            .common_features
+                            .as_ref()
+                            .unwrap()
+                            .tensors
+                            .as_ref()
+                            .unwrap()
+                            .get(&feature_id)
+                            .unwrap();
+                        let tensor = match source_tensor {
+                            GeneralTensor::FloatTensor(float_tensor) =>
+                            //Tensor::of_slice(
+                            {
+                                float_tensor
+                                    .floats
+                                    .iter()
+                                    .map(|x| x.into_inner() as f32)
+                                    .collect::<Vec<_>>()
+                            }
+                            _ => vec![0 as f32; cols],
+                        };
 
-                // since the tensor is found in common feature, add it in all batches
-                for row in bpr_start..bpr_end {
-                    for col in 0..cols {
-                        working_set[row * cols + col] = tensor[col];
+                        // since the tensor is found in common feature, add it in all batches
+                        for row in bpr_start..bpr_end {
+                            for col in 0..cols {
+                                working_set[row * cols + col] = tensor[col];
+                            }
+                        }
                     }
                 }
             }
@@ -298,9 +309,9 @@ impl BatchPredictionRequestToTorchTensorConverter {
     //           (INT64 --> INT64, DataRecord.discrete_feature)
     fn get_continuous(&self, bprs: &[BatchPredictionRequest], batch_ends: &[usize]) -> InputTensor {
         // These need to be part of model schema
-        let rows = batch_ends[batch_ends.len() - 1];
-        let cols = 5293;
-        let full_size = rows * cols;
+        let rows: usize = batch_ends[batch_ends.len() - 1];
+        let cols: usize = 5293;
+        let full_size: usize = rows * cols;
         let default_val = f32::NAN;
 
         let mut tensor = vec![default_val; full_size];
@@ -325,15 +336,18 @@ impl BatchPredictionRequestToTorchTensorConverter {
                     .unwrap();
 
                 for feature in common_features {
-                    if let Some(f_info) = self.feature_mapper.get(feature.0) {
-                        let idx = f_info.index_within_tensor as usize;
-                        if idx < cols {
-                            // Set value in each row
-                            for r in bpr_start..bpr_end {
-                                let flat_index = r * cols + idx;
-                                tensor[flat_index] = feature.1.into_inner() as f32;
+                    match self.feature_mapper.get(feature.0) {
+                        Some(f_info) => {
+                            let idx = f_info.index_within_tensor as usize;
+                            if idx < cols {
+                                // Set value in each row
+                                for r in bpr_start..bpr_end {
+                                    let flat_index: usize = r * cols + idx;
+                                    tensor[flat_index] = feature.1.into_inner() as f32;
+                                }
                             }
                         }
+                        None => (),
                     }
                     if self.continuous_features_to_report.contains(feature.0) {
                         self.continuous_feature_metrics
@@ -349,24 +363,28 @@ impl BatchPredictionRequestToTorchTensorConverter {
 
             // Process the batch of datarecords
             for r in bpr_start..bpr_end {
-                let dr: &DataRecord = &bpr.individual_features_list[r - bpr_start];
+                let dr: &DataRecord =
+                    &bpr.individual_features_list[usize::try_from(r - bpr_start).unwrap()];
                 if dr.continuous_features.is_some() {
                     for feature in dr.continuous_features.as_ref().unwrap() {
-                        if let Some(f_info) = self.feature_mapper.get(feature.0) {
-                            let idx = f_info.index_within_tensor as usize;
-                            let flat_index = r * cols + idx;
-                            if flat_index < tensor.len() && idx < cols {
-                                tensor[flat_index] = feature.1.into_inner() as f32;
+                        match self.feature_mapper.get(&feature.0) {
+                            Some(f_info) => {
+                                let idx = f_info.index_within_tensor as usize;
+                                let flat_index: usize = r * cols + idx;
+                                if flat_index < tensor.len() && idx < cols {
+                                    tensor[flat_index] = feature.1.into_inner() as f32;
+                                }
                             }
+                            None => (),
                         }
                         if self.continuous_features_to_report.contains(feature.0) {
                             self.continuous_feature_metrics
                                 .with_label_values(&[feature.0.to_string().as_str()])
-                                .observe(feature.1.into_inner())
+                                .observe(feature.1.into_inner() as f64)
                         } else if self.discrete_features_to_report.contains(feature.0) {
                             self.discrete_feature_metrics
                                 .with_label_values(&[feature.0.to_string().as_str()])
-                                .observe(feature.1.into_inner())
+                                .observe(feature.1.into_inner() as f64)
                         }
                     }
                 }
@@ -383,10 +401,10 @@ impl BatchPredictionRequestToTorchTensorConverter {
 
     fn get_binary(&self, bprs: &[BatchPredictionRequest], batch_ends: &[usize]) -> InputTensor {
         // These need to be part of model schema
-        let rows = batch_ends[batch_ends.len() - 1];
-        let cols = 149;
-        let full_size = rows * cols;
-        let default_val = 0;
+        let rows: usize = batch_ends[batch_ends.len() - 1];
+        let cols: usize = 149;
+        let full_size: usize = rows * cols;
+        let default_val: i64 = 0;
 
         let mut v = vec![default_val; full_size];
 
@@ -410,15 +428,18 @@ impl BatchPredictionRequestToTorchTensorConverter {
                     .unwrap();
 
                 for feature in common_features {
-                    if let Some(f_info) = self.feature_mapper.get(feature) {
-                        let idx = f_info.index_within_tensor as usize;
-                        if idx < cols {
-                            // Set value in each row
-                            for r in bpr_start..bpr_end {
-                                let flat_index = r * cols + idx;
-                                v[flat_index] = 1;
+                    match self.feature_mapper.get(feature) {
+                        Some(f_info) => {
+                            let idx = f_info.index_within_tensor as usize;
+                            if idx < cols {
+                                // Set value in each row
+                                for r in bpr_start..bpr_end {
+                                    let flat_index: usize = r * cols + idx;
+                                    v[flat_index] = 1;
+                                }
                             }
                         }
+                        None => (),
                     }
                 }
             }
@@ -428,10 +449,13 @@ impl BatchPredictionRequestToTorchTensorConverter {
                 let dr: &DataRecord = &bpr.individual_features_list[r - bpr_start];
                 if dr.binary_features.is_some() {
                     for feature in dr.binary_features.as_ref().unwrap() {
-                        if let Some(f_info) = self.feature_mapper.get(feature) {
-                            let idx = f_info.index_within_tensor as usize;
-                            let flat_index = r * cols + idx;
-                            v[flat_index] = 1;
+                        match self.feature_mapper.get(&feature) {
+                            Some(f_info) => {
+                                let idx = f_info.index_within_tensor as usize;
+                                let flat_index: usize = r * cols + idx;
+                                v[flat_index] = 1;
+                            }
+                            None => (),
                         }
                     }
                 }
@@ -448,10 +472,10 @@ impl BatchPredictionRequestToTorchTensorConverter {
     #[allow(dead_code)]
     fn get_discrete(&self, bprs: &[BatchPredictionRequest], batch_ends: &[usize]) -> InputTensor {
         // These need to be part of model schema
-        let rows = batch_ends[batch_ends.len() - 1];
-        let cols = 320;
-        let full_size = rows * cols;
-        let default_val = 0;
+        let rows: usize = batch_ends[batch_ends.len() - 1];
+        let cols: usize = 320;
+        let full_size: usize = rows * cols;
+        let default_val: i64 = 0;
 
         let mut v = vec![default_val; full_size];
 
@@ -475,15 +499,18 @@ impl BatchPredictionRequestToTorchTensorConverter {
                     .unwrap();
 
                 for feature in common_features {
-                    if let Some(f_info) = self.feature_mapper.get(feature.0) {
-                        let idx = f_info.index_within_tensor as usize;
-                        if idx < cols {
-                            // Set value in each row
-                            for r in bpr_start..bpr_end {
-                                let flat_index = r * cols + idx;
-                                v[flat_index] = *feature.1;
+                    match self.feature_mapper.get(feature.0) {
+                        Some(f_info) => {
+                            let idx = f_info.index_within_tensor as usize;
+                            if idx < cols {
+                                // Set value in each row
+                                for r in bpr_start..bpr_end {
+                                    let flat_index: usize = r * cols + idx;
+                                    v[flat_index] = *feature.1;
+                                }
                             }
                         }
+                        None => (),
                     }
                     if self.discrete_features_to_report.contains(feature.0) {
                         self.discrete_feature_metrics
@@ -495,15 +522,18 @@ impl BatchPredictionRequestToTorchTensorConverter {
 
             // Process the batch of datarecords
             for r in bpr_start..bpr_end {
-                let dr: &DataRecord = &bpr.individual_features_list[r];
+                let dr: &DataRecord = &bpr.individual_features_list[usize::try_from(r).unwrap()];
                 if dr.discrete_features.is_some() {
                     for feature in dr.discrete_features.as_ref().unwrap() {
-                        if let Some(f_info) = self.feature_mapper.get(feature.0) {
-                            let idx = f_info.index_within_tensor as usize;
-                            let flat_index = r * cols + idx;
-                            if flat_index < v.len() && idx < cols {
-                                v[flat_index] = *feature.1;
+                        match self.feature_mapper.get(&feature.0) {
+                            Some(f_info) => {
+                                let idx = f_info.index_within_tensor as usize;
+                                let flat_index: usize = r * cols + idx;
+                                if flat_index < v.len() && idx < cols {
+                                    v[flat_index] = *feature.1;
+                                }
                             }
+                            None => (),
                         }
                         if self.discrete_features_to_report.contains(feature.0) {
                             self.discrete_feature_metrics
@@ -569,7 +599,7 @@ impl Converter for BatchPredictionRequestToTorchTensorConverter {
             .map(|bpr| bpr.individual_features_list.len())
             .scan(0usize, |acc, e| {
                 //running total
-                *acc += e;
+                *acc = *acc + e;
                 Some(*acc)
             })
             .collect::<Vec<_>>();

navi/dr_transform/src/lib.rs

@@ -3,3 +3,4 @@ pub mod converter;
 #[cfg(test)]
 mod test;
 pub mod util;
+pub extern crate ort;

navi/navi/Cargo.toml

@@ -1,8 +1,7 @@
 [package]
 name = "navi"
-version = "2.0.42"
+version = "2.0.45"
 edition = "2021"
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [[bin]]
 name = "navi"
@@ -16,12 +15,19 @@ required-features=["torch"]
 name = "navi_onnx"
 path = "src/bin/navi_onnx.rs"
 required-features=["onnx"]
+[[bin]]
+name = "navi_onnx_test"
+path = "src/bin/bin_tests/navi_onnx_test.rs"
+[[bin]]
+name = "navi_torch_test"
+path = "src/bin/bin_tests/navi_torch_test.rs"
+required-features=["torch"]
 
 [features]
 default=[]
 navi_console=[]
 torch=["tch"]
-onnx=["ort"]
+onnx=[]
 tf=["tensorflow"]
 [dependencies]
 itertools = "0.10.5"
@@ -47,6 +53,7 @@ parking_lot = "0.12.1"
 rand = "0.8.5"
 rand_pcg = "0.3.1"
 random = "0.12.2"
+x509-parser = "0.15.0"
 sha256 = "1.0.3"
 tonic = { version = "0.6.2", features=['compression', 'tls'] }
 tokio = { version = "1.17.0", features = ["macros", "rt-multi-thread", "fs", "process"] }
@@ -55,16 +62,12 @@ npyz = "0.7.3"
 base64 = "0.21.0"
 histogram = "0.6.9"
 tch = {version = "0.10.3", optional = true}
-tensorflow = { version = "0.20.0", optional = true }
+tensorflow = { version = "0.18.0", optional = true }
 once_cell = {version = "1.17.1"}
 ndarray = "0.15"
 serde = "1.0.154"
 serde_json = "1.0.94"
 dr_transform = { path = "../dr_transform"}
-[target.'cfg(not(target_os="linux"))'.dependencies]
-ort = {git ="https://github.com/pykeio/ort.git", features=["profiling"], optional = true, tag="v1.14.2"}
-[target.'cfg(target_os="linux")'.dependencies]
-ort = {git ="https://github.com/pykeio/ort.git", features=["profiling", "tensorrt", "cuda", "copy-dylibs"], optional = true, tag="v1.14.2"}
 [build-dependencies]
 tonic-build = {version = "0.6.2", features=['prost', "compression"] }
 [profile.release]
@@ -74,3 +77,5 @@ ndarray-rand = "0.14.0"
 tokio-test = "*"
 assert_cmd = "2.0"
 criterion = "0.4.0"
+
+

navi/navi/proto/tensorflow/core/framework/full_type.proto

@@ -122,7 +122,7 @@ enum FullTypeId {
   //   TFT_TENSOR[TFT_INT32, TFT_UNKNOWN]
   //     is a Tensor of int32 element type and unknown shape.
   //
-  // TODO: Define TFT_SHAPE and add more examples.
+  // TODO(mdan): Define TFT_SHAPE and add more examples.
   TFT_TENSOR = 1000;
 
   // Array (or tensorflow::TensorList in the variant type registry).
@@ -178,7 +178,7 @@ enum FullTypeId {
   // object (for now).
 
   // The bool element type.
-  // TODO
+  // TODO(mdan): Quantized types, legacy representations (e.g. ref)
   TFT_BOOL = 200;
   // Integer element types.
   TFT_UINT8 = 201;
@@ -195,7 +195,7 @@ enum FullTypeId {
   TFT_DOUBLE = 211;
   TFT_BFLOAT16 = 215;
   // Complex element types.
-  // TODO: Represent as TFT_COMPLEX[TFT_DOUBLE] instead?
+  // TODO(mdan): Represent as TFT_COMPLEX[TFT_DOUBLE] instead?
   TFT_COMPLEX64 = 212;
   TFT_COMPLEX128 = 213;
   // The string element type.
@@ -240,7 +240,7 @@ enum FullTypeId {
   // ownership is in the true sense: "the op argument representing the lock is
   // available".
   // Mutex locks are the dynamic counterpart of control dependencies.
-  // TODO: Properly document this thing.
+  // TODO(mdan): Properly document this thing.
   //
   // Parametrization: TFT_MUTEX_LOCK[].
   TFT_MUTEX_LOCK = 10202;
@@ -271,6 +271,6 @@ message FullTypeDef {
   oneof attr {
     string s = 3;
     int64 i = 4;
-    // TODO: list/tensor, map? Need to reconcile with TFT_RECORD, etc.
+    // TODO(mdan): list/tensor, map? Need to reconcile with TFT_RECORD, etc.
   }
 }

navi/navi/proto/tensorflow/core/framework/function.proto

@@ -23,7 +23,7 @@ message FunctionDefLibrary {
 // with a value. When a GraphDef has a call to a function, it must
 // have binding for every attr defined in the signature.
 //
-// TODO:
+// TODO(zhifengc):
 //   * device spec, etc.
 message FunctionDef {
   // The definition of the function's name, arguments, return values,

navi/navi/proto/tensorflow/core/framework/node_def.proto

@@ -61,7 +61,7 @@ message NodeDef {
   // one of the names from the corresponding OpDef's attr field).
   // The values must have a type matching the corresponding OpDef
   // attr's type field.
-  // TODO: Add some examples here showing best practices.
+  // TODO(josh11b): Add some examples here showing best practices.
   map<string, AttrValue> attr = 5;
 
   message ExperimentalDebugInfo {

navi/navi/proto/tensorflow/core/framework/op_def.proto

@@ -96,7 +96,7 @@ message OpDef {
     // Human-readable description.
     string description = 4;
 
-    // TODO: bool is_optional?
+    // TODO(josh11b): bool is_optional?
 
     // --- Constraints ---
     // These constraints are only in effect if specified.  Default is no
@@ -139,7 +139,7 @@ message OpDef {
   // taking input from multiple devices with a tree of aggregate ops
   // that aggregate locally within each device (and possibly within
   // groups of nearby devices) before communicating.
-  // TODO: Implement that optimization.
+  // TODO(josh11b): Implement that optimization.
   bool is_aggregate = 16;  // for things like add
 
   // Other optimizations go here, like

navi/navi/proto/tensorflow/core/framework/step_stats.proto

@@ -53,7 +53,7 @@ message MemoryStats {
 
 // Time/size stats recorded for a single execution of a graph node.
 message NodeExecStats {
-  // TODO: Use some more compact form of node identity than
+  // TODO(tucker): Use some more compact form of node identity than
   // the full string name.  Either all processes should agree on a
   // global id (cost_id?) for each node, or we should use a hash of
   // the name.

navi/navi/proto/tensorflow/core/framework/tensor.proto

@@ -16,7 +16,7 @@ option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framewo
 message TensorProto {
   DataType dtype = 1;
 
-  // Shape of the tensor.  TODO: sort out the 0-rank issues.
+  // Shape of the tensor.  TODO(touts): sort out the 0-rank issues.
   TensorShapeProto tensor_shape = 2;
 
   // Only one of the representations below is set, one of "tensor_contents" and

navi/navi/proto/tensorflow/core/protobuf/config.proto

@@ -532,7 +532,7 @@ message ConfigProto {
 
     // We removed the flag client_handles_error_formatting. Marking the tag
     // number as reserved.
-    // TODO: Should we just remove this tag so that it can be
+    // TODO(shikharagarwal): Should we just remove this tag so that it can be
     // used in future for other purpose?
     reserved 2;
 
@@ -576,7 +576,7 @@ message ConfigProto {
     // - If isolate_session_state is true, session states are isolated.
     // - If isolate_session_state is false, session states are shared.
     //
-    // TODO: Add a single API that consistently treats
+    // TODO(b/129330037): Add a single API that consistently treats
     // isolate_session_state and ClusterSpec propagation.
     bool share_session_state_in_clusterspec_propagation = 8;
 
@@ -704,7 +704,7 @@ message ConfigProto {
 
 // Options for a single Run() call.
 message RunOptions {
-  // TODO Turn this into a TraceOptions proto which allows
+  // TODO(pbar) Turn this into a TraceOptions proto which allows
   // tracing to be controlled in a more orthogonal manner?
   enum TraceLevel {
     NO_TRACE = 0;
@@ -781,7 +781,7 @@ message RunMetadata {
   repeated GraphDef partition_graphs = 3;
 
   message FunctionGraphs {
-    // TODO: Include some sort of function/cache-key identifier?
+    // TODO(nareshmodi): Include some sort of function/cache-key identifier?
     repeated GraphDef partition_graphs = 1;
 
     GraphDef pre_optimization_graph = 2;

navi/navi/proto/tensorflow/core/protobuf/coordination_service.proto

@@ -194,7 +194,7 @@ service CoordinationService {
 
   // Report error to the task. RPC sets the receiving instance of coordination
   // service agent to error state permanently.
-  // TODO: Consider splitting this into a different RPC service.
+  // TODO(b/195990880): Consider splitting this into a different RPC service.
   rpc ReportErrorToAgent(ReportErrorToAgentRequest)
       returns (ReportErrorToAgentResponse);
 

navi/navi/proto/tensorflow/core/protobuf/debug.proto

@@ -46,7 +46,7 @@ message DebugTensorWatch {
   // are to be debugged, the callers of Session::Run() must use distinct
   // debug_urls to make sure that the streamed or dumped events do not overlap
   // among the invocations.
-  // TODO: More visible documentation of this in g3docs.
+  // TODO(cais): More visible documentation of this in g3docs.
   repeated string debug_urls = 4;
 
   // Do not error out if debug op creation fails (e.g., due to dtype

navi/navi/proto/tensorflow/core/protobuf/debug_event.proto

@@ -12,7 +12,7 @@ option java_package = "org.tensorflow.util";
 option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
 
 // Available modes for extracting debugging information from a Tensor.
-// TODO: Document the detailed column names and semantics in a separate
+// TODO(cais): Document the detailed column names and semantics in a separate
 // markdown file once the implementation settles.
 enum TensorDebugMode {
   UNSPECIFIED = 0;
@@ -223,7 +223,7 @@ message DebuggedDevice {
   // A debugger-generated ID for the device. Guaranteed to be unique within
   // the scope of the debugged TensorFlow program, including single-host and
   // multi-host settings.
-  // TODO: Test the uniqueness guarantee in multi-host settings.
+  // TODO(cais): Test the uniqueness guarantee in multi-host settings.
   int32 device_id = 2;
 }
 
@@ -264,7 +264,7 @@ message Execution {
   // field with the DebuggedDevice messages.
   repeated int32 output_tensor_device_ids = 9;
 
-  // TODO support, add more fields
+  // TODO(cais): When backporting to V1 Session.run() support, add more fields
   // such as fetches and feeds.
 }
 

navi/navi/proto/tensorflow/core/protobuf/distributed_runtime_payloads.proto

@@ -7,7 +7,7 @@ option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobu
 
 // Used to serialize and transmit tensorflow::Status payloads through
 // grpc::Status `error_details` since grpc::Status lacks payload API.
-// TODO: Use GRPC API once supported.
+// TODO(b/204231601): Use GRPC API once supported.
 message GrpcPayloadContainer {
   map<string, bytes> payloads = 1;
 }

navi/navi/proto/tensorflow/core/protobuf/eager_service.proto

@@ -172,7 +172,7 @@ message WaitQueueDoneRequest {
 }
 
 message WaitQueueDoneResponse {
-  // TODO: Consider adding NodeExecStats here to be able to
+  // TODO(nareshmodi): Consider adding NodeExecStats here to be able to
   // propagate some stats.
 }
 

navi/navi/proto/tensorflow/core/protobuf/master.proto

@@ -94,7 +94,7 @@ message ExtendSessionRequest {
 }
 
 message ExtendSessionResponse {
-  // TODO: Return something about the operation?
+  // TODO(mrry): Return something about the operation?
 
   // The new version number for the extended graph, to be used in the next call
   // to ExtendSession.

navi/navi/proto/tensorflow/core/protobuf/saved_object_graph.proto

@@ -176,7 +176,7 @@ message SavedBareConcreteFunction {
   // allows the ConcreteFunction to be called with nest structure inputs. This
   // field may not be populated. If this field is absent, the concrete function
   // can only be called with flat inputs.
-  // TODO: support calling saved ConcreteFunction with structured
+  // TODO(b/169361281): support calling saved ConcreteFunction with structured
   // inputs in C++ SavedModel API.
   FunctionSpec function_spec = 4;
 }

navi/navi/proto/tensorflow/core/protobuf/tensor_bundle.proto

@@ -17,7 +17,7 @@ option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobu
 
 // Special header that is associated with a bundle.
 //
-// TODO: maybe in the future, we can add information about
+// TODO(zongheng,zhifengc): maybe in the future, we can add information about
 // which binary produced this checkpoint, timestamp, etc. Sometime, these can be
 // valuable debugging information. And if needed, these can be used as defensive
 // information ensuring reader (binary version) of the checkpoint and the writer

navi/navi/proto/tensorflow/core/protobuf/worker.proto

@@ -188,7 +188,7 @@ message DeregisterGraphRequest {
 }
 
 message DeregisterGraphResponse {
-  // TODO: Optionally add summary stats for the graph.
+  // TODO(mrry): Optionally add summary stats for the graph.
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -294,7 +294,7 @@ message RunGraphResponse {
 
   // If the request asked for execution stats, the cost graph, or the partition
   // graphs, these are returned here.
-  // TODO: Package these in a RunMetadata instead.
+  // TODO(suharshs): Package these in a RunMetadata instead.
   StepStats step_stats = 2;
   CostGraphDef cost_graph = 3;
   repeated GraphDef partition_graph = 4;

navi/navi/proto/tensorflow_serving/apis/logging.proto

@@ -13,5 +13,5 @@ message LogMetadata {
   SamplingConfig sampling_config = 2;
   // List of tags used to load the relevant MetaGraphDef from SavedModel.
   repeated string saved_model_tags = 3;
-  // TODO: Add more metadata as mentioned in the bug.
+  // TODO(b/33279154): Add more metadata as mentioned in the bug.
 }

navi/navi/proto/tensorflow_serving/config/file_system_storage_path_source.proto

@@ -58,7 +58,7 @@ message FileSystemStoragePathSourceConfig {
 
   // A single servable name/base_path pair to monitor.
   // DEPRECATED: Use 'servables' instead.
-  // TODO: Stop using these fields, and ultimately remove them here.
+  // TODO(b/30898016): Stop using these fields, and ultimately remove them here.
   string servable_name = 1 [deprecated = true];
   string base_path = 2 [deprecated = true];
 
@@ -76,7 +76,7 @@ message FileSystemStoragePathSourceConfig {
   // check for a version to appear later.)
   // DEPRECATED: Use 'servable_versions_always_present' instead, which includes
   // this behavior.
-  // TODO: Remove 2019-10-31 or later.
+  // TODO(b/30898016): Remove 2019-10-31 or later.
   bool fail_if_zero_versions_at_startup = 4 [deprecated = true];
 
   // If true, the servable is always expected to exist on the underlying

navi/navi/proto/tensorflow_serving/config/model_server_config.proto

@@ -9,7 +9,7 @@ import "tensorflow_serving/config/logging_config.proto";
 option cc_enable_arenas = true;
 
 // The type of model.
-// TODO: DEPRECATED.
+// TODO(b/31336131): DEPRECATED.
 enum ModelType {
   MODEL_TYPE_UNSPECIFIED = 0 [deprecated = true];
   TENSORFLOW = 1 [deprecated = true];
@@ -31,7 +31,7 @@ message ModelConfig {
   string base_path = 2;
 
   // Type of model.
-  // TODO: DEPRECATED. Please use 'model_platform' instead.
+  // TODO(b/31336131): DEPRECATED. Please use 'model_platform' instead.
   ModelType model_type = 3 [deprecated = true];
 
   // Type of model (e.g. "tensorflow").

navi/navi/scripts/run_onnx.sh

@@ -1,10 +1,9 @@
 #!/bin/sh
 #RUST_LOG=debug LD_LIBRARY_PATH=so/onnx/lib target/release/navi_onnx --port 30 --num-worker-threads 8 --intra-op-parallelism 8 --inter-op-parallelism 8 \
 RUST_LOG=info LD_LIBRARY_PATH=so/onnx/lib cargo run --bin navi_onnx --features onnx -- \
-    --port 30 --num-worker-threads 8 --intra-op-parallelism 8 --inter-op-parallelism 8 \
+    --port 8030 --num-worker-threads 8 \
     --model-check-interval-secs 30 \
-    --model-dir models/int8 \
-    --output caligrated_probabilities \
-    --input "" \
     --modelsync-cli "echo" \
-    --onnx-ep-options use_arena=true
+    --onnx-ep-options use_arena=true \
+    --model-dir models/prod_home --output caligrated_probabilities  --input "" --intra-op-parallelism 8 --inter-op-parallelism 8 --max-batch-size 1 --batch-time-out-millis 1 \
+    --model-dir models/prod_home1 --output caligrated_probabilities  --input "" --intra-op-parallelism 8 --inter-op-parallelism 8 --max-batch-size 1 --batch-time-out-millis 1 \

navi/navi/src/bin/navi_onnx.rs

@@ -1,11 +1,24 @@
 use anyhow::Result;
+use log::info;
 use navi::cli_args::{ARGS, MODEL_SPECS};
 use navi::onnx_model::onnx::OnnxModel;
 use navi::{bootstrap, metrics};
 
 fn main() -> Result<()> {
     env_logger::init();
-    assert_eq!(MODEL_SPECS.len(), ARGS.inter_op_parallelism.len());
+    info!("global: {:?}", ARGS.onnx_global_thread_pool_options);
+    let assert_session_params = if ARGS.onnx_global_thread_pool_options.is_empty() {
+        // std::env::set_var("OMP_NUM_THREADS", "1");
+        info!("now we use per session thread pool");
+        MODEL_SPECS.len()
+    }
+    else {
+       info!("now we use global thread pool");
+        0
+    };
+    assert_eq!(assert_session_params, ARGS.inter_op_parallelism.len());
+    assert_eq!(assert_session_params, ARGS.inter_op_parallelism.len());
+
     metrics::register_custom_metrics();
     bootstrap::bootstrap(OnnxModel::new)
 }

navi/navi/src/bootstrap.rs

@@ -1,5 +1,6 @@
 use anyhow::Result;
 use log::{info, warn};
+use x509_parser::{prelude::{parse_x509_pem}, parse_x509_certificate};
 use std::collections::HashMap;
 use tokio::time::Instant;
 use tonic::{
@@ -27,6 +28,7 @@ use crate::cli_args::{ARGS, INPUTS, OUTPUTS};
 use crate::metrics::{
     NAVI_VERSION, NUM_PREDICTIONS, NUM_REQUESTS_FAILED, NUM_REQUESTS_FAILED_BY_MODEL,
     NUM_REQUESTS_RECEIVED, NUM_REQUESTS_RECEIVED_BY_MODEL, RESPONSE_TIME_COLLECTOR,
+    CERT_EXPIRY_EPOCH
 };
 use crate::predict_service::{Model, PredictService};
 use crate::tf_proto::tensorflow_serving::model_spec::VersionChoice::Version;
@@ -207,6 +209,9 @@ impl<T: Model> PredictionService for PredictService<T> {
                         PredictResult::DropDueToOverload => Err(Status::resource_exhausted("")),
                         PredictResult::ModelNotFound(idx) => {
                             Err(Status::not_found(format!("model index {}", idx)))
+                        },
+                        PredictResult::ModelNotReady(idx) => {
+                            Err(Status::unavailable(format!("model index {}", idx)))
                         }
                         PredictResult::ModelVersionNotFound(idx, version) => Err(
                             Status::not_found(format!("model index:{}, version {}", idx, version)),
@@ -230,6 +235,12 @@ impl<T: Model> PredictionService for PredictService<T> {
     }
 }
 
+// A function that takes a timestamp as input and returns a ticker stream
+fn report_expiry(expiry_time: i64) {
+    info!("Certificate expires at epoch: {:?}", expiry_time);
+    CERT_EXPIRY_EPOCH.set(expiry_time as i64);
+}
+
 pub fn bootstrap<T: Model>(model_factory: ModelFactory<T>) -> Result<()> {
     info!("package: {}, version: {}, args: {:?}", NAME, VERSION, *ARGS);
     //we follow SemVer. So here we assume MAJOR.MINOR.PATCH
@@ -246,6 +257,7 @@ pub fn bootstrap<T: Model>(model_factory: ModelFactory<T>) -> Result<()> {
         );
     }
 
+    
     tokio::runtime::Builder::new_multi_thread()
         .thread_name("async worker")
         .worker_threads(ARGS.num_worker_threads)
@@ -263,6 +275,21 @@ pub fn bootstrap<T: Model>(model_factory: ModelFactory<T>) -> Result<()> {
             let mut builder = if ARGS.ssl_dir.is_empty() {
                 Server::builder()
             } else {
+                // Read the pem file as a string
+                let pem_str = std::fs::read_to_string(format!("{}/server.crt", ARGS.ssl_dir)).unwrap();
+                let res = parse_x509_pem(&pem_str.as_bytes());
+                match res {
+                    Ok((rem, pem_2)) => {
+                        assert!(rem.is_empty());
+                        assert_eq!(pem_2.label, String::from("CERTIFICATE"));
+                        let res_x509 = parse_x509_certificate(&pem_2.contents);
+                        info!("Certificate label: {}", pem_2.label);
+                        assert!(res_x509.is_ok());
+                        report_expiry(res_x509.unwrap().1.validity().not_after.timestamp());
+                    },
+                    _ => panic!("PEM parsing failed: {:?}", res),
+                }
+
                 let key = tokio::fs::read(format!("{}/server.key", ARGS.ssl_dir))
                     .await
                     .expect("can't find key file");
@@ -278,7 +305,7 @@ pub fn bootstrap<T: Model>(model_factory: ModelFactory<T>) -> Result<()> {
                 let identity = Identity::from_pem(pem.clone(), key);
                 let client_ca_cert = Certificate::from_pem(pem.clone());
                 let tls = ServerTlsConfig::new()
-                    .identity(identity)
+                    .identity(identity) 
                     .client_ca_root(client_ca_cert);
                 Server::builder()
                     .tls_config(tls)

navi/navi/src/cli_args.rs

@@ -87,13 +87,11 @@ pub struct Args {
     pub intra_op_parallelism: Vec<String>,
     #[clap(
         long,
-        default_value = "14",
         help = "number of threads to parallelize computations of the graph"
     )]
     pub inter_op_parallelism: Vec<String>,
     #[clap(
         long,
-        default_value = "serving_default",
         help = "signature of a serving. only TF"
     )]
     pub serving_sig: Vec<String>,
@@ -107,10 +105,12 @@ pub struct Args {
         help = "max warmup records to use. warmup only implemented for TF"
     )]
     pub max_warmup_records: usize,
+    #[clap(long, value_parser = Args::parse_key_val::<String, String>, value_delimiter=',')]
+    pub onnx_global_thread_pool_options: Vec<(String, String)>,
     #[clap(
-        long,
-        default_value = "true",
-        help = "when to use graph parallelization. only for ONNX"
+    long,
+    default_value = "true",
+    help = "when to use graph parallelization. only for ONNX"
     )]
     pub onnx_use_parallel_mode: String,
     // #[clap(long, default_value = "false")]

navi/navi/src/lib.rs

@@ -146,6 +146,7 @@ pub enum PredictResult {
     Ok(Vec<TensorScores>, i64),
     DropDueToOverload,
     ModelNotFound(usize),
+    ModelNotReady(usize),
     ModelVersionNotFound(usize, i64),
 }
 

navi/navi/src/metrics.rs

@@ -171,6 +171,9 @@ lazy_static! {
         &["model_name"]
     )
     .expect("metric can be created");
+    pub static ref CERT_EXPIRY_EPOCH: IntGauge =
+        IntGauge::new(":navi:cert_expiry_epoch", "Timestamp when the current cert expires")
+            .expect("metric can be created");
 }
 
 pub fn register_custom_metrics() {
@@ -249,6 +252,10 @@ pub fn register_custom_metrics() {
     REGISTRY
         .register(Box::new(CONVERTER_TIME_COLLECTOR.clone()))
         .expect("collector can be registered");
+    REGISTRY
+    .register(Box::new(CERT_EXPIRY_EPOCH.clone()))
+    .expect("collector can be registered");
+
 }
 
 pub fn register_dynamic_metrics(c: &HistogramVec) {

navi/navi/src/onnx_model.rs

@@ -13,21 +13,22 @@ pub mod onnx {
     use dr_transform::converter::{BatchPredictionRequestToTorchTensorConverter, Converter};
     use itertools::Itertools;
     use log::{debug, info};
-    use ort::environment::Environment;
-    use ort::session::Session;
-    use ort::tensor::InputTensor;
-    use ort::{ExecutionProvider, GraphOptimizationLevel, SessionBuilder};
+    use dr_transform::ort::environment::Environment;
+    use dr_transform::ort::session::Session;
+    use dr_transform::ort::tensor::InputTensor;
+    use dr_transform::ort::{ExecutionProvider, GraphOptimizationLevel, SessionBuilder};
+    use dr_transform::ort::LoggingLevel;
     use serde_json::Value;
     use std::fmt::{Debug, Display};
     use std::sync::Arc;
     use std::{fmt, fs};
     use tokio::time::Instant;
-
     lazy_static! {
         pub static ref ENVIRONMENT: Arc<Environment> = Arc::new(
             Environment::builder()
                 .with_name("onnx home")
-                .with_log_level(ort::LoggingLevel::Error)
+                .with_log_level(LoggingLevel::Error)
+                .with_global_thread_pool(ARGS.onnx_global_thread_pool_options.clone())
                 .build()
                 .unwrap()
         );
@@ -101,23 +102,30 @@ pub mod onnx {
             let meta_info = format!("{}/{}/{}", ARGS.model_dir[idx], version, META_INFO);
             let mut builder = SessionBuilder::new(&ENVIRONMENT)?
                 .with_optimization_level(GraphOptimizationLevel::Level3)?
-                .with_parallel_execution(ARGS.onnx_use_parallel_mode == "true")?
-                .with_inter_threads(
-                    utils::get_config_or(
-                        model_config,
-                        "inter_op_parallelism",
-                        &ARGS.inter_op_parallelism[idx],
-                    )
-                    .parse()?,
-                )?
-                .with_intra_threads(
-                    utils::get_config_or(
-                        model_config,
-                        "intra_op_parallelism",
-                        &ARGS.intra_op_parallelism[idx],
-                    )
-                    .parse()?,
-                )?
+                .with_parallel_execution(ARGS.onnx_use_parallel_mode == "true")?;
+            if ARGS.onnx_global_thread_pool_options.is_empty() {
+                builder = builder
+                    .with_inter_threads(
+                        utils::get_config_or(
+                            model_config,
+                            "inter_op_parallelism",
+                            &ARGS.inter_op_parallelism[idx],
+                        )
+                            .parse()?,
+                    )?
+                    .with_intra_threads(
+                        utils::get_config_or(
+                            model_config,
+                            "intra_op_parallelism",
+                            &ARGS.intra_op_parallelism[idx],
+                        )
+                            .parse()?,
+                    )?;
+            }
+            else {
+                builder = builder.with_disable_per_session_threads()?;
+            }
+            builder = builder
                 .with_memory_pattern(ARGS.onnx_use_memory_pattern == "true")?
                 .with_execution_providers(&OnnxModel::ep_choices())?;
             match &ARGS.profiling {
@@ -181,7 +189,7 @@ pub mod onnx {
                     &version,
                     reporting_feature_ids,
                     Some(metrics::register_dynamic_metrics),
-                )),
+                )?),
             };
             onnx_model.warmup()?;
             Ok(onnx_model)

navi/navi/src/predict_service.rs

@@ -1,7 +1,7 @@
 use anyhow::{anyhow, Result};
 use arrayvec::ArrayVec;
 use itertools::Itertools;
-use log::{error, info, warn};
+use log::{error, info};
 use std::fmt::{Debug, Display};
 use std::string::String;
 use std::sync::Arc;
@@ -24,7 +24,7 @@ use serde_json::{self, Value};
 
 pub trait Model: Send + Sync + Display + Debug + 'static {
     fn warmup(&self) -> Result<()>;
-    //TODO: refactor this to return Vec<Vec<TensorScores>>, i.e.
+    //TODO: refactor this to return vec<vec<TensorScores>>, i.e.
     //we have the underlying runtime impl to split the response to each client.
     //It will eliminate some inefficient memory copy in onnx_model.rs as well as simplify code
     fn do_predict(
@@ -179,17 +179,17 @@ impl<T: Model> PredictService<T> {
         //initialize the latest version array
         let mut cur_versions = vec!["".to_owned(); MODEL_SPECS.len()];
         loop {
-            let config = utils::read_config(&meta_file).unwrap_or_else(|e| {
-                warn!("config file {} not found due to: {}", meta_file, e);
-                Value::Null
-            });
             info!("***polling for models***"); //nice deliminter
-            info!("config:{}", config);
             if let Some(ref cli) = ARGS.modelsync_cli {
                 if let Err(e) = call_external_modelsync(cli, &cur_versions).await {
                     error!("model sync cli running error:{}", e)
                 }
             }
+            let config = utils::read_config(&meta_file).unwrap_or_else(|e| {
+                info!("config file {} not found due to: {}", meta_file, e);
+                Value::Null
+            });
+            info!("config:{}", config);
             for (idx, cur_version) in cur_versions.iter_mut().enumerate() {
                 let model_dir = &ARGS.model_dir[idx];
                 PredictService::scan_load_latest_model_from_model_dir(
@@ -222,33 +222,39 @@ impl<T: Model> PredictService<T> {
             .map(|b| b.parse().unwrap())
             .collect::<Vec<u64>>();
         let no_msg_wait_millis = *batch_time_out_millis.iter().min().unwrap();
-        let mut all_model_predictors =
-            ArrayVec::<ArrayVec<BatchPredictor<T>, MAX_VERSIONS_PER_MODEL>, MAX_NUM_MODELS>::new();
+        let mut all_model_predictors: ArrayVec::<ArrayVec<BatchPredictor<T>, MAX_VERSIONS_PER_MODEL>, MAX_NUM_MODELS> =
+            (0 ..MAX_NUM_MODELS).map( |_| ArrayVec::<BatchPredictor<T>, MAX_VERSIONS_PER_MODEL>::new()).collect();
         loop {
             let msg = rx.try_recv();
             let no_more_msg = match msg {
                 Ok(PredictMessage::Predict(model_spec_at, version, val, resp, ts)) => {
                     if let Some(model_predictors) = all_model_predictors.get_mut(model_spec_at) {
-                        match version {
-                            None => model_predictors[0].push(val, resp, ts),
-                            Some(the_version) => match model_predictors
-                                .iter_mut()
-                                .find(|x| x.model.version() == the_version)
-                            {
-                                None => resp
-                                    .send(PredictResult::ModelVersionNotFound(
-                                        model_spec_at,
-                                        the_version,
-                                    ))
-                                    .unwrap_or_else(|e| {
-                                        error!("cannot send back version error: {:?}", e)
-                                    }),
-                                Some(predictor) => predictor.push(val, resp, ts),
-                            },
+                        if model_predictors.is_empty() {
+                            resp.send(PredictResult::ModelNotReady(model_spec_at))
+                                .unwrap_or_else(|e| error!("cannot send back model not ready error: {:?}", e));
+                        }
+                        else {
+                            match version {
+                                None => model_predictors[0].push(val, resp, ts),
+                                Some(the_version) => match model_predictors
+                                    .iter_mut()
+                                    .find(|x| x.model.version() == the_version)
+                                {
+                                    None => resp
+                                        .send(PredictResult::ModelVersionNotFound(
+                                            model_spec_at,
+                                            the_version,
+                                        ))
+                                        .unwrap_or_else(|e| {
+                                            error!("cannot send back version error: {:?}", e)
+                                        }),
+                                    Some(predictor) => predictor.push(val, resp, ts),
+                                },
+                            }
                         }
                     } else {
                         resp.send(PredictResult::ModelNotFound(model_spec_at))
-                            .unwrap_or_else(|e| error!("cannot send back model error: {:?}", e))
+                            .unwrap_or_else(|e| error!("cannot send back model not found error: {:?}", e))
                     }
                     MPSC_CHANNEL_SIZE.dec();
                     false
@@ -266,27 +272,23 @@ impl<T: Model> PredictService<T> {
                         queue_reset_ts: Instant::now(),
                         queue_earliest_rq_ts: Instant::now(),
                     };
-                    if idx < all_model_predictors.len() {
-                        metrics::NEW_MODEL_SNAPSHOT
-                            .with_label_values(&[&MODEL_SPECS[idx]])
-                            .inc();
+                    assert!(idx < all_model_predictors.len());
+                    metrics::NEW_MODEL_SNAPSHOT
+                        .with_label_values(&[&MODEL_SPECS[idx]])
+                        .inc();
 
-                        info!("now we serve updated model: {}", predictor.model);
-                        //we can do this since the vector is small
-                        let predictors = &mut all_model_predictors[idx];
-                        if predictors.len() == ARGS.versions_per_model {
-                            predictors.remove(predictors.len() - 1);
-                        }
-                        predictors.insert(0, predictor);
-                    } else {
-                        info!("now we serve new model: {:}", predictor.model);
-                        let mut predictors =
-                            ArrayVec::<BatchPredictor<T>, MAX_VERSIONS_PER_MODEL>::new();
-                        predictors.push(predictor);
-                        all_model_predictors.push(predictors);
-                        //check the invariant that we always push the last model to the end
-                        assert_eq!(all_model_predictors.len(), idx + 1)
+                    //we can do this since the vector is small
+                    let predictors = &mut all_model_predictors[idx];
+                    if predictors.len() == 0 {
+                        info!("now we serve new model: {}", predictor.model);
                     }
+                    else {
+                        info!("now we serve updated model: {}", predictor.model);
+                    }
+                    if predictors.len() == ARGS.versions_per_model {
+                        predictors.remove(predictors.len() - 1);
+                    }
+                    predictors.insert(0, predictor);
                     false
                 }
                 Err(TryRecvError::Empty) => true,

navi/segdense/Cargo.toml

@@ -3,9 +3,9 @@ name = "segdense"
 version = "0.1.0"
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+env_logger = "0.10.0"
 serde = { version = "1.0.104", features = ["derive"] }
 serde_json = "1.0.48"
 log = "0.4.17"

navi/segdense/src/error.rs

@@ -5,39 +5,49 @@ use std::fmt::Display;
  */
 #[derive(Debug)]
 pub enum SegDenseError {
-  IoError(std::io::Error),
-  Json(serde_json::Error),
-  JsonMissingRoot,
-  JsonMissingObject,
-  JsonMissingArray,
-  JsonArraySize,
-  JsonMissingInputFeature,
+    IoError(std::io::Error),
+    Json(serde_json::Error),
+    JsonMissingRoot,
+    JsonMissingObject,
+    JsonMissingArray,
+    JsonArraySize,
+    JsonMissingInputFeature,
 }
 
 impl Display for SegDenseError {
-  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-    match self {
-      SegDenseError::IoError(io_error) => write!(f, "{}", io_error),
-      SegDenseError::Json(serde_json) => write!(f, "{}", serde_json),
-      SegDenseError::JsonMissingRoot => write!(f, "{}", "SegDense JSON: Root Node note found!"),
-      SegDenseError::JsonMissingObject => write!(f, "{}", "SegDense JSON: Object note found!"),
-      SegDenseError::JsonMissingArray => write!(f, "{}", "SegDense JSON: Array Node note found!"),
-      SegDenseError::JsonArraySize => write!(f, "{}", "SegDense JSON: Array size not as expected!"),
-      SegDenseError::JsonMissingInputFeature => write!(f, "{}", "SegDense JSON: Missing input feature!"),
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            SegDenseError::IoError(io_error) => write!(f, "{}", io_error),
+            SegDenseError::Json(serde_json) => write!(f, "{}", serde_json),
+            SegDenseError::JsonMissingRoot => {
+                write!(f, "{}", "SegDense JSON: Root Node note found!")
+            }
+            SegDenseError::JsonMissingObject => {
+                write!(f, "{}", "SegDense JSON: Object note found!")
+            }
+            SegDenseError::JsonMissingArray => {
+                write!(f, "{}", "SegDense JSON: Array Node note found!")
+            }
+            SegDenseError::JsonArraySize => {
+                write!(f, "{}", "SegDense JSON: Array size not as expected!")
+            }
+            SegDenseError::JsonMissingInputFeature => {
+                write!(f, "{}", "SegDense JSON: Missing input feature!")
+            }
+        }
     }
-  }
 }
 
 impl std::error::Error for SegDenseError {}
 
 impl From<std::io::Error> for SegDenseError {
-  fn from(err: std::io::Error) -> Self {
-    SegDenseError::IoError(err)
-  }
+    fn from(err: std::io::Error) -> Self {
+        SegDenseError::IoError(err)
+    }
 }
 
 impl From<serde_json::Error> for SegDenseError {
-  fn from(err: serde_json::Error) -> Self {
-    SegDenseError::Json(err)
-  }
+    fn from(err: serde_json::Error) -> Self {
+        SegDenseError::Json(err)
+    }
 }

navi/segdense/src/lib.rs

@@ -1,4 +1,4 @@
 pub mod error;
-pub mod segdense_transform_spec_home_recap_2022; 
 pub mod mapper;
-pub mod util;
+pub mod segdense_transform_spec_home_recap_2022;
+pub mod util;

navi/segdense/src/main.rs

@@ -5,19 +5,18 @@ use segdense::error::SegDenseError;
 use segdense::util;
 
 fn main() -> Result<(), SegDenseError> {
-  env_logger::init();
-  let args: Vec<String> = env::args().collect();
-  
-  let schema_file_name: &str = if args.len() == 1 {
-    "json/compact.json"
-  } else {
-    &args[1]
-  };
+    env_logger::init();
+    let args: Vec<String> = env::args().collect();
 
-  let json_str = fs::read_to_string(schema_file_name)?;
+    let schema_file_name: &str = if args.len() == 1 {
+        "json/compact.json"
+    } else {
+        &args[1]
+    };
 
-  util::safe_load_config(&json_str)?;
+    let json_str = fs::read_to_string(schema_file_name)?;
 
-  Ok(())
+    util::safe_load_config(&json_str)?;
+
+    Ok(())
 }
-

navi/segdense/src/mapper.rs

@@ -19,13 +19,13 @@ pub struct FeatureMapper {
 impl FeatureMapper {
     pub fn new() -> FeatureMapper {
         FeatureMapper {
-            map: HashMap::new()
+            map: HashMap::new(),
         }
     }
 }
 
 pub trait MapWriter {
-    fn set(&mut self, feature_id: i64, info: FeatureInfo); 
+    fn set(&mut self, feature_id: i64, info: FeatureInfo);
 }
 
 pub trait MapReader {

navi/segdense/src/segdense_transform_spec_home_recap_2022.rs

@@ -164,7 +164,6 @@ pub struct ComplexFeatureTypeTransformSpec {
     pub tensor_shape: Vec<i64>,
 }
 
-
 #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct InputFeatureMapRecord {

navi/segdense/src/util.rs

@@ -1,23 +1,23 @@
+use log::debug;
 use std::fs;
-use log::{debug};
 
-use serde_json::{Value, Map};
+use serde_json::{Map, Value};
 
 use crate::error::SegDenseError;
-use crate::mapper::{FeatureMapper, FeatureInfo, MapWriter};
+use crate::mapper::{FeatureInfo, FeatureMapper, MapWriter};
 use crate::segdense_transform_spec_home_recap_2022::{self as seg_dense, InputFeature};
 
-pub fn load_config(file_name: &str) -> seg_dense::Root {
-    let json_str = fs::read_to_string(file_name).expect(
-        &format!("Unable to load segdense file {}", file_name));
-    let seg_dense_config = parse(&json_str).expect(
-        &format!("Unable to parse segdense file {}", file_name));
-    return seg_dense_config;
+pub fn load_config(file_name: &str) -> Result<seg_dense::Root, SegDenseError> {
+    let json_str = fs::read_to_string(file_name)?;
+    // &format!("Unable to load segdense file {}", file_name));
+    let seg_dense_config = parse(&json_str)?;
+    // &format!("Unable to parse segdense file {}", file_name));
+    Ok(seg_dense_config)
 }
 
 pub fn parse(json_str: &str) -> Result<seg_dense::Root, SegDenseError> {
     let root: seg_dense::Root = serde_json::from_str(json_str)?;
-    return Ok(root);
+    Ok(root)
 }
 
 /**
@@ -44,15 +44,8 @@ pub fn safe_load_config(json_str: &str) -> Result<FeatureMapper, SegDenseError>
     load_from_parsed_config(root)
 }
 
-pub fn load_from_parsed_config_ref(root: &seg_dense::Root) -> FeatureMapper {
-    load_from_parsed_config(root.clone()).unwrap_or_else(
-      |error| panic!("Error loading all_config.json - {}", error))
-}
-
 // Perf note : make 'root' un-owned
-pub fn load_from_parsed_config(root: seg_dense::Root) ->
-    Result<FeatureMapper, SegDenseError> {
-
+pub fn load_from_parsed_config(root: seg_dense::Root) -> Result<FeatureMapper, SegDenseError> {
     let v = root.input_features_map;
 
     // Do error check
@@ -86,7 +79,7 @@ pub fn load_from_parsed_config(root: seg_dense::Root) ->
             Some(info) => {
                 debug!("{:?}", info);
                 fm.set(feature_id, info)
-            },
+            }
             None => (),
         }
     }
@@ -94,19 +87,22 @@ pub fn load_from_parsed_config(root: seg_dense::Root) ->
     Ok(fm)
 }
 #[allow(dead_code)]
-fn add_feature_info_to_mapper(feature_mapper: &mut FeatureMapper, input_features: &Vec<InputFeature>) {
+fn add_feature_info_to_mapper(
+    feature_mapper: &mut FeatureMapper,
+    input_features: &Vec<InputFeature>,
+) {
     for input_feature in input_features.iter() {
-            let feature_id = input_feature.feature_id;
-            let feature_info = to_feature_info(input_feature);
-    
-            match feature_info {
-                Some(info) => {
-                    debug!("{:?}", info);
-                    feature_mapper.set(feature_id, info)
-                },
-                None => (),
+        let feature_id = input_feature.feature_id;
+        let feature_info = to_feature_info(input_feature);
+
+        match feature_info {
+            Some(info) => {
+                debug!("{:?}", info);
+                feature_mapper.set(feature_id, info)
             }
+            None => (),
         }
+    }
 }
 
 pub fn to_feature_info(input_feature: &seg_dense::InputFeature) -> Option<FeatureInfo> {
@@ -139,7 +135,7 @@ pub fn to_feature_info(input_feature: &seg_dense::InputFeature) -> Option<Featur
             2 => 0,
             3 => 2,
             _ => -1,
-        }
+        },
     };
 
     if input_feature.index < 0 {
@@ -156,4 +152,3 @@ pub fn to_feature_info(input_feature: &seg_dense::InputFeature) -> Option<Featur
         index_within_tensor: input_feature.index,
     })
 }
-

unified_user_actions/adapter/src/main/scala/com/twitter/unified_user_actions/adapter/behavioral_client_event/BUILD

@@ -1,13 +0,0 @@
-scala_library(
-    sources = [
-        "*.scala",
-    ],
-    tags = ["bazel-compatible"],
-    dependencies = [
-        "client-events/thrift/src/thrift/storage/twitter/behavioral_event:behavioral_event-scala",
-        "kafka/finagle-kafka/finatra-kafka/src/main/scala",
-        "unified_user_actions/adapter/src/main/scala/com/twitter/unified_user_actions/adapter:base",
-        "unified_user_actions/adapter/src/main/scala/com/twitter/unified_user_actions/adapter/common",
-        "unified_user_actions/thrift/src/main/thrift/com/twitter/unified_user_actions:unified_user_actions-scala",
-    ],
-)

unified_user_actions/adapter/src/main/scala/com/twitter/unified_user_actions/adapter/behavioral_client_event/BaseBCEAdapter.scala

@@ -1,96 +0,0 @@
-package com.twitter.unified_user_actions.adapter.behavioral_client_event
-
-import com.twitter.client_event_entities.serverside_context_key.latest.thriftscala.FlattenedServersideContextKey
-import com.twitter.storage.behavioral_event.thriftscala.EventLogContext
-import com.twitter.storage.behavioral_event.thriftscala.FlattenedEventLog
-import com.twitter.unified_user_actions.adapter.common.AdapterUtils
-import com.twitter.unified_user_actions.thriftscala.ActionType
-import com.twitter.unified_user_actions.thriftscala.BreadcrumbTweet
-import com.twitter.unified_user_actions.thriftscala.ClientEventNamespace
-import com.twitter.unified_user_actions.thriftscala.EventMetadata
-import com.twitter.unified_user_actions.thriftscala.Item
-import com.twitter.unified_user_actions.thriftscala.ProductSurface
-import com.twitter.unified_user_actions.thriftscala.ProductSurfaceInfo
-import com.twitter.unified_user_actions.thriftscala.SourceLineage
-import com.twitter.unified_user_actions.thriftscala.UnifiedUserAction
-import com.twitter.unified_user_actions.thriftscala.UserIdentifier
-
-case class ProductSurfaceRelated(
-  productSurface: Option[ProductSurface],
-  productSurfaceInfo: Option[ProductSurfaceInfo])
-
-trait BaseBCEAdapter {
-  def toUUA(e: FlattenedEventLog): Seq[UnifiedUserAction]
-
-  protected def getUserIdentifier(c: EventLogContext): UserIdentifier =
-    UserIdentifier(
-      userId = c.userId,
-      guestIdMarketing = c.guestIdMarketing
-    )
-
-  protected def getEventMetadata(e: FlattenedEventLog): EventMetadata =
-    EventMetadata(
-      sourceLineage = SourceLineage.BehavioralClientEvents,
-      sourceTimestampMs =
-        e.context.driftAdjustedEventCreatedAtMs.getOrElse(e.context.eventCreatedAtMs),
-      receivedTimestampMs = AdapterUtils.currentTimestampMs,
-      // Client UI language or from Gizmoduck which is what user set in Twitter App.
-      // Please see more at https://sourcegraph.twitter.biz/git.twitter.biz/source/-/blob/finatra-internal/international/src/main/scala/com/twitter/finatra/international/LanguageIdentifier.scala
-      // The format should be ISO 639-1.
-      language = e.context.languageCode.map(AdapterUtils.normalizeLanguageCode),
-      // Country code could be IP address (geoduck) or User registration country (gizmoduck) and the former takes precedence.
-      // We don’t know exactly which one is applied, unfortunately,
-      // see https://sourcegraph.twitter.biz/git.twitter.biz/source/-/blob/finatra-internal/international/src/main/scala/com/twitter/finatra/international/CountryIdentifier.scala
-      // The format should be ISO_3166-1_alpha-2.
-      countryCode = e.context.countryCode.map(AdapterUtils.normalizeCountryCode),
-      clientAppId = e.context.clientApplicationId,
-      clientVersion = e.context.clientVersion,
-      clientPlatform = e.context.clientPlatform,
-      viewHierarchy = e.v1ViewTypeHierarchy,
-      clientEventNamespace = Some(
-        ClientEventNamespace(
-          page = e.page,
-          section = e.section,
-          element = e.element,
-          action = e.actionName,
-          subsection = e.subsection
-        )),
-      breadcrumbViews = e.v1BreadcrumbViewTypeHierarchy,
-      breadcrumbTweets = e.v1BreadcrumbTweetIds.map { breadcrumbs =>
-        breadcrumbs.map { breadcrumb =>
-          BreadcrumbTweet(
-            tweetId = breadcrumb.serversideContextId.toLong,
-            sourceComponent = breadcrumb.sourceComponent)
-        }
-      }
-    )
-
-  protected def getBreadcrumbTweetIds(
-    breadcrumbTweetIds: Option[Seq[FlattenedServersideContextKey]]
-  ): Seq[BreadcrumbTweet] =
-    breadcrumbTweetIds
-      .getOrElse(Nil).map(breadcrumb => {
-        BreadcrumbTweet(
-          tweetId = breadcrumb.serversideContextId.toLong,
-          sourceComponent = breadcrumb.sourceComponent)
-      })
-
-  protected def getBreadcrumbViews(breadcrumbView: Option[Seq[String]]): Seq[String] =
-    breadcrumbView.getOrElse(Nil)
-
-  protected def getUnifiedUserAction(
-    event: FlattenedEventLog,
-    actionType: ActionType,
-    item: Item,
-    productSurface: Option[ProductSurface] = None,
-    productSurfaceInfo: Option[ProductSurfaceInfo] = None
-  ): UnifiedUserAction =
-    UnifiedUserAction(
-      userIdentifier = getUserIdentifier(event.context),
-      actionType = actionType,
-      item = item,
-      eventMetadata = getEventMetadata(event),
-      productSurface = productSurface,
-      productSurfaceInfo = productSurfaceInfo
-    )
-}

unified_user_actions/adapter/src/main/scala/com/twitter/unified_user_actions/adapter/behavioral_client_event/BehavioralClientEventAdapter.scala

@@ -1,39 +0,0 @@
-package com.twitter.unified_user_actions.adapter.behavioral_client_event
-
-import com.twitter.finagle.stats.NullStatsReceiver
-import com.twitter.finagle.stats.StatsReceiver
-import com.twitter.finatra.kafka.serde.UnKeyed
-import com.twitter.storage.behavioral_event.thriftscala.FlattenedEventLog
-import com.twitter.unified_user_actions.adapter.AbstractAdapter
-import com.twitter.unified_user_actions.thriftscala._
-
-class BehavioralClientEventAdapter
-    extends AbstractAdapter[FlattenedEventLog, UnKeyed, UnifiedUserAction] {
-
-  import BehavioralClientEventAdapter._
-
-  override def adaptOneToKeyedMany(
-    input: FlattenedEventLog,
-    statsReceiver: StatsReceiver = NullStatsReceiver
-  ): Seq[(UnKeyed, UnifiedUserAction)] =
-    adaptEvent(input).map { e => (UnKeyed, e) }
-}
-
-object BehavioralClientEventAdapter {
-  def adaptEvent(e: FlattenedEventLog): Seq[UnifiedUserAction] =
-    // See go/bcecoverage for event namespaces, usage and coverage details
-    Option(e)
-      .map { e =>
-        (e.page, e.actionName) match {
-          case (Some("tweet_details"), Some("impress")) =>
-            TweetImpressionBCEAdapter.TweetDetails.toUUA(e)
-          case (Some("fullscreen_video"), Some("impress")) =>
-            TweetImpressionBCEAdapter.FullscreenVideo.toUUA(e)
-          case (Some("fullscreen_image"), Some("impress")) =>
-            TweetImpressionBCEAdapter.FullscreenImage.toUUA(e)
-          case (Some("profile"), Some("impress")) =>
-            ProfileImpressionBCEAdapter.Profile.toUUA(e)
-          case _ => Nil
-        }
-      }.getOrElse(Nil)
-}

unified_user_actions/adapter/src/main/scala/com/twitter/unified_user_actions/adapter/behavioral_client_event/ImpressionBCEAdapter.scala

@@ -1,34 +0,0 @@
-package com.twitter.unified_user_actions.adapter.behavioral_client_event
-
-import com.twitter.client.behavioral_event.action.impress.latest.thriftscala.Impress
-import com.twitter.client_event_entities.serverside_context_key.latest.thriftscala.FlattenedServersideContextKey
-import com.twitter.unified_user_actions.thriftscala.Item
-
-trait ImpressionBCEAdapter extends BaseBCEAdapter {
-  type ImpressedItem <: Item
-
-  def getImpressedItem(
-    context: FlattenedServersideContextKey,
-    impression: Impress
-  ): ImpressedItem
-
-  /**
-   * The start time of an impression in milliseconds since epoch. In BCE, the impression
-   * tracking clock will start immediately after the page is visible with no initial delay.
-   */
-  def getImpressedStartTimestamp(impression: Impress): Long =
-    impression.visibilityPctDwellStartMs
-
-  /**
-   * The end time of an impression in milliseconds since epoch. In BCE, the impression
-   * tracking clock will end before the user exit the page.
-   */
-  def getImpressedEndTimestamp(impression: Impress): Long =
-    impression.visibilityPctDwellEndMs
-
-  /**
-   * The UI component that hosted the impressed item.
-   */
-  def getImpressedUISourceComponent(context: FlattenedServersideContextKey): String =
-    context.sourceComponent
-}

unified_user_actions/adapter/src/main/scala/com/twitter/unified_user_actions/adapter/behavioral_client_event/ProfileImpressionBCEAdapter.scala

@@ -1,52 +0,0 @@
-package com.twitter.unified_user_actions.adapter.behavioral_client_event
-
-import com.twitter.client.behavioral_event.action.impress.latest.thriftscala.Impress
-import com.twitter.client_event_entities.serverside_context_key.latest.thriftscala.FlattenedServersideContextKey
-import com.twitter.storage.behavioral_event.thriftscala.FlattenedEventLog
-import com.twitter.unified_user_actions.thriftscala.ActionType
-import com.twitter.unified_user_actions.thriftscala.ClientProfileV2Impression
-import com.twitter.unified_user_actions.thriftscala.Item
-import com.twitter.unified_user_actions.thriftscala.ProductSurface
-import com.twitter.unified_user_actions.thriftscala.ProfileActionInfo
-import com.twitter.unified_user_actions.thriftscala.ProfileInfo
-import com.twitter.unified_user_actions.thriftscala.UnifiedUserAction
-
-object ProfileImpressionBCEAdapter {
-  val Profile = new ProfileImpressionBCEAdapter()
-}
-
-class ProfileImpressionBCEAdapter extends ImpressionBCEAdapter {
-  override type ImpressedItem = Item.ProfileInfo
-
-  override def toUUA(e: FlattenedEventLog): Seq[UnifiedUserAction] =
-    (e.v2Impress, e.v1UserIds) match {
-      case (Some(v2Impress), Some(v1UserIds)) =>
-        v1UserIds.map { user =>
-          getUnifiedUserAction(
-            event = e,
-            actionType = ActionType.ClientProfileV2Impression,
-            item = getImpressedItem(user, v2Impress),
-            productSurface = Some(ProductSurface.ProfilePage)
-          )
-        }
-      case _ => Nil
-    }
-
-  override def getImpressedItem(
-    context: FlattenedServersideContextKey,
-    impression: Impress
-  ): ImpressedItem =
-    Item.ProfileInfo(
-      ProfileInfo(
-        actionProfileId = context.serversideContextId.toLong,
-        profileActionInfo = Some(
-          ProfileActionInfo.ClientProfileV2Impression(
-            ClientProfileV2Impression(
-              impressStartTimestampMs = getImpressedStartTimestamp(impression),
-              impressEndTimestampMs = getImpressedEndTimestamp(impression),
-              sourceComponent = getImpressedUISourceComponent(context)
-            )
-          )
-        )
-      ))
-}

unified_user_actions/adapter/src/main/scala/com/twitter/unified_user_actions/adapter/behavioral_client_event/TweetImpressionBCEAdapter.scala

@@ -1,84 +0,0 @@
-package com.twitter.unified_user_actions.adapter.behavioral_client_event
-
-import com.twitter.client.behavioral_event.action.impress.latest.thriftscala.Impress
-import com.twitter.client_event_entities.serverside_context_key.latest.thriftscala.FlattenedServersideContextKey
-import com.twitter.storage.behavioral_event.thriftscala.FlattenedEventLog
-import com.twitter.unified_user_actions.thriftscala.ActionType
-import com.twitter.unified_user_actions.thriftscala.ClientTweetV2Impression
-import com.twitter.unified_user_actions.thriftscala.Item
-import com.twitter.unified_user_actions.thriftscala.ProductSurface
-import com.twitter.unified_user_actions.thriftscala.TweetActionInfo
-import com.twitter.unified_user_actions.thriftscala.TweetInfo
-import com.twitter.unified_user_actions.thriftscala.UnifiedUserAction
-
-object TweetImpressionBCEAdapter {
-  val TweetDetails = new TweetImpressionBCEAdapter(ActionType.ClientTweetV2Impression)
-  val FullscreenVideo = new TweetImpressionBCEAdapter(
-    ActionType.ClientTweetVideoFullscreenV2Impression)
-  val FullscreenImage = new TweetImpressionBCEAdapter(
-    ActionType.ClientTweetImageFullscreenV2Impression)
-}
-
-class TweetImpressionBCEAdapter(actionType: ActionType) extends ImpressionBCEAdapter {
-  override type ImpressedItem = Item.TweetInfo
-
-  override def toUUA(e: FlattenedEventLog): Seq[UnifiedUserAction] =
-    (actionType, e.v2Impress, e.v1TweetIds, e.v1BreadcrumbTweetIds) match {
-      case (ActionType.ClientTweetV2Impression, Some(v2Impress), Some(v1TweetIds), _) =>
-        toUUAEvents(e, v2Impress, v1TweetIds)
-      case (
-            ActionType.ClientTweetVideoFullscreenV2Impression,
-            Some(v2Impress),
-            _,
-            Some(v1BreadcrumbTweetIds)) =>
-        toUUAEvents(e, v2Impress, v1BreadcrumbTweetIds)
-      case (
-            ActionType.ClientTweetImageFullscreenV2Impression,
-            Some(v2Impress),
-            _,
-            Some(v1BreadcrumbTweetIds)) =>
-        toUUAEvents(e, v2Impress, v1BreadcrumbTweetIds)
-      case _ => Nil
-    }
-
-  private def toUUAEvents(
-    e: FlattenedEventLog,
-    v2Impress: Impress,
-    v1TweetIds: Seq[FlattenedServersideContextKey]
-  ): Seq[UnifiedUserAction] =
-    v1TweetIds.map { tweet =>
-      getUnifiedUserAction(
-        event = e,
-        actionType = actionType,
-        item = getImpressedItem(tweet, v2Impress),
-        productSurface = getProductSurfaceRelated.productSurface,
-        productSurfaceInfo = getProductSurfaceRelated.productSurfaceInfo
-      )
-    }
-
-  override def getImpressedItem(
-    context: FlattenedServersideContextKey,
-    impression: Impress
-  ): ImpressedItem =
-    Item.TweetInfo(
-      TweetInfo(
-        actionTweetId = context.serversideContextId.toLong,
-        tweetActionInfo = Some(
-          TweetActionInfo.ClientTweetV2Impression(
-            ClientTweetV2Impression(
-              impressStartTimestampMs = getImpressedStartTimestamp(impression),
-              impressEndTimestampMs = getImpressedEndTimestamp(impression),
-              sourceComponent = getImpressedUISourceComponent(context)
-            )
-          ))
-      ))
-
-  private def getProductSurfaceRelated: ProductSurfaceRelated =
-    actionType match {
-      case ActionType.ClientTweetV2Impression =>
-        ProductSurfaceRelated(
-          productSurface = Some(ProductSurface.TweetDetailsPage),
-          productSurfaceInfo = None)
-      case _ => ProductSurfaceRelated(productSurface = None, productSurfaceInfo = None)
-    }
-}

unified_user_actions/adapter/src/test/scala/com/twitter/unified_user_actions/adapter/BehavioralClientEventAdapterSpec.scala

@@ -1,139 +0,0 @@
-package com.twitter.unified_user_actions.adapter
-
-import com.twitter.inject.Test
-import com.twitter.storage.behavioral_event.thriftscala.FlattenedEventLog
-import com.twitter.unified_user_actions.adapter.TestFixtures.BCEFixture
-import com.twitter.unified_user_actions.adapter.behavioral_client_event.BehavioralClientEventAdapter
-import com.twitter.unified_user_actions.thriftscala._
-import com.twitter.util.Time
-import org.scalatest.prop.TableDrivenPropertyChecks
-
-class BehavioralClientEventAdapterSpec extends Test with TableDrivenPropertyChecks {
-
-  test("basic event conversion should be correct") {
-    new BCEFixture {
-      Time.withTimeAt(frozenTime) { _ =>
-        val tests = Table(
-          ("event", "expected", "description"),
-          (
-            makeBCEEvent(),
-            makeUUAImpressEvent(productSurface = Some(ProductSurface.TweetDetailsPage)),
-            "tweet_details conversion"),
-          (makeBCEProfileImpressEvent(), makeUUAProfileImpressEvent(), "profile conversion"),
-          (
-            makeBCEVideoFullscreenImpressEvent(),
-            makeUUAVideoFullscreenImpressEvent(),
-            "fullscreen_video conversion"),
-          (
-            makeBCEImageFullscreenImpressEvent(),
-            makeUUAImageFullscreenImpressEvent(),
-            "fullscreen_image conversion"),
-        )
-        forEvery(tests) { (input: FlattenedEventLog, expected: UnifiedUserAction, desc: String) =>
-          assert(Seq(expected) === BehavioralClientEventAdapter.adaptEvent(input), desc)
-        }
-      }
-    }
-  }
-
-  test(
-    "tweet_details is NOT missing productSurface[Info] when empty breadcrumb components and breadcrumbs tweets id") {
-    new BCEFixture {
-      Time.withTimeAt(frozenTime) { _ =>
-        val input = makeBCEEvent(v1BreadcrumbViewTypeHierarchy = None, v1BreadcrumbTweetIds = None)
-        val expected =
-          makeUUAImpressEvent(
-            productSurface = Some(ProductSurface.TweetDetailsPage),
-            breadcrumbViews = None,
-            breadcrumbTweets = None)
-        val actual = BehavioralClientEventAdapter.adaptEvent(input)
-
-        assert(Seq(expected) === actual)
-      }
-    }
-  }
-
-  test("tweet_details is not missing productSurface[Info] when only breadcrumb tweets is empty") {
-    new BCEFixture {
-      Time.withTimeAt(frozenTime) { _ =>
-        val input = makeBCEEvent(v1BreadcrumbTweetIds = None)
-        val expected = makeUUAImpressEvent(
-          productSurface = Some(ProductSurface.TweetDetailsPage),
-          breadcrumbViews = Some(viewBreadcrumbs),
-          breadcrumbTweets = None
-        )
-        val actual = BehavioralClientEventAdapter.adaptEvent(input)
-
-        assert(Seq(expected) === actual)
-      }
-    }
-  }
-
-  test("unsupported events should be skipped") {
-    new BCEFixture {
-      val unsupportedPage = "unsupported_page"
-      val unsupportedAction = "unsupported_action"
-      val supportedNamespaces = Table(
-        ("page", "actions"),
-        ("tweet_details", Seq("impress")),
-        ("profile", Seq("impress")),
-      )
-
-      forAll(supportedNamespaces) { (page: String, actions: Seq[String]) =>
-        actions.foreach { supportedAction =>
-          assert(
-            BehavioralClientEventAdapter
-              .adaptEvent(
-                makeBCEEvent(
-                  currentPage = Some(unsupportedPage),
-                  actionName = Some(supportedAction))).isEmpty)
-
-          assert(BehavioralClientEventAdapter
-            .adaptEvent(
-              makeBCEEvent(currentPage = Some(page), actionName = Some(unsupportedAction))).isEmpty)
-        }
-      }
-    }
-  }
-
-  test("event w/ missing info should be skipped") {
-    new BCEFixture {
-      val eventsWithMissingInfo = Table(
-        ("event", "description"),
-        (null.asInstanceOf[FlattenedEventLog], "null event"),
-        (makeBCEEvent(v2Impress = None), "impression event missing v2Impress"),
-        (makeBCEEvent(v1TweetIds = None), "tweet event missing v1TweetIds"),
-        (makeBCEProfileImpressEvent(v1UserIds = None), "profile event missing v1UserIds"),
-        (
-          makeBCEVideoFullscreenImpressEvent(v1BreadcrumbTweetIds = None),
-          "fullscreen_video event missing v1BreadcrumbTweetIds"),
-        (
-          makeBCEImageFullscreenImpressEvent(v1BreadcrumbTweetIds = None),
-          "fullscreen_image event missing v1BreadcrumbTweetIds"),
-      )
-
-      forEvery(eventsWithMissingInfo) { (event: FlattenedEventLog, desc: String) =>
-        assert(
-          BehavioralClientEventAdapter
-            .adaptEvent(event).isEmpty,
-          desc)
-      }
-    }
-  }
-
-  test("use eventCreateAtMs when driftAdjustedTimetampMs is empty") {
-    new BCEFixture {
-      Time.withTimeAt(frozenTime) { _ =>
-        val input = makeBCEEvent(
-          context = makeBCEContext(driftAdjustedEventCreatedAtMs = None)
-        )
-        val expected = makeUUAImpressEvent(
-          createTs = eventCreatedTime,
-          productSurface = Some(ProductSurface.TweetDetailsPage))
-        val actual = BehavioralClientEventAdapter.adaptEvent(input)
-
-        assert(Seq(expected) === actual)
-      }
-    }
-  }
-}

unified_user_actions/service/src/main/scala/com/twitter/unified_user_actions/service/BehavioralClientEventService.scala

@@ -1,25 +0,0 @@
-package com.twitter.unified_user_actions.service;
-
-import com.twitter.finatra.decider.modules.DeciderModule
-import com.twitter.finatra.kafka.serde.UnKeyed
-import com.twitter.inject.server.TwitterServer
-import com.twitter.kafka.client.processor.AtLeastOnceProcessor
-import com.twitter.storage.behavioral_event.thriftscala.FlattenedEventLog
-import com.twitter.unified_user_actions.service.module.KafkaProcessorBehavioralClientEventModule
-
-object BehavioralClientEventServiceMain extends BehavioralClientEventService
-
-class BehavioralClientEventService extends TwitterServer {
-  override val modules = Seq(
-    KafkaProcessorBehavioralClientEventModule,
-    DeciderModule
-  )
-
-  override protected def setup(): Unit = {}
-
-  override protected def start(): Unit = {
-    val processor = injector.instance[AtLeastOnceProcessor[UnKeyed, FlattenedEventLog]]
-    closeOnExit(processor)
-    processor.start()
-  }
-}

unified_user_actions/service/src/main/scala/com/twitter/unified_user_actions/service/module/KafkaProcessorBehavioralClientEventModule.scala

@@ -1,87 +0,0 @@
-package com.twitter.unified_user_actions.service.module
-
-import com.google.inject.Provides
-import com.twitter.decider.Decider
-import com.twitter.finagle.stats.StatsReceiver
-import com.twitter.finatra.kafka.serde.UnKeyed
-import com.twitter.finatra.kafka.serde.UnKeyedSerde
-import com.twitter.inject.annotations.Flag
-import com.twitter.inject.TwitterModule
-import com.twitter.kafka.client.processor.AtLeastOnceProcessor
-import com.twitter.storage.behavioral_event.thriftscala.FlattenedEventLog
-import com.twitter.unified_user_actions.adapter.behavioral_client_event.BehavioralClientEventAdapter
-import com.twitter.unified_user_actions.kafka.CompressionTypeFlag
-import com.twitter.unified_user_actions.kafka.serde.NullableScalaSerdes
-import com.twitter.util.Duration
-import com.twitter.util.StorageUnit
-import com.twitter.util.logging.Logging
-import javax.inject.Singleton
-
-object KafkaProcessorBehavioralClientEventModule extends TwitterModule with Logging {
-  override def modules = Seq(FlagsModule)
-
-  private val adapter: BehavioralClientEventAdapter = new BehavioralClientEventAdapter
-  private final val processorName: String = "uuaProcessor"
-
-  @Provides
-  @Singleton
-  def providesKafkaProcessor(
-    decider: Decider,
-    @Flag(FlagsModule.cluster) cluster: String,
-    @Flag(FlagsModule.kafkaSourceCluster) kafkaSourceCluster: String,
-    @Flag(FlagsModule.kafkaDestCluster) kafkaDestCluster: String,
-    @Flag(FlagsModule.kafkaSourceTopic) kafkaSourceTopic: String,
-    @Flag(FlagsModule.kafkaSinkTopics) kafkaSinkTopics: Seq[String],
-    @Flag(FlagsModule.kafkaGroupId) kafkaGroupId: String,
-    @Flag(FlagsModule.kafkaProducerClientId) kafkaProducerClientId: String,
-    @Flag(FlagsModule.kafkaMaxPendingRequests) kafkaMaxPendingRequests: Int,
-    @Flag(FlagsModule.kafkaWorkerThreads) kafkaWorkerThreads: Int,
-    @Flag(FlagsModule.commitInterval) commitInterval: Duration,
-    @Flag(FlagsModule.maxPollRecords) maxPollRecords: Int,
-    @Flag(FlagsModule.maxPollInterval) maxPollInterval: Duration,
-    @Flag(FlagsModule.sessionTimeout) sessionTimeout: Duration,
-    @Flag(FlagsModule.fetchMax) fetchMax: StorageUnit,
-    @Flag(FlagsModule.batchSize) batchSize: StorageUnit,
-    @Flag(FlagsModule.linger) linger: Duration,
-    @Flag(FlagsModule.bufferMem) bufferMem: StorageUnit,
-    @Flag(FlagsModule.compressionType) compressionTypeFlag: CompressionTypeFlag,
-    @Flag(FlagsModule.retries) retries: Int,
-    @Flag(FlagsModule.retryBackoff) retryBackoff: Duration,
-    @Flag(FlagsModule.requestTimeout) requestTimeout: Duration,
-    @Flag(FlagsModule.enableTrustStore) enableTrustStore: Boolean,
-    @Flag(FlagsModule.trustStoreLocation) trustStoreLocation: String,
-    statsReceiver: StatsReceiver,
-  ): AtLeastOnceProcessor[UnKeyed, FlattenedEventLog] = {
-    KafkaProcessorProvider.provideDefaultAtLeastOnceProcessor(
-      name = processorName,
-      kafkaSourceCluster = kafkaSourceCluster,
-      kafkaGroupId = kafkaGroupId,
-      kafkaSourceTopic = kafkaSourceTopic,
-      sourceKeyDeserializer = UnKeyedSerde.deserializer,
-      sourceValueDeserializer = NullableScalaSerdes
-        .Thrift[FlattenedEventLog](statsReceiver.counter("deserializerErrors")).deserializer,
-      commitInterval = commitInterval,
-      maxPollRecords = maxPollRecords,
-      maxPollInterval = maxPollInterval,
-      sessionTimeout = sessionTimeout,
-      fetchMax = fetchMax,
-      processorMaxPendingRequests = kafkaMaxPendingRequests,
-      processorWorkerThreads = kafkaWorkerThreads,
-      adapter = adapter,
-      kafkaSinkTopics = kafkaSinkTopics,
-      kafkaDestCluster = kafkaDestCluster,
-      kafkaProducerClientId = kafkaProducerClientId,
-      batchSize = batchSize,
-      linger = linger,
-      bufferMem = bufferMem,
-      compressionType = compressionTypeFlag.compressionType,
-      retries = retries,
-      retryBackoff = retryBackoff,
-      requestTimeout = requestTimeout,
-      statsReceiver = statsReceiver,
-      trustStoreLocationOpt = if (enableTrustStore) Some(trustStoreLocation) else None,
-      decider = decider,
-      zone = ZoneFiltering.zoneMapping(cluster),
-    )
-  }
-}