Merge pull request #299 from mwvgroup/u/ch/lsst/main.py

Update LSST Cloud Run services

Author: hernandezc1

broker/cloud_run/lsst/classify_snn/deploy.sh

@@ -93,4 +93,7 @@ else
         --push-auth-service-account="${runinvoker_svcact}" \
         --dead-letter-topic="${ps_deadletter_topic}" \
         --max-delivery-attempts=5
+    gcloud pubsub subscriptions add-iam-policy-binding "${ps_input_subscrip}" \
+        --member="serviceAccount:${service_account}" \
+        --role="roles/pubsub.subscriber"
 fi

broker/cloud_run/lsst/classify_snn/main.py

@@ -99,6 +99,15 @@ def run():
 
 def _classify(alert_lite: pittgoogle.Alert) -> dict:
     """Classify the alert using SuperNNova."""
+
+    # check to see if the alert has a ssObjectId
+    if alert_lite.attributes.get("ssSource_ssObjectId"):
+        return {
+            "prob_class0": None,
+            "prob_class1": None,
+            "predicted_class": None,
+        }
+
     # init
     snn_df = _format_for_classifier(alert_lite)
     device = "cpu"

broker/cloud_run/lsst/classify_snn/requirements.txt

@@ -2,7 +2,7 @@ pandas==1.5.1
 numpy==1.23.3
 google-cloud-functions
 google-cloud-logging
-pittgoogle-client>=0.3.17
+pittgoogle-client>=0.3.18
 
 # for Cloud Run
 # https://cloud.google.com/run/docs/quickstarts/build-and-deploy/deploy-python-service

broker/cloud_run/lsst/classify_upsilon/deploy.sh

@@ -95,4 +95,7 @@ else
         --push-auth-service-account="${runinvoker_svcact}" \
         --dead-letter-topic="${ps_deadletter_topic}" \
         --max-delivery-attempts=5
+    gcloud pubsub subscriptions add-iam-policy-binding "${ps_input_subscrip}" \
+        --member="serviceAccount:${service_account}" \
+        --role="roles/pubsub.subscriber"
 fi

broker/cloud_run/lsst/classify_upsilon/main.py

@@ -56,16 +56,20 @@ def run() -> tuple[str, int]:
     # extract the envelope from the request that triggered the endpoint
     # this contains a single Pub/Sub message with the alert to be processed
     envelope = flask.request.get_json()
+
+    # unpack the alert. raises a `BadRequest` if the envelope does not contain a valid message
     try:
         alert_lite = pittgoogle.Alert.from_cloud_run(envelope, "default")
     except pittgoogle.exceptions.BadRequest as exc:
         return str(exc), HTTP_400
 
-    alert_lite_df = _create_lite_dataframe(alert_lite.dict["alert_lite"])
-    upsilon_dict = _classify_with_upsilon(alert_lite_df)
+    # classify
+    upsilon_dict = _classify(alert_lite)
     has_min_detections_in_any_band = any(
         upsilon_dict.get(f"n_data_points_{band}_band") >= 80 for band in SURVEY_BANDS
     )
+
+    # publish
     TOPIC.publish(
         pittgoogle.Alert.from_dict(
             {"alert_lite": alert_lite.dict["alert_lite"], "upsilon": upsilon_dict},
@@ -92,17 +96,21 @@ def run() -> tuple[str, int]:
     return "", HTTP_204
 
 
-def _classify_with_upsilon(alert_lite_df: pd.DataFrame) -> dict:
+def _classify(alert_lite: pittgoogle.Alert) -> dict:
     upsilon_dict = {}
+    alert_lite_df = _create_lite_dataframe(alert_lite.dict["alert_lite"])
+    is_ssobject = bool(alert_lite.attributes.get("ssSource_ssObjectId"))
+
     for band in SURVEY_BANDS:
         # ---Extract data
         filter_diaSources = alert_lite_df[alert_lite_df["band"] == band]
         flux_gt_zero = filter_diaSources["psfFlux"].to_numpy() > 0
         upsilon_dict[f"n_data_points_{band}_band"] = flux_gt_zero.sum().item()
+        # skip classification if the object has a ssObjectId
         # skip band if no detections or too few valid data points.
         # to avoid scipy's leastsq error: ("input vector length N=7 must not exceed output length M"), we require
         # that flux_gt_zero.sum() > 7
-        if filter_diaSources.empty or flux_gt_zero.sum() <= 7:
+        if is_ssobject or (filter_diaSources.empty or flux_gt_zero.sum() <= 7):
             upsilon_dict[f"{band}_label"] = None
             upsilon_dict[f"{band}_probability"] = None
             upsilon_dict[f"{band}_flag"] = None

broker/cloud_run/lsst/classify_upsilon/requirements.txt

@@ -4,7 +4,7 @@
 # file (or packaged with the function) in the same directory as `main.py`
 
 google-cloud-logging
-pittgoogle-client>=0.3.17
+pittgoogle-client>=0.3.18
 upsilon
 # required by upsilon
 # https://github.com/dwkim78/upsilon?tab=readme-ov-file#1-dependency

broker/cloud_run/lsst/ps_to_storage/deploy.sh

@@ -12,6 +12,7 @@ survey="${3:-lsst}"
 region="${4:-us-central1}"
 # get the environment variable
 PROJECT_ID=$GOOGLE_CLOUD_PROJECT
+PROJECT_NUMBER=$(gcloud projects describe "$PROJECT_ID" --format="value(projectNumber)")
 
 MODULE_NAME="alerts-to-storage"  # lower case required by cloud run
 ROUTE_RUN="/"  # url route that will trigger main.run()
@@ -37,7 +38,7 @@ ps_subscription_avro=$(define_GCP_resources "${survey}-alert_avros-counter")
 ps_topic_avro=$(define_GCP_resources "projects/${PROJECT_ID}/topics/${survey}-alert_avros")
 ps_trigger_topic=$(define_GCP_resources "${survey}-alerts_raw")
 runinvoker_svcact="cloud-run-invoker@${PROJECT_ID}.iam.gserviceaccount.com"
-
+service_account="service-${PROJECT_NUMBER}@gcp-sa-pubsub.iam.gserviceaccount.com"
 if [ "${teardown}" = "True" ]; then
     # ensure that we do not teardown production resources
     if [ "${testid}" != "False" ]; then
@@ -92,4 +93,7 @@ else
         --push-auth-service-account="${runinvoker_svcact}" \
         --dead-letter-topic="${ps_deadletter_topic}" \
         --max-delivery-attempts=5
+    gcloud pubsub subscriptions add-iam-policy-binding "${ps_input_subscrip}" \
+        --member="serviceAccount:${service_account}" \
+        --role="roles/pubsub.subscriber"
 fi

broker/cloud_run/lsst/ps_to_storage/main.py

@@ -106,5 +106,6 @@ def _create_file_metadata(alert: pittgoogle.Alert, event_id: str) -> dict:
     metadata["_".join(alert.get_key("sourceid"))] = alert.sourceid
     metadata["_".join(alert.get_key("ra"))] = alert.ra
     metadata["_".join(alert.get_key("dec"))] = alert.dec
+    metadata["kafka.timestamp"] = alert.attributes["kafka.timestamp"]
 
     return metadata

broker/cloud_run/lsst/ps_to_storage/requirements.txt

@@ -5,7 +5,7 @@
 
 google-cloud-logging
 google-cloud-storage
-pittgoogle-client>=0.3.17
+pittgoogle-client>=0.3.18
 
 # for Cloud Run
 # https://cloud.google.com/run/docs/quickstarts/build-and-deploy/deploy-python-service

broker/cloud_run/lsst/variability/deploy.sh

@@ -92,4 +92,7 @@ else
         --push-auth-service-account="${runinvoker_svcact}" \
         --dead-letter-topic="${ps_deadletter_topic}" \
         --max-delivery-attempts=5
+    gcloud pubsub subscriptions add-iam-policy-binding "${ps_input_subscrip}" \
+        --member="serviceAccount:${service_account}" \
+        --role="roles/pubsub.subscriber"
 fi