flyte 0.0.1b0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of flyte might be problematic. Click here for more details.

Files changed (390) hide show
  1. flyte/__init__.py +62 -0
  2. flyte/_api_commons.py +3 -0
  3. flyte/_bin/__init__.py +0 -0
  4. flyte/_bin/runtime.py +126 -0
  5. flyte/_build.py +25 -0
  6. flyte/_cache/__init__.py +12 -0
  7. flyte/_cache/cache.py +146 -0
  8. flyte/_cache/defaults.py +9 -0
  9. flyte/_cache/policy_function_body.py +42 -0
  10. flyte/_cli/__init__.py +0 -0
  11. flyte/_cli/_common.py +287 -0
  12. flyte/_cli/_create.py +42 -0
  13. flyte/_cli/_delete.py +23 -0
  14. flyte/_cli/_deploy.py +140 -0
  15. flyte/_cli/_get.py +235 -0
  16. flyte/_cli/_run.py +152 -0
  17. flyte/_cli/main.py +72 -0
  18. flyte/_code_bundle/__init__.py +8 -0
  19. flyte/_code_bundle/_ignore.py +113 -0
  20. flyte/_code_bundle/_packaging.py +187 -0
  21. flyte/_code_bundle/_utils.py +339 -0
  22. flyte/_code_bundle/bundle.py +178 -0
  23. flyte/_context.py +146 -0
  24. flyte/_datastructures.py +342 -0
  25. flyte/_deploy.py +202 -0
  26. flyte/_doc.py +29 -0
  27. flyte/_docstring.py +32 -0
  28. flyte/_environment.py +43 -0
  29. flyte/_group.py +31 -0
  30. flyte/_hash.py +23 -0
  31. flyte/_image.py +760 -0
  32. flyte/_initialize.py +634 -0
  33. flyte/_interface.py +84 -0
  34. flyte/_internal/__init__.py +3 -0
  35. flyte/_internal/controllers/__init__.py +115 -0
  36. flyte/_internal/controllers/_local_controller.py +118 -0
  37. flyte/_internal/controllers/_trace.py +40 -0
  38. flyte/_internal/controllers/pbhash.py +39 -0
  39. flyte/_internal/controllers/remote/__init__.py +40 -0
  40. flyte/_internal/controllers/remote/_action.py +141 -0
  41. flyte/_internal/controllers/remote/_client.py +43 -0
  42. flyte/_internal/controllers/remote/_controller.py +361 -0
  43. flyte/_internal/controllers/remote/_core.py +402 -0
  44. flyte/_internal/controllers/remote/_informer.py +361 -0
  45. flyte/_internal/controllers/remote/_service_protocol.py +50 -0
  46. flyte/_internal/imagebuild/__init__.py +11 -0
  47. flyte/_internal/imagebuild/docker_builder.py +416 -0
  48. flyte/_internal/imagebuild/image_builder.py +241 -0
  49. flyte/_internal/imagebuild/remote_builder.py +0 -0
  50. flyte/_internal/resolvers/__init__.py +0 -0
  51. flyte/_internal/resolvers/_task_module.py +54 -0
  52. flyte/_internal/resolvers/common.py +31 -0
  53. flyte/_internal/resolvers/default.py +28 -0
  54. flyte/_internal/runtime/__init__.py +0 -0
  55. flyte/_internal/runtime/convert.py +199 -0
  56. flyte/_internal/runtime/entrypoints.py +135 -0
  57. flyte/_internal/runtime/io.py +136 -0
  58. flyte/_internal/runtime/resources_serde.py +138 -0
  59. flyte/_internal/runtime/task_serde.py +210 -0
  60. flyte/_internal/runtime/taskrunner.py +190 -0
  61. flyte/_internal/runtime/types_serde.py +54 -0
  62. flyte/_logging.py +124 -0
  63. flyte/_protos/__init__.py +0 -0
  64. flyte/_protos/common/authorization_pb2.py +66 -0
  65. flyte/_protos/common/authorization_pb2.pyi +108 -0
  66. flyte/_protos/common/authorization_pb2_grpc.py +4 -0
  67. flyte/_protos/common/identifier_pb2.py +71 -0
  68. flyte/_protos/common/identifier_pb2.pyi +82 -0
  69. flyte/_protos/common/identifier_pb2_grpc.py +4 -0
  70. flyte/_protos/common/identity_pb2.py +48 -0
  71. flyte/_protos/common/identity_pb2.pyi +72 -0
  72. flyte/_protos/common/identity_pb2_grpc.py +4 -0
  73. flyte/_protos/common/list_pb2.py +36 -0
  74. flyte/_protos/common/list_pb2.pyi +69 -0
  75. flyte/_protos/common/list_pb2_grpc.py +4 -0
  76. flyte/_protos/common/policy_pb2.py +37 -0
  77. flyte/_protos/common/policy_pb2.pyi +27 -0
  78. flyte/_protos/common/policy_pb2_grpc.py +4 -0
  79. flyte/_protos/common/role_pb2.py +37 -0
  80. flyte/_protos/common/role_pb2.pyi +53 -0
  81. flyte/_protos/common/role_pb2_grpc.py +4 -0
  82. flyte/_protos/common/runtime_version_pb2.py +28 -0
  83. flyte/_protos/common/runtime_version_pb2.pyi +24 -0
  84. flyte/_protos/common/runtime_version_pb2_grpc.py +4 -0
  85. flyte/_protos/logs/dataplane/payload_pb2.py +96 -0
  86. flyte/_protos/logs/dataplane/payload_pb2.pyi +168 -0
  87. flyte/_protos/logs/dataplane/payload_pb2_grpc.py +4 -0
  88. flyte/_protos/secret/definition_pb2.py +49 -0
  89. flyte/_protos/secret/definition_pb2.pyi +93 -0
  90. flyte/_protos/secret/definition_pb2_grpc.py +4 -0
  91. flyte/_protos/secret/payload_pb2.py +62 -0
  92. flyte/_protos/secret/payload_pb2.pyi +94 -0
  93. flyte/_protos/secret/payload_pb2_grpc.py +4 -0
  94. flyte/_protos/secret/secret_pb2.py +38 -0
  95. flyte/_protos/secret/secret_pb2.pyi +6 -0
  96. flyte/_protos/secret/secret_pb2_grpc.py +198 -0
  97. flyte/_protos/secret/secret_pb2_grpc_grpc.py +198 -0
  98. flyte/_protos/validate/validate/validate_pb2.py +76 -0
  99. flyte/_protos/workflow/node_execution_service_pb2.py +26 -0
  100. flyte/_protos/workflow/node_execution_service_pb2.pyi +4 -0
  101. flyte/_protos/workflow/node_execution_service_pb2_grpc.py +32 -0
  102. flyte/_protos/workflow/queue_service_pb2.py +106 -0
  103. flyte/_protos/workflow/queue_service_pb2.pyi +141 -0
  104. flyte/_protos/workflow/queue_service_pb2_grpc.py +172 -0
  105. flyte/_protos/workflow/run_definition_pb2.py +128 -0
  106. flyte/_protos/workflow/run_definition_pb2.pyi +310 -0
  107. flyte/_protos/workflow/run_definition_pb2_grpc.py +4 -0
  108. flyte/_protos/workflow/run_logs_service_pb2.py +41 -0
  109. flyte/_protos/workflow/run_logs_service_pb2.pyi +28 -0
  110. flyte/_protos/workflow/run_logs_service_pb2_grpc.py +69 -0
  111. flyte/_protos/workflow/run_service_pb2.py +133 -0
  112. flyte/_protos/workflow/run_service_pb2.pyi +175 -0
  113. flyte/_protos/workflow/run_service_pb2_grpc.py +412 -0
  114. flyte/_protos/workflow/state_service_pb2.py +58 -0
  115. flyte/_protos/workflow/state_service_pb2.pyi +71 -0
  116. flyte/_protos/workflow/state_service_pb2_grpc.py +138 -0
  117. flyte/_protos/workflow/task_definition_pb2.py +72 -0
  118. flyte/_protos/workflow/task_definition_pb2.pyi +65 -0
  119. flyte/_protos/workflow/task_definition_pb2_grpc.py +4 -0
  120. flyte/_protos/workflow/task_service_pb2.py +44 -0
  121. flyte/_protos/workflow/task_service_pb2.pyi +31 -0
  122. flyte/_protos/workflow/task_service_pb2_grpc.py +104 -0
  123. flyte/_resources.py +226 -0
  124. flyte/_retry.py +32 -0
  125. flyte/_reusable_environment.py +25 -0
  126. flyte/_run.py +411 -0
  127. flyte/_secret.py +61 -0
  128. flyte/_task.py +367 -0
  129. flyte/_task_environment.py +200 -0
  130. flyte/_timeout.py +47 -0
  131. flyte/_tools.py +27 -0
  132. flyte/_trace.py +128 -0
  133. flyte/_utils/__init__.py +20 -0
  134. flyte/_utils/asyn.py +119 -0
  135. flyte/_utils/coro_management.py +25 -0
  136. flyte/_utils/file_handling.py +72 -0
  137. flyte/_utils/helpers.py +108 -0
  138. flyte/_utils/lazy_module.py +54 -0
  139. flyte/_utils/uv_script_parser.py +49 -0
  140. flyte/_version.py +21 -0
  141. flyte/connectors/__init__.py +0 -0
  142. flyte/errors.py +143 -0
  143. flyte/extras/__init__.py +5 -0
  144. flyte/extras/_container.py +273 -0
  145. flyte/io/__init__.py +11 -0
  146. flyte/io/_dataframe.py +0 -0
  147. flyte/io/_dir.py +448 -0
  148. flyte/io/_file.py +468 -0
  149. flyte/io/pickle/__init__.py +0 -0
  150. flyte/io/pickle/transformer.py +117 -0
  151. flyte/io/structured_dataset/__init__.py +129 -0
  152. flyte/io/structured_dataset/basic_dfs.py +219 -0
  153. flyte/io/structured_dataset/structured_dataset.py +1061 -0
  154. flyte/py.typed +0 -0
  155. flyte/remote/__init__.py +25 -0
  156. flyte/remote/_client/__init__.py +0 -0
  157. flyte/remote/_client/_protocols.py +131 -0
  158. flyte/remote/_client/auth/__init__.py +12 -0
  159. flyte/remote/_client/auth/_authenticators/__init__.py +0 -0
  160. flyte/remote/_client/auth/_authenticators/base.py +397 -0
  161. flyte/remote/_client/auth/_authenticators/client_credentials.py +73 -0
  162. flyte/remote/_client/auth/_authenticators/device_code.py +118 -0
  163. flyte/remote/_client/auth/_authenticators/external_command.py +79 -0
  164. flyte/remote/_client/auth/_authenticators/factory.py +200 -0
  165. flyte/remote/_client/auth/_authenticators/pkce.py +516 -0
  166. flyte/remote/_client/auth/_channel.py +184 -0
  167. flyte/remote/_client/auth/_client_config.py +83 -0
  168. flyte/remote/_client/auth/_default_html.py +32 -0
  169. flyte/remote/_client/auth/_grpc_utils/__init__.py +0 -0
  170. flyte/remote/_client/auth/_grpc_utils/auth_interceptor.py +288 -0
  171. flyte/remote/_client/auth/_grpc_utils/default_metadata_interceptor.py +151 -0
  172. flyte/remote/_client/auth/_keyring.py +143 -0
  173. flyte/remote/_client/auth/_token_client.py +260 -0
  174. flyte/remote/_client/auth/errors.py +16 -0
  175. flyte/remote/_client/controlplane.py +95 -0
  176. flyte/remote/_console.py +18 -0
  177. flyte/remote/_data.py +155 -0
  178. flyte/remote/_logs.py +116 -0
  179. flyte/remote/_project.py +86 -0
  180. flyte/remote/_run.py +873 -0
  181. flyte/remote/_secret.py +132 -0
  182. flyte/remote/_task.py +227 -0
  183. flyte/report/__init__.py +3 -0
  184. flyte/report/_report.py +178 -0
  185. flyte/report/_template.html +124 -0
  186. flyte/storage/__init__.py +24 -0
  187. flyte/storage/_remote_fs.py +34 -0
  188. flyte/storage/_storage.py +251 -0
  189. flyte/storage/_utils.py +5 -0
  190. flyte/types/__init__.py +13 -0
  191. flyte/types/_interface.py +25 -0
  192. flyte/types/_renderer.py +162 -0
  193. flyte/types/_string_literals.py +120 -0
  194. flyte/types/_type_engine.py +2210 -0
  195. flyte/types/_utils.py +80 -0
  196. flyte-0.0.1b0.dist-info/METADATA +179 -0
  197. flyte-0.0.1b0.dist-info/RECORD +390 -0
  198. flyte-0.0.1b0.dist-info/WHEEL +5 -0
  199. flyte-0.0.1b0.dist-info/entry_points.txt +3 -0
  200. flyte-0.0.1b0.dist-info/top_level.txt +1 -0
  201. union/__init__.py +54 -0
  202. union/_api_commons.py +3 -0
  203. union/_bin/__init__.py +0 -0
  204. union/_bin/runtime.py +113 -0
  205. union/_build.py +25 -0
  206. union/_cache/__init__.py +12 -0
  207. union/_cache/cache.py +141 -0
  208. union/_cache/defaults.py +9 -0
  209. union/_cache/policy_function_body.py +42 -0
  210. union/_cli/__init__.py +0 -0
  211. union/_cli/_common.py +263 -0
  212. union/_cli/_create.py +40 -0
  213. union/_cli/_delete.py +23 -0
  214. union/_cli/_deploy.py +120 -0
  215. union/_cli/_get.py +162 -0
  216. union/_cli/_params.py +579 -0
  217. union/_cli/_run.py +150 -0
  218. union/_cli/main.py +72 -0
  219. union/_code_bundle/__init__.py +8 -0
  220. union/_code_bundle/_ignore.py +113 -0
  221. union/_code_bundle/_packaging.py +187 -0
  222. union/_code_bundle/_utils.py +342 -0
  223. union/_code_bundle/bundle.py +176 -0
  224. union/_context.py +146 -0
  225. union/_datastructures.py +295 -0
  226. union/_deploy.py +185 -0
  227. union/_doc.py +29 -0
  228. union/_docstring.py +26 -0
  229. union/_environment.py +43 -0
  230. union/_group.py +31 -0
  231. union/_hash.py +23 -0
  232. union/_image.py +760 -0
  233. union/_initialize.py +585 -0
  234. union/_interface.py +84 -0
  235. union/_internal/__init__.py +3 -0
  236. union/_internal/controllers/__init__.py +77 -0
  237. union/_internal/controllers/_local_controller.py +77 -0
  238. union/_internal/controllers/pbhash.py +39 -0
  239. union/_internal/controllers/remote/__init__.py +40 -0
  240. union/_internal/controllers/remote/_action.py +131 -0
  241. union/_internal/controllers/remote/_client.py +43 -0
  242. union/_internal/controllers/remote/_controller.py +169 -0
  243. union/_internal/controllers/remote/_core.py +341 -0
  244. union/_internal/controllers/remote/_informer.py +260 -0
  245. union/_internal/controllers/remote/_service_protocol.py +44 -0
  246. union/_internal/imagebuild/__init__.py +11 -0
  247. union/_internal/imagebuild/docker_builder.py +416 -0
  248. union/_internal/imagebuild/image_builder.py +243 -0
  249. union/_internal/imagebuild/remote_builder.py +0 -0
  250. union/_internal/resolvers/__init__.py +0 -0
  251. union/_internal/resolvers/_task_module.py +31 -0
  252. union/_internal/resolvers/common.py +24 -0
  253. union/_internal/resolvers/default.py +27 -0
  254. union/_internal/runtime/__init__.py +0 -0
  255. union/_internal/runtime/convert.py +163 -0
  256. union/_internal/runtime/entrypoints.py +121 -0
  257. union/_internal/runtime/io.py +136 -0
  258. union/_internal/runtime/resources_serde.py +134 -0
  259. union/_internal/runtime/task_serde.py +202 -0
  260. union/_internal/runtime/taskrunner.py +179 -0
  261. union/_internal/runtime/types_serde.py +53 -0
  262. union/_logging.py +124 -0
  263. union/_protos/__init__.py +0 -0
  264. union/_protos/common/authorization_pb2.py +66 -0
  265. union/_protos/common/authorization_pb2.pyi +106 -0
  266. union/_protos/common/authorization_pb2_grpc.py +4 -0
  267. union/_protos/common/identifier_pb2.py +71 -0
  268. union/_protos/common/identifier_pb2.pyi +82 -0
  269. union/_protos/common/identifier_pb2_grpc.py +4 -0
  270. union/_protos/common/identity_pb2.py +48 -0
  271. union/_protos/common/identity_pb2.pyi +72 -0
  272. union/_protos/common/identity_pb2_grpc.py +4 -0
  273. union/_protos/common/list_pb2.py +36 -0
  274. union/_protos/common/list_pb2.pyi +69 -0
  275. union/_protos/common/list_pb2_grpc.py +4 -0
  276. union/_protos/common/policy_pb2.py +37 -0
  277. union/_protos/common/policy_pb2.pyi +27 -0
  278. union/_protos/common/policy_pb2_grpc.py +4 -0
  279. union/_protos/common/role_pb2.py +37 -0
  280. union/_protos/common/role_pb2.pyi +51 -0
  281. union/_protos/common/role_pb2_grpc.py +4 -0
  282. union/_protos/common/runtime_version_pb2.py +28 -0
  283. union/_protos/common/runtime_version_pb2.pyi +24 -0
  284. union/_protos/common/runtime_version_pb2_grpc.py +4 -0
  285. union/_protos/logs/dataplane/payload_pb2.py +96 -0
  286. union/_protos/logs/dataplane/payload_pb2.pyi +168 -0
  287. union/_protos/logs/dataplane/payload_pb2_grpc.py +4 -0
  288. union/_protos/secret/definition_pb2.py +49 -0
  289. union/_protos/secret/definition_pb2.pyi +93 -0
  290. union/_protos/secret/definition_pb2_grpc.py +4 -0
  291. union/_protos/secret/payload_pb2.py +62 -0
  292. union/_protos/secret/payload_pb2.pyi +94 -0
  293. union/_protos/secret/payload_pb2_grpc.py +4 -0
  294. union/_protos/secret/secret_pb2.py +38 -0
  295. union/_protos/secret/secret_pb2.pyi +6 -0
  296. union/_protos/secret/secret_pb2_grpc.py +198 -0
  297. union/_protos/validate/validate/validate_pb2.py +76 -0
  298. union/_protos/workflow/node_execution_service_pb2.py +26 -0
  299. union/_protos/workflow/node_execution_service_pb2.pyi +4 -0
  300. union/_protos/workflow/node_execution_service_pb2_grpc.py +32 -0
  301. union/_protos/workflow/queue_service_pb2.py +75 -0
  302. union/_protos/workflow/queue_service_pb2.pyi +103 -0
  303. union/_protos/workflow/queue_service_pb2_grpc.py +172 -0
  304. union/_protos/workflow/run_definition_pb2.py +100 -0
  305. union/_protos/workflow/run_definition_pb2.pyi +256 -0
  306. union/_protos/workflow/run_definition_pb2_grpc.py +4 -0
  307. union/_protos/workflow/run_logs_service_pb2.py +41 -0
  308. union/_protos/workflow/run_logs_service_pb2.pyi +28 -0
  309. union/_protos/workflow/run_logs_service_pb2_grpc.py +69 -0
  310. union/_protos/workflow/run_service_pb2.py +133 -0
  311. union/_protos/workflow/run_service_pb2.pyi +173 -0
  312. union/_protos/workflow/run_service_pb2_grpc.py +412 -0
  313. union/_protos/workflow/state_service_pb2.py +58 -0
  314. union/_protos/workflow/state_service_pb2.pyi +69 -0
  315. union/_protos/workflow/state_service_pb2_grpc.py +138 -0
  316. union/_protos/workflow/task_definition_pb2.py +72 -0
  317. union/_protos/workflow/task_definition_pb2.pyi +65 -0
  318. union/_protos/workflow/task_definition_pb2_grpc.py +4 -0
  319. union/_protos/workflow/task_service_pb2.py +44 -0
  320. union/_protos/workflow/task_service_pb2.pyi +31 -0
  321. union/_protos/workflow/task_service_pb2_grpc.py +104 -0
  322. union/_resources.py +226 -0
  323. union/_retry.py +32 -0
  324. union/_reusable_environment.py +25 -0
  325. union/_run.py +374 -0
  326. union/_secret.py +61 -0
  327. union/_task.py +354 -0
  328. union/_task_environment.py +186 -0
  329. union/_timeout.py +47 -0
  330. union/_tools.py +27 -0
  331. union/_utils/__init__.py +11 -0
  332. union/_utils/asyn.py +119 -0
  333. union/_utils/file_handling.py +71 -0
  334. union/_utils/helpers.py +46 -0
  335. union/_utils/lazy_module.py +54 -0
  336. union/_utils/uv_script_parser.py +49 -0
  337. union/_version.py +21 -0
  338. union/connectors/__init__.py +0 -0
  339. union/errors.py +128 -0
  340. union/extras/__init__.py +5 -0
  341. union/extras/_container.py +263 -0
  342. union/io/__init__.py +11 -0
  343. union/io/_dataframe.py +0 -0
  344. union/io/_dir.py +425 -0
  345. union/io/_file.py +418 -0
  346. union/io/pickle/__init__.py +0 -0
  347. union/io/pickle/transformer.py +117 -0
  348. union/io/structured_dataset/__init__.py +122 -0
  349. union/io/structured_dataset/basic_dfs.py +219 -0
  350. union/io/structured_dataset/structured_dataset.py +1057 -0
  351. union/py.typed +0 -0
  352. union/remote/__init__.py +23 -0
  353. union/remote/_client/__init__.py +0 -0
  354. union/remote/_client/_protocols.py +129 -0
  355. union/remote/_client/auth/__init__.py +12 -0
  356. union/remote/_client/auth/_authenticators/__init__.py +0 -0
  357. union/remote/_client/auth/_authenticators/base.py +391 -0
  358. union/remote/_client/auth/_authenticators/client_credentials.py +73 -0
  359. union/remote/_client/auth/_authenticators/device_code.py +120 -0
  360. union/remote/_client/auth/_authenticators/external_command.py +77 -0
  361. union/remote/_client/auth/_authenticators/factory.py +200 -0
  362. union/remote/_client/auth/_authenticators/pkce.py +515 -0
  363. union/remote/_client/auth/_channel.py +184 -0
  364. union/remote/_client/auth/_client_config.py +83 -0
  365. union/remote/_client/auth/_default_html.py +32 -0
  366. union/remote/_client/auth/_grpc_utils/__init__.py +0 -0
  367. union/remote/_client/auth/_grpc_utils/auth_interceptor.py +204 -0
  368. union/remote/_client/auth/_grpc_utils/default_metadata_interceptor.py +144 -0
  369. union/remote/_client/auth/_keyring.py +154 -0
  370. union/remote/_client/auth/_token_client.py +258 -0
  371. union/remote/_client/auth/errors.py +16 -0
  372. union/remote/_client/controlplane.py +86 -0
  373. union/remote/_data.py +149 -0
  374. union/remote/_logs.py +74 -0
  375. union/remote/_project.py +86 -0
  376. union/remote/_run.py +820 -0
  377. union/remote/_secret.py +132 -0
  378. union/remote/_task.py +193 -0
  379. union/report/__init__.py +3 -0
  380. union/report/_report.py +178 -0
  381. union/report/_template.html +124 -0
  382. union/storage/__init__.py +24 -0
  383. union/storage/_remote_fs.py +34 -0
  384. union/storage/_storage.py +247 -0
  385. union/storage/_utils.py +5 -0
  386. union/types/__init__.py +11 -0
  387. union/types/_renderer.py +162 -0
  388. union/types/_string_literals.py +120 -0
  389. union/types/_type_engine.py +2131 -0
  390. union/types/_utils.py +80 -0
flyte/py.typed ADDED
File without changes
@@ -0,0 +1,25 @@
1
+ """
2
+ Remote Entities that are accessible from the Union Server once deployed or created.
3
+ """
4
+
5
+ __all__ = [
6
+ "Action",
7
+ "ActionDetails",
8
+ "ActionInputs",
9
+ "ActionOutputs",
10
+ "Project",
11
+ "Run",
12
+ "RunDetails",
13
+ "Secret",
14
+ "Task",
15
+ "create_channel",
16
+ "upload_dir",
17
+ "upload_file",
18
+ ]
19
+
20
+ from ._client.auth import create_channel
21
+ from ._data import upload_dir, upload_file
22
+ from ._project import Project
23
+ from ._run import Action, ActionDetails, ActionInputs, ActionOutputs, Run, RunDetails
24
+ from ._secret import Secret
25
+ from ._task import Task
File without changes
@@ -0,0 +1,131 @@
1
+ from typing import AsyncIterator, Protocol
2
+
3
+ from flyteidl.admin import project_attributes_pb2, project_pb2, version_pb2
4
+ from flyteidl.service import dataproxy_pb2
5
+ from grpc.aio import UnaryStreamCall
6
+ from grpc.aio._typing import RequestType
7
+
8
+ from flyte._protos.secret import payload_pb2
9
+ from flyte._protos.workflow import run_logs_service_pb2, run_service_pb2, task_service_pb2
10
+
11
+
12
+ class MetadataServiceProtocol(Protocol):
13
+ async def GetVersion(self, request: version_pb2.GetVersionRequest) -> version_pb2.GetVersionResponse: ...
14
+
15
+
16
+ class ProjectDomainService(Protocol):
17
+ async def RegisterProject(
18
+ self, request: project_pb2.ProjectRegisterRequest
19
+ ) -> project_pb2.ProjectRegisterResponse: ...
20
+
21
+ async def UpdateProject(self, request: project_pb2.Project) -> project_pb2.ProjectUpdateResponse: ...
22
+
23
+ async def GetProject(self, request: project_pb2.ProjectGetRequest) -> project_pb2.Project: ...
24
+
25
+ async def ListProjects(self, request: project_pb2.ProjectListRequest) -> project_pb2.Projects: ...
26
+
27
+ async def GetDomains(self, request: project_pb2.GetDomainRequest) -> project_pb2.GetDomainsResponse: ...
28
+
29
+ async def UpdateProjectDomainAttributes(
30
+ self, request: project_attributes_pb2.ProjectAttributesUpdateRequest
31
+ ) -> project_pb2.ProjectUpdateResponse: ...
32
+
33
+ async def GetProjectDomainAttributes(
34
+ self, request: project_attributes_pb2.ProjectAttributesGetRequest
35
+ ) -> project_attributes_pb2.ProjectAttributes: ...
36
+
37
+ async def DeleteProjectDomainAttributes(
38
+ self, request: project_attributes_pb2.ProjectAttributesDeleteRequest
39
+ ) -> project_attributes_pb2.ProjectAttributesDeleteResponse: ...
40
+
41
+ async def UpdateProjectAttributes(
42
+ self, request: project_attributes_pb2.ProjectAttributesUpdateRequest
43
+ ) -> project_attributes_pb2.ProjectAttributesUpdateResponse: ...
44
+
45
+ async def GetProjectAttributes(
46
+ self, request: project_attributes_pb2.ProjectAttributesGetRequest
47
+ ) -> project_attributes_pb2.ProjectAttributes: ...
48
+
49
+ async def DeleteProjectAttributes(
50
+ self, request: project_attributes_pb2.ProjectAttributesDeleteRequest
51
+ ) -> project_attributes_pb2.ProjectAttributesDeleteResponse: ...
52
+
53
+
54
+ class TaskService(Protocol):
55
+ async def DeployTask(self, request: task_service_pb2.DeployTaskRequest) -> task_service_pb2.DeployTaskResponse: ...
56
+
57
+ async def GetTaskDetails(
58
+ self, request: task_service_pb2.GetTaskDetailsRequest
59
+ ) -> task_service_pb2.GetTaskDetailsResponse: ...
60
+
61
+
62
+ class RunService(Protocol):
63
+ async def CreateRun(self, request: run_service_pb2.CreateRunRequest) -> run_service_pb2.CreateRunResponse: ...
64
+
65
+ async def AbortRun(self, request: run_service_pb2.AbortRunRequest) -> run_service_pb2.AbortRunResponse: ...
66
+
67
+ async def GetRunDetails(
68
+ self, request: run_service_pb2.GetRunDetailsRequest
69
+ ) -> run_service_pb2.GetRunDetailsResponse: ...
70
+
71
+ async def WatchRunDetails(
72
+ self, request: run_service_pb2.WatchRunDetailsRequest
73
+ ) -> AsyncIterator[run_service_pb2.WatchRunDetailsResponse]: ...
74
+
75
+ async def GetActionDetails(
76
+ self, request: run_service_pb2.GetActionDetailsRequest
77
+ ) -> run_service_pb2.GetActionDetailsResponse: ...
78
+
79
+ async def WatchActionDetails(
80
+ self, request: run_service_pb2.WatchActionDetailsRequest
81
+ ) -> AsyncIterator[run_service_pb2.WatchActionDetailsResponse]: ...
82
+
83
+ async def GetActionData(
84
+ self, request: run_service_pb2.GetActionDataRequest
85
+ ) -> run_service_pb2.GetActionDataResponse: ...
86
+
87
+ async def ListRuns(self, request: run_service_pb2.ListRunsRequest) -> run_service_pb2.ListRunsResponse: ...
88
+
89
+ async def WatchRuns(
90
+ self, request: run_service_pb2.WatchRunsRequest
91
+ ) -> AsyncIterator[run_service_pb2.WatchRunsResponse]: ...
92
+
93
+ async def ListActions(self, request: run_service_pb2.ListActionsRequest) -> run_service_pb2.ListActionsResponse: ...
94
+
95
+ async def WatchActions(
96
+ self, request: run_service_pb2.WatchActionsRequest
97
+ ) -> AsyncIterator[run_service_pb2.WatchActionsResponse]: ...
98
+
99
+
100
+ class DataProxyService(Protocol):
101
+ async def CreateUploadLocation(
102
+ self, request: dataproxy_pb2.CreateUploadLocationRequest
103
+ ) -> dataproxy_pb2.CreateUploadLocationResponse: ...
104
+
105
+ async def CreateDownloadLocation(
106
+ self, request: dataproxy_pb2.CreateDownloadLocationRequest
107
+ ) -> dataproxy_pb2.CreateDownloadLocationResponse: ...
108
+
109
+ async def CreateDownloadLink(
110
+ self, request: dataproxy_pb2.CreateDownloadLinkRequest
111
+ ) -> dataproxy_pb2.CreateDownloadLinkResponse: ...
112
+
113
+ async def GetData(self, request: dataproxy_pb2.GetDataRequest) -> dataproxy_pb2.GetDataResponse: ...
114
+
115
+
116
+ class RunLogsService(Protocol):
117
+ def TailLogs(
118
+ self, request: run_logs_service_pb2.TailLogsRequest
119
+ ) -> UnaryStreamCall[RequestType, run_logs_service_pb2.TailLogsResponse]: ...
120
+
121
+
122
+ class SecretService(Protocol):
123
+ async def CreateSecret(self, request: payload_pb2.CreateSecretRequest) -> payload_pb2.CreateSecretResponse: ...
124
+
125
+ async def UpdateSecret(self, request: payload_pb2.UpdateSecretRequest) -> payload_pb2.UpdateSecretResponse: ...
126
+
127
+ async def GetSecret(self, request: payload_pb2.GetSecretRequest) -> payload_pb2.GetSecretResponse: ...
128
+
129
+ async def ListSecrets(self, request: payload_pb2.ListSecretsRequest) -> payload_pb2.ListSecretsResponse: ...
130
+
131
+ async def DeleteSecret(self, request: payload_pb2.DeleteSecretRequest) -> payload_pb2.DeleteSecretResponse: ...
@@ -0,0 +1,12 @@
1
+ from flyte.remote._client.auth._channel import create_channel
2
+ from flyte.remote._client.auth._client_config import AuthType, ClientConfig
3
+ from flyte.remote._client.auth.errors import AccessTokenNotFoundError, AuthenticationError, AuthenticationPending
4
+
5
+ __all__ = [
6
+ "AccessTokenNotFoundError",
7
+ "AuthType",
8
+ "AuthenticationError",
9
+ "AuthenticationPending",
10
+ "ClientConfig",
11
+ "create_channel",
12
+ ]
File without changes
@@ -0,0 +1,397 @@
1
+ import asyncio
2
+ import dataclasses
3
+ import ssl
4
+ import typing
5
+ from abc import abstractmethod
6
+ from http import HTTPStatus
7
+
8
+ import httpx
9
+ from grpc.aio import Metadata
10
+
11
+ from flyte.remote._client.auth._client_config import ClientConfig, ClientConfigStore
12
+ from flyte.remote._client.auth._keyring import Credentials, KeyringStore
13
+
14
+
15
+ @dataclasses.dataclass
16
+ class GrpcAuthMetadata:
17
+ creds_id: str
18
+ pairs: Metadata
19
+
20
+
21
+ class Authenticator(object):
22
+ """
23
+ Base authenticator for all authentication flows
24
+ """
25
+
26
+ def __init__(
27
+ self,
28
+ endpoint: str,
29
+ *,
30
+ cfg_store: typing.Optional[ClientConfigStore] = None,
31
+ client_config: typing.Optional[ClientConfig] = None,
32
+ credentials: typing.Optional[Credentials] = None,
33
+ http_session: typing.Optional[httpx.AsyncClient] = None,
34
+ http_proxy_url: typing.Optional[str] = None,
35
+ verify: bool = True,
36
+ ca_cert_path: typing.Optional[str] = None,
37
+ **kwargs,
38
+ ):
39
+ """
40
+ Initialize the base authenticator.
41
+
42
+ :param endpoint: The endpoint URL for authentication
43
+ :param cfg_store: Optional client configuration store for retrieving remote configuration
44
+ :param client_config: Optional client configuration containing authentication settings
45
+ :param credentials: Optional credentials to use for authentication
46
+ :param http_session: Optional HTTP session to use for requests
47
+ :param http_proxy_url: Optional HTTP proxy URL
48
+ :param verify: Whether to verify SSL certificates
49
+ :param ca_cert_path: Optional path to CA certificate file
50
+ :param kwargs: Additional keyword arguments passed to get_async_session, which may include:
51
+ - auth: Authentication implementation to use
52
+ - params: Query parameters to include in request URLs
53
+ - headers: HTTP headers to include in requests
54
+ - cookies: Cookies to include in requests
55
+ - cert: SSL client certificate (path or tuple)
56
+ - http1: Whether to enable HTTP/1.1 support
57
+ - http2: Whether to enable HTTP/2 support
58
+ - proxies: Proxy configuration mapping
59
+ - mounts: Mounted transports for specific URL patterns
60
+ - timeout: Request timeout configuration
61
+ - follow_redirects: Whether to follow redirects
62
+ - limits: Connection pool limits
63
+ - max_redirects: Maximum number of redirects to follow
64
+ - event_hooks: Event hooks for request/response lifecycle
65
+ - base_url: Base URL to join with relative URLs
66
+ - transport: Transport implementation to use
67
+ - app: ASGI application to handle requests
68
+ """
69
+ self._endpoint = endpoint
70
+ self._creds = credentials or KeyringStore.retrieve(endpoint)
71
+ self._http_proxy_url = http_proxy_url
72
+ self._verify = verify
73
+ self._ca_cert_path = ca_cert_path
74
+ self._client_config = client_config
75
+ self._cfg_store = cfg_store
76
+ # Will be populated by _ensure_remote_config
77
+ self._resolved_config: ClientConfig | None = None
78
+ # Lock for coroutine safety
79
+ self._async_lock = asyncio.Lock()
80
+ self._http_session = http_session or get_async_session(**kwargs)
81
+ # Id for tracking credential refresh state
82
+ self._creds_id = self._creds.id if self._creds else None
83
+
84
+ async def _resolve_config(self) -> ClientConfig:
85
+ """
86
+ Resolves and merges client configuration with remote configuration.
87
+
88
+ This method fetches the remote configuration from the cfg_store and merges it with
89
+ the local client_config, prioritizing local settings over remote ones.
90
+
91
+ This method is thread-safe and coroutine-safe, ensuring the remote config is fetched
92
+ only once regardless of concurrent access from multiple threads or coroutines.
93
+
94
+ :return: A merged ClientConfig object containing resolved configuration settings
95
+ """
96
+ # First check without locks for performance
97
+ if self._resolved_config is not None:
98
+ return self._resolved_config
99
+
100
+ if self._cfg_store is None:
101
+ raise ValueError("ClientConfigStore is not set. Cannot resolve configuration.")
102
+
103
+ remote_config = await self._cfg_store.get_client_config()
104
+ self._resolved_config = (
105
+ remote_config.with_override(self._client_config) if self._client_config else remote_config
106
+ )
107
+
108
+ return self._resolved_config
109
+
110
+ def get_credentials(self) -> typing.Optional[Credentials]:
111
+ """
112
+ Get the current credentials.
113
+
114
+ :return: The current credentials or None if not set
115
+ """
116
+ return self._creds
117
+
118
+ def _set_credentials(self, creds: Credentials):
119
+ """
120
+ Set the credentials.
121
+
122
+ :param creds: The credentials to set
123
+ """
124
+ self._creds = creds
125
+
126
+ async def get_grpc_call_auth_metadata(self) -> typing.Optional[GrpcAuthMetadata]:
127
+ """
128
+ Fetch the authentication metadata for gRPC calls.
129
+
130
+ :return: A tuple of (header_key, header_value) or None if no credentials are available
131
+ """
132
+ creds = self.get_credentials()
133
+ if creds:
134
+ cfg = await self._resolve_config()
135
+ return GrpcAuthMetadata(
136
+ creds_id=creds.id,
137
+ pairs=Metadata((cfg.header_key, f"Bearer {creds.access_token}")),
138
+ )
139
+ return None
140
+
141
+ async def refresh_credentials(self, creds_id: str | None = None):
142
+ """
143
+ Refresh the credentials asynchronously with thread and asyncio safety.
144
+
145
+ This method implements a thread-safe and coroutine-safe credential refresh mechanism.
146
+ It uses a timestamp-based approach to prevent redundant credential refreshes when
147
+ multiple threads or coroutines attempt to refresh credentials simultaneously.
148
+
149
+ The caller should capture the current _creds_timestamp before attempting to use credentials.
150
+ If credential usage fails, the caller can pass that timestamp to this method.
151
+ If the timestamp matches the current value, a refresh is needed; otherwise,
152
+ another thread has already refreshed the credentials.
153
+
154
+ :param creds_id: The id of credentials when they were last accessed by the caller.
155
+ If None, force a refresh regardless of id.
156
+ :raises: May raise authentication-related exceptions if the refresh fails
157
+ """
158
+ # If creds_id is None, force refresh
159
+ # If creds_id matches current value, credentials need refresh
160
+ # If creds_id doesn't match, another thread already refreshed credentials
161
+ if creds_id and creds_id != self._creds_id:
162
+ # Credentials have been refreshed by another thread/coroutine since caller read them
163
+ return
164
+
165
+ # Use the async lock to ensure coroutine safety
166
+ async with self._async_lock:
167
+ # Double-check pattern to avoid unnecessary work
168
+ if creds_id and creds_id != self._creds_id:
169
+ # Another thread/coroutine refreshed credentials while we were waiting for the lock
170
+ return
171
+
172
+ # Perform the actual credential refresh
173
+ try:
174
+ self._creds = await self._do_refresh_credentials()
175
+ KeyringStore.store(self._creds)
176
+ except Exception:
177
+ KeyringStore.delete(self._endpoint)
178
+ raise
179
+
180
+ # Update the timestamp to indicate credentials have been refreshed
181
+ self._creds_id = self._creds.id
182
+
183
+ @abstractmethod
184
+ async def _do_refresh_credentials(self) -> Credentials:
185
+ """
186
+ Perform the actual credential refresh operation.
187
+
188
+ This method must be implemented by subclasses to handle the specific authentication flow.
189
+ It should update the internal credentials object (_creds) with a new access token.
190
+
191
+ Implementations typically use the resolved configuration from _resolve_config() to
192
+ determine authentication endpoints, scopes, audience, and other parameters needed for
193
+ the specific authentication flow.
194
+
195
+ :raises: May raise authentication-related exceptions if the refresh fails
196
+ """
197
+ ...
198
+
199
+
200
+ class AsyncAuthenticatedClient(httpx.AsyncClient):
201
+ """
202
+ An httpx.AsyncClient that automatically adds authentication headers to requests.
203
+ This class extends httpx.AsyncClient which is inherently async for network operations.
204
+ """
205
+
206
+ def __init__(self, authenticator: Authenticator, **kwargs):
207
+ """
208
+ Initialize the authenticated client.
209
+
210
+ :param authenticator: The authenticator to use for authentication
211
+ :param kwargs: Additional arguments passed to the httpx.AsyncClient constructor
212
+ """
213
+ super().__init__(**kwargs)
214
+ self.auth_adapter = AsyncAuthenticationHTTPAdapter(authenticator)
215
+ self.authenticator = authenticator
216
+
217
+ async def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
218
+ """
219
+ Sends the request with added authentication headers.
220
+ Must be async because it performs network IO operations and may need to refresh credentials.
221
+ If the response returns a 401 status code, refreshes the credentials and retries the request.
222
+
223
+ :param request: The request object to send.
224
+ :param kwargs: Additional keyword arguments passed to the parent httpx.AsyncClient.send method, which may
225
+ include:
226
+ - auth: Authentication implementation to use for this request
227
+ - follow_redirects: Whether to follow redirects for this request
228
+ - timeout: Request timeout configuration for this request
229
+ :return: The response object.
230
+ """
231
+
232
+ creds_id = await self.auth_adapter.add_auth_header(request)
233
+ response = await super().send(request, **kwargs)
234
+
235
+ if response.status_code == HTTPStatus.UNAUTHORIZED:
236
+ await self.authenticator.refresh_credentials(creds_id=creds_id)
237
+ await self.auth_adapter.add_auth_header(request)
238
+ response = await super().send(request, **kwargs)
239
+
240
+ return response
241
+
242
+
243
+ class AsyncAuthenticationHTTPAdapter:
244
+ """
245
+ A custom async HTTP adapter that adds authentication headers to requests of an httpx.AsyncClient.
246
+ This is the async equivalent of AuthenticationHTTPAdapter for requests.
247
+ """
248
+
249
+ def __init__(self, authenticator: Authenticator):
250
+ """
251
+ Initialize the authentication HTTP adapter.
252
+
253
+ :param authenticator: The authenticator to use for authentication
254
+ """
255
+ self.authenticator = authenticator
256
+
257
+ async def add_auth_header(self, request: httpx.Request) -> typing.Optional[str]:
258
+ """
259
+ Adds authentication headers to the request.
260
+ Must be async because it may call refresh_credentials which performs IO operations.
261
+
262
+ :param request: The request object to add headers to.
263
+ :return: The credentials ID (creds_id) used for tracking credential refresh state
264
+ """
265
+ if self.authenticator.get_credentials() is None:
266
+ await self.authenticator.refresh_credentials()
267
+
268
+ metadata = await self.authenticator.get_grpc_call_auth_metadata()
269
+ if metadata is None:
270
+ return None
271
+ for key, value in metadata.pairs.keys():
272
+ request.headers[key] = value
273
+ return metadata.creds_id
274
+
275
+
276
+ def upgrade_async_session_to_proxy_authenticated(
277
+ http_session: httpx.AsyncClient, proxy_authenticator: typing.Optional[Authenticator] = None, **kwargs
278
+ ) -> httpx.AsyncClient:
279
+ """
280
+ Given an httpx.AsyncClient, it returns a new session that uses AsyncAuthenticationHTTPAdapter
281
+ to perform authentication with a proxy in front of Flyte
282
+
283
+ :param http_session: httpx.AsyncClient Precreated session
284
+ :param proxy_authenticator: Optional authenticator for proxy authentication
285
+ :param kwargs: Additional arguments passed to AsyncAuthenticatedClient, which may include:
286
+ - auth: Authentication implementation to use
287
+ - params: Query parameters to include in request URLs
288
+ - headers: HTTP headers to include in requests
289
+ - cookies: Cookies to include in requests
290
+ - verify: SSL verification mode (True/False/path to certificate)
291
+ - cert: SSL client certificate (path or tuple)
292
+ - http1: Whether to enable HTTP/1.1 support
293
+ - http2: Whether to enable HTTP/2 support
294
+ - proxies: Proxy configuration mapping
295
+ - mounts: Mounted transports for specific URL patterns
296
+ - timeout: Request timeout configuration
297
+ - follow_redirects: Whether to follow redirects
298
+ - limits: Connection pool limits
299
+ - max_redirects: Maximum number of redirects to follow
300
+ - event_hooks: Event hooks for request/response lifecycle
301
+ - base_url: Base URL to join with relative URLs
302
+ - transport: Transport implementation to use
303
+ - app: ASGI application to handle requests
304
+ :return: httpx.AsyncClient with authentication
305
+ """
306
+ if proxy_authenticator:
307
+ return AsyncAuthenticatedClient(proxy_authenticator, **kwargs)
308
+ else:
309
+ return http_session
310
+
311
+
312
+ def get_async_session(
313
+ proxy_authenticator: Authenticator | None = None,
314
+ ca_cert_path: str | None = None,
315
+ verify: bool | None = None,
316
+ **kwargs,
317
+ ) -> httpx.AsyncClient:
318
+ """
319
+ Returns a new httpx.AsyncClient with proxy authentication if proxy_authenticator is provided.
320
+
321
+ This function creates a new httpx.AsyncClient and optionally configures it with proxy authentication
322
+ if a proxy authenticator is provided.
323
+
324
+ :param proxy_authenticator: Optional authenticator for proxy authentication
325
+ :param ca_cert_path: Optional path to CA certificate file for SSL verification
326
+ :param verify: Optional SSL verification mode (True/False/path to certificate)
327
+ :param kwargs: Additional keyword arguments passed to httpx.AsyncClient constructor and AsyncAuthenticatedClient,
328
+ which may include:
329
+ - auth: Authentication implementation to use
330
+ - params: Query parameters to include in request URLs
331
+ - headers: HTTP headers to include in requests
332
+ - cookies: Cookies to include in requests
333
+ - cert: SSL client certificate (path or tuple)
334
+ - http1: Whether to enable HTTP/1.1 support
335
+ - http2: Whether to enable HTTP/2 support
336
+ - proxies: Proxy configuration mapping
337
+ - mounts: Mounted transports for specific URL patterns
338
+ - timeout: Request timeout configuration
339
+ - follow_redirects: Whether to follow redirects
340
+ - limits: Connection pool limits
341
+ - max_redirects: Maximum number of redirects to follow
342
+ - event_hooks: Event hooks for request/response lifecycle
343
+ - base_url: Base URL to join with relative URLs
344
+ - transport: Transport implementation to use
345
+ - app: ASGI application to handle requests
346
+ - proxy_env: Environment variables for proxy command
347
+ - proxy_timeout: Timeout for proxy command execution
348
+ - header_key: Header key to use for authentication
349
+ - endpoint: The endpoint URL for authentication
350
+ - client_id: Client ID for authentication
351
+ - client_secret: Client secret for authentication
352
+ - scopes: List of scopes to request during authentication
353
+ - audience: Audience for the token
354
+ - http_proxy_url: HTTP proxy URL
355
+ :return: An httpx.AsyncClient instance, optionally configured with proxy authentication
356
+ """
357
+
358
+ # Extract known httpx.AsyncClient parameters from kwargs
359
+ client_kwargs = {
360
+ k: v
361
+ for k, v in kwargs.items()
362
+ if k
363
+ in [
364
+ "auth",
365
+ "params",
366
+ "headers",
367
+ "cookies",
368
+ "verify",
369
+ "cert",
370
+ "http1",
371
+ "http2",
372
+ "proxies",
373
+ "mounts",
374
+ "timeout",
375
+ "follow_redirects",
376
+ "limits",
377
+ "max_redirects",
378
+ "event_hooks",
379
+ "base_url",
380
+ "transport",
381
+ "app",
382
+ ]
383
+ }
384
+
385
+ if ca_cert_path:
386
+ context = ssl.create_default_context(capath=ca_cert_path)
387
+ verify = True if context is not None else False
388
+
389
+ if verify is not None:
390
+ client_kwargs["verify"] = verify
391
+
392
+ http_session = httpx.AsyncClient(**client_kwargs)
393
+ if proxy_authenticator:
394
+ http_session = upgrade_async_session_to_proxy_authenticated(
395
+ http_session, proxy_authenticator=proxy_authenticator, **kwargs
396
+ )
397
+ return http_session
@@ -0,0 +1,73 @@
1
+ from flyte._logging import logger
2
+ from flyte.remote._client.auth import _token_client as token_client
3
+ from flyte.remote._client.auth._authenticators.base import Authenticator
4
+ from flyte.remote._client.auth._keyring import Credentials
5
+
6
+
7
+ class ClientCredentialsAuthenticator(Authenticator):
8
+ """
9
+ This Authenticator uses ClientId and ClientSecret to authenticate
10
+ """
11
+
12
+ def __init__(
13
+ self,
14
+ client_id: str,
15
+ client_credentials_secret: str,
16
+ **kwargs,
17
+ ):
18
+ """
19
+ Initialize the client credentials authenticator.
20
+
21
+ :param client_id: The client ID for authentication
22
+ :param client_credentials_secret: The client secret for authentication
23
+ :param kwargs: Additional keyword arguments passed to the base Authenticator
24
+
25
+ **Keyword Arguments passed to base Authenticator**:
26
+ :param endpoint: The endpoint URL for authentication (required)
27
+ :param cfg_store: Optional client configuration store for retrieving remote configuration
28
+ :param client_config: Optional client configuration containing authentication settings
29
+ :param credentials: Optional credentials to use for authentication
30
+ :param http_session: Optional HTTP session to use for requests
31
+ :param http_proxy_url: Optional HTTP proxy URL
32
+ :param verify: Whether to verify SSL certificates (default: True)
33
+ :param ca_cert_path: Optional path to CA certificate file
34
+ :param scopes: List of scopes to request during authentication
35
+ :param audience: Audience for the token
36
+ """
37
+ if not client_id or not client_credentials_secret:
38
+ raise ValueError("both client_id and client_credentials_secret are required.")
39
+ self._client_id = client_id
40
+ self._client_credentials_secret = client_credentials_secret
41
+ super().__init__(**kwargs)
42
+
43
+ async def _do_refresh_credentials(self) -> Credentials:
44
+ """
45
+ Refreshes the authentication credentials using client credentials flow.
46
+
47
+ This function is used by the _handle_rpc_error() decorator, depending on the AUTH_MODE config object.
48
+ This handler is meant for SDK use-cases of auth (like pyflyte, or when users call SDK functions that require
49
+ access to Admin, like when waiting for another workflow to complete from within a task). This function uses
50
+ basic auth, which means the credentials for basic auth must be present from wherever this code is running.
51
+ """
52
+ cfg = await self._resolve_config()
53
+
54
+ # Note that unlike the Pkce flow, the client ID does not come from Admin.
55
+ logger.debug(f"Basic authorization flow with client id {self._client_id} scope {cfg.scopes}")
56
+ authorization_header = token_client.get_basic_authorization_header(
57
+ self._client_id, self._client_credentials_secret
58
+ )
59
+
60
+ token, refresh_token, expires_in = await token_client.get_token(
61
+ token_endpoint=cfg.token_endpoint,
62
+ authorization_header=authorization_header,
63
+ http_proxy_url=self._http_proxy_url,
64
+ verify=self._verify,
65
+ scopes=cfg.scopes,
66
+ audience=cfg.audience,
67
+ http_session=self._http_session,
68
+ )
69
+
70
+ logger.info("Retrieved new token, expires in {}".format(expires_in))
71
+ return Credentials(
72
+ for_endpoint=self._endpoint, access_token=token, refresh_token=refresh_token, expires_in=expires_in
73
+ )