chore: Support Token Provider Mode #2160

lhrotk · 2024-01-22T16:43:31Z

What changes are proposed in this pull request?

Support token provider auth mode, auto refresh token while job is running.

How is this patch tested?

I have written tests (not required for typo or doc fix) and confirmed the proposed feature/bug-fix/change works.

Does this PR change any dependencies?

No. You can skip this section.
[] Yes. Make sure the dependencies are resolved correctly, and list changes here.

Does this PR add a new feature? If so, have you added samples on website?

No. You can skip this section.
Yes. Make sure you have added samples following below steps.

lhrotk · 2024-01-22T16:43:41Z

/azp run

github-actions · 2024-01-22T16:43:44Z

Hey @lhrotk 👋!
Thank you so much for contributing to our repository 🙌.
Someone from SynapseML Team will be reviewing this pull request soon.

We use semantic commit messages to streamline the release process.
Before your pull request can be merged, you should make sure your first commit and PR title start with a semantic prefix.
This helps us to create release messages and credit you for your hard work!

Examples of commit messages with semantic prefixes:

fix: Fix LightGBM crashes with empty partitions
feat: Make HTTP on Spark back-offs configurable
docs: Update Spark Serving usage
build: Add codecov support
perf: improve LightGBM memory usage
refactor: make python code generation rely on classes
style: Remove nulls from CNTKModel
test: Add test coverage for CNTKModel

To test your commit locally, please follow our guild on building from source.
Check out the developer guide for additional guidance on testing your change.

azure-pipelines · 2024-01-22T16:43:53Z

Azure Pipelines successfully started running 1 pipeline(s).

lhrotk · 2024-01-23T08:48:17Z

/azp run

azure-pipelines · 2024-01-23T08:48:31Z

Azure Pipelines successfully started running 1 pipeline(s).

codecov-commenter · 2024-01-23T09:02:33Z

Codecov Report

Attention: Patch coverage is 38.06452% with 96 lines in your changes are missing coverage. Please review.

Project coverage is 84.16%. Comparing base (345bad9) to head (66fb110).

Files	Patch %	Lines
...crosoft/azure/synapse/ml/fabric/FabricClient.scala	35.71%	36 Missing ⚠️
...t/azure/synapse/ml/fabric/OpenAITokenLibrary.scala	16.66%	25 Missing ⚠️
...ft/azure/synapse/ml/fabric/FabricTokenParser.scala	0.00%	14 Missing ⚠️
...re/synapse/ml/logging/common/PlatformDetails.scala	36.36%	7 Missing ⚠️
...azure/synapse/ml/services/text/TextAnalytics.scala	50.00%	5 Missing ⚠️
...ure/synapse/ml/services/CognitiveServiceBase.scala	66.66%	3 Missing ⚠️
...soft/azure/synapse/ml/services/openai/OpenAI.scala	60.00%	2 Missing ⚠️
...napse/ml/logging/fabric/CertifiedEventClient.scala	33.33%	2 Missing ⚠️
.../microsoft/azure/synapse/ml/fabric/RESTUtils.scala	0.00%	1 Missing ⚠️
...crosoft/azure/synapse/ml/fabric/TokenLibrary.scala	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2160      +/-   ##
==========================================
- Coverage   84.47%   84.16%   -0.31%     
==========================================
  Files         325      328       +3     
  Lines       16959    17067     +108     
  Branches     1524     1548      +24     
==========================================
+ Hits        14326    14365      +39     
- Misses       2633     2702      +69

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

lhrotk · 2024-01-23T11:59:49Z

/azp run

azure-pipelines · 2024-01-23T11:59:59Z

Azure Pipelines successfully started running 1 pipeline(s).

lhrotk · 2024-01-24T09:22:35Z

/azp run

azure-pipelines · 2024-01-24T09:22:45Z

Azure Pipelines successfully started running 1 pipeline(s).

mhamilton723 · 2024-01-24T18:58:37Z

build.sbt

@@ -32,6 +32,8 @@ val extraDependencies = Seq(
  "org.scalactic" %% "scalactic" % "3.2.14",
  "io.spray" %% "spray-json" % "1.3.5",
  "com.jcraft" % "jsch" % "0.1.54",
+  "com.pauldijou" %% "jwt-core" % "3.0.0",
+  "org.json" % "json" % "20180130",


Use existing spray parsing dep please

mhamilton723 · 2024-01-24T18:58:58Z

build.sbt

@@ -32,6 +32,8 @@ val extraDependencies = Seq(
  "org.scalactic" %% "scalactic" % "3.2.14",
  "io.spray" %% "spray-json" % "1.3.5",
  "com.jcraft" % "jsch" % "0.1.54",
+  "com.pauldijou" %% "jwt-core" % "3.0.0",


we have some code that parses jwt, check the certified events logging code

I didn't find any jwt related code, I am reading this: certified_event, can you point me to the right place

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/OpenAITokenLibrary.scala

lhrotk · 2024-01-25T10:44:40Z

/azp run

azure-pipelines · 2024-01-25T10:44:52Z

Azure Pipelines successfully started running 1 pipeline(s).

lhrotk · 2024-01-26T07:25:18Z

/azp run

azure-pipelines · 2024-01-26T07:25:28Z

Azure Pipelines successfully started running 1 pipeline(s).

lhrotk · 2024-02-07T08:19:23Z

/azp run

azure-pipelines · 2024-02-07T08:19:33Z

Azure Pipelines successfully started running 1 pipeline(s).

lhrotk · 2024-02-21T12:09:42Z

/azp run

azure-pipelines · 2024-02-21T12:09:54Z

Azure Pipelines successfully started running 1 pipeline(s).

lhrotk · 2024-02-22T05:06:38Z

/azp run

azure-pipelines · 2024-02-22T05:06:49Z

Azure Pipelines successfully started running 1 pipeline(s).

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala

mhamilton723 · 2024-02-22T19:33:00Z

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/OpenAITokenLibrary.scala

+import spray.json.DefaultJsonProtocol.StringJsonFormat
+
+object OpenAITokenLibrary extends SynapseMLLogging with AuthHeaderProvider {
+  var MLMWCToken = "";


can we try to eliminate this var by making it a def or something?

I need another var to tell it is mwc token or aad token and I think it is strange to change value of another var in def for this var. And also this var is a catched token and will expire sometime while def is called only once, wdyt?

core/src/main/scala/com/microsoft/azure/synapse/ml/logging/fabric/CertifiedEventClient.scala

lhrotk · 2024-02-23T14:35:20Z

/azp run

azure-pipelines · 2024-02-23T14:35:31Z

Azure Pipelines successfully started running 1 pipeline(s).

mhamilton723 · 2024-02-23T23:23:41Z

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala

@@ -244,6 +247,19 @@ trait HasOpenAITextParams extends HasOpenAISharedParams {
  }
 }

+trait HasOpenAICognitiveServiceInput extends HasCognitiveServiceInput {
+  override protected def getCustomAuthHeader(row: Row): Option[String] = {
+    var providedCustomHeader = getValueOpt(row, CustomAuthHeader)


mhamilton723 · 2024-02-23T23:24:50Z

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala

+          case _: Exception =>
+            // Handle MalformedURLException or other exceptions
+            None
+        }
+      case None =>
+        None


Do we expect these branches to be triggered ever?

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/OpenAITokenLibrary.scala

mhamilton723 · 2024-02-23T23:27:34Z

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricTokenParser.scala

+
+class JwtTokenExpiryMissingException(message: String) extends Exception(message)
+
+class FabricTokenParser(JWToken: String) {


make all the APIs private or private[ml] if they arent supposed to be used by a end user

lhrotk · 2024-02-26T08:26:33Z

/azp run

azure-pipelines · 2024-02-26T08:26:44Z

Azure Pipelines successfully started running 1 pipeline(s).

mhamilton723 · 2024-02-27T19:44:24Z

/azp run

azure-pipelines · 2024-02-27T19:44:36Z

Azure Pipelines successfully started running 1 pipeline(s).

mslhrotk added 2 commits January 22, 2024 12:33

add header provider

e379af7

add trident mwc token provider

10d8675

lhrotk requested a review from mhamilton723 as a code owner January 22, 2024 16:43

lhrotk force-pushed the cl/add-header-provider branch from a7db643 to ffa21ac Compare January 23, 2024 08:37

bug fix

8122f86

lhrotk force-pushed the cl/add-header-provider branch from ffa21ac to 8122f86 Compare January 23, 2024 11:59

mhamilton723 reviewed Jan 24, 2024

View reviewed changes

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/OpenAITokenLibrary.scala Show resolved Hide resolved

add fabric token provider in cog service

aeb302d

lhrotk force-pushed the cl/add-header-provider branch from a739978 to aeb302d Compare January 25, 2024 10:43

make platform details work on executor

dc1cffa

use tokenservice endpoint instead of resolve at runtime

3c2dd56

fix style

f111d8f

fix fabric token issue for text analysis

3299f3b

fix style

fd7bd6b

mhamilton723 reviewed Feb 22, 2024

View reviewed changes

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala Outdated Show resolved Hide resolved

mhamilton723 reviewed Feb 22, 2024

View reviewed changes

cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/CognitiveServiceBase.scala Outdated Show resolved Hide resolved

mhamilton723 reviewed Feb 22, 2024

View reviewed changes

core/src/main/scala/com/microsoft/azure/synapse/ml/fabric/FabricClient.scala Outdated Show resolved Hide resolved