ob-metaflow-stubs 6.0.4.4rc0__py2.py3-none-any.whl → 6.0.4.6__py2.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 (261) hide show
  1. metaflow-stubs/__init__.pyi +732 -738
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -3
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +7 -6
  8. metaflow-stubs/client/filecache.pyi +3 -3
  9. metaflow-stubs/events.pyi +3 -3
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +9 -9
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +5 -5
  14. metaflow-stubs/{meta_files.pyi → info_file.pyi} +6 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +3 -3
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -6
  20. metaflow-stubs/metaflow_current.pyi +34 -34
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +3 -3
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +2 -2
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +2 -2
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +4 -4
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +2 -2
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +5 -5
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +2 -2
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +4 -4
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +3 -3
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +3 -3
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +3 -3
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +3 -3
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +4 -4
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +3 -3
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +3 -3
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/__init__.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/__init__.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_state_machine.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/__init__.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/__init__.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/spinners.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_cli.pyi +3 -3
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_config.pyi +3 -3
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/capsule.pyi +4 -4
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/click_importer.pyi +2 -2
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/__init__.pyi +2 -2
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/code_packager.pyi +2 -2
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/__init__.pyi +2 -2
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/cli_generator.pyi +2 -2
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/config_utils.pyi +5 -5
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/schema_export.pyi +2 -2
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/typed_configs.pyi +3 -3
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +3 -3
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +3 -3
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +5 -5
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/experimental/__init__.pyi +2 -2
  87. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/perimeters.pyi +2 -2
  88. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/utils.pyi +4 -4
  89. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +2 -2
  90. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +8 -12
  91. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  92. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +2 -2
  93. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  94. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +2 -2
  95. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +11 -13
  96. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +11 -13
  97. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +2 -2
  98. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +4 -4
  99. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +3 -3
  100. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +2 -2
  101. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +2 -2
  102. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +2 -2
  103. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  104. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +2 -2
  105. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +2 -2
  106. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  107. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +2 -2
  108. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  109. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  110. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  111. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  112. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  113. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  114. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  115. metaflow-stubs/multicore_utils.pyi +2 -2
  116. metaflow-stubs/ob_internal.pyi +2 -2
  117. metaflow-stubs/parameters.pyi +5 -5
  118. metaflow-stubs/plugins/__init__.pyi +14 -14
  119. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  120. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  121. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  122. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  123. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  124. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  125. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  126. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  127. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  128. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  129. metaflow-stubs/plugins/argo/argo_workflows.pyi +6 -6
  130. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  131. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  132. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  133. metaflow-stubs/plugins/argo/exit_hooks.pyi +4 -4
  134. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  135. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  136. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  137. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  138. metaflow-stubs/plugins/aws/batch/batch.pyi +4 -4
  139. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  140. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -4
  141. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  142. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  143. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  144. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  145. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  146. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +3 -3
  147. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  148. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  149. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  150. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  151. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  152. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  153. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  154. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  155. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  156. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  157. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  158. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  159. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  160. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  161. metaflow-stubs/plugins/cards/card_decorator.pyi +5 -3
  162. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  163. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  164. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  165. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  166. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  167. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  168. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  169. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  170. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  171. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  172. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  173. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  174. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  175. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  176. metaflow-stubs/plugins/datatools/s3/s3.pyi +4 -4
  177. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  178. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  179. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  180. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  181. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  182. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  183. metaflow-stubs/plugins/exit_hook/__init__.pyi +2 -2
  184. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +2 -2
  185. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  186. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  187. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  188. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  189. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  190. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  191. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  192. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  193. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  194. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +3 -3
  195. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +4 -4
  196. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  197. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -4
  198. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  199. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  200. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  201. metaflow-stubs/plugins/parallel_decorator.pyi +3 -3
  202. metaflow-stubs/plugins/perimeters.pyi +2 -2
  203. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  204. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  205. metaflow-stubs/plugins/pypi/conda_decorator.pyi +8 -5
  206. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -6
  207. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  208. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +4 -4
  209. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  210. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  211. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  212. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  213. metaflow-stubs/plugins/secrets/__init__.pyi +3 -3
  214. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  215. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  216. metaflow-stubs/plugins/secrets/secrets_func.pyi +2 -2
  217. metaflow-stubs/plugins/secrets/secrets_spec.pyi +2 -2
  218. metaflow-stubs/plugins/secrets/utils.pyi +2 -2
  219. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  220. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  221. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  222. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  223. metaflow-stubs/plugins/torchtune/__init__.pyi +2 -2
  224. metaflow-stubs/plugins/uv/__init__.pyi +2 -2
  225. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -3
  226. metaflow-stubs/profilers/__init__.pyi +2 -2
  227. metaflow-stubs/pylint_wrapper.pyi +2 -2
  228. metaflow-stubs/runner/__init__.pyi +2 -2
  229. metaflow-stubs/runner/deployer.pyi +30 -30
  230. metaflow-stubs/runner/deployer_impl.pyi +3 -3
  231. metaflow-stubs/runner/metaflow_runner.pyi +5 -5
  232. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  233. metaflow-stubs/runner/nbrun.pyi +2 -2
  234. metaflow-stubs/runner/subprocess_manager.pyi +2 -3
  235. metaflow-stubs/runner/utils.pyi +3 -3
  236. metaflow-stubs/system/__init__.pyi +2 -2
  237. metaflow-stubs/system/system_logger.pyi +3 -3
  238. metaflow-stubs/system/system_monitor.pyi +2 -2
  239. metaflow-stubs/tagging_util.pyi +2 -2
  240. metaflow-stubs/tuple_util.pyi +2 -2
  241. metaflow-stubs/user_configs/__init__.pyi +3 -2
  242. metaflow-stubs/user_configs/config_decorators.pyi +251 -0
  243. metaflow-stubs/user_configs/config_options.pyi +6 -5
  244. metaflow-stubs/user_configs/config_parameters.pyi +9 -7
  245. {ob_metaflow_stubs-6.0.4.4rc0.dist-info → ob_metaflow_stubs-6.0.4.6.dist-info}/METADATA +1 -1
  246. ob_metaflow_stubs-6.0.4.6.dist-info/RECORD +249 -0
  247. metaflow-stubs/packaging_sys/__init__.pyi +0 -430
  248. metaflow-stubs/packaging_sys/backend.pyi +0 -73
  249. metaflow-stubs/packaging_sys/distribution_support.pyi +0 -57
  250. metaflow-stubs/packaging_sys/tar_backend.pyi +0 -53
  251. metaflow-stubs/packaging_sys/utils.pyi +0 -26
  252. metaflow-stubs/packaging_sys/v1.pyi +0 -145
  253. metaflow-stubs/user_decorators/__init__.pyi +0 -15
  254. metaflow-stubs/user_decorators/common.pyi +0 -38
  255. metaflow-stubs/user_decorators/mutable_flow.pyi +0 -223
  256. metaflow-stubs/user_decorators/mutable_step.pyi +0 -152
  257. metaflow-stubs/user_decorators/user_flow_decorator.pyi +0 -137
  258. metaflow-stubs/user_decorators/user_step_decorator.pyi +0 -323
  259. ob_metaflow_stubs-6.0.4.4rc0.dist-info/RECORD +0 -260
  260. {ob_metaflow_stubs-6.0.4.4rc0.dist-info → ob_metaflow_stubs-6.0.4.6.dist-info}/WHEEL +0 -0
  261. {ob_metaflow_stubs-6.0.4.4rc0.dist-info → ob_metaflow_stubs-6.0.4.6.dist-info}/top_level.txt +0 -0
@@ -1,20 +1,19 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.16.1.1+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-07-15T21:03:18.366830 #
3
+ # MF version: 2.15.21.3+obcheckpoint(0.2.4);ob(v1) #
4
+ # Generated on 2025-07-21T18:19:17.308440 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
16
- from . import meta_files as meta_files
17
- from . import packaging_sys as packaging_sys
16
+ from . import info_file as info_file
18
17
  from . import exception as exception
19
18
  from . import metaflow_config as metaflow_config
20
19
  from . import multicore_utils as multicore_utils
@@ -24,7 +23,6 @@ from . import metaflow_current as metaflow_current
24
23
  from .metaflow_current import current as current
25
24
  from . import parameters as parameters
26
25
  from . import user_configs as user_configs
27
- from . import user_decorators as user_decorators
28
26
  from . import tagging_util as tagging_util
29
27
  from . import metadata_provider as metadata_provider
30
28
  from . import flowspec as flowspec
@@ -35,14 +33,12 @@ from .parameters import JSONType as JSONType
35
33
  from .user_configs.config_parameters import Config as Config
36
34
  from .user_configs.config_parameters import ConfigValue as ConfigValue
37
35
  from .user_configs.config_parameters import config_expr as config_expr
38
- from .user_decorators.user_step_decorator import UserStepDecorator as UserStepDecorator
39
- from .user_decorators.user_step_decorator import StepMutator as StepMutator
40
- from .user_decorators.user_step_decorator import user_step_decorator as user_step_decorator
41
- from .user_decorators.user_flow_decorator import FlowMutator as FlowMutator
36
+ from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
+ from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
42
38
  from . import cards as cards
43
- from . import metaflow_git as metaflow_git
44
39
  from . import tuple_util as tuple_util
45
40
  from . import events as events
41
+ from . import metaflow_git as metaflow_git
46
42
  from . import runner as runner
47
43
  from . import plugins as plugins
48
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
@@ -87,8 +83,6 @@ from . import ob_internal as ob_internal
87
83
 
88
84
  EXT_PKG: str
89
85
 
90
- USER_SKIP_STEP: dict
91
-
92
86
  @typing.overload
93
87
  def step(f: typing.Callable[[FlowSpecDerived], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
94
88
  """
@@ -162,46 +156,61 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
162
156
  """
163
157
  ...
164
158
 
165
- def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
159
+ @typing.overload
160
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
166
161
  """
167
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
168
-
169
- User code call
170
- --------------
171
- @ollama(
172
- models=[...],
173
- ...
174
- )
162
+ Specifies secrets to be retrieved and injected as environment variables prior to
163
+ the execution of a step.
175
164
 
176
- Valid backend options
177
- ---------------------
178
- - 'local': Run as a separate process on the local task machine.
179
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
180
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
181
165
 
182
- Valid model options
183
- -------------------
184
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
166
+ Parameters
167
+ ----------
168
+ sources : List[Union[str, Dict[str, Any]]], default: []
169
+ List of secret specs, defining how the secrets are to be retrieved
170
+ role : str, optional, default: None
171
+ Role to use for fetching secrets
172
+ """
173
+ ...
174
+
175
+ @typing.overload
176
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
177
+ ...
178
+
179
+ @typing.overload
180
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
181
+ ...
182
+
183
+ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: typing.Optional[str] = None):
184
+ """
185
+ Specifies secrets to be retrieved and injected as environment variables prior to
186
+ the execution of a step.
185
187
 
186
188
 
187
189
  Parameters
188
190
  ----------
189
- models: list[str]
190
- List of Ollama containers running models in sidecars.
191
- backend: str
192
- Determines where and how to run the Ollama process.
193
- force_pull: bool
194
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
195
- cache_update_policy: str
196
- Cache update policy: "auto", "force", or "never".
197
- force_cache_update: bool
198
- Simple override for "force" cache update policy.
199
- debug: bool
200
- Whether to turn on verbose debugging logs.
201
- circuit_breaker_config: dict
202
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
203
- timeout_config: dict
204
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
191
+ sources : List[Union[str, Dict[str, Any]]], default: []
192
+ List of secret specs, defining how the secrets are to be retrieved
193
+ role : str, optional, default: None
194
+ Role to use for fetching secrets
195
+ """
196
+ ...
197
+
198
+ @typing.overload
199
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
200
+ """
201
+ Decorator prototype for all step decorators. This function gets specialized
202
+ and imported for all decorators types by _import_plugin_decorators().
203
+ """
204
+ ...
205
+
206
+ @typing.overload
207
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
208
+ ...
209
+
210
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
211
+ """
212
+ Decorator prototype for all step decorators. This function gets specialized
213
+ and imported for all decorators types by _import_plugin_decorators().
205
214
  """
206
215
  ...
207
216
 
@@ -260,6 +269,49 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
260
269
  """
261
270
  ...
262
271
 
272
+ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
273
+ """
274
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
275
+
276
+ User code call
277
+ --------------
278
+ @ollama(
279
+ models=[...],
280
+ ...
281
+ )
282
+
283
+ Valid backend options
284
+ ---------------------
285
+ - 'local': Run as a separate process on the local task machine.
286
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
287
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
288
+
289
+ Valid model options
290
+ -------------------
291
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
292
+
293
+
294
+ Parameters
295
+ ----------
296
+ models: list[str]
297
+ List of Ollama containers running models in sidecars.
298
+ backend: str
299
+ Determines where and how to run the Ollama process.
300
+ force_pull: bool
301
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
302
+ cache_update_policy: str
303
+ Cache update policy: "auto", "force", or "never".
304
+ force_cache_update: bool
305
+ Simple override for "force" cache update policy.
306
+ debug: bool
307
+ Whether to turn on verbose debugging logs.
308
+ circuit_breaker_config: dict
309
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
310
+ timeout_config: dict
311
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
312
+ """
313
+ ...
314
+
263
315
  def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card_refresh_interval: int, max_retries: int, retry_alert_frequency: int, engine_args: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
264
316
  """
265
317
  This decorator is used to run vllm APIs as Metaflow task sidecars.
@@ -310,84 +362,6 @@ def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card
310
362
  """
311
363
  ...
312
364
 
313
- @typing.overload
314
- def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
315
- """
316
- Specifies the Conda environment for the step.
317
-
318
- Information in this decorator will augment any
319
- attributes set in the `@conda_base` flow-level decorator. Hence,
320
- you can use `@conda_base` to set packages required by all
321
- steps and use `@conda` to specify step-specific overrides.
322
-
323
-
324
- Parameters
325
- ----------
326
- packages : Dict[str, str], default {}
327
- Packages to use for this step. The key is the name of the package
328
- and the value is the version to use.
329
- libraries : Dict[str, str], default {}
330
- Supported for backward compatibility. When used with packages, packages will take precedence.
331
- python : str, optional, default None
332
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
333
- that the version used will correspond to the version of the Python interpreter used to start the run.
334
- disabled : bool, default False
335
- If set to True, disables @conda.
336
- """
337
- ...
338
-
339
- @typing.overload
340
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
341
- ...
342
-
343
- @typing.overload
344
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
345
- ...
346
-
347
- def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
348
- """
349
- Specifies the Conda environment for the step.
350
-
351
- Information in this decorator will augment any
352
- attributes set in the `@conda_base` flow-level decorator. Hence,
353
- you can use `@conda_base` to set packages required by all
354
- steps and use `@conda` to specify step-specific overrides.
355
-
356
-
357
- Parameters
358
- ----------
359
- packages : Dict[str, str], default {}
360
- Packages to use for this step. The key is the name of the package
361
- and the value is the version to use.
362
- libraries : Dict[str, str], default {}
363
- Supported for backward compatibility. When used with packages, packages will take precedence.
364
- python : str, optional, default None
365
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
366
- that the version used will correspond to the version of the Python interpreter used to start the run.
367
- disabled : bool, default False
368
- If set to True, disables @conda.
369
- """
370
- ...
371
-
372
- @typing.overload
373
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
374
- """
375
- Decorator prototype for all step decorators. This function gets specialized
376
- and imported for all decorators types by _import_plugin_decorators().
377
- """
378
- ...
379
-
380
- @typing.overload
381
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
382
- ...
383
-
384
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
385
- """
386
- Decorator prototype for all step decorators. This function gets specialized
387
- and imported for all decorators types by _import_plugin_decorators().
388
- """
389
- ...
390
-
391
365
  @typing.overload
392
366
  def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
393
367
  """
@@ -518,104 +492,176 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
518
492
  ...
519
493
 
520
494
  @typing.overload
521
- def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
495
+ def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
522
496
  """
523
- Specifies that the step will success under all circumstances.
524
-
525
- The decorator will create an optional artifact, specified by `var`, which
526
- contains the exception raised. You can use it to detect the presence
527
- of errors, indicating that all happy-path artifacts produced by the step
528
- are missing.
497
+ Specifies environment variables to be set prior to the execution of a step.
529
498
 
530
499
 
531
500
  Parameters
532
501
  ----------
533
- var : str, optional, default None
534
- Name of the artifact in which to store the caught exception.
535
- If not specified, the exception is not stored.
536
- print_exception : bool, default True
537
- Determines whether or not the exception is printed to
538
- stdout when caught.
502
+ vars : Dict[str, str], default {}
503
+ Dictionary of environment variables to set.
539
504
  """
540
505
  ...
541
506
 
542
507
  @typing.overload
543
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
508
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
544
509
  ...
545
510
 
546
511
  @typing.overload
547
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
512
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
548
513
  ...
549
514
 
550
- def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
515
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
551
516
  """
552
- Specifies that the step will success under all circumstances.
553
-
554
- The decorator will create an optional artifact, specified by `var`, which
555
- contains the exception raised. You can use it to detect the presence
556
- of errors, indicating that all happy-path artifacts produced by the step
557
- are missing.
517
+ Specifies environment variables to be set prior to the execution of a step.
558
518
 
559
519
 
560
520
  Parameters
561
521
  ----------
562
- var : str, optional, default None
563
- Name of the artifact in which to store the caught exception.
564
- If not specified, the exception is not stored.
565
- print_exception : bool, default True
566
- Determines whether or not the exception is printed to
567
- stdout when caught.
522
+ vars : Dict[str, str], default {}
523
+ Dictionary of environment variables to set.
568
524
  """
569
525
  ...
570
526
 
571
527
  @typing.overload
572
- def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
528
+ def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
573
529
  """
574
- Specifies the PyPI packages for the step.
530
+ Creates a human-readable report, a Metaflow Card, after this step completes.
575
531
 
576
- Information in this decorator will augment any
577
- attributes set in the `@pyi_base` flow-level decorator. Hence,
578
- you can use `@pypi_base` to set packages required by all
579
- steps and use `@pypi` to specify step-specific overrides.
532
+ Note that you may add multiple `@card` decorators in a step with different parameters.
580
533
 
581
534
 
582
535
  Parameters
583
536
  ----------
584
- packages : Dict[str, str], default: {}
585
- Packages to use for this step. The key is the name of the package
537
+ type : str, default 'default'
538
+ Card type.
539
+ id : str, optional, default None
540
+ If multiple cards are present, use this id to identify this card.
541
+ options : Dict[str, Any], default {}
542
+ Options passed to the card. The contents depend on the card type.
543
+ timeout : int, default 45
544
+ Interrupt reporting if it takes more than this many seconds.
545
+ """
546
+ ...
547
+
548
+ @typing.overload
549
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
550
+ ...
551
+
552
+ @typing.overload
553
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
554
+ ...
555
+
556
+ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
557
+ """
558
+ Creates a human-readable report, a Metaflow Card, after this step completes.
559
+
560
+ Note that you may add multiple `@card` decorators in a step with different parameters.
561
+
562
+
563
+ Parameters
564
+ ----------
565
+ type : str, default 'default'
566
+ Card type.
567
+ id : str, optional, default None
568
+ If multiple cards are present, use this id to identify this card.
569
+ options : Dict[str, Any], default {}
570
+ Options passed to the card. The contents depend on the card type.
571
+ timeout : int, default 45
572
+ Interrupt reporting if it takes more than this many seconds.
573
+ """
574
+ ...
575
+
576
+ @typing.overload
577
+ def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
578
+ """
579
+ Specifies the Conda environment for the step.
580
+
581
+ Information in this decorator will augment any
582
+ attributes set in the `@conda_base` flow-level decorator. Hence,
583
+ you can use `@conda_base` to set packages required by all
584
+ steps and use `@conda` to specify step-specific overrides.
585
+
586
+
587
+ Parameters
588
+ ----------
589
+ packages : Dict[str, str], default {}
590
+ Packages to use for this step. The key is the name of the package
586
591
  and the value is the version to use.
587
- python : str, optional, default: None
592
+ libraries : Dict[str, str], default {}
593
+ Supported for backward compatibility. When used with packages, packages will take precedence.
594
+ python : str, optional, default None
588
595
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
589
596
  that the version used will correspond to the version of the Python interpreter used to start the run.
597
+ disabled : bool, default False
598
+ If set to True, disables @conda.
590
599
  """
591
600
  ...
592
601
 
593
602
  @typing.overload
594
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
603
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
595
604
  ...
596
605
 
597
606
  @typing.overload
598
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
607
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
599
608
  ...
600
609
 
601
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
610
+ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
602
611
  """
603
- Specifies the PyPI packages for the step.
612
+ Specifies the Conda environment for the step.
604
613
 
605
614
  Information in this decorator will augment any
606
- attributes set in the `@pyi_base` flow-level decorator. Hence,
607
- you can use `@pypi_base` to set packages required by all
608
- steps and use `@pypi` to specify step-specific overrides.
615
+ attributes set in the `@conda_base` flow-level decorator. Hence,
616
+ you can use `@conda_base` to set packages required by all
617
+ steps and use `@conda` to specify step-specific overrides.
609
618
 
610
619
 
611
620
  Parameters
612
621
  ----------
613
- packages : Dict[str, str], default: {}
622
+ packages : Dict[str, str], default {}
614
623
  Packages to use for this step. The key is the name of the package
615
624
  and the value is the version to use.
616
- python : str, optional, default: None
625
+ libraries : Dict[str, str], default {}
626
+ Supported for backward compatibility. When used with packages, packages will take precedence.
627
+ python : str, optional, default None
617
628
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
618
629
  that the version used will correspond to the version of the Python interpreter used to start the run.
630
+ disabled : bool, default False
631
+ If set to True, disables @conda.
632
+ """
633
+ ...
634
+
635
+ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
636
+ """
637
+ Specifies that this step should execute on DGX cloud.
638
+
639
+
640
+ Parameters
641
+ ----------
642
+ gpu : int
643
+ Number of GPUs to use.
644
+ gpu_type : str
645
+ Type of Nvidia GPU to use.
646
+ """
647
+ ...
648
+
649
+ @typing.overload
650
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
651
+ """
652
+ Decorator prototype for all step decorators. This function gets specialized
653
+ and imported for all decorators types by _import_plugin_decorators().
654
+ """
655
+ ...
656
+
657
+ @typing.overload
658
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
659
+ ...
660
+
661
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
662
+ """
663
+ Decorator prototype for all step decorators. This function gets specialized
664
+ and imported for all decorators types by _import_plugin_decorators().
619
665
  """
620
666
  ...
621
667
 
@@ -709,21 +755,53 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
709
755
  ...
710
756
 
711
757
  @typing.overload
712
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
758
+ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
713
759
  """
714
- Decorator prototype for all step decorators. This function gets specialized
715
- and imported for all decorators types by _import_plugin_decorators().
760
+ Specifies the PyPI packages for the step.
761
+
762
+ Information in this decorator will augment any
763
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
764
+ you can use `@pypi_base` to set packages required by all
765
+ steps and use `@pypi` to specify step-specific overrides.
766
+
767
+
768
+ Parameters
769
+ ----------
770
+ packages : Dict[str, str], default: {}
771
+ Packages to use for this step. The key is the name of the package
772
+ and the value is the version to use.
773
+ python : str, optional, default: None
774
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
775
+ that the version used will correspond to the version of the Python interpreter used to start the run.
716
776
  """
717
777
  ...
718
778
 
719
779
  @typing.overload
720
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
780
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
721
781
  ...
722
782
 
723
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
783
+ @typing.overload
784
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
785
+ ...
786
+
787
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
724
788
  """
725
- Decorator prototype for all step decorators. This function gets specialized
726
- and imported for all decorators types by _import_plugin_decorators().
789
+ Specifies the PyPI packages for the step.
790
+
791
+ Information in this decorator will augment any
792
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
793
+ you can use `@pypi_base` to set packages required by all
794
+ steps and use `@pypi` to specify step-specific overrides.
795
+
796
+
797
+ Parameters
798
+ ----------
799
+ packages : Dict[str, str], default: {}
800
+ Packages to use for this step. The key is the name of the package
801
+ and the value is the version to use.
802
+ python : str, optional, default: None
803
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
804
+ that the version used will correspond to the version of the Python interpreter used to start the run.
727
805
  """
728
806
  ...
729
807
 
@@ -744,250 +822,110 @@ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[t
744
822
  ...
745
823
 
746
824
  @typing.overload
747
- def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
825
+ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
748
826
  """
749
- Specifies environment variables to be set prior to the execution of a step.
827
+ Specifies the resources needed when executing this step.
828
+
829
+ Use `@resources` to specify the resource requirements
830
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
831
+
832
+ You can choose the compute layer on the command line by executing e.g.
833
+ ```
834
+ python myflow.py run --with batch
835
+ ```
836
+ or
837
+ ```
838
+ python myflow.py run --with kubernetes
839
+ ```
840
+ which executes the flow on the desired system using the
841
+ requirements specified in `@resources`.
750
842
 
751
843
 
752
844
  Parameters
753
845
  ----------
754
- vars : Dict[str, str], default {}
755
- Dictionary of environment variables to set.
846
+ cpu : int, default 1
847
+ Number of CPUs required for this step.
848
+ gpu : int, optional, default None
849
+ Number of GPUs required for this step.
850
+ disk : int, optional, default None
851
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
852
+ memory : int, default 4096
853
+ Memory size (in MB) required for this step.
854
+ shared_memory : int, optional, default None
855
+ The value for the size (in MiB) of the /dev/shm volume for this step.
856
+ This parameter maps to the `--shm-size` option in Docker.
756
857
  """
757
858
  ...
758
859
 
759
860
  @typing.overload
760
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
861
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
761
862
  ...
762
863
 
763
864
  @typing.overload
764
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
865
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
765
866
  ...
766
867
 
767
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
868
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
768
869
  """
769
- Specifies environment variables to be set prior to the execution of a step.
870
+ Specifies the resources needed when executing this step.
871
+
872
+ Use `@resources` to specify the resource requirements
873
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
874
+
875
+ You can choose the compute layer on the command line by executing e.g.
876
+ ```
877
+ python myflow.py run --with batch
878
+ ```
879
+ or
880
+ ```
881
+ python myflow.py run --with kubernetes
882
+ ```
883
+ which executes the flow on the desired system using the
884
+ requirements specified in `@resources`.
770
885
 
771
886
 
772
887
  Parameters
773
888
  ----------
774
- vars : Dict[str, str], default {}
775
- Dictionary of environment variables to set.
889
+ cpu : int, default 1
890
+ Number of CPUs required for this step.
891
+ gpu : int, optional, default None
892
+ Number of GPUs required for this step.
893
+ disk : int, optional, default None
894
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
895
+ memory : int, default 4096
896
+ Memory size (in MB) required for this step.
897
+ shared_memory : int, optional, default None
898
+ The value for the size (in MiB) of the /dev/shm volume for this step.
899
+ This parameter maps to the `--shm-size` option in Docker.
776
900
  """
777
901
  ...
778
902
 
779
- @typing.overload
780
- def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
903
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
781
904
  """
782
- Enables checkpointing for a step.
905
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
783
906
 
784
907
  > Examples
785
908
 
786
- - Saving Checkpoints
909
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
910
+ ```python
911
+ @huggingface_hub
912
+ @step
913
+ def pull_model_from_huggingface(self):
914
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
915
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
916
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
917
+ # value of the function is a reference to the model in the backend storage.
918
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
919
+
920
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
921
+ self.llama_model = current.huggingface_hub.snapshot_download(
922
+ repo_id=self.model_id,
923
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
924
+ )
925
+ self.next(self.train)
926
+ ```
787
927
 
788
- ```python
789
- @checkpoint
790
- @step
791
- def train(self):
792
- model = create_model(self.parameters, checkpoint_path = None)
793
- for i in range(self.epochs):
794
- # some training logic
795
- loss = model.train(self.dataset)
796
- if i % 10 == 0:
797
- model.save(
798
- current.checkpoint.directory,
799
- )
800
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
801
- # and returns a reference dictionary to the checkpoint saved in the datastore
802
- self.latest_checkpoint = current.checkpoint.save(
803
- name="epoch_checkpoint",
804
- metadata={
805
- "epoch": i,
806
- "loss": loss,
807
- }
808
- )
809
- ```
810
-
811
- - Using Loaded Checkpoints
812
-
813
- ```python
814
- @retry(times=3)
815
- @checkpoint
816
- @step
817
- def train(self):
818
- # Assume that the task has restarted and the previous attempt of the task
819
- # saved a checkpoint
820
- checkpoint_path = None
821
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
822
- print("Loaded checkpoint from the previous attempt")
823
- checkpoint_path = current.checkpoint.directory
824
-
825
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
826
- for i in range(self.epochs):
827
- ...
828
- ```
829
-
830
-
831
- Parameters
832
- ----------
833
- load_policy : str, default: "fresh"
834
- The policy for loading the checkpoint. The following policies are supported:
835
- - "eager": Loads the the latest available checkpoint within the namespace.
836
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
837
- will be loaded at the start of the task.
838
- - "none": Do not load any checkpoint
839
- - "fresh": Loads the lastest checkpoint created within the running Task.
840
- This mode helps loading checkpoints across various retry attempts of the same task.
841
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
842
- created within the task will be loaded when the task is retries execution on failure.
843
-
844
- temp_dir_root : str, default: None
845
- The root directory under which `current.checkpoint.directory` will be created.
846
- """
847
- ...
848
-
849
- @typing.overload
850
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
851
- ...
852
-
853
- @typing.overload
854
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
855
- ...
856
-
857
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
858
- """
859
- Enables checkpointing for a step.
860
-
861
- > Examples
862
-
863
- - Saving Checkpoints
864
-
865
- ```python
866
- @checkpoint
867
- @step
868
- def train(self):
869
- model = create_model(self.parameters, checkpoint_path = None)
870
- for i in range(self.epochs):
871
- # some training logic
872
- loss = model.train(self.dataset)
873
- if i % 10 == 0:
874
- model.save(
875
- current.checkpoint.directory,
876
- )
877
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
878
- # and returns a reference dictionary to the checkpoint saved in the datastore
879
- self.latest_checkpoint = current.checkpoint.save(
880
- name="epoch_checkpoint",
881
- metadata={
882
- "epoch": i,
883
- "loss": loss,
884
- }
885
- )
886
- ```
887
-
888
- - Using Loaded Checkpoints
889
-
890
- ```python
891
- @retry(times=3)
892
- @checkpoint
893
- @step
894
- def train(self):
895
- # Assume that the task has restarted and the previous attempt of the task
896
- # saved a checkpoint
897
- checkpoint_path = None
898
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
899
- print("Loaded checkpoint from the previous attempt")
900
- checkpoint_path = current.checkpoint.directory
901
-
902
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
903
- for i in range(self.epochs):
904
- ...
905
- ```
906
-
907
-
908
- Parameters
909
- ----------
910
- load_policy : str, default: "fresh"
911
- The policy for loading the checkpoint. The following policies are supported:
912
- - "eager": Loads the the latest available checkpoint within the namespace.
913
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
914
- will be loaded at the start of the task.
915
- - "none": Do not load any checkpoint
916
- - "fresh": Loads the lastest checkpoint created within the running Task.
917
- This mode helps loading checkpoints across various retry attempts of the same task.
918
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
919
- created within the task will be loaded when the task is retries execution on failure.
920
-
921
- temp_dir_root : str, default: None
922
- The root directory under which `current.checkpoint.directory` will be created.
923
- """
924
- ...
925
-
926
- @typing.overload
927
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
928
- """
929
- Specifies secrets to be retrieved and injected as environment variables prior to
930
- the execution of a step.
931
-
932
-
933
- Parameters
934
- ----------
935
- sources : List[Union[str, Dict[str, Any]]], default: []
936
- List of secret specs, defining how the secrets are to be retrieved
937
- role : str, optional, default: None
938
- Role to use for fetching secrets
939
- """
940
- ...
941
-
942
- @typing.overload
943
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
944
- ...
945
-
946
- @typing.overload
947
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
948
- ...
949
-
950
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: typing.Optional[str] = None):
951
- """
952
- Specifies secrets to be retrieved and injected as environment variables prior to
953
- the execution of a step.
954
-
955
-
956
- Parameters
957
- ----------
958
- sources : List[Union[str, Dict[str, Any]]], default: []
959
- List of secret specs, defining how the secrets are to be retrieved
960
- role : str, optional, default: None
961
- Role to use for fetching secrets
962
- """
963
- ...
964
-
965
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
966
- """
967
- Decorator that helps cache, version and store models/datasets from huggingface hub.
968
-
969
- > Examples
970
-
971
- **Usage: creating references of models from huggingface that may be loaded in downstream steps**
972
- ```python
973
- @huggingface_hub
974
- @step
975
- def pull_model_from_huggingface(self):
976
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
977
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
978
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
979
- # value of the function is a reference to the model in the backend storage.
980
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
981
-
982
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
983
- self.llama_model = current.huggingface_hub.snapshot_download(
984
- repo_id=self.model_id,
985
- allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
986
- )
987
- self.next(self.train)
988
- ```
989
-
990
- **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
928
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
991
929
  ```python
992
930
  @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
993
931
  @step
@@ -1042,20 +980,6 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.
1042
980
  """
1043
981
  ...
1044
982
 
1045
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1046
- """
1047
- Specifies that this step should execute on DGX cloud.
1048
-
1049
-
1050
- Parameters
1051
- ----------
1052
- gpu : int
1053
- Number of GPUs to use.
1054
- gpu_type : str
1055
- Type of Nvidia GPU to use.
1056
- """
1057
- ...
1058
-
1059
983
  @typing.overload
1060
984
  def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1061
985
  """
@@ -1116,309 +1040,275 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
1116
1040
  ...
1117
1041
 
1118
1042
  @typing.overload
1119
- def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1043
+ def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1120
1044
  """
1121
- Creates a human-readable report, a Metaflow Card, after this step completes.
1045
+ Specifies that the step will success under all circumstances.
1122
1046
 
1123
- Note that you may add multiple `@card` decorators in a step with different parameters.
1047
+ The decorator will create an optional artifact, specified by `var`, which
1048
+ contains the exception raised. You can use it to detect the presence
1049
+ of errors, indicating that all happy-path artifacts produced by the step
1050
+ are missing.
1124
1051
 
1125
1052
 
1126
1053
  Parameters
1127
1054
  ----------
1128
- type : str, default 'default'
1129
- Card type.
1130
- id : str, optional, default None
1131
- If multiple cards are present, use this id to identify this card.
1132
- options : Dict[str, Any], default {}
1133
- Options passed to the card. The contents depend on the card type.
1134
- timeout : int, default 45
1135
- Interrupt reporting if it takes more than this many seconds.
1055
+ var : str, optional, default None
1056
+ Name of the artifact in which to store the caught exception.
1057
+ If not specified, the exception is not stored.
1058
+ print_exception : bool, default True
1059
+ Determines whether or not the exception is printed to
1060
+ stdout when caught.
1136
1061
  """
1137
1062
  ...
1138
1063
 
1139
1064
  @typing.overload
1140
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1065
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1141
1066
  ...
1142
1067
 
1143
1068
  @typing.overload
1144
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1069
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1145
1070
  ...
1146
1071
 
1147
- def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
1072
+ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
1148
1073
  """
1149
- Creates a human-readable report, a Metaflow Card, after this step completes.
1074
+ Specifies that the step will success under all circumstances.
1150
1075
 
1151
- Note that you may add multiple `@card` decorators in a step with different parameters.
1076
+ The decorator will create an optional artifact, specified by `var`, which
1077
+ contains the exception raised. You can use it to detect the presence
1078
+ of errors, indicating that all happy-path artifacts produced by the step
1079
+ are missing.
1152
1080
 
1153
1081
 
1154
1082
  Parameters
1155
1083
  ----------
1156
- type : str, default 'default'
1157
- Card type.
1158
- id : str, optional, default None
1159
- If multiple cards are present, use this id to identify this card.
1160
- options : Dict[str, Any], default {}
1161
- Options passed to the card. The contents depend on the card type.
1162
- timeout : int, default 45
1163
- Interrupt reporting if it takes more than this many seconds.
1164
- """
1165
- ...
1166
-
1167
- @typing.overload
1168
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1169
- """
1170
- Internal decorator to support Fast bakery
1171
- """
1172
- ...
1173
-
1174
- @typing.overload
1175
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1176
- ...
1177
-
1178
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1179
- """
1180
- Internal decorator to support Fast bakery
1084
+ var : str, optional, default None
1085
+ Name of the artifact in which to store the caught exception.
1086
+ If not specified, the exception is not stored.
1087
+ print_exception : bool, default True
1088
+ Determines whether or not the exception is printed to
1089
+ stdout when caught.
1181
1090
  """
1182
1091
  ...
1183
1092
 
1184
1093
  @typing.overload
1185
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1094
+ def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1186
1095
  """
1187
- Specifies the resources needed when executing this step.
1096
+ Enables checkpointing for a step.
1188
1097
 
1189
- Use `@resources` to specify the resource requirements
1190
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1098
+ > Examples
1191
1099
 
1192
- You can choose the compute layer on the command line by executing e.g.
1193
- ```
1194
- python myflow.py run --with batch
1195
- ```
1196
- or
1197
- ```
1198
- python myflow.py run --with kubernetes
1199
- ```
1200
- which executes the flow on the desired system using the
1201
- requirements specified in `@resources`.
1100
+ - Saving Checkpoints
1202
1101
 
1102
+ ```python
1103
+ @checkpoint
1104
+ @step
1105
+ def train(self):
1106
+ model = create_model(self.parameters, checkpoint_path = None)
1107
+ for i in range(self.epochs):
1108
+ # some training logic
1109
+ loss = model.train(self.dataset)
1110
+ if i % 10 == 0:
1111
+ model.save(
1112
+ current.checkpoint.directory,
1113
+ )
1114
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
1115
+ # and returns a reference dictionary to the checkpoint saved in the datastore
1116
+ self.latest_checkpoint = current.checkpoint.save(
1117
+ name="epoch_checkpoint",
1118
+ metadata={
1119
+ "epoch": i,
1120
+ "loss": loss,
1121
+ }
1122
+ )
1123
+ ```
1203
1124
 
1204
- Parameters
1205
- ----------
1206
- cpu : int, default 1
1207
- Number of CPUs required for this step.
1208
- gpu : int, optional, default None
1209
- Number of GPUs required for this step.
1210
- disk : int, optional, default None
1211
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1212
- memory : int, default 4096
1213
- Memory size (in MB) required for this step.
1214
- shared_memory : int, optional, default None
1215
- The value for the size (in MiB) of the /dev/shm volume for this step.
1216
- This parameter maps to the `--shm-size` option in Docker.
1217
- """
1218
- ...
1219
-
1220
- @typing.overload
1221
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1222
- ...
1223
-
1224
- @typing.overload
1225
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1226
- ...
1227
-
1228
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
1229
- """
1230
- Specifies the resources needed when executing this step.
1125
+ - Using Loaded Checkpoints
1231
1126
 
1232
- Use `@resources` to specify the resource requirements
1233
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1127
+ ```python
1128
+ @retry(times=3)
1129
+ @checkpoint
1130
+ @step
1131
+ def train(self):
1132
+ # Assume that the task has restarted and the previous attempt of the task
1133
+ # saved a checkpoint
1134
+ checkpoint_path = None
1135
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1136
+ print("Loaded checkpoint from the previous attempt")
1137
+ checkpoint_path = current.checkpoint.directory
1234
1138
 
1235
- You can choose the compute layer on the command line by executing e.g.
1236
- ```
1237
- python myflow.py run --with batch
1238
- ```
1239
- or
1240
- ```
1241
- python myflow.py run --with kubernetes
1139
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1140
+ for i in range(self.epochs):
1141
+ ...
1242
1142
  ```
1243
- which executes the flow on the desired system using the
1244
- requirements specified in `@resources`.
1245
1143
 
1246
1144
 
1247
1145
  Parameters
1248
1146
  ----------
1249
- cpu : int, default 1
1250
- Number of CPUs required for this step.
1251
- gpu : int, optional, default None
1252
- Number of GPUs required for this step.
1253
- disk : int, optional, default None
1254
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1255
- memory : int, default 4096
1256
- Memory size (in MB) required for this step.
1257
- shared_memory : int, optional, default None
1258
- The value for the size (in MiB) of the /dev/shm volume for this step.
1259
- This parameter maps to the `--shm-size` option in Docker.
1260
- """
1261
- ...
1262
-
1263
- @typing.overload
1264
- def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1265
- """
1266
- Specifies the flow(s) that this flow depends on.
1267
-
1268
- ```
1269
- @trigger_on_finish(flow='FooFlow')
1270
- ```
1271
- or
1272
- ```
1273
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1274
- ```
1275
- This decorator respects the @project decorator and triggers the flow
1276
- when upstream runs within the same namespace complete successfully
1277
-
1278
- Additionally, you can specify project aware upstream flow dependencies
1279
- by specifying the fully qualified project_flow_name.
1280
- ```
1281
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1282
- ```
1283
- or
1284
- ```
1285
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1286
- ```
1287
-
1288
- You can also specify just the project or project branch (other values will be
1289
- inferred from the current project or project branch):
1290
- ```
1291
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1292
- ```
1293
-
1294
- Note that `branch` is typically one of:
1295
- - `prod`
1296
- - `user.bob`
1297
- - `test.my_experiment`
1298
- - `prod.staging`
1299
-
1147
+ load_policy : str, default: "fresh"
1148
+ The policy for loading the checkpoint. The following policies are supported:
1149
+ - "eager": Loads the the latest available checkpoint within the namespace.
1150
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1151
+ will be loaded at the start of the task.
1152
+ - "none": Do not load any checkpoint
1153
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1154
+ This mode helps loading checkpoints across various retry attempts of the same task.
1155
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1156
+ created within the task will be loaded when the task is retries execution on failure.
1300
1157
 
1301
- Parameters
1302
- ----------
1303
- flow : Union[str, Dict[str, str]], optional, default None
1304
- Upstream flow dependency for this flow.
1305
- flows : List[Union[str, Dict[str, str]]], default []
1306
- Upstream flow dependencies for this flow.
1307
- options : Dict[str, Any], default {}
1308
- Backend-specific configuration for tuning eventing behavior.
1158
+ temp_dir_root : str, default: None
1159
+ The root directory under which `current.checkpoint.directory` will be created.
1309
1160
  """
1310
1161
  ...
1311
1162
 
1312
1163
  @typing.overload
1313
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1164
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1314
1165
  ...
1315
1166
 
1316
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1167
+ @typing.overload
1168
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1169
+ ...
1170
+
1171
+ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
1317
1172
  """
1318
- Specifies the flow(s) that this flow depends on.
1173
+ Enables checkpointing for a step.
1319
1174
 
1320
- ```
1321
- @trigger_on_finish(flow='FooFlow')
1322
- ```
1323
- or
1324
- ```
1325
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1326
- ```
1327
- This decorator respects the @project decorator and triggers the flow
1328
- when upstream runs within the same namespace complete successfully
1175
+ > Examples
1329
1176
 
1330
- Additionally, you can specify project aware upstream flow dependencies
1331
- by specifying the fully qualified project_flow_name.
1332
- ```
1333
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1334
- ```
1335
- or
1336
- ```
1337
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1338
- ```
1177
+ - Saving Checkpoints
1339
1178
 
1340
- You can also specify just the project or project branch (other values will be
1341
- inferred from the current project or project branch):
1342
- ```
1343
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1179
+ ```python
1180
+ @checkpoint
1181
+ @step
1182
+ def train(self):
1183
+ model = create_model(self.parameters, checkpoint_path = None)
1184
+ for i in range(self.epochs):
1185
+ # some training logic
1186
+ loss = model.train(self.dataset)
1187
+ if i % 10 == 0:
1188
+ model.save(
1189
+ current.checkpoint.directory,
1190
+ )
1191
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
1192
+ # and returns a reference dictionary to the checkpoint saved in the datastore
1193
+ self.latest_checkpoint = current.checkpoint.save(
1194
+ name="epoch_checkpoint",
1195
+ metadata={
1196
+ "epoch": i,
1197
+ "loss": loss,
1198
+ }
1199
+ )
1344
1200
  ```
1345
1201
 
1346
- Note that `branch` is typically one of:
1347
- - `prod`
1348
- - `user.bob`
1349
- - `test.my_experiment`
1350
- - `prod.staging`
1202
+ - Using Loaded Checkpoints
1203
+
1204
+ ```python
1205
+ @retry(times=3)
1206
+ @checkpoint
1207
+ @step
1208
+ def train(self):
1209
+ # Assume that the task has restarted and the previous attempt of the task
1210
+ # saved a checkpoint
1211
+ checkpoint_path = None
1212
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1213
+ print("Loaded checkpoint from the previous attempt")
1214
+ checkpoint_path = current.checkpoint.directory
1215
+
1216
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1217
+ for i in range(self.epochs):
1218
+ ...
1219
+ ```
1351
1220
 
1352
1221
 
1353
1222
  Parameters
1354
1223
  ----------
1355
- flow : Union[str, Dict[str, str]], optional, default None
1356
- Upstream flow dependency for this flow.
1357
- flows : List[Union[str, Dict[str, str]]], default []
1358
- Upstream flow dependencies for this flow.
1359
- options : Dict[str, Any], default {}
1360
- Backend-specific configuration for tuning eventing behavior.
1224
+ load_policy : str, default: "fresh"
1225
+ The policy for loading the checkpoint. The following policies are supported:
1226
+ - "eager": Loads the the latest available checkpoint within the namespace.
1227
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1228
+ will be loaded at the start of the task.
1229
+ - "none": Do not load any checkpoint
1230
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1231
+ This mode helps loading checkpoints across various retry attempts of the same task.
1232
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1233
+ created within the task will be loaded when the task is retries execution on failure.
1234
+
1235
+ temp_dir_root : str, default: None
1236
+ The root directory under which `current.checkpoint.directory` will be created.
1361
1237
  """
1362
1238
  ...
1363
1239
 
1364
1240
  @typing.overload
1365
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1241
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1366
1242
  """
1367
- Specifies the times when the flow should be run when running on a
1368
- production scheduler.
1243
+ Internal decorator to support Fast bakery
1244
+ """
1245
+ ...
1246
+
1247
+ @typing.overload
1248
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1249
+ ...
1250
+
1251
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1252
+ """
1253
+ Internal decorator to support Fast bakery
1254
+ """
1255
+ ...
1256
+
1257
+ @typing.overload
1258
+ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1259
+ """
1260
+ Specifies the Conda environment for all steps of the flow.
1261
+
1262
+ Use `@conda_base` to set common libraries required by all
1263
+ steps and use `@conda` to specify step-specific additions.
1369
1264
 
1370
1265
 
1371
1266
  Parameters
1372
1267
  ----------
1373
- hourly : bool, default False
1374
- Run the workflow hourly.
1375
- daily : bool, default True
1376
- Run the workflow daily.
1377
- weekly : bool, default False
1378
- Run the workflow weekly.
1379
- cron : str, optional, default None
1380
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1381
- specified by this expression.
1382
- timezone : str, optional, default None
1383
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1384
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1268
+ packages : Dict[str, str], default {}
1269
+ Packages to use for this flow. The key is the name of the package
1270
+ and the value is the version to use.
1271
+ libraries : Dict[str, str], default {}
1272
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1273
+ python : str, optional, default None
1274
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1275
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1276
+ disabled : bool, default False
1277
+ If set to True, disables Conda.
1385
1278
  """
1386
1279
  ...
1387
1280
 
1388
1281
  @typing.overload
1389
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1282
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1390
1283
  ...
1391
1284
 
1392
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1285
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1393
1286
  """
1394
- Specifies the times when the flow should be run when running on a
1395
- production scheduler.
1287
+ Specifies the Conda environment for all steps of the flow.
1288
+
1289
+ Use `@conda_base` to set common libraries required by all
1290
+ steps and use `@conda` to specify step-specific additions.
1396
1291
 
1397
1292
 
1398
1293
  Parameters
1399
1294
  ----------
1400
- hourly : bool, default False
1401
- Run the workflow hourly.
1402
- daily : bool, default True
1403
- Run the workflow daily.
1404
- weekly : bool, default False
1405
- Run the workflow weekly.
1406
- cron : str, optional, default None
1407
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1408
- specified by this expression.
1409
- timezone : str, optional, default None
1410
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1411
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1295
+ packages : Dict[str, str], default {}
1296
+ Packages to use for this flow. The key is the name of the package
1297
+ and the value is the version to use.
1298
+ libraries : Dict[str, str], default {}
1299
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1300
+ python : str, optional, default None
1301
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1302
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1303
+ disabled : bool, default False
1304
+ If set to True, disables Conda.
1412
1305
  """
1413
1306
  ...
1414
1307
 
1415
- def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1308
+ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1416
1309
  """
1417
- The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1418
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1419
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1420
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1421
- starts only after all sensors finish.
1310
+ The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1311
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1422
1312
 
1423
1313
 
1424
1314
  Parameters
@@ -1440,18 +1330,21 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1440
1330
  Name of the sensor on Airflow
1441
1331
  description : str
1442
1332
  Description of sensor in the Airflow UI
1443
- bucket_key : Union[str, List[str]]
1444
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1445
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1446
- bucket_name : str
1447
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1448
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1449
- wildcard_match : bool
1450
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1451
- aws_conn_id : str
1452
- a reference to the s3 connection on Airflow. (Default: None)
1453
- verify : bool
1454
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1333
+ external_dag_id : str
1334
+ The dag_id that contains the task you want to wait for.
1335
+ external_task_ids : List[str]
1336
+ The list of task_ids that you want to wait for.
1337
+ If None (default value) the sensor waits for the DAG. (Default: None)
1338
+ allowed_states : List[str]
1339
+ Iterable of allowed states, (Default: ['success'])
1340
+ failed_states : List[str]
1341
+ Iterable of failed or dis-allowed states. (Default: None)
1342
+ execution_delta : datetime.timedelta
1343
+ time difference with the previous execution to look at,
1344
+ the default is the same logical date as the current task or DAG. (Default: None)
1345
+ check_existence: bool
1346
+ Set to True to check if the external task exists or check if
1347
+ the DAG to wait for exists. (Default: True)
1455
1348
  """
1456
1349
  ...
1457
1350
 
@@ -1569,6 +1462,57 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1569
1462
  """
1570
1463
  ...
1571
1464
 
1465
+ @typing.overload
1466
+ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1467
+ """
1468
+ Specifies the times when the flow should be run when running on a
1469
+ production scheduler.
1470
+
1471
+
1472
+ Parameters
1473
+ ----------
1474
+ hourly : bool, default False
1475
+ Run the workflow hourly.
1476
+ daily : bool, default True
1477
+ Run the workflow daily.
1478
+ weekly : bool, default False
1479
+ Run the workflow weekly.
1480
+ cron : str, optional, default None
1481
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1482
+ specified by this expression.
1483
+ timezone : str, optional, default None
1484
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1485
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1486
+ """
1487
+ ...
1488
+
1489
+ @typing.overload
1490
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1491
+ ...
1492
+
1493
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1494
+ """
1495
+ Specifies the times when the flow should be run when running on a
1496
+ production scheduler.
1497
+
1498
+
1499
+ Parameters
1500
+ ----------
1501
+ hourly : bool, default False
1502
+ Run the workflow hourly.
1503
+ daily : bool, default True
1504
+ Run the workflow daily.
1505
+ weekly : bool, default False
1506
+ Run the workflow weekly.
1507
+ cron : str, optional, default None
1508
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1509
+ specified by this expression.
1510
+ timezone : str, optional, default None
1511
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1512
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1513
+ """
1514
+ ...
1515
+
1572
1516
  @typing.overload
1573
1517
  def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1574
1518
  """
@@ -1610,45 +1554,13 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1610
1554
  """
1611
1555
  ...
1612
1556
 
1613
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1614
- """
1615
- Specifies what flows belong to the same project.
1616
-
1617
- A project-specific namespace is created for all flows that
1618
- use the same `@project(name)`.
1619
-
1620
-
1621
- Parameters
1622
- ----------
1623
- name : str
1624
- Project name. Make sure that the name is unique amongst all
1625
- projects that use the same production scheduler. The name may
1626
- contain only lowercase alphanumeric characters and underscores.
1627
-
1628
- branch : Optional[str], default None
1629
- The branch to use. If not specified, the branch is set to
1630
- `user.<username>` unless `production` is set to `True`. This can
1631
- also be set on the command line using `--branch` as a top-level option.
1632
- It is an error to specify `branch` in the decorator and on the command line.
1633
-
1634
- production : bool, default False
1635
- Whether or not the branch is the production branch. This can also be set on the
1636
- command line using `--production` as a top-level option. It is an error to specify
1637
- `production` in the decorator and on the command line.
1638
- The project branch name will be:
1639
- - if `branch` is specified:
1640
- - if `production` is True: `prod.<branch>`
1641
- - if `production` is False: `test.<branch>`
1642
- - if `branch` is not specified:
1643
- - if `production` is True: `prod`
1644
- - if `production` is False: `user.<username>`
1645
- """
1646
- ...
1647
-
1648
- def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1557
+ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1649
1558
  """
1650
- The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1651
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1559
+ The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1560
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1561
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1562
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1563
+ starts only after all sensors finish.
1652
1564
 
1653
1565
 
1654
1566
  Parameters
@@ -1670,21 +1582,119 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1670
1582
  Name of the sensor on Airflow
1671
1583
  description : str
1672
1584
  Description of sensor in the Airflow UI
1673
- external_dag_id : str
1674
- The dag_id that contains the task you want to wait for.
1675
- external_task_ids : List[str]
1676
- The list of task_ids that you want to wait for.
1677
- If None (default value) the sensor waits for the DAG. (Default: None)
1678
- allowed_states : List[str]
1679
- Iterable of allowed states, (Default: ['success'])
1680
- failed_states : List[str]
1681
- Iterable of failed or dis-allowed states. (Default: None)
1682
- execution_delta : datetime.timedelta
1683
- time difference with the previous execution to look at,
1684
- the default is the same logical date as the current task or DAG. (Default: None)
1685
- check_existence: bool
1686
- Set to True to check if the external task exists or check if
1687
- the DAG to wait for exists. (Default: True)
1585
+ bucket_key : Union[str, List[str]]
1586
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1587
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1588
+ bucket_name : str
1589
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1590
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1591
+ wildcard_match : bool
1592
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1593
+ aws_conn_id : str
1594
+ a reference to the s3 connection on Airflow. (Default: None)
1595
+ verify : bool
1596
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1597
+ """
1598
+ ...
1599
+
1600
+ @typing.overload
1601
+ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1602
+ """
1603
+ Specifies the flow(s) that this flow depends on.
1604
+
1605
+ ```
1606
+ @trigger_on_finish(flow='FooFlow')
1607
+ ```
1608
+ or
1609
+ ```
1610
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1611
+ ```
1612
+ This decorator respects the @project decorator and triggers the flow
1613
+ when upstream runs within the same namespace complete successfully
1614
+
1615
+ Additionally, you can specify project aware upstream flow dependencies
1616
+ by specifying the fully qualified project_flow_name.
1617
+ ```
1618
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1619
+ ```
1620
+ or
1621
+ ```
1622
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1623
+ ```
1624
+
1625
+ You can also specify just the project or project branch (other values will be
1626
+ inferred from the current project or project branch):
1627
+ ```
1628
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1629
+ ```
1630
+
1631
+ Note that `branch` is typically one of:
1632
+ - `prod`
1633
+ - `user.bob`
1634
+ - `test.my_experiment`
1635
+ - `prod.staging`
1636
+
1637
+
1638
+ Parameters
1639
+ ----------
1640
+ flow : Union[str, Dict[str, str]], optional, default None
1641
+ Upstream flow dependency for this flow.
1642
+ flows : List[Union[str, Dict[str, str]]], default []
1643
+ Upstream flow dependencies for this flow.
1644
+ options : Dict[str, Any], default {}
1645
+ Backend-specific configuration for tuning eventing behavior.
1646
+ """
1647
+ ...
1648
+
1649
+ @typing.overload
1650
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1651
+ ...
1652
+
1653
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1654
+ """
1655
+ Specifies the flow(s) that this flow depends on.
1656
+
1657
+ ```
1658
+ @trigger_on_finish(flow='FooFlow')
1659
+ ```
1660
+ or
1661
+ ```
1662
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1663
+ ```
1664
+ This decorator respects the @project decorator and triggers the flow
1665
+ when upstream runs within the same namespace complete successfully
1666
+
1667
+ Additionally, you can specify project aware upstream flow dependencies
1668
+ by specifying the fully qualified project_flow_name.
1669
+ ```
1670
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1671
+ ```
1672
+ or
1673
+ ```
1674
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1675
+ ```
1676
+
1677
+ You can also specify just the project or project branch (other values will be
1678
+ inferred from the current project or project branch):
1679
+ ```
1680
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1681
+ ```
1682
+
1683
+ Note that `branch` is typically one of:
1684
+ - `prod`
1685
+ - `user.bob`
1686
+ - `test.my_experiment`
1687
+ - `prod.staging`
1688
+
1689
+
1690
+ Parameters
1691
+ ----------
1692
+ flow : Union[str, Dict[str, str]], optional, default None
1693
+ Upstream flow dependency for this flow.
1694
+ flows : List[Union[str, Dict[str, str]]], default []
1695
+ Upstream flow dependencies for this flow.
1696
+ options : Dict[str, Any], default {}
1697
+ Backend-specific configuration for tuning eventing behavior.
1688
1698
  """
1689
1699
  ...
1690
1700
 
@@ -1781,54 +1791,38 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1781
1791
  """
1782
1792
  ...
1783
1793
 
1784
- @typing.overload
1785
- def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1794
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1786
1795
  """
1787
- Specifies the Conda environment for all steps of the flow.
1796
+ Specifies what flows belong to the same project.
1788
1797
 
1789
- Use `@conda_base` to set common libraries required by all
1790
- steps and use `@conda` to specify step-specific additions.
1798
+ A project-specific namespace is created for all flows that
1799
+ use the same `@project(name)`.
1791
1800
 
1792
1801
 
1793
1802
  Parameters
1794
1803
  ----------
1795
- packages : Dict[str, str], default {}
1796
- Packages to use for this flow. The key is the name of the package
1797
- and the value is the version to use.
1798
- libraries : Dict[str, str], default {}
1799
- Supported for backward compatibility. When used with packages, packages will take precedence.
1800
- python : str, optional, default None
1801
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1802
- that the version used will correspond to the version of the Python interpreter used to start the run.
1803
- disabled : bool, default False
1804
- If set to True, disables Conda.
1805
- """
1806
- ...
1807
-
1808
- @typing.overload
1809
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1810
- ...
1811
-
1812
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1813
- """
1814
- Specifies the Conda environment for all steps of the flow.
1815
-
1816
- Use `@conda_base` to set common libraries required by all
1817
- steps and use `@conda` to specify step-specific additions.
1804
+ name : str
1805
+ Project name. Make sure that the name is unique amongst all
1806
+ projects that use the same production scheduler. The name may
1807
+ contain only lowercase alphanumeric characters and underscores.
1818
1808
 
1809
+ branch : Optional[str], default None
1810
+ The branch to use. If not specified, the branch is set to
1811
+ `user.<username>` unless `production` is set to `True`. This can
1812
+ also be set on the command line using `--branch` as a top-level option.
1813
+ It is an error to specify `branch` in the decorator and on the command line.
1819
1814
 
1820
- Parameters
1821
- ----------
1822
- packages : Dict[str, str], default {}
1823
- Packages to use for this flow. The key is the name of the package
1824
- and the value is the version to use.
1825
- libraries : Dict[str, str], default {}
1826
- Supported for backward compatibility. When used with packages, packages will take precedence.
1827
- python : str, optional, default None
1828
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1829
- that the version used will correspond to the version of the Python interpreter used to start the run.
1830
- disabled : bool, default False
1831
- If set to True, disables Conda.
1815
+ production : bool, default False
1816
+ Whether or not the branch is the production branch. This can also be set on the
1817
+ command line using `--production` as a top-level option. It is an error to specify
1818
+ `production` in the decorator and on the command line.
1819
+ The project branch name will be:
1820
+ - if `branch` is specified:
1821
+ - if `production` is True: `prod.<branch>`
1822
+ - if `production` is False: `test.<branch>`
1823
+ - if `branch` is not specified:
1824
+ - if `production` is True: `prod`
1825
+ - if `production` is False: `user.<username>`
1832
1826
  """
1833
1827
  ...
1834
1828