chore: fix propagation of fabric telemetry (#2403)

mhamilton723 · web-flow · commit 85a66877da13 · 2025-07-29T16:15:47.000-04:00
* chore: fix propegation of fabric telemetry

* chore: refine names
diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala
@@ -10,19 +10,19 @@ import com.microsoft.azure.synapse.ml.fabric.FabricClient
 import com.microsoft.azure.synapse.ml.io.http._
 import com.microsoft.azure.synapse.ml.logging.SynapseMLLogging
 import com.microsoft.azure.synapse.ml.logging.common.PlatformDetails
-import com.microsoft.azure.synapse.ml.param.{ GlobalKey, GlobalParams, HasGlobalParams, ServiceParam }
-import com.microsoft.azure.synapse.ml.stages.{ DropColumns, Lambda }
+import com.microsoft.azure.synapse.ml.param.{GlobalKey, GlobalParams, HasGlobalParams, ServiceParam, TypedArrayParam}
+import com.microsoft.azure.synapse.ml.stages.{DropColumns, Lambda}
 import org.apache.commons.lang.StringUtils
 import org.apache.http.NameValuePair
-import org.apache.http.client.methods.{ HttpEntityEnclosingRequestBase, HttpPost, HttpRequestBase }
+import org.apache.http.client.methods.{HttpEntityEnclosingRequestBase, HttpPost, HttpRequestBase}
 import org.apache.http.client.utils.URLEncodedUtils
 import org.apache.http.entity.AbstractHttpEntity
 import org.apache.http.impl.client.CloseableHttpClient
 import org.apache.spark.ml.param._
-import org.apache.spark.ml.{ ComplexParamsWritable, NamespaceInjections, PipelineModel, Transformer }
-import org.apache.spark.sql.functions.{ col, lit, struct }
+import org.apache.spark.ml.{ComplexParamsWritable, NamespaceInjections, PipelineModel, Transformer}
+import org.apache.spark.sql.functions.{col, lit, struct}
 import org.apache.spark.sql.types._
-import org.apache.spark.sql.{ DataFrame, Dataset, Row }
+import org.apache.spark.sql.{DataFrame, Dataset, Row}
 import spray.json.DefaultJsonProtocol._
 
 import java.net.URI
@@ -206,13 +206,37 @@ trait HasCustomHeaders extends HasServiceParams {
 
   // For Pyspark compatability accept Java HashMap as input to parameter
   // py4J only natively supports conversions from Python Dict to Java HashMap
-  def setCustomHeaders(v: java.util.HashMap[String,String]): this.type = {
+  def setCustomHeaders(v: java.util.HashMap[String, String]): this.type = {
     setCustomHeaders(v.asScala.toMap)
   }
+}
+
+trait HasTelemHeaders extends HasServiceParams {
+
+  private[ml] val telemHeaders = new ServiceParam[Map[String, String]](
+    this, "telemHeaders", "Map of Custom Header Key-Value Tuples."
+  )
+
+  private[ml] def setTelemHeaders(v: Map[String, String]): this.type = {
+    setScalarParam(telemHeaders, v)
+  }
+
+  // For Pyspark compatability accept Java HashMap as input to parameter
+  // py4J only natively supports conversions from Python Dict to Java HashMap
+  private[ml] def setTelemHeaders(v: java.util.HashMap[String, String]): this.type = {
+    setTelemHeaders(v.asScala.toMap)
+  }
+
+  setDefault(telemHeaders -> Left(Map("x-ai-telemetry-properties"->
+    s"""{
+       |"OriginatingService": "SynapseML",
+       |"ClientArtifactType": "Spark",
+       |"OperationName": "${this.getClass.getName}"
+       |}""".stripMargin.replaceAll("\n", ""))))
 
-  def getCustomHeaders: Map[String, String] = getScalarParam(customHeaders)
 }
 
+
 trait HasCustomCogServiceDomain extends Wrappable with HasURL with HasUrlPath {
   def setCustomServiceName(v: String): this.type = {
     setUrl(s"https://$v.cognitiveservices.azure.com/" + urlPath.stripPrefix("/"))
@@ -281,7 +305,7 @@ object URLEncodingUtils {
 }
 
 trait HasCognitiveServiceInput extends HasURL with HasSubscriptionKey with HasAADToken with HasCustomAuthHeader
-  with HasCustomHeaders with SynapseMLLogging {
+  with HasCustomHeaders with HasTelemHeaders with SynapseMLLogging {
 
   val customUrlRoot: Param[String] = new Param[String](
     this, "customUrlRoot", "The custom URL root for the service. " +
@@ -334,7 +358,7 @@ trait HasCognitiveServiceInput extends HasURL with HasSubscriptionKey with HasAA
 
   protected def getCustomAuthHeader(row: Row): Option[String] = {
     val providedCustomAuthHeader = getValueOpt(row, CustomAuthHeader)
-    if (providedCustomAuthHeader .isEmpty && PlatformDetails.runningOnFabric()) {
+    if (providedCustomAuthHeader.isEmpty && PlatformDetails.runningOnFabric()) {
       logInfo("Using Default AAD Token On Fabric")
       Option(FabricClient.getCognitiveMWCTokenAuthHeader)
     } else {
@@ -362,6 +386,7 @@ trait HasCognitiveServiceInput extends HasURL with HasSubscriptionKey with HasAA
     val contentTypeValue = contentType(row)
     val customAuthHeaderOpt = getCustomAuthHeader(row)
     val customHeadersOpt = getCustomHeaders(row)
+    val telemHeadersOpt =  getValueOpt(row, telemHeaders)
 
     if (subscriptionKeyOpt.nonEmpty) {
       headers += (subscriptionKeyHeaderName -> getValue(row, subscriptionKey))
@@ -376,6 +401,7 @@ trait HasCognitiveServiceInput extends HasURL with HasSubscriptionKey with HasAA
         headers += ("x-ms-workload-resource-moniker" -> UUID.randomUUID().toString)
       }
     }
+
     if (customHeadersOpt.nonEmpty) {
       customHeadersOpt.foreach { m =>
         m.foreach { case (headerName, headerValue) =>
@@ -384,10 +410,17 @@ trait HasCognitiveServiceInput extends HasURL with HasSubscriptionKey with HasAA
       }
     }
 
+    if (telemHeadersOpt.nonEmpty) {
+      telemHeadersOpt.foreach { m =>
+        m.foreach { case (headerName, headerValue) =>
+          headers += (headerName -> headerValue)
+        }
+      }
+    }
+
     if (addContentType && !StringUtils.isEmpty(contentTypeValue)) {
       headers += ("Content-Type" -> contentTypeValue)
     }
-
     new scala.collection.immutable.TreeMap[String, String]() ++ headers
   }
 
@@ -514,7 +547,7 @@ abstract class CognitiveServicesBaseNoHandler(val uid: String) extends Transform
     errorCol -> (this.uid + "_error")
   )
 
-  if(PlatformDetails.runningOnFabric()) {
+  if (PlatformDetails.runningOnFabric()) {
     setDefaultInternalEndpoint(FabricClient.MLWorkloadEndpointML)
   }
 
diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPSchema.scala
@@ -199,10 +199,7 @@ case class HTTPRequestData(requestLine: RequestLineData,
       request.setProtocolVersion(pv.toHTTPCore))
     request.setHeaders(headers.map(_.toHTTPCore) ++
       Array(new BasicHeader(
-        "User-Agent", s"synapseml/${BuildInfo.version}${HeaderValues.PlatformInfo}"),
-        new BasicHeader(
-          "x-ai-telemetry-properties", "{\"ClientArtifactType\": \"AIFunctionsSpark\"}"
-        )))
+        "User-Agent", s"synapseml/${BuildInfo.version}${HeaderValues.PlatformInfo}")))
     request
   }
   //scalastyle:on cyclomatic.complexity
