flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__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.
Files changed (455) hide show
  1. flyte/__init__.py +83 -30
  2. flyte/_bin/connect.py +61 -0
  3. flyte/_bin/debug.py +38 -0
  4. flyte/_bin/runtime.py +87 -19
  5. flyte/_bin/serve.py +351 -0
  6. flyte/_build.py +3 -2
  7. flyte/_cache/cache.py +6 -5
  8. flyte/_cache/local_cache.py +216 -0
  9. flyte/_code_bundle/_ignore.py +31 -5
  10. flyte/_code_bundle/_packaging.py +42 -11
  11. flyte/_code_bundle/_utils.py +57 -34
  12. flyte/_code_bundle/bundle.py +130 -27
  13. flyte/_constants.py +1 -0
  14. flyte/_context.py +21 -5
  15. flyte/_custom_context.py +73 -0
  16. flyte/_debug/constants.py +37 -0
  17. flyte/_debug/utils.py +17 -0
  18. flyte/_debug/vscode.py +315 -0
  19. flyte/_deploy.py +396 -75
  20. flyte/_deployer.py +109 -0
  21. flyte/_environment.py +94 -11
  22. flyte/_excepthook.py +37 -0
  23. flyte/_group.py +2 -1
  24. flyte/_hash.py +1 -16
  25. flyte/_image.py +544 -234
  26. flyte/_initialize.py +443 -294
  27. flyte/_interface.py +40 -5
  28. flyte/_internal/controllers/__init__.py +22 -8
  29. flyte/_internal/controllers/_local_controller.py +159 -35
  30. flyte/_internal/controllers/_trace.py +18 -10
  31. flyte/_internal/controllers/remote/__init__.py +38 -9
  32. flyte/_internal/controllers/remote/_action.py +82 -12
  33. flyte/_internal/controllers/remote/_client.py +6 -2
  34. flyte/_internal/controllers/remote/_controller.py +290 -64
  35. flyte/_internal/controllers/remote/_core.py +155 -95
  36. flyte/_internal/controllers/remote/_informer.py +40 -20
  37. flyte/_internal/controllers/remote/_service_protocol.py +2 -2
  38. flyte/_internal/imagebuild/__init__.py +2 -10
  39. flyte/_internal/imagebuild/docker_builder.py +391 -84
  40. flyte/_internal/imagebuild/image_builder.py +111 -55
  41. flyte/_internal/imagebuild/remote_builder.py +409 -0
  42. flyte/_internal/imagebuild/utils.py +79 -0
  43. flyte/_internal/resolvers/_app_env_module.py +92 -0
  44. flyte/_internal/resolvers/_task_module.py +5 -38
  45. flyte/_internal/resolvers/app_env.py +26 -0
  46. flyte/_internal/resolvers/common.py +8 -1
  47. flyte/_internal/resolvers/default.py +2 -2
  48. flyte/_internal/runtime/convert.py +322 -33
  49. flyte/_internal/runtime/entrypoints.py +106 -18
  50. flyte/_internal/runtime/io.py +71 -23
  51. flyte/_internal/runtime/resources_serde.py +21 -7
  52. flyte/_internal/runtime/reuse.py +125 -0
  53. flyte/_internal/runtime/rusty.py +196 -0
  54. flyte/_internal/runtime/task_serde.py +239 -66
  55. flyte/_internal/runtime/taskrunner.py +48 -8
  56. flyte/_internal/runtime/trigger_serde.py +162 -0
  57. flyte/_internal/runtime/types_serde.py +7 -16
  58. flyte/_keyring/file.py +115 -0
  59. flyte/_link.py +30 -0
  60. flyte/_logging.py +241 -42
  61. flyte/_map.py +312 -0
  62. flyte/_metrics.py +59 -0
  63. flyte/_module.py +74 -0
  64. flyte/_pod.py +30 -0
  65. flyte/_resources.py +296 -33
  66. flyte/_retry.py +1 -7
  67. flyte/_reusable_environment.py +72 -7
  68. flyte/_run.py +461 -132
  69. flyte/_secret.py +47 -11
  70. flyte/_serve.py +333 -0
  71. flyte/_task.py +245 -56
  72. flyte/_task_environment.py +219 -97
  73. flyte/_task_plugins.py +47 -0
  74. flyte/_tools.py +8 -8
  75. flyte/_trace.py +15 -24
  76. flyte/_trigger.py +1027 -0
  77. flyte/_utils/__init__.py +12 -1
  78. flyte/_utils/asyn.py +3 -1
  79. flyte/_utils/async_cache.py +139 -0
  80. flyte/_utils/coro_management.py +5 -4
  81. flyte/_utils/description_parser.py +19 -0
  82. flyte/_utils/docker_credentials.py +173 -0
  83. flyte/_utils/helpers.py +45 -19
  84. flyte/_utils/module_loader.py +123 -0
  85. flyte/_utils/org_discovery.py +57 -0
  86. flyte/_utils/uv_script_parser.py +8 -1
  87. flyte/_version.py +16 -3
  88. flyte/app/__init__.py +27 -0
  89. flyte/app/_app_environment.py +362 -0
  90. flyte/app/_connector_environment.py +40 -0
  91. flyte/app/_deploy.py +130 -0
  92. flyte/app/_parameter.py +343 -0
  93. flyte/app/_runtime/__init__.py +3 -0
  94. flyte/app/_runtime/app_serde.py +383 -0
  95. flyte/app/_types.py +113 -0
  96. flyte/app/extras/__init__.py +9 -0
  97. flyte/app/extras/_auth_middleware.py +217 -0
  98. flyte/app/extras/_fastapi.py +93 -0
  99. flyte/app/extras/_model_loader/__init__.py +3 -0
  100. flyte/app/extras/_model_loader/config.py +7 -0
  101. flyte/app/extras/_model_loader/loader.py +288 -0
  102. flyte/cli/__init__.py +12 -0
  103. flyte/cli/_abort.py +28 -0
  104. flyte/cli/_build.py +114 -0
  105. flyte/cli/_common.py +493 -0
  106. flyte/cli/_create.py +371 -0
  107. flyte/cli/_delete.py +45 -0
  108. flyte/cli/_deploy.py +401 -0
  109. flyte/cli/_gen.py +316 -0
  110. flyte/cli/_get.py +446 -0
  111. flyte/cli/_option.py +33 -0
  112. {union/_cli → flyte/cli}/_params.py +152 -153
  113. flyte/cli/_plugins.py +209 -0
  114. flyte/cli/_prefetch.py +292 -0
  115. flyte/cli/_run.py +690 -0
  116. flyte/cli/_serve.py +338 -0
  117. flyte/cli/_update.py +86 -0
  118. flyte/cli/_user.py +20 -0
  119. flyte/cli/main.py +246 -0
  120. flyte/config/__init__.py +3 -0
  121. flyte/config/_config.py +248 -0
  122. flyte/config/_internal.py +73 -0
  123. flyte/config/_reader.py +225 -0
  124. flyte/connectors/__init__.py +11 -0
  125. flyte/connectors/_connector.py +330 -0
  126. flyte/connectors/_server.py +194 -0
  127. flyte/connectors/utils.py +159 -0
  128. flyte/errors.py +134 -2
  129. flyte/extend.py +24 -0
  130. flyte/extras/_container.py +69 -56
  131. flyte/git/__init__.py +3 -0
  132. flyte/git/_config.py +279 -0
  133. flyte/io/__init__.py +8 -1
  134. flyte/io/{structured_dataset → _dataframe}/__init__.py +32 -30
  135. flyte/io/{structured_dataset → _dataframe}/basic_dfs.py +75 -68
  136. flyte/io/{structured_dataset/structured_dataset.py → _dataframe/dataframe.py} +207 -242
  137. flyte/io/_dir.py +575 -113
  138. flyte/io/_file.py +587 -141
  139. flyte/io/_hashing_io.py +342 -0
  140. flyte/io/extend.py +7 -0
  141. flyte/models.py +635 -0
  142. flyte/prefetch/__init__.py +22 -0
  143. flyte/prefetch/_hf_model.py +563 -0
  144. flyte/remote/__init__.py +14 -3
  145. flyte/remote/_action.py +879 -0
  146. flyte/remote/_app.py +346 -0
  147. flyte/remote/_auth_metadata.py +42 -0
  148. flyte/remote/_client/_protocols.py +62 -4
  149. flyte/remote/_client/auth/_auth_utils.py +19 -0
  150. flyte/remote/_client/auth/_authenticators/base.py +8 -2
  151. flyte/remote/_client/auth/_authenticators/device_code.py +4 -5
  152. flyte/remote/_client/auth/_authenticators/factory.py +4 -0
  153. flyte/remote/_client/auth/_authenticators/passthrough.py +79 -0
  154. flyte/remote/_client/auth/_authenticators/pkce.py +17 -18
  155. flyte/remote/_client/auth/_channel.py +47 -18
  156. flyte/remote/_client/auth/_client_config.py +5 -3
  157. flyte/remote/_client/auth/_keyring.py +15 -2
  158. flyte/remote/_client/auth/_token_client.py +3 -3
  159. flyte/remote/_client/controlplane.py +206 -18
  160. flyte/remote/_common.py +66 -0
  161. flyte/remote/_data.py +107 -22
  162. flyte/remote/_logs.py +116 -33
  163. flyte/remote/_project.py +21 -19
  164. flyte/remote/_run.py +164 -631
  165. flyte/remote/_secret.py +72 -29
  166. flyte/remote/_task.py +387 -46
  167. flyte/remote/_trigger.py +368 -0
  168. flyte/remote/_user.py +43 -0
  169. flyte/report/_report.py +10 -6
  170. flyte/storage/__init__.py +13 -1
  171. flyte/storage/_config.py +237 -0
  172. flyte/storage/_parallel_reader.py +289 -0
  173. flyte/storage/_storage.py +268 -59
  174. flyte/syncify/__init__.py +56 -0
  175. flyte/syncify/_api.py +414 -0
  176. flyte/types/__init__.py +39 -0
  177. flyte/types/_interface.py +22 -7
  178. flyte/{io/pickle/transformer.py → types/_pickle.py} +37 -9
  179. flyte/types/_string_literals.py +8 -9
  180. flyte/types/_type_engine.py +230 -129
  181. flyte/types/_utils.py +1 -1
  182. flyte-2.0.0b46.data/scripts/debug.py +38 -0
  183. flyte-2.0.0b46.data/scripts/runtime.py +194 -0
  184. flyte-2.0.0b46.dist-info/METADATA +352 -0
  185. flyte-2.0.0b46.dist-info/RECORD +221 -0
  186. flyte-2.0.0b46.dist-info/entry_points.txt +8 -0
  187. flyte-2.0.0b46.dist-info/licenses/LICENSE +201 -0
  188. flyte/_api_commons.py +0 -3
  189. flyte/_cli/_common.py +0 -287
  190. flyte/_cli/_create.py +0 -42
  191. flyte/_cli/_delete.py +0 -23
  192. flyte/_cli/_deploy.py +0 -140
  193. flyte/_cli/_get.py +0 -235
  194. flyte/_cli/_run.py +0 -152
  195. flyte/_cli/main.py +0 -72
  196. flyte/_datastructures.py +0 -342
  197. flyte/_internal/controllers/pbhash.py +0 -39
  198. flyte/_protos/common/authorization_pb2.py +0 -66
  199. flyte/_protos/common/authorization_pb2.pyi +0 -108
  200. flyte/_protos/common/authorization_pb2_grpc.py +0 -4
  201. flyte/_protos/common/identifier_pb2.py +0 -71
  202. flyte/_protos/common/identifier_pb2.pyi +0 -82
  203. flyte/_protos/common/identifier_pb2_grpc.py +0 -4
  204. flyte/_protos/common/identity_pb2.py +0 -48
  205. flyte/_protos/common/identity_pb2.pyi +0 -72
  206. flyte/_protos/common/identity_pb2_grpc.py +0 -4
  207. flyte/_protos/common/list_pb2.py +0 -36
  208. flyte/_protos/common/list_pb2.pyi +0 -69
  209. flyte/_protos/common/list_pb2_grpc.py +0 -4
  210. flyte/_protos/common/policy_pb2.py +0 -37
  211. flyte/_protos/common/policy_pb2.pyi +0 -27
  212. flyte/_protos/common/policy_pb2_grpc.py +0 -4
  213. flyte/_protos/common/role_pb2.py +0 -37
  214. flyte/_protos/common/role_pb2.pyi +0 -53
  215. flyte/_protos/common/role_pb2_grpc.py +0 -4
  216. flyte/_protos/common/runtime_version_pb2.py +0 -28
  217. flyte/_protos/common/runtime_version_pb2.pyi +0 -24
  218. flyte/_protos/common/runtime_version_pb2_grpc.py +0 -4
  219. flyte/_protos/logs/dataplane/payload_pb2.py +0 -96
  220. flyte/_protos/logs/dataplane/payload_pb2.pyi +0 -168
  221. flyte/_protos/logs/dataplane/payload_pb2_grpc.py +0 -4
  222. flyte/_protos/secret/definition_pb2.py +0 -49
  223. flyte/_protos/secret/definition_pb2.pyi +0 -93
  224. flyte/_protos/secret/definition_pb2_grpc.py +0 -4
  225. flyte/_protos/secret/payload_pb2.py +0 -62
  226. flyte/_protos/secret/payload_pb2.pyi +0 -94
  227. flyte/_protos/secret/payload_pb2_grpc.py +0 -4
  228. flyte/_protos/secret/secret_pb2.py +0 -38
  229. flyte/_protos/secret/secret_pb2.pyi +0 -6
  230. flyte/_protos/secret/secret_pb2_grpc.py +0 -198
  231. flyte/_protos/secret/secret_pb2_grpc_grpc.py +0 -198
  232. flyte/_protos/validate/validate/validate_pb2.py +0 -76
  233. flyte/_protos/workflow/node_execution_service_pb2.py +0 -26
  234. flyte/_protos/workflow/node_execution_service_pb2.pyi +0 -4
  235. flyte/_protos/workflow/node_execution_service_pb2_grpc.py +0 -32
  236. flyte/_protos/workflow/queue_service_pb2.py +0 -106
  237. flyte/_protos/workflow/queue_service_pb2.pyi +0 -141
  238. flyte/_protos/workflow/queue_service_pb2_grpc.py +0 -172
  239. flyte/_protos/workflow/run_definition_pb2.py +0 -128
  240. flyte/_protos/workflow/run_definition_pb2.pyi +0 -310
  241. flyte/_protos/workflow/run_definition_pb2_grpc.py +0 -4
  242. flyte/_protos/workflow/run_logs_service_pb2.py +0 -41
  243. flyte/_protos/workflow/run_logs_service_pb2.pyi +0 -28
  244. flyte/_protos/workflow/run_logs_service_pb2_grpc.py +0 -69
  245. flyte/_protos/workflow/run_service_pb2.py +0 -133
  246. flyte/_protos/workflow/run_service_pb2.pyi +0 -175
  247. flyte/_protos/workflow/run_service_pb2_grpc.py +0 -412
  248. flyte/_protos/workflow/state_service_pb2.py +0 -58
  249. flyte/_protos/workflow/state_service_pb2.pyi +0 -71
  250. flyte/_protos/workflow/state_service_pb2_grpc.py +0 -138
  251. flyte/_protos/workflow/task_definition_pb2.py +0 -72
  252. flyte/_protos/workflow/task_definition_pb2.pyi +0 -65
  253. flyte/_protos/workflow/task_definition_pb2_grpc.py +0 -4
  254. flyte/_protos/workflow/task_service_pb2.py +0 -44
  255. flyte/_protos/workflow/task_service_pb2.pyi +0 -31
  256. flyte/_protos/workflow/task_service_pb2_grpc.py +0 -104
  257. flyte/io/_dataframe.py +0 -0
  258. flyte/io/pickle/__init__.py +0 -0
  259. flyte/remote/_console.py +0 -18
  260. flyte-0.0.1b0.dist-info/METADATA +0 -179
  261. flyte-0.0.1b0.dist-info/RECORD +0 -390
  262. flyte-0.0.1b0.dist-info/entry_points.txt +0 -3
  263. union/__init__.py +0 -54
  264. union/_api_commons.py +0 -3
  265. union/_bin/__init__.py +0 -0
  266. union/_bin/runtime.py +0 -113
  267. union/_build.py +0 -25
  268. union/_cache/__init__.py +0 -12
  269. union/_cache/cache.py +0 -141
  270. union/_cache/defaults.py +0 -9
  271. union/_cache/policy_function_body.py +0 -42
  272. union/_cli/__init__.py +0 -0
  273. union/_cli/_common.py +0 -263
  274. union/_cli/_create.py +0 -40
  275. union/_cli/_delete.py +0 -23
  276. union/_cli/_deploy.py +0 -120
  277. union/_cli/_get.py +0 -162
  278. union/_cli/_run.py +0 -150
  279. union/_cli/main.py +0 -72
  280. union/_code_bundle/__init__.py +0 -8
  281. union/_code_bundle/_ignore.py +0 -113
  282. union/_code_bundle/_packaging.py +0 -187
  283. union/_code_bundle/_utils.py +0 -342
  284. union/_code_bundle/bundle.py +0 -176
  285. union/_context.py +0 -146
  286. union/_datastructures.py +0 -295
  287. union/_deploy.py +0 -185
  288. union/_doc.py +0 -29
  289. union/_docstring.py +0 -26
  290. union/_environment.py +0 -43
  291. union/_group.py +0 -31
  292. union/_hash.py +0 -23
  293. union/_image.py +0 -760
  294. union/_initialize.py +0 -585
  295. union/_interface.py +0 -84
  296. union/_internal/__init__.py +0 -3
  297. union/_internal/controllers/__init__.py +0 -77
  298. union/_internal/controllers/_local_controller.py +0 -77
  299. union/_internal/controllers/pbhash.py +0 -39
  300. union/_internal/controllers/remote/__init__.py +0 -40
  301. union/_internal/controllers/remote/_action.py +0 -131
  302. union/_internal/controllers/remote/_client.py +0 -43
  303. union/_internal/controllers/remote/_controller.py +0 -169
  304. union/_internal/controllers/remote/_core.py +0 -341
  305. union/_internal/controllers/remote/_informer.py +0 -260
  306. union/_internal/controllers/remote/_service_protocol.py +0 -44
  307. union/_internal/imagebuild/__init__.py +0 -11
  308. union/_internal/imagebuild/docker_builder.py +0 -416
  309. union/_internal/imagebuild/image_builder.py +0 -243
  310. union/_internal/imagebuild/remote_builder.py +0 -0
  311. union/_internal/resolvers/__init__.py +0 -0
  312. union/_internal/resolvers/_task_module.py +0 -31
  313. union/_internal/resolvers/common.py +0 -24
  314. union/_internal/resolvers/default.py +0 -27
  315. union/_internal/runtime/__init__.py +0 -0
  316. union/_internal/runtime/convert.py +0 -163
  317. union/_internal/runtime/entrypoints.py +0 -121
  318. union/_internal/runtime/io.py +0 -136
  319. union/_internal/runtime/resources_serde.py +0 -134
  320. union/_internal/runtime/task_serde.py +0 -202
  321. union/_internal/runtime/taskrunner.py +0 -179
  322. union/_internal/runtime/types_serde.py +0 -53
  323. union/_logging.py +0 -124
  324. union/_protos/__init__.py +0 -0
  325. union/_protos/common/authorization_pb2.py +0 -66
  326. union/_protos/common/authorization_pb2.pyi +0 -106
  327. union/_protos/common/authorization_pb2_grpc.py +0 -4
  328. union/_protos/common/identifier_pb2.py +0 -71
  329. union/_protos/common/identifier_pb2.pyi +0 -82
  330. union/_protos/common/identifier_pb2_grpc.py +0 -4
  331. union/_protos/common/identity_pb2.py +0 -48
  332. union/_protos/common/identity_pb2.pyi +0 -72
  333. union/_protos/common/identity_pb2_grpc.py +0 -4
  334. union/_protos/common/list_pb2.py +0 -36
  335. union/_protos/common/list_pb2.pyi +0 -69
  336. union/_protos/common/list_pb2_grpc.py +0 -4
  337. union/_protos/common/policy_pb2.py +0 -37
  338. union/_protos/common/policy_pb2.pyi +0 -27
  339. union/_protos/common/policy_pb2_grpc.py +0 -4
  340. union/_protos/common/role_pb2.py +0 -37
  341. union/_protos/common/role_pb2.pyi +0 -51
  342. union/_protos/common/role_pb2_grpc.py +0 -4
  343. union/_protos/common/runtime_version_pb2.py +0 -28
  344. union/_protos/common/runtime_version_pb2.pyi +0 -24
  345. union/_protos/common/runtime_version_pb2_grpc.py +0 -4
  346. union/_protos/logs/dataplane/payload_pb2.py +0 -96
  347. union/_protos/logs/dataplane/payload_pb2.pyi +0 -168
  348. union/_protos/logs/dataplane/payload_pb2_grpc.py +0 -4
  349. union/_protos/secret/definition_pb2.py +0 -49
  350. union/_protos/secret/definition_pb2.pyi +0 -93
  351. union/_protos/secret/definition_pb2_grpc.py +0 -4
  352. union/_protos/secret/payload_pb2.py +0 -62
  353. union/_protos/secret/payload_pb2.pyi +0 -94
  354. union/_protos/secret/payload_pb2_grpc.py +0 -4
  355. union/_protos/secret/secret_pb2.py +0 -38
  356. union/_protos/secret/secret_pb2.pyi +0 -6
  357. union/_protos/secret/secret_pb2_grpc.py +0 -198
  358. union/_protos/validate/validate/validate_pb2.py +0 -76
  359. union/_protos/workflow/node_execution_service_pb2.py +0 -26
  360. union/_protos/workflow/node_execution_service_pb2.pyi +0 -4
  361. union/_protos/workflow/node_execution_service_pb2_grpc.py +0 -32
  362. union/_protos/workflow/queue_service_pb2.py +0 -75
  363. union/_protos/workflow/queue_service_pb2.pyi +0 -103
  364. union/_protos/workflow/queue_service_pb2_grpc.py +0 -172
  365. union/_protos/workflow/run_definition_pb2.py +0 -100
  366. union/_protos/workflow/run_definition_pb2.pyi +0 -256
  367. union/_protos/workflow/run_definition_pb2_grpc.py +0 -4
  368. union/_protos/workflow/run_logs_service_pb2.py +0 -41
  369. union/_protos/workflow/run_logs_service_pb2.pyi +0 -28
  370. union/_protos/workflow/run_logs_service_pb2_grpc.py +0 -69
  371. union/_protos/workflow/run_service_pb2.py +0 -133
  372. union/_protos/workflow/run_service_pb2.pyi +0 -173
  373. union/_protos/workflow/run_service_pb2_grpc.py +0 -412
  374. union/_protos/workflow/state_service_pb2.py +0 -58
  375. union/_protos/workflow/state_service_pb2.pyi +0 -69
  376. union/_protos/workflow/state_service_pb2_grpc.py +0 -138
  377. union/_protos/workflow/task_definition_pb2.py +0 -72
  378. union/_protos/workflow/task_definition_pb2.pyi +0 -65
  379. union/_protos/workflow/task_definition_pb2_grpc.py +0 -4
  380. union/_protos/workflow/task_service_pb2.py +0 -44
  381. union/_protos/workflow/task_service_pb2.pyi +0 -31
  382. union/_protos/workflow/task_service_pb2_grpc.py +0 -104
  383. union/_resources.py +0 -226
  384. union/_retry.py +0 -32
  385. union/_reusable_environment.py +0 -25
  386. union/_run.py +0 -374
  387. union/_secret.py +0 -61
  388. union/_task.py +0 -354
  389. union/_task_environment.py +0 -186
  390. union/_timeout.py +0 -47
  391. union/_tools.py +0 -27
  392. union/_utils/__init__.py +0 -11
  393. union/_utils/asyn.py +0 -119
  394. union/_utils/file_handling.py +0 -71
  395. union/_utils/helpers.py +0 -46
  396. union/_utils/lazy_module.py +0 -54
  397. union/_utils/uv_script_parser.py +0 -49
  398. union/_version.py +0 -21
  399. union/connectors/__init__.py +0 -0
  400. union/errors.py +0 -128
  401. union/extras/__init__.py +0 -5
  402. union/extras/_container.py +0 -263
  403. union/io/__init__.py +0 -11
  404. union/io/_dataframe.py +0 -0
  405. union/io/_dir.py +0 -425
  406. union/io/_file.py +0 -418
  407. union/io/pickle/__init__.py +0 -0
  408. union/io/pickle/transformer.py +0 -117
  409. union/io/structured_dataset/__init__.py +0 -122
  410. union/io/structured_dataset/basic_dfs.py +0 -219
  411. union/io/structured_dataset/structured_dataset.py +0 -1057
  412. union/py.typed +0 -0
  413. union/remote/__init__.py +0 -23
  414. union/remote/_client/__init__.py +0 -0
  415. union/remote/_client/_protocols.py +0 -129
  416. union/remote/_client/auth/__init__.py +0 -12
  417. union/remote/_client/auth/_authenticators/__init__.py +0 -0
  418. union/remote/_client/auth/_authenticators/base.py +0 -391
  419. union/remote/_client/auth/_authenticators/client_credentials.py +0 -73
  420. union/remote/_client/auth/_authenticators/device_code.py +0 -120
  421. union/remote/_client/auth/_authenticators/external_command.py +0 -77
  422. union/remote/_client/auth/_authenticators/factory.py +0 -200
  423. union/remote/_client/auth/_authenticators/pkce.py +0 -515
  424. union/remote/_client/auth/_channel.py +0 -184
  425. union/remote/_client/auth/_client_config.py +0 -83
  426. union/remote/_client/auth/_default_html.py +0 -32
  427. union/remote/_client/auth/_grpc_utils/__init__.py +0 -0
  428. union/remote/_client/auth/_grpc_utils/auth_interceptor.py +0 -204
  429. union/remote/_client/auth/_grpc_utils/default_metadata_interceptor.py +0 -144
  430. union/remote/_client/auth/_keyring.py +0 -154
  431. union/remote/_client/auth/_token_client.py +0 -258
  432. union/remote/_client/auth/errors.py +0 -16
  433. union/remote/_client/controlplane.py +0 -86
  434. union/remote/_data.py +0 -149
  435. union/remote/_logs.py +0 -74
  436. union/remote/_project.py +0 -86
  437. union/remote/_run.py +0 -820
  438. union/remote/_secret.py +0 -132
  439. union/remote/_task.py +0 -193
  440. union/report/__init__.py +0 -3
  441. union/report/_report.py +0 -178
  442. union/report/_template.html +0 -124
  443. union/storage/__init__.py +0 -24
  444. union/storage/_remote_fs.py +0 -34
  445. union/storage/_storage.py +0 -247
  446. union/storage/_utils.py +0 -5
  447. union/types/__init__.py +0 -11
  448. union/types/_renderer.py +0 -162
  449. union/types/_string_literals.py +0 -120
  450. union/types/_type_engine.py +0 -2131
  451. union/types/_utils.py +0 -80
  452. /flyte/{_cli → _debug}/__init__.py +0 -0
  453. /flyte/{_protos → _keyring}/__init__.py +0 -0
  454. {flyte-0.0.1b0.dist-info → flyte-2.0.0b46.dist-info}/WHEEL +0 -0
  455. {flyte-0.0.1b0.dist-info → flyte-2.0.0b46.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,879 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ from collections import UserDict
5
+ from dataclasses import dataclass
6
+ from datetime import datetime, timedelta, timezone
7
+ from functools import cached_property
8
+ from typing import (
9
+ Any,
10
+ AsyncGenerator,
11
+ AsyncIterator,
12
+ Dict,
13
+ Iterator,
14
+ List,
15
+ Literal,
16
+ Tuple,
17
+ Union,
18
+ cast,
19
+ )
20
+
21
+ import grpc
22
+ import rich.pretty
23
+ import rich.repr
24
+ from flyteidl2.common import identifier_pb2, list_pb2, phase_pb2
25
+ from flyteidl2.task import common_pb2
26
+ from flyteidl2.workflow import run_definition_pb2, run_service_pb2
27
+ from flyteidl2.workflow.run_service_pb2 import WatchActionDetailsResponse
28
+ from rich.console import Console
29
+ from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
30
+
31
+ from flyte import types
32
+ from flyte._initialize import ensure_client, get_client, get_init_config
33
+ from flyte._interface import default_output_name
34
+ from flyte.models import ActionPhase
35
+ from flyte.remote._common import ToJSONMixin
36
+ from flyte.remote._logs import Logs
37
+ from flyte.syncify import syncify
38
+
39
+ WaitFor = Literal["terminal", "running", "logs-ready"]
40
+
41
+
42
+ def _action_time_phase(
43
+ action: run_definition_pb2.Action | run_definition_pb2.ActionDetails,
44
+ ) -> rich.repr.Result:
45
+ """
46
+ Rich representation of the action time and phase.
47
+ """
48
+ start_time = action.status.start_time.ToDatetime().replace(tzinfo=timezone.utc)
49
+ yield "start_time", start_time.isoformat()
50
+ if action.status.phase in [
51
+ phase_pb2.ACTION_PHASE_FAILED,
52
+ phase_pb2.ACTION_PHASE_SUCCEEDED,
53
+ phase_pb2.ACTION_PHASE_ABORTED,
54
+ phase_pb2.ACTION_PHASE_TIMED_OUT,
55
+ ]:
56
+ end_time = action.status.end_time.ToDatetime().replace(tzinfo=timezone.utc)
57
+ yield "end_time", end_time.isoformat()
58
+ yield "run_time", f"{(end_time - start_time).seconds} secs"
59
+ else:
60
+ yield "end_time", None
61
+ yield "run_time", f"{(datetime.now(timezone.utc) - start_time).seconds} secs"
62
+ yield "phase", phase_pb2.ActionPhase.Name(action.status.phase)
63
+ if isinstance(action, run_definition_pb2.ActionDetails):
64
+ yield (
65
+ "error",
66
+ (f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA"),
67
+ )
68
+
69
+
70
+ def _action_rich_repr(action: run_definition_pb2.Action) -> rich.repr.Result:
71
+ """
72
+ Rich representation of the action.
73
+ """
74
+ yield "name", action.id.run.name
75
+ if action.metadata.HasField("task"):
76
+ yield "task name", action.metadata.task.id.name
77
+ yield "type", action.metadata.task.task_type
78
+ elif action.metadata.HasField("trace"):
79
+ yield "trace", action.metadata.trace.name
80
+ yield "type", "trace"
81
+ yield "action name", action.id.name
82
+ yield from _action_time_phase(action)
83
+ yield "group", action.metadata.group
84
+ yield "parent", action.metadata.parent
85
+ yield "attempts", action.status.attempts
86
+
87
+
88
+ def _attempt_rich_repr(
89
+ action: List[run_definition_pb2.ActionAttempt],
90
+ ) -> rich.repr.Result:
91
+ for attempt in action:
92
+ yield "attempt", attempt.attempt
93
+ yield "phase", phase_pb2.ActionPhase.Name(attempt.phase)
94
+ yield "logs_available", attempt.logs_available
95
+
96
+
97
+ def _action_details_rich_repr(
98
+ action: run_definition_pb2.ActionDetails,
99
+ ) -> rich.repr.Result:
100
+ """
101
+ Rich representation of the action details.
102
+ """
103
+ yield "name", action.id.run.name
104
+ yield from _action_time_phase(action)
105
+ if action.HasField("task"):
106
+ yield "task", action.task.task_template.id.name
107
+ yield "task_type", action.task.task_template.type
108
+ yield "task_version", action.task.task_template.id.version
109
+ yield "attempts", action.attempts
110
+ yield "error", (f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA")
111
+ yield "phase", phase_pb2.ActionPhase.Name(action.status.phase)
112
+ yield "group", action.metadata.group
113
+ yield "parent", action.metadata.parent
114
+
115
+
116
+ def _action_done_check(phase: phase_pb2.ActionPhase) -> bool:
117
+ """
118
+ Check if the action is done.
119
+ """
120
+ return phase in [
121
+ phase_pb2.ACTION_PHASE_FAILED,
122
+ phase_pb2.ACTION_PHASE_SUCCEEDED,
123
+ phase_pb2.ACTION_PHASE_ABORTED,
124
+ phase_pb2.ACTION_PHASE_TIMED_OUT,
125
+ ]
126
+
127
+
128
+ @dataclass
129
+ class Action(ToJSONMixin):
130
+ """
131
+ A class representing an action. It is used to manage the "execution" of a task and its state on the remote API.
132
+
133
+ From a datamodel perspective, a Run consists of actions. All actions are linearly nested under a parent action.
134
+ Actions have unique auto-generated identifiers, that are unique within a parent action.
135
+
136
+ <pre>
137
+ run
138
+ - a0
139
+ - action1 under a0
140
+ - action2 under a0
141
+ - action1 under action2 under a0
142
+ - action2 under action1 under action2 under a0
143
+ - ...
144
+ - ...
145
+ </pre>
146
+ """
147
+
148
+ pb2: run_definition_pb2.Action
149
+ _details: ActionDetails | None = None
150
+
151
+ @syncify
152
+ @classmethod
153
+ async def listall(
154
+ cls,
155
+ for_run_name: str,
156
+ in_phase: Tuple[ActionPhase | str, ...] | None = None,
157
+ filters: str | None = None,
158
+ sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
159
+ ) -> Union[Iterator[Action], AsyncIterator[Action]]:
160
+ """
161
+ Get all actions for a given run.
162
+
163
+ :param for_run_name: The name of the run.
164
+ :param in_phase: Filter actions by one or more phases.
165
+ :param filters: The filters to apply to the project list.
166
+ :param sort_by: The sorting criteria for the project list, in the format (field, order).
167
+ :return: An iterator of actions.
168
+ """
169
+ ensure_client()
170
+ token = None
171
+ sort_by = sort_by or ("created_at", "asc")
172
+ sort_pb2 = list_pb2.Sort(
173
+ key=sort_by[0],
174
+ direction=(list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING),
175
+ )
176
+
177
+ # Build filters for phase
178
+ filter_list = []
179
+ if in_phase:
180
+ from flyteidl2.common import phase_pb2
181
+
182
+ from flyte._logging import logger
183
+
184
+ phases = [
185
+ str(p.to_protobuf_value())
186
+ if isinstance(p, ActionPhase)
187
+ else str(phase_pb2.ActionPhase.Value(f"ACTION_PHASE_{p.upper()}"))
188
+ for p in in_phase
189
+ ]
190
+ logger.debug(f"Fetching action phases: {phases}")
191
+ if len(phases) > 1:
192
+ filter_list.append(
193
+ list_pb2.Filter(
194
+ function=list_pb2.Filter.Function.VALUE_IN,
195
+ field="phase",
196
+ values=phases,
197
+ ),
198
+ )
199
+ else:
200
+ filter_list.append(
201
+ list_pb2.Filter(
202
+ function=list_pb2.Filter.Function.EQUAL,
203
+ field="phase",
204
+ values=phases[0],
205
+ ),
206
+ )
207
+
208
+ cfg = get_init_config()
209
+ while True:
210
+ req = list_pb2.ListRequest(
211
+ limit=100,
212
+ token=token,
213
+ sort_by=sort_pb2,
214
+ filters=filter_list if filter_list else None,
215
+ )
216
+ resp = await get_client().run_service.ListActions(
217
+ run_service_pb2.ListActionsRequest(
218
+ request=req,
219
+ run_id=identifier_pb2.RunIdentifier(
220
+ org=cfg.org,
221
+ project=cfg.project,
222
+ domain=cfg.domain,
223
+ name=for_run_name,
224
+ ),
225
+ )
226
+ )
227
+ token = resp.token
228
+ for r in resp.actions:
229
+ yield cls(r)
230
+ if not token:
231
+ break
232
+
233
+ @syncify
234
+ @classmethod
235
+ async def get(
236
+ cls,
237
+ uri: str | None = None,
238
+ /,
239
+ run_name: str | None = None,
240
+ name: str | None = None,
241
+ ) -> Action:
242
+ """
243
+ Get a run by its ID or name. If both are provided, the ID will take precedence.
244
+
245
+ :param uri: The URI of the action.
246
+ :param run_name: The name of the action.
247
+ :param name: The name of the action.
248
+ """
249
+ ensure_client()
250
+ cfg = get_init_config()
251
+ details: ActionDetails = await ActionDetails.get_details.aio(
252
+ identifier_pb2.ActionIdentifier(
253
+ run=identifier_pb2.RunIdentifier(
254
+ org=cfg.org,
255
+ project=cfg.project,
256
+ domain=cfg.domain,
257
+ name=run_name,
258
+ ),
259
+ name=name,
260
+ ),
261
+ )
262
+ return cls(
263
+ pb2=run_definition_pb2.Action(
264
+ id=details.action_id,
265
+ metadata=details.pb2.metadata,
266
+ status=details.pb2.status,
267
+ ),
268
+ _details=details,
269
+ )
270
+
271
+ @property
272
+ def phase(self) -> ActionPhase:
273
+ """
274
+ Get the phase of the action.
275
+
276
+ Returns:
277
+ The current execution phase as an ActionPhase enum
278
+ """
279
+ return ActionPhase.from_protobuf(self.pb2.status.phase)
280
+
281
+ @property
282
+ def raw_phase(self) -> phase_pb2.ActionPhase:
283
+ """
284
+ Get the raw phase of the action.
285
+ """
286
+ return self.pb2.status.phase
287
+
288
+ @property
289
+ def name(self) -> str:
290
+ """
291
+ Get the name of the action.
292
+ """
293
+ return self.action_id.name
294
+
295
+ @property
296
+ def run_name(self) -> str:
297
+ """
298
+ Get the name of the run.
299
+ """
300
+ return self.action_id.run.name
301
+
302
+ @property
303
+ def task_name(self) -> str | None:
304
+ """
305
+ Get the name of the task.
306
+ """
307
+ if self.pb2.metadata.HasField("task") and self.pb2.metadata.task.HasField("id"):
308
+ return self.pb2.metadata.task.id.name
309
+ return None
310
+
311
+ @property
312
+ def action_id(self) -> identifier_pb2.ActionIdentifier:
313
+ """
314
+ Get the action ID.
315
+ """
316
+ return self.pb2.id
317
+
318
+ @property
319
+ def start_time(self) -> datetime:
320
+ """
321
+ Get the start time of the action.
322
+ """
323
+ return self.pb2.status.start_time.ToDatetime().replace(tzinfo=timezone.utc)
324
+
325
+ @syncify
326
+ async def show_logs(
327
+ self,
328
+ attempt: int | None = None,
329
+ max_lines: int = 30,
330
+ show_ts: bool = False,
331
+ raw: bool = False,
332
+ filter_system: bool = False,
333
+ ):
334
+ """
335
+ Display logs for the action.
336
+
337
+ :param attempt: The attempt number to show logs for (defaults to latest attempt).
338
+ :param max_lines: Maximum number of log lines to display in the viewer.
339
+ :param show_ts: Whether to show timestamps with each log line.
340
+ :param raw: If True, print logs directly without the interactive viewer.
341
+ :param filter_system: If True, filter out system-generated log lines.
342
+ """
343
+ details = await self.details()
344
+ if not details.is_running and not details.done():
345
+ # TODO we can short circuit here if the attempt is not the last one and it is done!
346
+ await self.wait(wait_for="logs-ready")
347
+ details = await self.details()
348
+ if not attempt:
349
+ attempt = details.attempts
350
+ return await Logs.create_viewer(
351
+ action_id=self.action_id,
352
+ attempt=attempt,
353
+ max_lines=max_lines,
354
+ show_ts=show_ts,
355
+ raw=raw,
356
+ filter_system=filter_system,
357
+ )
358
+
359
+ async def details(self) -> ActionDetails:
360
+ """
361
+ Get the details of the action. This is a placeholder for getting the action details.
362
+ """
363
+ if not self._details:
364
+ self._details = await ActionDetails.get_details.aio(self.action_id)
365
+ return cast(ActionDetails, self._details)
366
+
367
+ async def watch(
368
+ self, cache_data_on_done: bool = False, wait_for: WaitFor = "terminal"
369
+ ) -> AsyncGenerator[ActionDetails, None]:
370
+ """
371
+ Watch the action for updates, updating the internal Action state with latest details.
372
+
373
+ This method updates both the cached details and the protobuf representation,
374
+ ensuring that properties like `phase` reflect the current state.
375
+ """
376
+ ad = None
377
+ async for ad in ActionDetails.watch.aio(self.action_id):
378
+ if ad is None:
379
+ return
380
+ self._details = ad
381
+ # Update the protobuf with the latest status and metadata
382
+ self.pb2.status.CopyFrom(ad.pb2.status)
383
+ self.pb2.metadata.CopyFrom(ad.pb2.metadata)
384
+ yield ad
385
+ if wait_for == "running" and ad.is_running:
386
+ break
387
+ elif wait_for == "logs-ready" and ad.logs_available():
388
+ break
389
+ if ad.done():
390
+ break
391
+ if cache_data_on_done and ad and ad.done():
392
+ await cast(ActionDetails, self._details).outputs()
393
+
394
+ async def wait(self, quiet: bool = False, wait_for: WaitFor = "terminal") -> None:
395
+ """
396
+ Wait for the run to complete, displaying a rich progress panel with status transitions,
397
+ time elapsed, and error details in case of failure.
398
+ """
399
+ console = Console()
400
+ if self.done():
401
+ if not quiet:
402
+ if self.pb2.status.phase == phase_pb2.ACTION_PHASE_SUCCEEDED:
403
+ console.print(
404
+ f"[bold green]Action '{self.name}' in Run '{self.run_name}'"
405
+ f" completed successfully.[/bold green]"
406
+ )
407
+ else:
408
+ details = await self.details()
409
+ error_message = details.error_info.message if details.error_info else ""
410
+ console.print(
411
+ f"[bold red]Action '{self.name}' in Run '{self.run_name}'"
412
+ f" exited unsuccessfully in state {self.phase} with error: {error_message}[/bold red]"
413
+ )
414
+ return
415
+
416
+ try:
417
+ with Progress(
418
+ SpinnerColumn(),
419
+ TextColumn("[progress.description]{task.description}"),
420
+ TimeElapsedColumn(),
421
+ console=console,
422
+ transient=True,
423
+ disable=quiet,
424
+ ) as progress:
425
+ task_id = progress.add_task(f"Waiting for run '{self.name}'...", start=False)
426
+ progress.start_task(task_id)
427
+
428
+ async for ad in self.watch(cache_data_on_done=True, wait_for=wait_for):
429
+ if ad is None:
430
+ progress.stop_task(task_id)
431
+ break
432
+
433
+ if ad.is_running and wait_for == "running":
434
+ progress.start_task(task_id)
435
+ break
436
+
437
+ if ad.logs_available() and wait_for == "logs-ready":
438
+ progress.start_task(task_id)
439
+ break
440
+
441
+ # Update progress description with the current phase
442
+ progress.update(
443
+ task_id,
444
+ description=f"Run: {self.run_name} in {ad.phase}, Runtime: {ad.runtime} secs "
445
+ f"Attempts[{ad.attempts}]",
446
+ )
447
+
448
+ # If the action is done, handle the final state
449
+ if ad.done():
450
+ progress.stop_task(task_id)
451
+ if not quiet:
452
+ if ad.pb2.status.phase == phase_pb2.ACTION_PHASE_SUCCEEDED:
453
+ console.print(f"[bold green]Run '{self.run_name}' completed successfully.[/bold green]")
454
+ else:
455
+ error_message = ad.error_info.message if ad.error_info else ""
456
+ console.print(
457
+ f"[bold red]Run '{self.run_name}' exited unsuccessfully in state {ad.phase}"
458
+ f" with error: {error_message}[/bold red]"
459
+ )
460
+ break
461
+ except asyncio.CancelledError:
462
+ # Handle cancellation gracefully
463
+ pass
464
+ except KeyboardInterrupt:
465
+ # Handle keyboard interrupt gracefully
466
+ pass
467
+
468
+ def done(self) -> bool:
469
+ """
470
+ Check if the action is done.
471
+ """
472
+ return _action_done_check(self.raw_phase)
473
+
474
+ async def sync(self) -> Action:
475
+ """
476
+ Sync the action with the remote server. This is a placeholder for syncing the action.
477
+ """
478
+ return self
479
+
480
+ def __rich_repr__(self) -> rich.repr.Result:
481
+ """
482
+ Rich representation of the Action object.
483
+ """
484
+ yield from _action_rich_repr(self.pb2)
485
+ if self._details:
486
+ yield from self._details.__rich_repr__()
487
+
488
+ def __repr__(self) -> str:
489
+ """
490
+ String representation of the Action object.
491
+ """
492
+ import rich.pretty
493
+
494
+ return rich.pretty.pretty_repr(self)
495
+
496
+
497
+ @dataclass
498
+ class ActionDetails(ToJSONMixin):
499
+ """
500
+ A class representing an action. It is used to manage the run of a task and its state on the remote Union API.
501
+ """
502
+
503
+ pb2: run_definition_pb2.ActionDetails
504
+ _inputs: ActionInputs | None = None
505
+ _outputs: ActionOutputs | None = None
506
+
507
+ @syncify
508
+ @classmethod
509
+ async def get_details(cls, action_id: identifier_pb2.ActionIdentifier) -> ActionDetails:
510
+ """
511
+ Get the details of the action. This is a placeholder for getting the action details.
512
+ """
513
+ ensure_client()
514
+ resp = await get_client().run_service.GetActionDetails(
515
+ run_service_pb2.GetActionDetailsRequest(
516
+ action_id=action_id,
517
+ )
518
+ )
519
+ return ActionDetails(resp.details)
520
+
521
+ @syncify
522
+ @classmethod
523
+ async def get(
524
+ cls,
525
+ uri: str | None = None,
526
+ /,
527
+ run_name: str | None = None,
528
+ name: str | None = None,
529
+ ) -> ActionDetails:
530
+ """
531
+ Get a run by its ID or name. If both are provided, the ID will take precedence.
532
+
533
+ :param uri: The URI of the action.
534
+ :param name: The name of the action.
535
+ :param run_name: The name of the run.
536
+ """
537
+ ensure_client()
538
+ if not uri:
539
+ assert name is not None and run_name is not None, "Either uri or name and run_name must be provided"
540
+ cfg = get_init_config()
541
+ return await cls.get_details.aio(
542
+ identifier_pb2.ActionIdentifier(
543
+ run=identifier_pb2.RunIdentifier(
544
+ org=cfg.org,
545
+ project=cfg.project,
546
+ domain=cfg.domain,
547
+ name=run_name,
548
+ ),
549
+ name=name,
550
+ ),
551
+ )
552
+
553
+ @syncify
554
+ @classmethod
555
+ async def watch(cls, action_id: identifier_pb2.ActionIdentifier) -> AsyncIterator[ActionDetails]:
556
+ """
557
+ Watch the action for updates. This is a placeholder for watching the action.
558
+ """
559
+ ensure_client()
560
+ if not action_id:
561
+ raise ValueError("Action ID is required")
562
+
563
+ call = cast(
564
+ AsyncIterator[WatchActionDetailsResponse],
565
+ get_client().run_service.WatchActionDetails(
566
+ request=run_service_pb2.WatchActionDetailsRequest(
567
+ action_id=action_id,
568
+ )
569
+ ),
570
+ )
571
+ try:
572
+ async for resp in call:
573
+ v = cls(resp.details)
574
+ yield v
575
+ if v.done():
576
+ return
577
+ except grpc.aio.AioRpcError as e:
578
+ if e.code() == grpc.StatusCode.CANCELLED:
579
+ pass
580
+ else:
581
+ raise e
582
+
583
+ async def watch_updates(self, cache_data_on_done: bool = False) -> AsyncGenerator[ActionDetails, None]:
584
+ """
585
+ Watch for updates to the action details, yielding each update until the action is done.
586
+
587
+ :param cache_data_on_done: If True, cache inputs and outputs when the action completes.
588
+ """
589
+ async for d in self.watch.aio(action_id=self.pb2.id):
590
+ yield d
591
+ if d.done():
592
+ self.pb2 = d.pb2
593
+ break
594
+
595
+ if cache_data_on_done and self.done():
596
+ await self._cache_data.aio()
597
+
598
+ @property
599
+ def phase(self) -> ActionPhase:
600
+ """
601
+ Get the phase of the action.
602
+
603
+ Returns:
604
+ The current execution phase as an ActionPhase enum
605
+ """
606
+ return ActionPhase.from_protobuf(self.status.phase)
607
+
608
+ @property
609
+ def raw_phase(self) -> phase_pb2.ActionPhase:
610
+ """
611
+ Get the raw phase of the action.
612
+ """
613
+ return self.status.phase
614
+
615
+ @property
616
+ def is_running(self) -> bool:
617
+ """
618
+ Check if the action is currently running.
619
+ """
620
+ return self.status.phase == phase_pb2.ACTION_PHASE_RUNNING
621
+
622
+ @property
623
+ def name(self) -> str:
624
+ """
625
+ Get the name of the action.
626
+ """
627
+ return self.action_id.name
628
+
629
+ @property
630
+ def run_name(self) -> str:
631
+ """
632
+ Get the name of the run.
633
+ """
634
+ return self.action_id.run.name
635
+
636
+ @property
637
+ def task_name(self) -> str | None:
638
+ """
639
+ Get the name of the task.
640
+ """
641
+ if self.pb2.metadata.HasField("task") and self.pb2.metadata.task.HasField("id"):
642
+ return self.pb2.metadata.task.id.name
643
+ return None
644
+
645
+ @property
646
+ def action_id(self) -> identifier_pb2.ActionIdentifier:
647
+ """
648
+ Get the action ID.
649
+ """
650
+ return self.pb2.id
651
+
652
+ @property
653
+ def metadata(self) -> run_definition_pb2.ActionMetadata:
654
+ """
655
+ Get the metadata of the action.
656
+ """
657
+ return self.pb2.metadata
658
+
659
+ @property
660
+ def status(self) -> run_definition_pb2.ActionStatus:
661
+ """
662
+ Get the status of the action.
663
+ """
664
+ return self.pb2.status
665
+
666
+ @property
667
+ def error_info(self) -> run_definition_pb2.ErrorInfo | None:
668
+ """
669
+ Get the error information if the action failed, otherwise returns None.
670
+ """
671
+ if self.pb2.HasField("error_info"):
672
+ return self.pb2.error_info
673
+ return None
674
+
675
+ @property
676
+ def abort_info(self) -> run_definition_pb2.AbortInfo | None:
677
+ """
678
+ Get the abort information if the action was aborted, otherwise returns None.
679
+ """
680
+ if self.pb2.HasField("abort_info"):
681
+ return self.pb2.abort_info
682
+ return None
683
+
684
+ @property
685
+ def runtime(self) -> timedelta:
686
+ """
687
+ Get the runtime of the action.
688
+ """
689
+ start_time = self.pb2.status.start_time.ToDatetime().replace(tzinfo=timezone.utc)
690
+ if self.pb2.status.HasField("end_time"):
691
+ end_time = self.pb2.status.end_time.ToDatetime().replace(tzinfo=timezone.utc)
692
+ return end_time - start_time
693
+ return datetime.now(timezone.utc) - start_time
694
+
695
+ @property
696
+ def attempts(self) -> int:
697
+ """
698
+ Get the number of attempts of the action.
699
+ """
700
+ return self.pb2.status.attempts
701
+
702
+ def logs_available(self, attempt: int | None = None) -> bool:
703
+ """
704
+ Check if logs are available for the action, optionally for a specific attempt.
705
+ If attempt is None, it checks for the latest attempt.
706
+ """
707
+ if attempt is None:
708
+ attempt = self.pb2.status.attempts
709
+ attempts = self.pb2.attempts
710
+ if attempts and len(attempts) >= attempt:
711
+ return attempts[attempt - 1].logs_available
712
+ return False
713
+
714
+ @syncify
715
+ async def _cache_data(self) -> bool:
716
+ """
717
+ Cache the inputs and outputs of the action.
718
+ :return: Returns True if Action is terminal and all data is cached else False.
719
+ """
720
+ from flyte._internal.runtime import convert
721
+
722
+ if self._inputs and self._outputs:
723
+ return True
724
+ if self._inputs and not self.done():
725
+ return False
726
+ resp = await get_client().run_service.GetActionData(
727
+ request=run_service_pb2.GetActionDataRequest(
728
+ action_id=self.pb2.id,
729
+ )
730
+ )
731
+ native_iface = None
732
+ if self.pb2.HasField("task"):
733
+ iface = self.pb2.task.task_template.interface
734
+ native_iface = types.guess_interface(iface)
735
+ elif self.pb2.HasField("trace"):
736
+ iface = self.pb2.trace.interface
737
+ native_iface = types.guess_interface(iface)
738
+
739
+ if resp.inputs:
740
+ data_dict = (
741
+ await convert.convert_from_inputs_to_native(native_iface, convert.Inputs(resp.inputs))
742
+ if native_iface
743
+ else {}
744
+ )
745
+ self._inputs = ActionInputs(pb2=resp.inputs, data=data_dict)
746
+
747
+ if resp.outputs:
748
+ data_tuple = (
749
+ await convert.convert_outputs_to_native(native_iface, convert.Outputs(resp.outputs))
750
+ if native_iface
751
+ else ()
752
+ )
753
+ if not isinstance(data_tuple, tuple):
754
+ data_tuple = (data_tuple,)
755
+ self._outputs = ActionOutputs(pb2=resp.outputs, data=data_tuple)
756
+
757
+ return self._outputs is not None
758
+
759
+ async def inputs(self) -> ActionInputs:
760
+ """
761
+ Return the inputs of the action.
762
+ Will return instantly if inputs are available else will fetch and return.
763
+ """
764
+ if not self._inputs:
765
+ await self._cache_data.aio()
766
+ return cast(ActionInputs, self._inputs)
767
+
768
+ async def outputs(self) -> ActionOutputs:
769
+ """
770
+ Returns the outputs of the action, returns instantly if outputs are already cached, else fetches them and
771
+ returns. If Action is not in a terminal state, raise a RuntimeError.
772
+
773
+ :return: ActionOutputs
774
+ """
775
+ if not self._outputs:
776
+ if not await self._cache_data.aio():
777
+ raise RuntimeError(
778
+ "Action is not in a terminal state, outputs are not available. "
779
+ "Please wait for the action to complete."
780
+ )
781
+ return cast(ActionOutputs, self._outputs)
782
+
783
+ def done(self) -> bool:
784
+ """
785
+ Check if the action is in a terminal state (completed or failed). This is a placeholder for checking the
786
+ action state.
787
+ """
788
+ return _action_done_check(self.raw_phase)
789
+
790
+ def __rich_repr__(self) -> rich.repr.Result:
791
+ """
792
+ Rich representation of the Action object.
793
+ """
794
+ yield from _action_details_rich_repr(self.pb2)
795
+
796
+ def __repr__(self) -> str:
797
+ """
798
+ String representation of the Action object.
799
+ """
800
+ import rich.pretty
801
+
802
+ return rich.pretty.pretty_repr(self)
803
+
804
+
805
+ @dataclass
806
+ class ActionInputs(UserDict, ToJSONMixin):
807
+ """
808
+ A class representing the inputs of an action. It is used to manage the inputs of a task and its state on the
809
+ remote Union API.
810
+
811
+ ActionInputs extends from a `UserDict` and hence is accessible like a dictionary
812
+
813
+ Example Usage:
814
+ ```python
815
+ action = Action.get(...)
816
+ print(action.inputs())
817
+ ```
818
+ Output:
819
+ ```bash
820
+ {
821
+ "x": ...,
822
+ "y": ...,
823
+ }
824
+ ```
825
+ """
826
+
827
+ pb2: common_pb2.Inputs
828
+ data: Dict[str, Any]
829
+
830
+ def __repr__(self):
831
+ import rich.pretty
832
+
833
+ import flyte.types as types
834
+
835
+ return rich.pretty.pretty_repr(types.literal_string_repr(self.pb2))
836
+
837
+
838
+ class ActionOutputs(tuple, ToJSONMixin):
839
+ """
840
+ A class representing the outputs of an action. The outputs are by default represented as a Tuple. To access them,
841
+ you can simply read them as a tuple (assign to individual variables, use index to access) or you can use the
842
+ property `named_outputs` to retrieve a dictionary of outputs with keys that represent output names
843
+ which are usually auto-generated `o0, o1, o2, o3, ...`.
844
+
845
+ Example Usage:
846
+ ```python
847
+ action = Action.get(...)
848
+ print(action.outputs())
849
+ ```
850
+ Output:
851
+ ```python
852
+ ("val1", "val2", ...)
853
+ ```
854
+ OR
855
+ ```python
856
+ action = Action.get(...)
857
+ print(action.outputs().named_outputs)
858
+ ```
859
+ Output:
860
+ ```bash
861
+ {"o0": "val1", "o1": "val2", ...}
862
+ ```
863
+ """
864
+
865
+ def __new__(cls, pb2: common_pb2.Outputs, data: Tuple[Any, ...]):
866
+ # Create the tuple part
867
+ obj = super().__new__(cls, data)
868
+ # Store extra data (you can't do this here directly since it's immutable)
869
+ obj.pb2 = pb2
870
+ return obj
871
+
872
+ def __init__(self, pb2: common_pb2.Outputs, data: Tuple[Any, ...]):
873
+ # Normally you'd set instance attributes here,
874
+ # but we've already set `pb2` in `__new__`
875
+ self.pb2 = pb2
876
+
877
+ @cached_property
878
+ def named_outputs(self) -> dict[str, Any]:
879
+ return {default_output_name(i): x for i, x in enumerate(self)}