ob-metaflow-stubs 6.0.4.4rc0__py2.py3-none-any.whl → 6.0.4.5__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 +800 -806
  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 +8 -7
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  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 +4 -4
  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 +49 -49
  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 +3 -3
  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 +4 -4
  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 +3 -3
  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 +4 -4
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +5 -5
  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 +3 -3
  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 +3 -3
  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 +4 -4
  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 +3 -3
  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 +2 -2
  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 +3 -3
  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 +3 -3
  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 +4 -4
  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 +4 -4
  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 +2 -2
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +6 -6
  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 +5 -5
  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 +5 -5
  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 +4 -4
  118. metaflow-stubs/plugins/__init__.pyi +12 -12
  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 +7 -7
  130. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  131. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  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 +2 -2
  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 +2 -2
  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 +2 -2
  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 +5 -5
  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 +3 -3
  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 +2 -2
  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 +4 -5
  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 +2 -2
  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 +4 -4
  222. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  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 +6 -6
  230. metaflow-stubs/runner/deployer_impl.pyi +3 -3
  231. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  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 +8 -6
  245. {ob_metaflow_stubs-6.0.4.4rc0.dist-info → ob_metaflow_stubs-6.0.4.5.dist-info}/METADATA +1 -1
  246. ob_metaflow_stubs-6.0.4.5.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.5.dist-info}/WHEEL +0 -0
  261. {ob_metaflow_stubs-6.0.4.4rc0.dist-info → ob_metaflow_stubs-6.0.4.5.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.2+obcheckpoint(0.2.4);ob(v1) #
4
+ # Generated on 2025-07-16T21:13:36.500276 #
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,22 +33,20 @@ 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
40
+ from . import metaflow_git as metaflow_git
45
41
  from . import events as events
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
49
45
  from . import includefile as includefile
50
46
  from .includefile import IncludeFile as IncludeFile
51
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
52
47
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
53
48
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
49
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
54
50
  from . import client as client
55
51
  from .client.core import namespace as namespace
56
52
  from .client.core import get_namespace as get_namespace
@@ -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,460 +156,525 @@ 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 parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, 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
- )
175
-
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
-
182
- Valid model options
183
- -------------------
184
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
185
-
186
-
187
- Parameters
188
- ----------
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.
162
+ Decorator prototype for all step decorators. This function gets specialized
163
+ and imported for all decorators types by _import_plugin_decorators().
205
164
  """
206
165
  ...
207
166
 
208
167
  @typing.overload
209
- def retry(*, times: int = 3, minutes_between_retries: int = 2) -> 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]]]:
168
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
169
+ ...
170
+
171
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
210
172
  """
211
- Specifies the number of times the task corresponding
212
- to a step needs to be retried.
173
+ Decorator prototype for all step decorators. This function gets specialized
174
+ and imported for all decorators types by _import_plugin_decorators().
175
+ """
176
+ ...
177
+
178
+ @typing.overload
179
+ 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]]]:
180
+ """
181
+ Specifies a timeout for your step.
213
182
 
214
- This decorator is useful for handling transient errors, such as networking issues.
215
- If your task contains operations that can't be retried safely, e.g. database updates,
216
- it is advisable to annotate it with `@retry(times=0)`.
183
+ This decorator is useful if this step may hang indefinitely.
217
184
 
218
- This can be used in conjunction with the `@catch` decorator. The `@catch`
219
- decorator will execute a no-op task after all retries have been exhausted,
220
- ensuring that the flow execution can continue.
185
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
186
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
187
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
188
+
189
+ Note that all the values specified in parameters are added together so if you specify
190
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
221
191
 
222
192
 
223
193
  Parameters
224
194
  ----------
225
- times : int, default 3
226
- Number of times to retry this task.
227
- minutes_between_retries : int, default 2
228
- Number of minutes between retries.
195
+ seconds : int, default 0
196
+ Number of seconds to wait prior to timing out.
197
+ minutes : int, default 0
198
+ Number of minutes to wait prior to timing out.
199
+ hours : int, default 0
200
+ Number of hours to wait prior to timing out.
229
201
  """
230
202
  ...
231
203
 
232
204
  @typing.overload
233
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
205
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
234
206
  ...
235
207
 
236
208
  @typing.overload
237
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
238
- ...
239
-
240
- def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
241
- """
242
- Specifies the number of times the task corresponding
243
- to a step needs to be retried.
244
-
245
- This decorator is useful for handling transient errors, such as networking issues.
246
- If your task contains operations that can't be retried safely, e.g. database updates,
247
- it is advisable to annotate it with `@retry(times=0)`.
248
-
249
- This can be used in conjunction with the `@catch` decorator. The `@catch`
250
- decorator will execute a no-op task after all retries have been exhausted,
251
- ensuring that the flow execution can continue.
252
-
253
-
254
- Parameters
255
- ----------
256
- times : int, default 3
257
- Number of times to retry this task.
258
- minutes_between_retries : int, default 2
259
- Number of minutes between retries.
260
- """
209
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
261
210
  ...
262
211
 
263
- 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]]]:
212
+ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
264
213
  """
265
- This decorator is used to run vllm APIs as Metaflow task sidecars.
266
-
267
- User code call
268
- --------------
269
- @vllm(
270
- model="...",
271
- ...
272
- )
214
+ Specifies a timeout for your step.
273
215
 
274
- Valid backend options
275
- ---------------------
276
- - 'local': Run as a separate process on the local task machine.
216
+ This decorator is useful if this step may hang indefinitely.
277
217
 
278
- Valid model options
279
- -------------------
280
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
218
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
219
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
220
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
281
221
 
282
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
283
- If you need multiple models, you must create multiple @vllm decorators.
222
+ Note that all the values specified in parameters are added together so if you specify
223
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
284
224
 
285
225
 
286
226
  Parameters
287
227
  ----------
288
- model: str
289
- HuggingFace model identifier to be served by vLLM.
290
- backend: str
291
- Determines where and how to run the vLLM process.
292
- openai_api_server: bool
293
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
294
- Default is False (uses native engine).
295
- Set to True for backward compatibility with existing code.
296
- debug: bool
297
- Whether to turn on verbose debugging logs.
298
- card_refresh_interval: int
299
- Interval in seconds for refreshing the vLLM status card.
300
- Only used when openai_api_server=True.
301
- max_retries: int
302
- Maximum number of retries checking for vLLM server startup.
303
- Only used when openai_api_server=True.
304
- retry_alert_frequency: int
305
- Frequency of alert logs for vLLM server startup retries.
306
- Only used when openai_api_server=True.
307
- engine_args : dict
308
- Additional keyword arguments to pass to the vLLM engine.
309
- For example, `tensor_parallel_size=2`.
228
+ seconds : int, default 0
229
+ Number of seconds to wait prior to timing out.
230
+ minutes : int, default 0
231
+ Number of minutes to wait prior to timing out.
232
+ hours : int, default 0
233
+ Number of hours to wait prior to timing out.
310
234
  """
311
235
  ...
312
236
 
313
237
  @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]]]:
238
+ 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]]]:
315
239
  """
316
- Specifies the Conda environment for the step.
240
+ Specifies the resources needed when executing this step.
317
241
 
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.
242
+ Use `@resources` to specify the resource requirements
243
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
244
+
245
+ You can choose the compute layer on the command line by executing e.g.
246
+ ```
247
+ python myflow.py run --with batch
248
+ ```
249
+ or
250
+ ```
251
+ python myflow.py run --with kubernetes
252
+ ```
253
+ which executes the flow on the desired system using the
254
+ requirements specified in `@resources`.
322
255
 
323
256
 
324
257
  Parameters
325
258
  ----------
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.
259
+ cpu : int, default 1
260
+ Number of CPUs required for this step.
261
+ gpu : int, optional, default None
262
+ Number of GPUs required for this step.
263
+ disk : int, optional, default None
264
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
265
+ memory : int, default 4096
266
+ Memory size (in MB) required for this step.
267
+ shared_memory : int, optional, default None
268
+ The value for the size (in MiB) of the /dev/shm volume for this step.
269
+ This parameter maps to the `--shm-size` option in Docker.
336
270
  """
337
271
  ...
338
272
 
339
273
  @typing.overload
340
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
274
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
341
275
  ...
342
276
 
343
277
  @typing.overload
344
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
278
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
345
279
  ...
346
280
 
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):
281
+ 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):
348
282
  """
349
- Specifies the Conda environment for the step.
283
+ Specifies the resources needed when executing this step.
350
284
 
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.
285
+ Use `@resources` to specify the resource requirements
286
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
287
+
288
+ You can choose the compute layer on the command line by executing e.g.
289
+ ```
290
+ python myflow.py run --with batch
291
+ ```
292
+ or
293
+ ```
294
+ python myflow.py run --with kubernetes
295
+ ```
296
+ which executes the flow on the desired system using the
297
+ requirements specified in `@resources`.
355
298
 
356
299
 
357
300
  Parameters
358
301
  ----------
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().
302
+ cpu : int, default 1
303
+ Number of CPUs required for this step.
304
+ gpu : int, optional, default None
305
+ Number of GPUs required for this step.
306
+ disk : int, optional, default None
307
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
308
+ memory : int, default 4096
309
+ Memory size (in MB) required for this step.
310
+ shared_memory : int, optional, default None
311
+ The value for the size (in MiB) of the /dev/shm volume for this step.
312
+ This parameter maps to the `--shm-size` option in Docker.
388
313
  """
389
314
  ...
390
315
 
391
- @typing.overload
392
- 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]]]:
316
+ 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]]]:
393
317
  """
394
- Enables loading / saving of models within a step.
318
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
395
319
 
396
320
  > Examples
397
- - Saving Models
321
+
322
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
398
323
  ```python
399
- @model
400
- @step
401
- def train(self):
402
- # current.model.save returns a dictionary reference to the model saved
403
- self.my_model = current.model.save(
404
- path_to_my_model,
405
- label="my_model",
406
- metadata={
407
- "epochs": 10,
408
- "batch-size": 32,
409
- "learning-rate": 0.001,
410
- }
411
- )
412
- self.next(self.test)
324
+ @huggingface_hub
325
+ @step
326
+ def pull_model_from_huggingface(self):
327
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
328
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
329
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
330
+ # value of the function is a reference to the model in the backend storage.
331
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
413
332
 
414
- @model(load="my_model")
415
- @step
416
- def test(self):
417
- # `current.model.loaded` returns a dictionary of the loaded models
418
- # where the key is the name of the artifact and the value is the path to the model
419
- print(os.listdir(current.model.loaded["my_model"]))
420
- self.next(self.end)
333
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
334
+ self.llama_model = current.huggingface_hub.snapshot_download(
335
+ repo_id=self.model_id,
336
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
337
+ )
338
+ self.next(self.train)
421
339
  ```
422
340
 
423
- - Loading models
341
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
424
342
  ```python
425
- @step
426
- def train(self):
427
- # current.model.load returns the path to the model loaded
428
- checkpoint_path = current.model.load(
429
- self.checkpoint_key,
430
- )
431
- model_path = current.model.load(
432
- self.model,
433
- )
434
- self.next(self.test)
343
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
344
+ @step
345
+ def pull_model_from_huggingface(self):
346
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
347
+ ```
348
+
349
+ ```python
350
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
351
+ @step
352
+ def finetune_model(self):
353
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
354
+ # path_to_model will be /my-directory
355
+ ```
356
+
357
+ ```python
358
+ # Takes all the arguments passed to `snapshot_download`
359
+ # except for `local_dir`
360
+ @huggingface_hub(load=[
361
+ {
362
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
363
+ },
364
+ {
365
+ "repo_id": "myorg/mistral-lora",
366
+ "repo_type": "model",
367
+ },
368
+ ])
369
+ @step
370
+ def finetune_model(self):
371
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
372
+ # path_to_model will be /my-directory
435
373
  ```
436
374
 
437
375
 
438
376
  Parameters
439
377
  ----------
440
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
441
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
442
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
443
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
444
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
445
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
378
+ temp_dir_root : str, optional
379
+ The root directory that will hold the temporary directory where objects will be downloaded.
446
380
 
447
- temp_dir_root : str, default: None
448
- The root directory under which `current.model.loaded` will store loaded models
381
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
382
+ The list of repos (models/datasets) to load.
383
+
384
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
385
+
386
+ - If repo (model/dataset) is not found in the datastore:
387
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
388
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
389
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
390
+
391
+ - If repo is found in the datastore:
392
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
449
393
  """
450
394
  ...
451
395
 
452
396
  @typing.overload
453
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
397
+ 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]]]:
398
+ """
399
+ Creates a human-readable report, a Metaflow Card, after this step completes.
400
+
401
+ Note that you may add multiple `@card` decorators in a step with different parameters.
402
+
403
+
404
+ Parameters
405
+ ----------
406
+ type : str, default 'default'
407
+ Card type.
408
+ id : str, optional, default None
409
+ If multiple cards are present, use this id to identify this card.
410
+ options : Dict[str, Any], default {}
411
+ Options passed to the card. The contents depend on the card type.
412
+ timeout : int, default 45
413
+ Interrupt reporting if it takes more than this many seconds.
414
+ """
454
415
  ...
455
416
 
456
417
  @typing.overload
457
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
418
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
458
419
  ...
459
420
 
460
- def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
421
+ @typing.overload
422
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
423
+ ...
424
+
425
+ 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):
461
426
  """
462
- Enables loading / saving of models within a step.
427
+ Creates a human-readable report, a Metaflow Card, after this step completes.
463
428
 
464
- > Examples
465
- - Saving Models
466
- ```python
467
- @model
468
- @step
469
- def train(self):
470
- # current.model.save returns a dictionary reference to the model saved
471
- self.my_model = current.model.save(
472
- path_to_my_model,
473
- label="my_model",
474
- metadata={
475
- "epochs": 10,
476
- "batch-size": 32,
477
- "learning-rate": 0.001,
478
- }
479
- )
480
- self.next(self.test)
429
+ Note that you may add multiple `@card` decorators in a step with different parameters.
481
430
 
482
- @model(load="my_model")
483
- @step
484
- def test(self):
485
- # `current.model.loaded` returns a dictionary of the loaded models
486
- # where the key is the name of the artifact and the value is the path to the model
487
- print(os.listdir(current.model.loaded["my_model"]))
488
- self.next(self.end)
489
- ```
490
431
 
491
- - Loading models
492
- ```python
493
- @step
494
- def train(self):
495
- # current.model.load returns the path to the model loaded
496
- checkpoint_path = current.model.load(
497
- self.checkpoint_key,
498
- )
499
- model_path = current.model.load(
500
- self.model,
501
- )
502
- self.next(self.test)
503
- ```
432
+ Parameters
433
+ ----------
434
+ type : str, default 'default'
435
+ Card type.
436
+ id : str, optional, default None
437
+ If multiple cards are present, use this id to identify this card.
438
+ options : Dict[str, Any], default {}
439
+ Options passed to the card. The contents depend on the card type.
440
+ timeout : int, default 45
441
+ Interrupt reporting if it takes more than this many seconds.
442
+ """
443
+ ...
444
+
445
+ 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]]]:
446
+ """
447
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
448
+
449
+ User code call
450
+ --------------
451
+ @vllm(
452
+ model="...",
453
+ ...
454
+ )
455
+
456
+ Valid backend options
457
+ ---------------------
458
+ - 'local': Run as a separate process on the local task machine.
459
+
460
+ Valid model options
461
+ -------------------
462
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
463
+
464
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
465
+ If you need multiple models, you must create multiple @vllm decorators.
504
466
 
505
467
 
506
468
  Parameters
507
469
  ----------
508
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
509
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
510
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
511
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
512
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
513
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
514
-
515
- temp_dir_root : str, default: None
516
- The root directory under which `current.model.loaded` will store loaded models
470
+ model: str
471
+ HuggingFace model identifier to be served by vLLM.
472
+ backend: str
473
+ Determines where and how to run the vLLM process.
474
+ openai_api_server: bool
475
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
476
+ Default is False (uses native engine).
477
+ Set to True for backward compatibility with existing code.
478
+ debug: bool
479
+ Whether to turn on verbose debugging logs.
480
+ card_refresh_interval: int
481
+ Interval in seconds for refreshing the vLLM status card.
482
+ Only used when openai_api_server=True.
483
+ max_retries: int
484
+ Maximum number of retries checking for vLLM server startup.
485
+ Only used when openai_api_server=True.
486
+ retry_alert_frequency: int
487
+ Frequency of alert logs for vLLM server startup retries.
488
+ Only used when openai_api_server=True.
489
+ engine_args : dict
490
+ Additional keyword arguments to pass to the vLLM engine.
491
+ For example, `tensor_parallel_size=2`.
517
492
  """
518
493
  ...
519
494
 
520
495
  @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]]]:
496
+ 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]]]:
522
497
  """
523
- Specifies that the step will success under all circumstances.
498
+ Enables checkpointing for a step.
524
499
 
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.
500
+ > Examples
501
+
502
+ - Saving Checkpoints
503
+
504
+ ```python
505
+ @checkpoint
506
+ @step
507
+ def train(self):
508
+ model = create_model(self.parameters, checkpoint_path = None)
509
+ for i in range(self.epochs):
510
+ # some training logic
511
+ loss = model.train(self.dataset)
512
+ if i % 10 == 0:
513
+ model.save(
514
+ current.checkpoint.directory,
515
+ )
516
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
517
+ # and returns a reference dictionary to the checkpoint saved in the datastore
518
+ self.latest_checkpoint = current.checkpoint.save(
519
+ name="epoch_checkpoint",
520
+ metadata={
521
+ "epoch": i,
522
+ "loss": loss,
523
+ }
524
+ )
525
+ ```
526
+
527
+ - Using Loaded Checkpoints
528
+
529
+ ```python
530
+ @retry(times=3)
531
+ @checkpoint
532
+ @step
533
+ def train(self):
534
+ # Assume that the task has restarted and the previous attempt of the task
535
+ # saved a checkpoint
536
+ checkpoint_path = None
537
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
538
+ print("Loaded checkpoint from the previous attempt")
539
+ checkpoint_path = current.checkpoint.directory
540
+
541
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
542
+ for i in range(self.epochs):
543
+ ...
544
+ ```
529
545
 
530
546
 
531
547
  Parameters
532
548
  ----------
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.
549
+ load_policy : str, default: "fresh"
550
+ The policy for loading the checkpoint. The following policies are supported:
551
+ - "eager": Loads the the latest available checkpoint within the namespace.
552
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
553
+ will be loaded at the start of the task.
554
+ - "none": Do not load any checkpoint
555
+ - "fresh": Loads the lastest checkpoint created within the running Task.
556
+ This mode helps loading checkpoints across various retry attempts of the same task.
557
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
558
+ created within the task will be loaded when the task is retries execution on failure.
559
+
560
+ temp_dir_root : str, default: None
561
+ The root directory under which `current.checkpoint.directory` will be created.
539
562
  """
540
563
  ...
541
564
 
542
565
  @typing.overload
543
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
566
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
544
567
  ...
545
568
 
546
569
  @typing.overload
547
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
570
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
548
571
  ...
549
572
 
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):
573
+ 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):
551
574
  """
552
- Specifies that the step will success under all circumstances.
575
+ Enables checkpointing for a step.
553
576
 
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.
577
+ > Examples
578
+
579
+ - Saving Checkpoints
580
+
581
+ ```python
582
+ @checkpoint
583
+ @step
584
+ def train(self):
585
+ model = create_model(self.parameters, checkpoint_path = None)
586
+ for i in range(self.epochs):
587
+ # some training logic
588
+ loss = model.train(self.dataset)
589
+ if i % 10 == 0:
590
+ model.save(
591
+ current.checkpoint.directory,
592
+ )
593
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
594
+ # and returns a reference dictionary to the checkpoint saved in the datastore
595
+ self.latest_checkpoint = current.checkpoint.save(
596
+ name="epoch_checkpoint",
597
+ metadata={
598
+ "epoch": i,
599
+ "loss": loss,
600
+ }
601
+ )
602
+ ```
603
+
604
+ - Using Loaded Checkpoints
605
+
606
+ ```python
607
+ @retry(times=3)
608
+ @checkpoint
609
+ @step
610
+ def train(self):
611
+ # Assume that the task has restarted and the previous attempt of the task
612
+ # saved a checkpoint
613
+ checkpoint_path = None
614
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
615
+ print("Loaded checkpoint from the previous attempt")
616
+ checkpoint_path = current.checkpoint.directory
617
+
618
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
619
+ for i in range(self.epochs):
620
+ ...
621
+ ```
558
622
 
559
623
 
560
624
  Parameters
561
625
  ----------
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.
626
+ load_policy : str, default: "fresh"
627
+ The policy for loading the checkpoint. The following policies are supported:
628
+ - "eager": Loads the the latest available checkpoint within the namespace.
629
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
630
+ will be loaded at the start of the task.
631
+ - "none": Do not load any checkpoint
632
+ - "fresh": Loads the lastest checkpoint created within the running Task.
633
+ This mode helps loading checkpoints across various retry attempts of the same task.
634
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
635
+ created within the task will be loaded when the task is retries execution on failure.
636
+
637
+ temp_dir_root : str, default: None
638
+ The root directory under which `current.checkpoint.directory` will be created.
568
639
  """
569
640
  ...
570
641
 
571
642
  @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]]]:
643
+ 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]]]:
573
644
  """
574
- Specifies the PyPI packages for the step.
575
-
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.
645
+ Specifies secrets to be retrieved and injected as environment variables prior to
646
+ the execution of a step.
580
647
 
581
648
 
582
649
  Parameters
583
650
  ----------
584
- packages : Dict[str, str], default: {}
585
- Packages to use for this step. The key is the name of the package
586
- and the value is the version to use.
587
- python : str, optional, default: None
588
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
589
- that the version used will correspond to the version of the Python interpreter used to start the run.
651
+ sources : List[Union[str, Dict[str, Any]]], default: []
652
+ List of secret specs, defining how the secrets are to be retrieved
653
+ role : str, optional, default: None
654
+ Role to use for fetching secrets
590
655
  """
591
656
  ...
592
657
 
593
658
  @typing.overload
594
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
659
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
595
660
  ...
596
661
 
597
662
  @typing.overload
598
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
663
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
599
664
  ...
600
665
 
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):
666
+ 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):
602
667
  """
603
- Specifies the PyPI packages for the step.
604
-
605
- 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.
668
+ Specifies secrets to be retrieved and injected as environment variables prior to
669
+ the execution of a step.
609
670
 
610
671
 
611
672
  Parameters
612
673
  ----------
613
- packages : Dict[str, str], default: {}
614
- Packages to use for this step. The key is the name of the package
615
- and the value is the version to use.
616
- python : str, optional, default: None
617
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
618
- that the version used will correspond to the version of the Python interpreter used to start the run.
674
+ sources : List[Union[str, Dict[str, Any]]], default: []
675
+ List of secret specs, defining how the secrets are to be retrieved
676
+ role : str, optional, default: None
677
+ Role to use for fetching secrets
619
678
  """
620
679
  ...
621
680
 
@@ -709,554 +768,489 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
709
768
  ...
710
769
 
711
770
  @typing.overload
712
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
713
- """
714
- Decorator prototype for all step decorators. This function gets specialized
715
- and imported for all decorators types by _import_plugin_decorators().
716
- """
717
- ...
718
-
719
- @typing.overload
720
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
721
- ...
722
-
723
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
724
- """
725
- Decorator prototype for all step decorators. This function gets specialized
726
- and imported for all decorators types by _import_plugin_decorators().
727
- """
728
- ...
729
-
730
- def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> 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]]]:
771
+ 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]]]:
731
772
  """
732
- Specifies that this step should execute on DGX cloud.
733
-
773
+ Specifies the Conda environment for the step.
734
774
 
735
- Parameters
736
- ----------
737
- gpu : int
738
- Number of GPUs to use.
739
- gpu_type : str
740
- Type of Nvidia GPU to use.
741
- queue_timeout : int
742
- Time to keep the job in NVCF's queue.
743
- """
744
- ...
745
-
746
- @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]]]:
748
- """
749
- Specifies environment variables to be set prior to the execution of a step.
775
+ Information in this decorator will augment any
776
+ attributes set in the `@conda_base` flow-level decorator. Hence,
777
+ you can use `@conda_base` to set packages required by all
778
+ steps and use `@conda` to specify step-specific overrides.
750
779
 
751
780
 
752
781
  Parameters
753
782
  ----------
754
- vars : Dict[str, str], default {}
755
- Dictionary of environment variables to set.
783
+ packages : Dict[str, str], default {}
784
+ Packages to use for this step. The key is the name of the package
785
+ and the value is the version to use.
786
+ libraries : Dict[str, str], default {}
787
+ Supported for backward compatibility. When used with packages, packages will take precedence.
788
+ python : str, optional, default None
789
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
790
+ that the version used will correspond to the version of the Python interpreter used to start the run.
791
+ disabled : bool, default False
792
+ If set to True, disables @conda.
756
793
  """
757
794
  ...
758
795
 
759
796
  @typing.overload
760
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
797
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
761
798
  ...
762
799
 
763
800
  @typing.overload
764
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
801
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
765
802
  ...
766
803
 
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] = {}):
804
+ 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):
768
805
  """
769
- Specifies environment variables to be set prior to the execution of a step.
806
+ Specifies the Conda environment for the step.
807
+
808
+ Information in this decorator will augment any
809
+ attributes set in the `@conda_base` flow-level decorator. Hence,
810
+ you can use `@conda_base` to set packages required by all
811
+ steps and use `@conda` to specify step-specific overrides.
770
812
 
771
813
 
772
814
  Parameters
773
815
  ----------
774
- vars : Dict[str, str], default {}
775
- Dictionary of environment variables to set.
816
+ packages : Dict[str, str], default {}
817
+ Packages to use for this step. The key is the name of the package
818
+ and the value is the version to use.
819
+ libraries : Dict[str, str], default {}
820
+ Supported for backward compatibility. When used with packages, packages will take precedence.
821
+ python : str, optional, default None
822
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
823
+ that the version used will correspond to the version of the Python interpreter used to start the run.
824
+ disabled : bool, default False
825
+ If set to True, disables @conda.
776
826
  """
777
827
  ...
778
828
 
779
829
  @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]]]:
830
+ 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]]]:
781
831
  """
782
- Enables checkpointing for a step.
832
+ Enables loading / saving of models within a step.
783
833
 
784
834
  > Examples
785
-
786
- - Saving Checkpoints
787
-
835
+ - Saving Models
788
836
  ```python
789
- @checkpoint
837
+ @model
790
838
  @step
791
839
  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
- ```
840
+ # current.model.save returns a dictionary reference to the model saved
841
+ self.my_model = current.model.save(
842
+ path_to_my_model,
843
+ label="my_model",
844
+ metadata={
845
+ "epochs": 10,
846
+ "batch-size": 32,
847
+ "learning-rate": 0.001,
848
+ }
849
+ )
850
+ self.next(self.test)
810
851
 
811
- - Using Loaded Checkpoints
852
+ @model(load="my_model")
853
+ @step
854
+ def test(self):
855
+ # `current.model.loaded` returns a dictionary of the loaded models
856
+ # where the key is the name of the artifact and the value is the path to the model
857
+ print(os.listdir(current.model.loaded["my_model"]))
858
+ self.next(self.end)
859
+ ```
812
860
 
861
+ - Loading models
813
862
  ```python
814
- @retry(times=3)
815
- @checkpoint
816
863
  @step
817
864
  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
- ...
865
+ # current.model.load returns the path to the model loaded
866
+ checkpoint_path = current.model.load(
867
+ self.checkpoint_key,
868
+ )
869
+ model_path = current.model.load(
870
+ self.model,
871
+ )
872
+ self.next(self.test)
828
873
  ```
829
874
 
830
875
 
831
876
  Parameters
832
877
  ----------
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.
878
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
879
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
880
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
881
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
882
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
883
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
843
884
 
844
885
  temp_dir_root : str, default: None
845
- The root directory under which `current.checkpoint.directory` will be created.
886
+ The root directory under which `current.model.loaded` will store loaded models
846
887
  """
847
888
  ...
848
889
 
849
890
  @typing.overload
850
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
891
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
851
892
  ...
852
893
 
853
894
  @typing.overload
854
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
895
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
855
896
  ...
856
897
 
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):
898
+ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
858
899
  """
859
- Enables checkpointing for a step.
900
+ Enables loading / saving of models within a step.
860
901
 
861
902
  > Examples
862
-
863
- - Saving Checkpoints
864
-
903
+ - Saving Models
865
904
  ```python
866
- @checkpoint
905
+ @model
867
906
  @step
868
907
  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
- ```
908
+ # current.model.save returns a dictionary reference to the model saved
909
+ self.my_model = current.model.save(
910
+ path_to_my_model,
911
+ label="my_model",
912
+ metadata={
913
+ "epochs": 10,
914
+ "batch-size": 32,
915
+ "learning-rate": 0.001,
916
+ }
917
+ )
918
+ self.next(self.test)
887
919
 
888
- - Using Loaded Checkpoints
920
+ @model(load="my_model")
921
+ @step
922
+ def test(self):
923
+ # `current.model.loaded` returns a dictionary of the loaded models
924
+ # where the key is the name of the artifact and the value is the path to the model
925
+ print(os.listdir(current.model.loaded["my_model"]))
926
+ self.next(self.end)
927
+ ```
889
928
 
929
+ - Loading models
890
930
  ```python
891
- @retry(times=3)
892
- @checkpoint
893
931
  @step
894
932
  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
- ...
933
+ # current.model.load returns the path to the model loaded
934
+ checkpoint_path = current.model.load(
935
+ self.checkpoint_key,
936
+ )
937
+ model_path = current.model.load(
938
+ self.model,
939
+ )
940
+ self.next(self.test)
905
941
  ```
906
942
 
907
943
 
908
944
  Parameters
909
945
  ----------
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.
946
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
947
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
948
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
949
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
950
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
951
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
920
952
 
921
953
  temp_dir_root : str, default: None
922
- The root directory under which `current.checkpoint.directory` will be created.
954
+ The root directory under which `current.model.loaded` will store loaded models
923
955
  """
924
956
  ...
925
957
 
926
958
  @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]]]:
959
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
928
960
  """
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
961
+ Internal decorator to support Fast bakery
939
962
  """
940
963
  ...
941
964
 
942
965
  @typing.overload
943
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
966
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
944
967
  ...
945
968
 
946
- @typing.overload
947
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
969
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
970
+ """
971
+ Internal decorator to support Fast bakery
972
+ """
948
973
  ...
949
974
 
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):
975
+ 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]]]:
951
976
  """
952
- Specifies secrets to be retrieved and injected as environment variables prior to
953
- the execution of a step.
977
+ Specifies that this step should execute on DGX cloud.
954
978
 
955
979
 
956
980
  Parameters
957
981
  ----------
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
982
+ gpu : int
983
+ Number of GPUs to use.
984
+ gpu_type : str
985
+ Type of Nvidia GPU to use.
962
986
  """
963
987
  ...
964
988
 
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]]]:
989
+ @typing.overload
990
+ 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]]]:
966
991
  """
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)**
991
- ```python
992
- @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
993
- @step
994
- def pull_model_from_huggingface(self):
995
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
996
- ```
997
-
998
- ```python
999
- @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
1000
- @step
1001
- def finetune_model(self):
1002
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1003
- # path_to_model will be /my-directory
1004
- ```
992
+ Specifies the PyPI packages for the step.
1005
993
 
1006
- ```python
1007
- # Takes all the arguments passed to `snapshot_download`
1008
- # except for `local_dir`
1009
- @huggingface_hub(load=[
1010
- {
1011
- "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
1012
- },
1013
- {
1014
- "repo_id": "myorg/mistral-lora",
1015
- "repo_type": "model",
1016
- },
1017
- ])
1018
- @step
1019
- def finetune_model(self):
1020
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1021
- # path_to_model will be /my-directory
1022
- ```
994
+ Information in this decorator will augment any
995
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
996
+ you can use `@pypi_base` to set packages required by all
997
+ steps and use `@pypi` to specify step-specific overrides.
1023
998
 
1024
999
 
1025
1000
  Parameters
1026
1001
  ----------
1027
- temp_dir_root : str, optional
1028
- The root directory that will hold the temporary directory where objects will be downloaded.
1029
-
1030
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
1031
- The list of repos (models/datasets) to load.
1032
-
1033
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
1034
-
1035
- - If repo (model/dataset) is not found in the datastore:
1036
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
1037
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
1038
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
1039
-
1040
- - If repo is found in the datastore:
1041
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
1002
+ packages : Dict[str, str], default: {}
1003
+ Packages to use for this step. The key is the name of the package
1004
+ and the value is the version to use.
1005
+ python : str, optional, default: None
1006
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1007
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1042
1008
  """
1043
1009
  ...
1044
1010
 
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]]]:
1011
+ @typing.overload
1012
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1013
+ ...
1014
+
1015
+ @typing.overload
1016
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1017
+ ...
1018
+
1019
+ 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):
1046
1020
  """
1047
- Specifies that this step should execute on DGX cloud.
1021
+ Specifies the PyPI packages for the step.
1022
+
1023
+ Information in this decorator will augment any
1024
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1025
+ you can use `@pypi_base` to set packages required by all
1026
+ steps and use `@pypi` to specify step-specific overrides.
1048
1027
 
1049
1028
 
1050
1029
  Parameters
1051
1030
  ----------
1052
- gpu : int
1053
- Number of GPUs to use.
1054
- gpu_type : str
1055
- Type of Nvidia GPU to use.
1031
+ packages : Dict[str, str], default: {}
1032
+ Packages to use for this step. The key is the name of the package
1033
+ and the value is the version to use.
1034
+ python : str, optional, default: None
1035
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1036
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1056
1037
  """
1057
1038
  ...
1058
1039
 
1059
1040
  @typing.overload
1060
- 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]]]:
1041
+ def retry(*, times: int = 3, minutes_between_retries: int = 2) -> 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
1042
  """
1062
- Specifies a timeout for your step.
1063
-
1064
- This decorator is useful if this step may hang indefinitely.
1043
+ Specifies the number of times the task corresponding
1044
+ to a step needs to be retried.
1065
1045
 
1066
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1067
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1068
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1046
+ This decorator is useful for handling transient errors, such as networking issues.
1047
+ If your task contains operations that can't be retried safely, e.g. database updates,
1048
+ it is advisable to annotate it with `@retry(times=0)`.
1069
1049
 
1070
- Note that all the values specified in parameters are added together so if you specify
1071
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1050
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1051
+ decorator will execute a no-op task after all retries have been exhausted,
1052
+ ensuring that the flow execution can continue.
1072
1053
 
1073
1054
 
1074
1055
  Parameters
1075
1056
  ----------
1076
- seconds : int, default 0
1077
- Number of seconds to wait prior to timing out.
1078
- minutes : int, default 0
1079
- Number of minutes to wait prior to timing out.
1080
- hours : int, default 0
1081
- Number of hours to wait prior to timing out.
1057
+ times : int, default 3
1058
+ Number of times to retry this task.
1059
+ minutes_between_retries : int, default 2
1060
+ Number of minutes between retries.
1082
1061
  """
1083
1062
  ...
1084
1063
 
1085
1064
  @typing.overload
1086
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1065
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1087
1066
  ...
1088
1067
 
1089
1068
  @typing.overload
1090
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1069
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1091
1070
  ...
1092
1071
 
1093
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
1072
+ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
1094
1073
  """
1095
- Specifies a timeout for your step.
1096
-
1097
- This decorator is useful if this step may hang indefinitely.
1074
+ Specifies the number of times the task corresponding
1075
+ to a step needs to be retried.
1098
1076
 
1099
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1100
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1101
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1077
+ This decorator is useful for handling transient errors, such as networking issues.
1078
+ If your task contains operations that can't be retried safely, e.g. database updates,
1079
+ it is advisable to annotate it with `@retry(times=0)`.
1102
1080
 
1103
- Note that all the values specified in parameters are added together so if you specify
1104
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1081
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1082
+ decorator will execute a no-op task after all retries have been exhausted,
1083
+ ensuring that the flow execution can continue.
1105
1084
 
1106
1085
 
1107
1086
  Parameters
1108
1087
  ----------
1109
- seconds : int, default 0
1110
- Number of seconds to wait prior to timing out.
1111
- minutes : int, default 0
1112
- Number of minutes to wait prior to timing out.
1113
- hours : int, default 0
1114
- Number of hours to wait prior to timing out.
1088
+ times : int, default 3
1089
+ Number of times to retry this task.
1090
+ minutes_between_retries : int, default 2
1091
+ Number of minutes between retries.
1115
1092
  """
1116
1093
  ...
1117
1094
 
1118
1095
  @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]]]:
1096
+ 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]]]:
1120
1097
  """
1121
- Creates a human-readable report, a Metaflow Card, after this step completes.
1122
-
1123
- Note that you may add multiple `@card` decorators in a step with different parameters.
1098
+ Specifies environment variables to be set prior to the execution of a step.
1124
1099
 
1125
1100
 
1126
1101
  Parameters
1127
1102
  ----------
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.
1103
+ vars : Dict[str, str], default {}
1104
+ Dictionary of environment variables to set.
1136
1105
  """
1137
1106
  ...
1138
1107
 
1139
1108
  @typing.overload
1140
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1109
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1141
1110
  ...
1142
1111
 
1143
1112
  @typing.overload
1144
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1113
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1145
1114
  ...
1146
1115
 
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):
1116
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1148
1117
  """
1149
- Creates a human-readable report, a Metaflow Card, after this step completes.
1118
+ Specifies environment variables to be set prior to the execution of a step.
1150
1119
 
1151
- Note that you may add multiple `@card` decorators in a step with different parameters.
1120
+
1121
+ Parameters
1122
+ ----------
1123
+ vars : Dict[str, str], default {}
1124
+ Dictionary of environment variables to set.
1125
+ """
1126
+ ...
1127
+
1128
+ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> 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]]]:
1129
+ """
1130
+ Specifies that this step should execute on DGX cloud.
1152
1131
 
1153
1132
 
1154
1133
  Parameters
1155
1134
  ----------
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.
1135
+ gpu : int
1136
+ Number of GPUs to use.
1137
+ gpu_type : str
1138
+ Type of Nvidia GPU to use.
1139
+ queue_timeout : int
1140
+ Time to keep the job in NVCF's queue.
1141
+ """
1142
+ ...
1143
+
1144
+ 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]]]:
1145
+ """
1146
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
1147
+
1148
+ User code call
1149
+ --------------
1150
+ @ollama(
1151
+ models=[...],
1152
+ ...
1153
+ )
1154
+
1155
+ Valid backend options
1156
+ ---------------------
1157
+ - 'local': Run as a separate process on the local task machine.
1158
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
1159
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
1160
+
1161
+ Valid model options
1162
+ -------------------
1163
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
1164
+
1165
+
1166
+ Parameters
1167
+ ----------
1168
+ models: list[str]
1169
+ List of Ollama containers running models in sidecars.
1170
+ backend: str
1171
+ Determines where and how to run the Ollama process.
1172
+ force_pull: bool
1173
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
1174
+ cache_update_policy: str
1175
+ Cache update policy: "auto", "force", or "never".
1176
+ force_cache_update: bool
1177
+ Simple override for "force" cache update policy.
1178
+ debug: bool
1179
+ Whether to turn on verbose debugging logs.
1180
+ circuit_breaker_config: dict
1181
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
1182
+ timeout_config: dict
1183
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
1164
1184
  """
1165
1185
  ...
1166
1186
 
1167
1187
  @typing.overload
1168
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1188
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1169
1189
  """
1170
- Internal decorator to support Fast bakery
1190
+ Decorator prototype for all step decorators. This function gets specialized
1191
+ and imported for all decorators types by _import_plugin_decorators().
1171
1192
  """
1172
1193
  ...
1173
1194
 
1174
1195
  @typing.overload
1175
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1196
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1176
1197
  ...
1177
1198
 
1178
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1199
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1179
1200
  """
1180
- Internal decorator to support Fast bakery
1201
+ Decorator prototype for all step decorators. This function gets specialized
1202
+ and imported for all decorators types by _import_plugin_decorators().
1181
1203
  """
1182
1204
  ...
1183
1205
 
1184
1206
  @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]]]:
1207
+ 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]]]:
1186
1208
  """
1187
- Specifies the resources needed when executing this step.
1188
-
1189
- Use `@resources` to specify the resource requirements
1190
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1209
+ Specifies that the step will success under all circumstances.
1191
1210
 
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`.
1211
+ The decorator will create an optional artifact, specified by `var`, which
1212
+ contains the exception raised. You can use it to detect the presence
1213
+ of errors, indicating that all happy-path artifacts produced by the step
1214
+ are missing.
1202
1215
 
1203
1216
 
1204
1217
  Parameters
1205
1218
  ----------
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.
1219
+ var : str, optional, default None
1220
+ Name of the artifact in which to store the caught exception.
1221
+ If not specified, the exception is not stored.
1222
+ print_exception : bool, default True
1223
+ Determines whether or not the exception is printed to
1224
+ stdout when caught.
1217
1225
  """
1218
1226
  ...
1219
1227
 
1220
1228
  @typing.overload
1221
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1229
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1222
1230
  ...
1223
1231
 
1224
1232
  @typing.overload
1225
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1233
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1226
1234
  ...
1227
1235
 
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):
1236
+ 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):
1229
1237
  """
1230
- Specifies the resources needed when executing this step.
1231
-
1232
- Use `@resources` to specify the resource requirements
1233
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1238
+ Specifies that the step will success under all circumstances.
1234
1239
 
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
1242
- ```
1243
- which executes the flow on the desired system using the
1244
- requirements specified in `@resources`.
1240
+ The decorator will create an optional artifact, specified by `var`, which
1241
+ contains the exception raised. You can use it to detect the presence
1242
+ of errors, indicating that all happy-path artifacts produced by the step
1243
+ are missing.
1245
1244
 
1246
1245
 
1247
1246
  Parameters
1248
1247
  ----------
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.
1248
+ var : str, optional, default None
1249
+ Name of the artifact in which to store the caught exception.
1250
+ If not specified, the exception is not stored.
1251
+ print_exception : bool, default True
1252
+ Determines whether or not the exception is printed to
1253
+ stdout when caught.
1260
1254
  """
1261
1255
  ...
1262
1256
 
@@ -1412,49 +1406,6 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1412
1406
  """
1413
1407
  ...
1414
1408
 
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]]:
1416
- """
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.
1422
-
1423
-
1424
- Parameters
1425
- ----------
1426
- timeout : int
1427
- Time, in seconds before the task times out and fails. (Default: 3600)
1428
- poke_interval : int
1429
- Time in seconds that the job should wait in between each try. (Default: 60)
1430
- mode : str
1431
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1432
- exponential_backoff : bool
1433
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1434
- pool : str
1435
- the slot pool this task should run in,
1436
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1437
- soft_fail : bool
1438
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1439
- name : str
1440
- Name of the sensor on Airflow
1441
- description : str
1442
- 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)
1455
- """
1456
- ...
1457
-
1458
1409
  def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1459
1410
  """
1460
1411
  Allows setting external datastores to save data for the
@@ -1610,41 +1561,6 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1610
1561
  """
1611
1562
  ...
1612
1563
 
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
1564
  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]]:
1649
1565
  """
1650
1566
  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.
@@ -1688,6 +1604,100 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1688
1604
  """
1689
1605
  ...
1690
1606
 
1607
+ 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]]:
1608
+ """
1609
+ 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)
1610
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1611
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1612
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1613
+ starts only after all sensors finish.
1614
+
1615
+
1616
+ Parameters
1617
+ ----------
1618
+ timeout : int
1619
+ Time, in seconds before the task times out and fails. (Default: 3600)
1620
+ poke_interval : int
1621
+ Time in seconds that the job should wait in between each try. (Default: 60)
1622
+ mode : str
1623
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1624
+ exponential_backoff : bool
1625
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1626
+ pool : str
1627
+ the slot pool this task should run in,
1628
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1629
+ soft_fail : bool
1630
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1631
+ name : str
1632
+ Name of the sensor on Airflow
1633
+ description : str
1634
+ Description of sensor in the Airflow UI
1635
+ bucket_key : Union[str, List[str]]
1636
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1637
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1638
+ bucket_name : str
1639
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1640
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1641
+ wildcard_match : bool
1642
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1643
+ aws_conn_id : str
1644
+ a reference to the s3 connection on Airflow. (Default: None)
1645
+ verify : bool
1646
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1647
+ """
1648
+ ...
1649
+
1650
+ @typing.overload
1651
+ 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]]:
1652
+ """
1653
+ Specifies the Conda environment for all steps of the flow.
1654
+
1655
+ Use `@conda_base` to set common libraries required by all
1656
+ steps and use `@conda` to specify step-specific additions.
1657
+
1658
+
1659
+ Parameters
1660
+ ----------
1661
+ packages : Dict[str, str], default {}
1662
+ Packages to use for this flow. The key is the name of the package
1663
+ and the value is the version to use.
1664
+ libraries : Dict[str, str], default {}
1665
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1666
+ python : str, optional, default None
1667
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1668
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1669
+ disabled : bool, default False
1670
+ If set to True, disables Conda.
1671
+ """
1672
+ ...
1673
+
1674
+ @typing.overload
1675
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1676
+ ...
1677
+
1678
+ 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):
1679
+ """
1680
+ Specifies the Conda environment for all steps of the flow.
1681
+
1682
+ Use `@conda_base` to set common libraries required by all
1683
+ steps and use `@conda` to specify step-specific additions.
1684
+
1685
+
1686
+ Parameters
1687
+ ----------
1688
+ packages : Dict[str, str], default {}
1689
+ Packages to use for this flow. The key is the name of the package
1690
+ and the value is the version to use.
1691
+ libraries : Dict[str, str], default {}
1692
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1693
+ python : str, optional, default None
1694
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1695
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1696
+ disabled : bool, default False
1697
+ If set to True, disables Conda.
1698
+ """
1699
+ ...
1700
+
1691
1701
  @typing.overload
1692
1702
  def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1693
1703
  """
@@ -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