ob-metaflow-stubs 6.0.7.3__py2.py3-none-any.whl → 6.0.7.4__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 (262) hide show
  1. metaflow-stubs/__init__.pyi +984 -984
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +2 -2
  14. metaflow-stubs/meta_files.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +57 -57
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +1 -1
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +3 -3
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +2 -2
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +4 -4
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +2 -2
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +3 -3
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +3 -3
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +1 -1
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/__init__.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_state_machine.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/__init__.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/__init__.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/spinners.pyi +1 -1
  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 +2 -2
  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 +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/__init__.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/code_packager.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/__init__.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/cli_generator.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/config_utils.pyi +2 -2
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/schema_export.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/typed_configs.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +2 -2
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +2 -2
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +3 -3
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/experimental/__init__.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/perimeters.pyi +1 -1
  88. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/utils.pyi +3 -3
  89. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +1 -1
  90. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +2 -2
  91. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  92. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +1 -1
  93. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  94. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  95. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +2 -2
  96. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +2 -2
  97. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +1 -1
  98. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +3 -3
  99. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +1 -1
  100. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +1 -1
  101. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  102. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  103. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  104. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  105. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  106. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  107. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  108. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  109. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  110. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  111. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  112. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  113. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  114. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  115. metaflow-stubs/mf_extensions/outerbounds/toplevel/s3_proxy.pyi +1 -1
  116. metaflow-stubs/multicore_utils.pyi +1 -1
  117. metaflow-stubs/ob_internal.pyi +1 -1
  118. metaflow-stubs/packaging_sys/__init__.pyi +5 -5
  119. metaflow-stubs/packaging_sys/backend.pyi +3 -3
  120. metaflow-stubs/packaging_sys/distribution_support.pyi +4 -4
  121. metaflow-stubs/packaging_sys/tar_backend.pyi +5 -5
  122. metaflow-stubs/packaging_sys/utils.pyi +1 -1
  123. metaflow-stubs/packaging_sys/v1.pyi +3 -3
  124. metaflow-stubs/parameters.pyi +2 -2
  125. metaflow-stubs/plugins/__init__.pyi +12 -12
  126. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  127. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  128. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  129. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  130. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  131. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  132. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  133. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  134. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  135. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  136. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  137. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  138. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  139. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  140. metaflow-stubs/plugins/argo/exit_hooks.pyi +1 -1
  141. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  142. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  143. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  144. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  145. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  146. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  147. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  148. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  149. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  150. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  151. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  152. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  153. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  154. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  155. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +4 -4
  156. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  157. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  158. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  159. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  160. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  161. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  162. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  163. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  164. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  165. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  166. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  167. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  168. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  169. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  170. metaflow-stubs/plugins/cards/card_modules/basic.pyi +1 -1
  171. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  172. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  173. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  174. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  175. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  176. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  177. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  178. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  179. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  180. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  181. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  182. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  183. metaflow-stubs/plugins/datatools/s3/s3.pyi +2 -2
  184. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  185. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  186. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  187. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  188. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  189. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  190. metaflow-stubs/plugins/exit_hook/__init__.pyi +1 -1
  191. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +1 -1
  192. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  193. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  194. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  195. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  196. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  197. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  198. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  199. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  200. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  201. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  202. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  203. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  204. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  205. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  206. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  207. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  208. metaflow-stubs/plugins/optuna/__init__.pyi +1 -1
  209. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  210. metaflow-stubs/plugins/perimeters.pyi +1 -1
  211. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  212. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  213. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  214. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  215. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  216. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  217. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  218. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  219. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  220. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  221. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  222. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  223. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  224. metaflow-stubs/plugins/secrets/secrets_func.pyi +1 -1
  225. metaflow-stubs/plugins/secrets/secrets_spec.pyi +1 -1
  226. metaflow-stubs/plugins/secrets/utils.pyi +1 -1
  227. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  228. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  229. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  230. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  231. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  232. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  233. metaflow-stubs/plugins/uv/uv_environment.pyi +1 -1
  234. metaflow-stubs/profilers/__init__.pyi +1 -1
  235. metaflow-stubs/pylint_wrapper.pyi +1 -1
  236. metaflow-stubs/runner/__init__.pyi +1 -1
  237. metaflow-stubs/runner/deployer.pyi +3 -3
  238. metaflow-stubs/runner/deployer_impl.pyi +1 -1
  239. metaflow-stubs/runner/metaflow_runner.pyi +2 -2
  240. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  241. metaflow-stubs/runner/nbrun.pyi +1 -1
  242. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  243. metaflow-stubs/runner/utils.pyi +3 -3
  244. metaflow-stubs/system/__init__.pyi +1 -1
  245. metaflow-stubs/system/system_logger.pyi +1 -1
  246. metaflow-stubs/system/system_monitor.pyi +1 -1
  247. metaflow-stubs/tagging_util.pyi +1 -1
  248. metaflow-stubs/tuple_util.pyi +1 -1
  249. metaflow-stubs/user_configs/__init__.pyi +1 -1
  250. metaflow-stubs/user_configs/config_options.pyi +2 -2
  251. metaflow-stubs/user_configs/config_parameters.pyi +5 -5
  252. metaflow-stubs/user_decorators/__init__.pyi +1 -1
  253. metaflow-stubs/user_decorators/common.pyi +1 -1
  254. metaflow-stubs/user_decorators/mutable_flow.pyi +3 -3
  255. metaflow-stubs/user_decorators/mutable_step.pyi +4 -4
  256. metaflow-stubs/user_decorators/user_flow_decorator.pyi +3 -3
  257. metaflow-stubs/user_decorators/user_step_decorator.pyi +5 -5
  258. {ob_metaflow_stubs-6.0.7.3.dist-info → ob_metaflow_stubs-6.0.7.4.dist-info}/METADATA +1 -1
  259. ob_metaflow_stubs-6.0.7.4.dist-info/RECORD +262 -0
  260. ob_metaflow_stubs-6.0.7.3.dist-info/RECORD +0 -262
  261. {ob_metaflow_stubs-6.0.7.3.dist-info → ob_metaflow_stubs-6.0.7.4.dist-info}/WHEEL +0 -0
  262. {ob_metaflow_stubs-6.0.7.3.dist-info → ob_metaflow_stubs-6.0.7.4.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.17.1.0+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-08-20T21:57:30.472423 #
4
+ # Generated on 2025-08-21T17:03:57.569240 #
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
 
@@ -40,16 +40,16 @@ from .user_decorators.user_step_decorator import StepMutator as StepMutator
40
40
  from .user_decorators.user_step_decorator import user_step_decorator as user_step_decorator
41
41
  from .user_decorators.user_flow_decorator import FlowMutator as FlowMutator
42
42
  from . import events as events
43
- from . import tuple_util as tuple_util
44
- from . import cards as cards
45
43
  from . import metaflow_git as metaflow_git
44
+ from . import cards as cards
45
+ from . import tuple_util as tuple_util
46
46
  from . import runner as runner
47
47
  from . import plugins as plugins
48
48
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
49
49
  from . import includefile as includefile
50
50
  from .includefile import IncludeFile as IncludeFile
51
- from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
52
51
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
52
+ from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
53
53
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
54
54
  from . import client as client
55
55
  from .client.core import namespace as namespace
@@ -167,145 +167,471 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
167
167
  """
168
168
  ...
169
169
 
170
- 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]]]:
170
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = 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]]]:
171
171
  """
172
- Specifies that this step should execute on DGX cloud.
172
+ Specifies that this step should execute on Kubernetes.
173
173
 
174
174
 
175
175
  Parameters
176
176
  ----------
177
- gpu : int
178
- Number of GPUs to use.
179
- gpu_type : str
180
- Type of Nvidia GPU to use.
181
- queue_timeout : int
182
- Time to keep the job in NVCF's queue.
177
+ cpu : int, default 1
178
+ Number of CPUs required for this step. If `@resources` is
179
+ also present, the maximum value from all decorators is used.
180
+ memory : int, default 4096
181
+ Memory size (in MB) required for this step. If
182
+ `@resources` is also present, the maximum value from all decorators is
183
+ used.
184
+ disk : int, default 10240
185
+ Disk size (in MB) required for this step. If
186
+ `@resources` is also present, the maximum value from all decorators is
187
+ used.
188
+ image : str, optional, default None
189
+ Docker image to use when launching on Kubernetes. If not specified, and
190
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
191
+ not, a default Docker image mapping to the current version of Python is used.
192
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
193
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
194
+ image_pull_secrets: List[str], default []
195
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
196
+ Kubernetes image pull secrets to use when pulling container images
197
+ in Kubernetes.
198
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
199
+ Kubernetes service account to use when launching pod in Kubernetes.
200
+ secrets : List[str], optional, default None
201
+ Kubernetes secrets to use when launching pod in Kubernetes. These
202
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
203
+ in Metaflow configuration.
204
+ node_selector: Union[Dict[str,str], str], optional, default None
205
+ Kubernetes node selector(s) to apply to the pod running the task.
206
+ Can be passed in as a comma separated string of values e.g.
207
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
208
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
209
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
210
+ Kubernetes namespace to use when launching pod in Kubernetes.
211
+ gpu : int, optional, default None
212
+ Number of GPUs required for this step. A value of zero implies that
213
+ the scheduled node should not have GPUs.
214
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
215
+ The vendor of the GPUs to be used for this step.
216
+ tolerations : List[Dict[str,str]], default []
217
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
218
+ Kubernetes tolerations to use when launching pod in Kubernetes.
219
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
220
+ Kubernetes labels to use when launching pod in Kubernetes.
221
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
222
+ Kubernetes annotations to use when launching pod in Kubernetes.
223
+ use_tmpfs : bool, default False
224
+ This enables an explicit tmpfs mount for this step.
225
+ tmpfs_tempdir : bool, default True
226
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
227
+ tmpfs_size : int, optional, default: None
228
+ The value for the size (in MiB) of the tmpfs mount for this step.
229
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
230
+ memory allocated for this step.
231
+ tmpfs_path : str, optional, default /metaflow_temp
232
+ Path to tmpfs mount for this step.
233
+ persistent_volume_claims : Dict[str, str], optional, default None
234
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
235
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
236
+ shared_memory: int, optional
237
+ Shared memory size (in MiB) required for this step
238
+ port: int, optional
239
+ Port number to specify in the Kubernetes job object
240
+ compute_pool : str, optional, default None
241
+ Compute pool to be used for for this step.
242
+ If not specified, any accessible compute pool within the perimeter is used.
243
+ hostname_resolution_timeout: int, default 10 * 60
244
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
245
+ Only applicable when @parallel is used.
246
+ qos: str, default: Burstable
247
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
248
+
249
+ security_context: Dict[str, Any], optional, default None
250
+ Container security context. Applies to the task container. Allows the following keys:
251
+ - privileged: bool, optional, default None
252
+ - allow_privilege_escalation: bool, optional, default None
253
+ - run_as_user: int, optional, default None
254
+ - run_as_group: int, optional, default None
255
+ - run_as_non_root: bool, optional, default None
183
256
  """
184
257
  ...
185
258
 
186
259
  @typing.overload
187
- 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]]]:
260
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
188
261
  """
189
- Specifies environment variables to be set prior to the execution of a step.
262
+ Internal decorator to support Fast bakery
263
+ """
264
+ ...
265
+
266
+ @typing.overload
267
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
268
+ ...
269
+
270
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
271
+ """
272
+ Internal decorator to support Fast bakery
273
+ """
274
+ ...
275
+
276
+ @typing.overload
277
+ 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]]]:
278
+ """
279
+ Specifies the PyPI packages for the step.
280
+
281
+ Information in this decorator will augment any
282
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
283
+ you can use `@pypi_base` to set packages required by all
284
+ steps and use `@pypi` to specify step-specific overrides.
190
285
 
191
286
 
192
287
  Parameters
193
288
  ----------
194
- vars : Dict[str, str], default {}
195
- Dictionary of environment variables to set.
289
+ packages : Dict[str, str], default: {}
290
+ Packages to use for this step. The key is the name of the package
291
+ and the value is the version to use.
292
+ python : str, optional, default: None
293
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
294
+ that the version used will correspond to the version of the Python interpreter used to start the run.
196
295
  """
197
296
  ...
198
297
 
199
298
  @typing.overload
200
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
299
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
201
300
  ...
202
301
 
203
302
  @typing.overload
204
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
303
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
205
304
  ...
206
305
 
207
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
306
+ 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):
208
307
  """
209
- Specifies environment variables to be set prior to the execution of a step.
308
+ Specifies the PyPI packages for the step.
309
+
310
+ Information in this decorator will augment any
311
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
312
+ you can use `@pypi_base` to set packages required by all
313
+ steps and use `@pypi` to specify step-specific overrides.
210
314
 
211
315
 
212
316
  Parameters
213
317
  ----------
214
- vars : Dict[str, str], default {}
215
- Dictionary of environment variables to set.
318
+ packages : Dict[str, str], default: {}
319
+ Packages to use for this step. The key is the name of the package
320
+ and the value is the version to use.
321
+ python : str, optional, default: None
322
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
323
+ that the version used will correspond to the version of the Python interpreter used to start the run.
216
324
  """
217
325
  ...
218
326
 
219
327
  @typing.overload
220
- 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]]]:
328
+ 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]]]:
221
329
  """
222
- Specifies secrets to be retrieved and injected as environment variables prior to
223
- the execution of a step.
330
+ Specifies that the step will success under all circumstances.
331
+
332
+ The decorator will create an optional artifact, specified by `var`, which
333
+ contains the exception raised. You can use it to detect the presence
334
+ of errors, indicating that all happy-path artifacts produced by the step
335
+ are missing.
224
336
 
225
337
 
226
338
  Parameters
227
339
  ----------
228
- sources : List[Union[str, Dict[str, Any]]], default: []
229
- List of secret specs, defining how the secrets are to be retrieved
230
- role : str, optional, default: None
231
- Role to use for fetching secrets
340
+ var : str, optional, default None
341
+ Name of the artifact in which to store the caught exception.
342
+ If not specified, the exception is not stored.
343
+ print_exception : bool, default True
344
+ Determines whether or not the exception is printed to
345
+ stdout when caught.
232
346
  """
233
347
  ...
234
348
 
235
349
  @typing.overload
236
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
350
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
237
351
  ...
238
352
 
239
353
  @typing.overload
240
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
354
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
241
355
  ...
242
356
 
243
- 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):
357
+ 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):
244
358
  """
245
- Specifies secrets to be retrieved and injected as environment variables prior to
246
- the execution of a step.
359
+ Specifies that the step will success under all circumstances.
360
+
361
+ The decorator will create an optional artifact, specified by `var`, which
362
+ contains the exception raised. You can use it to detect the presence
363
+ of errors, indicating that all happy-path artifacts produced by the step
364
+ are missing.
247
365
 
248
366
 
249
367
  Parameters
250
368
  ----------
251
- sources : List[Union[str, Dict[str, Any]]], default: []
252
- List of secret specs, defining how the secrets are to be retrieved
253
- role : str, optional, default: None
254
- Role to use for fetching secrets
369
+ var : str, optional, default None
370
+ Name of the artifact in which to store the caught exception.
371
+ If not specified, the exception is not stored.
372
+ print_exception : bool, default True
373
+ Determines whether or not the exception is printed to
374
+ stdout when caught.
255
375
  """
256
376
  ...
257
377
 
258
- @typing.overload
259
- 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]]]:
378
+ 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]]]:
260
379
  """
261
- Enables loading / saving of models within a step.
380
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
262
381
 
263
- > Examples
264
- - Saving Models
265
- ```python
266
- @model
267
- @step
268
- def train(self):
269
- # current.model.save returns a dictionary reference to the model saved
270
- self.my_model = current.model.save(
271
- path_to_my_model,
272
- label="my_model",
273
- metadata={
274
- "epochs": 10,
275
- "batch-size": 32,
276
- "learning-rate": 0.001,
277
- }
278
- )
279
- self.next(self.test)
382
+ User code call
383
+ --------------
384
+ @vllm(
385
+ model="...",
386
+ ...
387
+ )
280
388
 
281
- @model(load="my_model")
282
- @step
283
- def test(self):
284
- # `current.model.loaded` returns a dictionary of the loaded models
285
- # where the key is the name of the artifact and the value is the path to the model
286
- print(os.listdir(current.model.loaded["my_model"]))
287
- self.next(self.end)
288
- ```
389
+ Valid backend options
390
+ ---------------------
391
+ - 'local': Run as a separate process on the local task machine.
289
392
 
290
- - Loading models
291
- ```python
292
- @step
293
- def train(self):
294
- # current.model.load returns the path to the model loaded
295
- checkpoint_path = current.model.load(
296
- self.checkpoint_key,
297
- )
298
- model_path = current.model.load(
299
- self.model,
300
- )
301
- self.next(self.test)
302
- ```
393
+ Valid model options
394
+ -------------------
395
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
396
+
397
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
398
+ If you need multiple models, you must create multiple @vllm decorators.
303
399
 
304
400
 
305
401
  Parameters
306
402
  ----------
307
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
308
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
403
+ model: str
404
+ HuggingFace model identifier to be served by vLLM.
405
+ backend: str
406
+ Determines where and how to run the vLLM process.
407
+ openai_api_server: bool
408
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
409
+ Default is False (uses native engine).
410
+ Set to True for backward compatibility with existing code.
411
+ debug: bool
412
+ Whether to turn on verbose debugging logs.
413
+ card_refresh_interval: int
414
+ Interval in seconds for refreshing the vLLM status card.
415
+ Only used when openai_api_server=True.
416
+ max_retries: int
417
+ Maximum number of retries checking for vLLM server startup.
418
+ Only used when openai_api_server=True.
419
+ retry_alert_frequency: int
420
+ Frequency of alert logs for vLLM server startup retries.
421
+ Only used when openai_api_server=True.
422
+ engine_args : dict
423
+ Additional keyword arguments to pass to the vLLM engine.
424
+ For example, `tensor_parallel_size=2`.
425
+ """
426
+ ...
427
+
428
+ @typing.overload
429
+ 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]]]:
430
+ """
431
+ Specifies the number of times the task corresponding
432
+ to a step needs to be retried.
433
+
434
+ This decorator is useful for handling transient errors, such as networking issues.
435
+ If your task contains operations that can't be retried safely, e.g. database updates,
436
+ it is advisable to annotate it with `@retry(times=0)`.
437
+
438
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
439
+ decorator will execute a no-op task after all retries have been exhausted,
440
+ ensuring that the flow execution can continue.
441
+
442
+
443
+ Parameters
444
+ ----------
445
+ times : int, default 3
446
+ Number of times to retry this task.
447
+ minutes_between_retries : int, default 2
448
+ Number of minutes between retries.
449
+ """
450
+ ...
451
+
452
+ @typing.overload
453
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
454
+ ...
455
+
456
+ @typing.overload
457
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
458
+ ...
459
+
460
+ 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):
461
+ """
462
+ Specifies the number of times the task corresponding
463
+ to a step needs to be retried.
464
+
465
+ This decorator is useful for handling transient errors, such as networking issues.
466
+ If your task contains operations that can't be retried safely, e.g. database updates,
467
+ it is advisable to annotate it with `@retry(times=0)`.
468
+
469
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
470
+ decorator will execute a no-op task after all retries have been exhausted,
471
+ ensuring that the flow execution can continue.
472
+
473
+
474
+ Parameters
475
+ ----------
476
+ times : int, default 3
477
+ Number of times to retry this task.
478
+ minutes_between_retries : int, default 2
479
+ Number of minutes between retries.
480
+ """
481
+ ...
482
+
483
+ @typing.overload
484
+ 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]]]:
485
+ """
486
+ Specifies environment variables to be set prior to the execution of a step.
487
+
488
+
489
+ Parameters
490
+ ----------
491
+ vars : Dict[str, str], default {}
492
+ Dictionary of environment variables to set.
493
+ """
494
+ ...
495
+
496
+ @typing.overload
497
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
498
+ ...
499
+
500
+ @typing.overload
501
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
502
+ ...
503
+
504
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
505
+ """
506
+ Specifies environment variables to be set prior to the execution of a step.
507
+
508
+
509
+ Parameters
510
+ ----------
511
+ vars : Dict[str, str], default {}
512
+ Dictionary of environment variables to set.
513
+ """
514
+ ...
515
+
516
+ @typing.overload
517
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
518
+ """
519
+ Decorator prototype for all step decorators. This function gets specialized
520
+ and imported for all decorators types by _import_plugin_decorators().
521
+ """
522
+ ...
523
+
524
+ @typing.overload
525
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
526
+ ...
527
+
528
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
529
+ """
530
+ Decorator prototype for all step decorators. This function gets specialized
531
+ and imported for all decorators types by _import_plugin_decorators().
532
+ """
533
+ ...
534
+
535
+ @typing.overload
536
+ 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]]]:
537
+ """
538
+ Creates a human-readable report, a Metaflow Card, after this step completes.
539
+
540
+ Note that you may add multiple `@card` decorators in a step with different parameters.
541
+
542
+
543
+ Parameters
544
+ ----------
545
+ type : str, default 'default'
546
+ Card type.
547
+ id : str, optional, default None
548
+ If multiple cards are present, use this id to identify this card.
549
+ options : Dict[str, Any], default {}
550
+ Options passed to the card. The contents depend on the card type.
551
+ timeout : int, default 45
552
+ Interrupt reporting if it takes more than this many seconds.
553
+ """
554
+ ...
555
+
556
+ @typing.overload
557
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
558
+ ...
559
+
560
+ @typing.overload
561
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
562
+ ...
563
+
564
+ 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):
565
+ """
566
+ Creates a human-readable report, a Metaflow Card, after this step completes.
567
+
568
+ Note that you may add multiple `@card` decorators in a step with different parameters.
569
+
570
+
571
+ Parameters
572
+ ----------
573
+ type : str, default 'default'
574
+ Card type.
575
+ id : str, optional, default None
576
+ If multiple cards are present, use this id to identify this card.
577
+ options : Dict[str, Any], default {}
578
+ Options passed to the card. The contents depend on the card type.
579
+ timeout : int, default 45
580
+ Interrupt reporting if it takes more than this many seconds.
581
+ """
582
+ ...
583
+
584
+ @typing.overload
585
+ 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]]]:
586
+ """
587
+ Enables loading / saving of models within a step.
588
+
589
+ > Examples
590
+ - Saving Models
591
+ ```python
592
+ @model
593
+ @step
594
+ def train(self):
595
+ # current.model.save returns a dictionary reference to the model saved
596
+ self.my_model = current.model.save(
597
+ path_to_my_model,
598
+ label="my_model",
599
+ metadata={
600
+ "epochs": 10,
601
+ "batch-size": 32,
602
+ "learning-rate": 0.001,
603
+ }
604
+ )
605
+ self.next(self.test)
606
+
607
+ @model(load="my_model")
608
+ @step
609
+ def test(self):
610
+ # `current.model.loaded` returns a dictionary of the loaded models
611
+ # where the key is the name of the artifact and the value is the path to the model
612
+ print(os.listdir(current.model.loaded["my_model"]))
613
+ self.next(self.end)
614
+ ```
615
+
616
+ - Loading models
617
+ ```python
618
+ @step
619
+ def train(self):
620
+ # current.model.load returns the path to the model loaded
621
+ checkpoint_path = current.model.load(
622
+ self.checkpoint_key,
623
+ )
624
+ model_path = current.model.load(
625
+ self.model,
626
+ )
627
+ self.next(self.test)
628
+ ```
629
+
630
+
631
+ Parameters
632
+ ----------
633
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
634
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
309
635
  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`.
310
636
  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
311
637
  the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
@@ -384,156 +710,103 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
384
710
  """
385
711
  ...
386
712
 
387
- 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]]]:
713
+ @typing.overload
714
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
388
715
  """
389
- Specifies that this step should execute on DGX cloud.
390
-
391
-
392
- Parameters
393
- ----------
394
- gpu : int
395
- Number of GPUs to use.
396
- gpu_type : str
397
- Type of Nvidia GPU to use.
716
+ Decorator prototype for all step decorators. This function gets specialized
717
+ and imported for all decorators types by _import_plugin_decorators().
398
718
  """
399
719
  ...
400
720
 
401
- 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]]]:
721
+ @typing.overload
722
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
723
+ ...
724
+
725
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
402
726
  """
403
- Decorator that helps cache, version and store models/datasets from huggingface hub.
404
-
405
- > Examples
406
-
407
- **Usage: creating references of models from huggingface that may be loaded in downstream steps**
408
- ```python
409
- @huggingface_hub
410
- @step
411
- def pull_model_from_huggingface(self):
412
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
413
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
414
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
415
- # value of the function is a reference to the model in the backend storage.
416
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
417
-
418
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
419
- self.llama_model = current.huggingface_hub.snapshot_download(
420
- repo_id=self.model_id,
421
- allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
422
- )
423
- self.next(self.train)
424
- ```
425
-
426
- **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
427
- ```python
428
- @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
429
- @step
430
- def pull_model_from_huggingface(self):
431
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
432
- ```
433
-
434
- ```python
435
- @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
436
- @step
437
- def finetune_model(self):
438
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
439
- # path_to_model will be /my-directory
440
- ```
441
-
442
- ```python
443
- # Takes all the arguments passed to `snapshot_download`
444
- # except for `local_dir`
445
- @huggingface_hub(load=[
446
- {
447
- "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
448
- },
449
- {
450
- "repo_id": "myorg/mistral-lora",
451
- "repo_type": "model",
452
- },
453
- ])
454
- @step
455
- def finetune_model(self):
456
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
457
- # path_to_model will be /my-directory
458
- ```
727
+ Decorator prototype for all step decorators. This function gets specialized
728
+ and imported for all decorators types by _import_plugin_decorators().
729
+ """
730
+ ...
731
+
732
+ def s3_proxy(*, integration_name: typing.Optional[str] = None, write_mode: typing.Optional[str] = None, debug: typing.Optional[bool] = 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]]]:
733
+ """
734
+ S3 Proxy decorator for routing S3 requests through a local proxy service.
459
735
 
460
736
 
461
737
  Parameters
462
738
  ----------
463
- temp_dir_root : str, optional
464
- The root directory that will hold the temporary directory where objects will be downloaded.
465
-
466
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
467
- The list of repos (models/datasets) to load.
468
-
469
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
470
-
471
- - If repo (model/dataset) is not found in the datastore:
472
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
473
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
474
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
475
-
476
- - If repo is found in the datastore:
477
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
739
+ integration_name : str, optional
740
+ Name of the S3 proxy integration. If not specified, will use the only
741
+ available S3 proxy integration in the namespace (fails if multiple exist).
742
+ write_mode : str, optional
743
+ The desired behavior during write operations to target (origin) S3 bucket.
744
+ allowed options are:
745
+ "origin-and-cache" -> write to both the target S3 bucket and local object
746
+ storage
747
+ "origin" -> only write to the target S3 bucket
748
+ "cache" -> only write to the object storage service used for caching
749
+ debug : bool, optional
750
+ Enable debug logging for proxy operations.
478
751
  """
479
752
  ...
480
753
 
481
754
  @typing.overload
482
- 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]]]:
755
+ 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]]]:
483
756
  """
484
- Specifies a timeout for your step.
485
-
486
- This decorator is useful if this step may hang indefinitely.
487
-
488
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
489
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
490
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
757
+ Specifies the Conda environment for the step.
491
758
 
492
- Note that all the values specified in parameters are added together so if you specify
493
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
759
+ Information in this decorator will augment any
760
+ attributes set in the `@conda_base` flow-level decorator. Hence,
761
+ you can use `@conda_base` to set packages required by all
762
+ steps and use `@conda` to specify step-specific overrides.
494
763
 
495
764
 
496
765
  Parameters
497
766
  ----------
498
- seconds : int, default 0
499
- Number of seconds to wait prior to timing out.
500
- minutes : int, default 0
501
- Number of minutes to wait prior to timing out.
502
- hours : int, default 0
503
- Number of hours to wait prior to timing out.
767
+ packages : Dict[str, str], default {}
768
+ Packages to use for this step. The key is the name of the package
769
+ and the value is the version to use.
770
+ libraries : Dict[str, str], default {}
771
+ Supported for backward compatibility. When used with packages, packages will take precedence.
772
+ python : str, optional, default None
773
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
774
+ that the version used will correspond to the version of the Python interpreter used to start the run.
775
+ disabled : bool, default False
776
+ If set to True, disables @conda.
504
777
  """
505
778
  ...
506
779
 
507
780
  @typing.overload
508
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
781
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
509
782
  ...
510
783
 
511
784
  @typing.overload
512
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
785
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
513
786
  ...
514
787
 
515
- 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):
788
+ 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):
516
789
  """
517
- Specifies a timeout for your step.
518
-
519
- This decorator is useful if this step may hang indefinitely.
520
-
521
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
522
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
523
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
790
+ Specifies the Conda environment for the step.
524
791
 
525
- Note that all the values specified in parameters are added together so if you specify
526
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
792
+ Information in this decorator will augment any
793
+ attributes set in the `@conda_base` flow-level decorator. Hence,
794
+ you can use `@conda_base` to set packages required by all
795
+ steps and use `@conda` to specify step-specific overrides.
527
796
 
528
797
 
529
798
  Parameters
530
799
  ----------
531
- seconds : int, default 0
532
- Number of seconds to wait prior to timing out.
533
- minutes : int, default 0
534
- Number of minutes to wait prior to timing out.
535
- hours : int, default 0
536
- Number of hours to wait prior to timing out.
800
+ packages : Dict[str, str], default {}
801
+ Packages to use for this step. The key is the name of the package
802
+ and the value is the version to use.
803
+ libraries : Dict[str, str], default {}
804
+ Supported for backward compatibility. When used with packages, packages will take precedence.
805
+ python : str, optional, default None
806
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
807
+ that the version used will correspond to the version of the Python interpreter used to start the run.
808
+ disabled : bool, default False
809
+ If set to True, disables @conda.
537
810
  """
538
811
  ...
539
812
 
@@ -580,346 +853,194 @@ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy:
580
853
  """
581
854
  ...
582
855
 
583
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = 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]]]:
856
+ @typing.overload
857
+ def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
584
858
  """
585
- Specifies that this step should execute on Kubernetes.
586
-
587
-
588
- Parameters
589
- ----------
590
- cpu : int, default 1
591
- Number of CPUs required for this step. If `@resources` is
592
- also present, the maximum value from all decorators is used.
593
- memory : int, default 4096
594
- Memory size (in MB) required for this step. If
595
- `@resources` is also present, the maximum value from all decorators is
596
- used.
597
- disk : int, default 10240
598
- Disk size (in MB) required for this step. If
599
- `@resources` is also present, the maximum value from all decorators is
600
- used.
601
- image : str, optional, default None
602
- Docker image to use when launching on Kubernetes. If not specified, and
603
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
604
- not, a default Docker image mapping to the current version of Python is used.
605
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
606
- If given, the imagePullPolicy to be applied to the Docker image of the step.
607
- image_pull_secrets: List[str], default []
608
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
609
- Kubernetes image pull secrets to use when pulling container images
610
- in Kubernetes.
611
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
612
- Kubernetes service account to use when launching pod in Kubernetes.
613
- secrets : List[str], optional, default None
614
- Kubernetes secrets to use when launching pod in Kubernetes. These
615
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
616
- in Metaflow configuration.
617
- node_selector: Union[Dict[str,str], str], optional, default None
618
- Kubernetes node selector(s) to apply to the pod running the task.
619
- Can be passed in as a comma separated string of values e.g.
620
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
621
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
622
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
623
- Kubernetes namespace to use when launching pod in Kubernetes.
624
- gpu : int, optional, default None
625
- Number of GPUs required for this step. A value of zero implies that
626
- the scheduled node should not have GPUs.
627
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
628
- The vendor of the GPUs to be used for this step.
629
- tolerations : List[Dict[str,str]], default []
630
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
631
- Kubernetes tolerations to use when launching pod in Kubernetes.
632
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
633
- Kubernetes labels to use when launching pod in Kubernetes.
634
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
635
- Kubernetes annotations to use when launching pod in Kubernetes.
636
- use_tmpfs : bool, default False
637
- This enables an explicit tmpfs mount for this step.
638
- tmpfs_tempdir : bool, default True
639
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
640
- tmpfs_size : int, optional, default: None
641
- The value for the size (in MiB) of the tmpfs mount for this step.
642
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
643
- memory allocated for this step.
644
- tmpfs_path : str, optional, default /metaflow_temp
645
- Path to tmpfs mount for this step.
646
- persistent_volume_claims : Dict[str, str], optional, default None
647
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
648
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
649
- shared_memory: int, optional
650
- Shared memory size (in MiB) required for this step
651
- port: int, optional
652
- Port number to specify in the Kubernetes job object
653
- compute_pool : str, optional, default None
654
- Compute pool to be used for for this step.
655
- If not specified, any accessible compute pool within the perimeter is used.
656
- hostname_resolution_timeout: int, default 10 * 60
657
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
658
- Only applicable when @parallel is used.
659
- qos: str, default: Burstable
660
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
661
-
662
- security_context: Dict[str, Any], optional, default None
663
- Container security context. Applies to the task container. Allows the following keys:
664
- - privileged: bool, optional, default None
665
- - allow_privilege_escalation: bool, optional, default None
666
- - run_as_user: int, optional, default None
667
- - run_as_group: int, optional, default None
668
- - run_as_non_root: bool, optional, default None
859
+ Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
860
+ It exists to make it easier for users to know that this decorator should only be used with
861
+ a Neo Cloud like Nebius.
669
862
  """
670
863
  ...
671
864
 
672
865
  @typing.overload
673
- 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]]]:
866
+ def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
867
+ ...
868
+
869
+ def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
674
870
  """
675
- Specifies the PyPI packages for the step.
676
-
677
- Information in this decorator will augment any
678
- attributes set in the `@pyi_base` flow-level decorator. Hence,
679
- you can use `@pypi_base` to set packages required by all
680
- steps and use `@pypi` to specify step-specific overrides.
681
-
682
-
683
- Parameters
684
- ----------
685
- packages : Dict[str, str], default: {}
686
- Packages to use for this step. The key is the name of the package
687
- and the value is the version to use.
688
- python : str, optional, default: None
689
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
690
- that the version used will correspond to the version of the Python interpreter used to start the run.
871
+ Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
872
+ It exists to make it easier for users to know that this decorator should only be used with
873
+ a Neo Cloud like Nebius.
691
874
  """
692
875
  ...
693
876
 
694
877
  @typing.overload
695
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
696
- ...
697
-
698
- @typing.overload
699
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
700
- ...
701
-
702
- 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):
878
+ 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]]]:
703
879
  """
704
- Specifies the PyPI packages for the step.
880
+ Specifies a timeout for your step.
705
881
 
706
- Information in this decorator will augment any
707
- attributes set in the `@pyi_base` flow-level decorator. Hence,
708
- you can use `@pypi_base` to set packages required by all
709
- steps and use `@pypi` to specify step-specific overrides.
882
+ This decorator is useful if this step may hang indefinitely.
883
+
884
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
885
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
886
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
887
+
888
+ Note that all the values specified in parameters are added together so if you specify
889
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
710
890
 
711
891
 
712
892
  Parameters
713
893
  ----------
714
- packages : Dict[str, str], default: {}
715
- Packages to use for this step. The key is the name of the package
716
- and the value is the version to use.
717
- python : str, optional, default: None
718
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
719
- that the version used will correspond to the version of the Python interpreter used to start the run.
720
- """
721
- ...
722
-
723
- @typing.overload
724
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
725
- """
726
- Decorator prototype for all step decorators. This function gets specialized
727
- and imported for all decorators types by _import_plugin_decorators().
728
- """
729
- ...
730
-
731
- @typing.overload
732
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
733
- ...
734
-
735
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
736
- """
737
- Decorator prototype for all step decorators. This function gets specialized
738
- and imported for all decorators types by _import_plugin_decorators().
894
+ seconds : int, default 0
895
+ Number of seconds to wait prior to timing out.
896
+ minutes : int, default 0
897
+ Number of minutes to wait prior to timing out.
898
+ hours : int, default 0
899
+ Number of hours to wait prior to timing out.
739
900
  """
740
901
  ...
741
902
 
742
903
  @typing.overload
743
- def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
744
- """
745
- A simple decorator that demonstrates using CardDecoratorInjector
746
- to inject a card and render simple markdown content.
747
- """
904
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
748
905
  ...
749
906
 
750
907
  @typing.overload
751
- def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
908
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
752
909
  ...
753
910
 
754
- def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
911
+ 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):
755
912
  """
756
- A simple decorator that demonstrates using CardDecoratorInjector
757
- to inject a card and render simple markdown content.
913
+ Specifies a timeout for your step.
914
+
915
+ This decorator is useful if this step may hang indefinitely.
916
+
917
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
918
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
919
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
920
+
921
+ Note that all the values specified in parameters are added together so if you specify
922
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
923
+
924
+
925
+ Parameters
926
+ ----------
927
+ seconds : int, default 0
928
+ Number of seconds to wait prior to timing out.
929
+ minutes : int, default 0
930
+ Number of minutes to wait prior to timing out.
931
+ hours : int, default 0
932
+ Number of hours to wait prior to timing out.
758
933
  """
759
934
  ...
760
935
 
761
936
  @typing.overload
762
- def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
937
+ def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
763
938
  """
764
- Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
939
+ CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
765
940
  It exists to make it easier for users to know that this decorator should only be used with
766
- a Neo Cloud like Nebius.
941
+ a Neo Cloud like CoreWeave.
767
942
  """
768
943
  ...
769
944
 
770
945
  @typing.overload
771
- def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
946
+ def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
772
947
  ...
773
948
 
774
- def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
949
+ def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
775
950
  """
776
- Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
951
+ CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
777
952
  It exists to make it easier for users to know that this decorator should only be used with
778
- a Neo Cloud like Nebius.
953
+ a Neo Cloud like CoreWeave.
779
954
  """
780
955
  ...
781
956
 
782
957
  @typing.overload
783
- 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]]]:
958
+ 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]]]:
784
959
  """
785
- Specifies the number of times the task corresponding
786
- to a step needs to be retried.
960
+ Specifies the resources needed when executing this step.
787
961
 
788
- This decorator is useful for handling transient errors, such as networking issues.
789
- If your task contains operations that can't be retried safely, e.g. database updates,
790
- it is advisable to annotate it with `@retry(times=0)`.
962
+ Use `@resources` to specify the resource requirements
963
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
791
964
 
792
- This can be used in conjunction with the `@catch` decorator. The `@catch`
793
- decorator will execute a no-op task after all retries have been exhausted,
794
- ensuring that the flow execution can continue.
965
+ You can choose the compute layer on the command line by executing e.g.
966
+ ```
967
+ python myflow.py run --with batch
968
+ ```
969
+ or
970
+ ```
971
+ python myflow.py run --with kubernetes
972
+ ```
973
+ which executes the flow on the desired system using the
974
+ requirements specified in `@resources`.
795
975
 
796
976
 
797
977
  Parameters
798
978
  ----------
799
- times : int, default 3
800
- Number of times to retry this task.
801
- minutes_between_retries : int, default 2
802
- Number of minutes between retries.
979
+ cpu : int, default 1
980
+ Number of CPUs required for this step.
981
+ gpu : int, optional, default None
982
+ Number of GPUs required for this step.
983
+ disk : int, optional, default None
984
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
985
+ memory : int, default 4096
986
+ Memory size (in MB) required for this step.
987
+ shared_memory : int, optional, default None
988
+ The value for the size (in MiB) of the /dev/shm volume for this step.
989
+ This parameter maps to the `--shm-size` option in Docker.
803
990
  """
804
991
  ...
805
992
 
806
993
  @typing.overload
807
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
994
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
808
995
  ...
809
996
 
810
997
  @typing.overload
811
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
998
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
812
999
  ...
813
1000
 
814
- 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):
1001
+ 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):
815
1002
  """
816
- Specifies the number of times the task corresponding
817
- to a step needs to be retried.
818
-
819
- This decorator is useful for handling transient errors, such as networking issues.
820
- If your task contains operations that can't be retried safely, e.g. database updates,
821
- it is advisable to annotate it with `@retry(times=0)`.
822
-
823
- This can be used in conjunction with the `@catch` decorator. The `@catch`
824
- decorator will execute a no-op task after all retries have been exhausted,
825
- ensuring that the flow execution can continue.
826
-
1003
+ Specifies the resources needed when executing this step.
827
1004
 
828
- Parameters
829
- ----------
830
- times : int, default 3
831
- Number of times to retry this task.
832
- minutes_between_retries : int, default 2
833
- Number of minutes between retries.
834
- """
835
- ...
836
-
837
- @typing.overload
838
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
839
- """
840
- Decorator prototype for all step decorators. This function gets specialized
841
- and imported for all decorators types by _import_plugin_decorators().
842
- """
843
- ...
844
-
845
- @typing.overload
846
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
847
- ...
848
-
849
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
850
- """
851
- Decorator prototype for all step decorators. This function gets specialized
852
- and imported for all decorators types by _import_plugin_decorators().
853
- """
854
- ...
855
-
856
- @typing.overload
857
- 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]]]:
858
- """
859
- Specifies the Conda environment for the step.
1005
+ Use `@resources` to specify the resource requirements
1006
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
860
1007
 
861
- Information in this decorator will augment any
862
- attributes set in the `@conda_base` flow-level decorator. Hence,
863
- you can use `@conda_base` to set packages required by all
864
- steps and use `@conda` to specify step-specific overrides.
1008
+ You can choose the compute layer on the command line by executing e.g.
1009
+ ```
1010
+ python myflow.py run --with batch
1011
+ ```
1012
+ or
1013
+ ```
1014
+ python myflow.py run --with kubernetes
1015
+ ```
1016
+ which executes the flow on the desired system using the
1017
+ requirements specified in `@resources`.
865
1018
 
866
1019
 
867
1020
  Parameters
868
1021
  ----------
869
- packages : Dict[str, str], default {}
870
- Packages to use for this step. The key is the name of the package
871
- and the value is the version to use.
872
- libraries : Dict[str, str], default {}
873
- Supported for backward compatibility. When used with packages, packages will take precedence.
874
- python : str, optional, default None
875
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
876
- that the version used will correspond to the version of the Python interpreter used to start the run.
877
- disabled : bool, default False
878
- If set to True, disables @conda.
1022
+ cpu : int, default 1
1023
+ Number of CPUs required for this step.
1024
+ gpu : int, optional, default None
1025
+ Number of GPUs required for this step.
1026
+ disk : int, optional, default None
1027
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1028
+ memory : int, default 4096
1029
+ Memory size (in MB) required for this step.
1030
+ shared_memory : int, optional, default None
1031
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1032
+ This parameter maps to the `--shm-size` option in Docker.
879
1033
  """
880
1034
  ...
881
1035
 
882
1036
  @typing.overload
883
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
884
- ...
885
-
886
- @typing.overload
887
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
888
- ...
889
-
890
- 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):
1037
+ 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]]]:
891
1038
  """
892
- Specifies the Conda environment for the step.
1039
+ Enables checkpointing for a step.
893
1040
 
894
- Information in this decorator will augment any
895
- attributes set in the `@conda_base` flow-level decorator. Hence,
896
- you can use `@conda_base` to set packages required by all
897
- steps and use `@conda` to specify step-specific overrides.
1041
+ > Examples
898
1042
 
899
-
900
- Parameters
901
- ----------
902
- packages : Dict[str, str], default {}
903
- Packages to use for this step. The key is the name of the package
904
- and the value is the version to use.
905
- libraries : Dict[str, str], default {}
906
- Supported for backward compatibility. When used with packages, packages will take precedence.
907
- python : str, optional, default None
908
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
909
- that the version used will correspond to the version of the Python interpreter used to start the run.
910
- disabled : bool, default False
911
- If set to True, disables @conda.
912
- """
913
- ...
914
-
915
- @typing.overload
916
- 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]]]:
917
- """
918
- Enables checkpointing for a step.
919
-
920
- > Examples
921
-
922
- - Saving Checkpoints
1043
+ - Saving Checkpoints
923
1044
 
924
1045
  ```python
925
1046
  @checkpoint
@@ -1059,477 +1180,214 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
1059
1180
  """
1060
1181
  ...
1061
1182
 
1062
- def s3_proxy(*, integration_name: typing.Optional[str] = None, write_mode: typing.Optional[str] = None, debug: typing.Optional[bool] = 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]]]:
1063
- """
1064
- S3 Proxy decorator for routing S3 requests through a local proxy service.
1065
-
1066
-
1067
- Parameters
1068
- ----------
1069
- integration_name : str, optional
1070
- Name of the S3 proxy integration. If not specified, will use the only
1071
- available S3 proxy integration in the namespace (fails if multiple exist).
1072
- write_mode : str, optional
1073
- The desired behavior during write operations to target (origin) S3 bucket.
1074
- allowed options are:
1075
- "origin-and-cache" -> write to both the target S3 bucket and local object
1076
- storage
1077
- "origin" -> only write to the target S3 bucket
1078
- "cache" -> only write to the object storage service used for caching
1079
- debug : bool, optional
1080
- Enable debug logging for proxy operations.
1081
- """
1082
- ...
1083
-
1084
- @typing.overload
1085
- 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]]]:
1086
- """
1087
- Specifies that the step will success under all circumstances.
1088
-
1089
- The decorator will create an optional artifact, specified by `var`, which
1090
- contains the exception raised. You can use it to detect the presence
1091
- of errors, indicating that all happy-path artifacts produced by the step
1092
- are missing.
1093
-
1094
-
1095
- Parameters
1096
- ----------
1097
- var : str, optional, default None
1098
- Name of the artifact in which to store the caught exception.
1099
- If not specified, the exception is not stored.
1100
- print_exception : bool, default True
1101
- Determines whether or not the exception is printed to
1102
- stdout when caught.
1103
- """
1104
- ...
1105
-
1106
- @typing.overload
1107
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1108
- ...
1109
-
1110
- @typing.overload
1111
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1112
- ...
1113
-
1114
- 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):
1115
- """
1116
- Specifies that the step will success under all circumstances.
1117
-
1118
- The decorator will create an optional artifact, specified by `var`, which
1119
- contains the exception raised. You can use it to detect the presence
1120
- of errors, indicating that all happy-path artifacts produced by the step
1121
- are missing.
1122
-
1123
-
1124
- Parameters
1125
- ----------
1126
- var : str, optional, default None
1127
- Name of the artifact in which to store the caught exception.
1128
- If not specified, the exception is not stored.
1129
- print_exception : bool, default True
1130
- Determines whether or not the exception is printed to
1131
- stdout when caught.
1132
- """
1133
- ...
1134
-
1135
1183
  @typing.overload
1136
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1137
- """
1138
- Internal decorator to support Fast bakery
1139
- """
1140
- ...
1141
-
1142
- @typing.overload
1143
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1144
- ...
1145
-
1146
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1147
- """
1148
- Internal decorator to support Fast bakery
1149
- """
1150
- ...
1151
-
1152
- @typing.overload
1153
- 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]]]:
1154
- """
1155
- Specifies the resources needed when executing this step.
1156
-
1157
- Use `@resources` to specify the resource requirements
1158
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1159
-
1160
- You can choose the compute layer on the command line by executing e.g.
1161
- ```
1162
- python myflow.py run --with batch
1163
- ```
1164
- or
1165
- ```
1166
- python myflow.py run --with kubernetes
1167
- ```
1168
- which executes the flow on the desired system using the
1169
- requirements specified in `@resources`.
1170
-
1171
-
1172
- Parameters
1173
- ----------
1174
- cpu : int, default 1
1175
- Number of CPUs required for this step.
1176
- gpu : int, optional, default None
1177
- Number of GPUs required for this step.
1178
- disk : int, optional, default None
1179
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1180
- memory : int, default 4096
1181
- Memory size (in MB) required for this step.
1182
- shared_memory : int, optional, default None
1183
- The value for the size (in MiB) of the /dev/shm volume for this step.
1184
- This parameter maps to the `--shm-size` option in Docker.
1185
- """
1186
- ...
1187
-
1188
- @typing.overload
1189
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1190
- ...
1191
-
1192
- @typing.overload
1193
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1194
- ...
1195
-
1196
- 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):
1197
- """
1198
- Specifies the resources needed when executing this step.
1199
-
1200
- Use `@resources` to specify the resource requirements
1201
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1202
-
1203
- You can choose the compute layer on the command line by executing e.g.
1204
- ```
1205
- python myflow.py run --with batch
1206
- ```
1207
- or
1208
- ```
1209
- python myflow.py run --with kubernetes
1210
- ```
1211
- which executes the flow on the desired system using the
1212
- requirements specified in `@resources`.
1213
-
1214
-
1215
- Parameters
1216
- ----------
1217
- cpu : int, default 1
1218
- Number of CPUs required for this step.
1219
- gpu : int, optional, default None
1220
- Number of GPUs required for this step.
1221
- disk : int, optional, default None
1222
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1223
- memory : int, default 4096
1224
- Memory size (in MB) required for this step.
1225
- shared_memory : int, optional, default None
1226
- The value for the size (in MiB) of the /dev/shm volume for this step.
1227
- This parameter maps to the `--shm-size` option in Docker.
1228
- """
1229
- ...
1230
-
1231
- @typing.overload
1232
- 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]]]:
1184
+ 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]]]:
1233
1185
  """
1234
- Creates a human-readable report, a Metaflow Card, after this step completes.
1235
-
1236
- Note that you may add multiple `@card` decorators in a step with different parameters.
1186
+ Specifies secrets to be retrieved and injected as environment variables prior to
1187
+ the execution of a step.
1237
1188
 
1238
1189
 
1239
1190
  Parameters
1240
1191
  ----------
1241
- type : str, default 'default'
1242
- Card type.
1243
- id : str, optional, default None
1244
- If multiple cards are present, use this id to identify this card.
1245
- options : Dict[str, Any], default {}
1246
- Options passed to the card. The contents depend on the card type.
1247
- timeout : int, default 45
1248
- Interrupt reporting if it takes more than this many seconds.
1192
+ sources : List[Union[str, Dict[str, Any]]], default: []
1193
+ List of secret specs, defining how the secrets are to be retrieved
1194
+ role : str, optional, default: None
1195
+ Role to use for fetching secrets
1249
1196
  """
1250
1197
  ...
1251
1198
 
1252
1199
  @typing.overload
1253
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1200
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1254
1201
  ...
1255
1202
 
1256
1203
  @typing.overload
1257
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1204
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1258
1205
  ...
1259
1206
 
1260
- 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):
1207
+ 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):
1261
1208
  """
1262
- Creates a human-readable report, a Metaflow Card, after this step completes.
1263
-
1264
- Note that you may add multiple `@card` decorators in a step with different parameters.
1209
+ Specifies secrets to be retrieved and injected as environment variables prior to
1210
+ the execution of a step.
1265
1211
 
1266
1212
 
1267
1213
  Parameters
1268
1214
  ----------
1269
- type : str, default 'default'
1270
- Card type.
1271
- id : str, optional, default None
1272
- If multiple cards are present, use this id to identify this card.
1273
- options : Dict[str, Any], default {}
1274
- Options passed to the card. The contents depend on the card type.
1275
- timeout : int, default 45
1276
- Interrupt reporting if it takes more than this many seconds.
1215
+ sources : List[Union[str, Dict[str, Any]]], default: []
1216
+ List of secret specs, defining how the secrets are to be retrieved
1217
+ role : str, optional, default: None
1218
+ Role to use for fetching secrets
1277
1219
  """
1278
1220
  ...
1279
1221
 
1280
1222
  @typing.overload
1281
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1223
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1282
1224
  """
1283
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1284
- It exists to make it easier for users to know that this decorator should only be used with
1285
- a Neo Cloud like CoreWeave.
1225
+ A simple decorator that demonstrates using CardDecoratorInjector
1226
+ to inject a card and render simple markdown content.
1286
1227
  """
1287
1228
  ...
1288
1229
 
1289
1230
  @typing.overload
1290
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1231
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1291
1232
  ...
1292
1233
 
1293
- def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1234
+ def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1294
1235
  """
1295
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1296
- It exists to make it easier for users to know that this decorator should only be used with
1297
- a Neo Cloud like CoreWeave.
1236
+ A simple decorator that demonstrates using CardDecoratorInjector
1237
+ to inject a card and render simple markdown content.
1298
1238
  """
1299
1239
  ...
1300
1240
 
1301
- 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]]]:
1241
+ 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]]]:
1302
1242
  """
1303
- This decorator is used to run vllm APIs as Metaflow task sidecars.
1304
-
1305
- User code call
1306
- --------------
1307
- @vllm(
1308
- model="...",
1309
- ...
1310
- )
1311
-
1312
- Valid backend options
1313
- ---------------------
1314
- - 'local': Run as a separate process on the local task machine.
1315
-
1316
- Valid model options
1317
- -------------------
1318
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
1319
-
1320
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
1321
- If you need multiple models, you must create multiple @vllm decorators.
1243
+ Specifies that this step should execute on DGX cloud.
1322
1244
 
1323
1245
 
1324
1246
  Parameters
1325
1247
  ----------
1326
- model: str
1327
- HuggingFace model identifier to be served by vLLM.
1328
- backend: str
1329
- Determines where and how to run the vLLM process.
1330
- openai_api_server: bool
1331
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
1332
- Default is False (uses native engine).
1333
- Set to True for backward compatibility with existing code.
1334
- debug: bool
1335
- Whether to turn on verbose debugging logs.
1336
- card_refresh_interval: int
1337
- Interval in seconds for refreshing the vLLM status card.
1338
- Only used when openai_api_server=True.
1339
- max_retries: int
1340
- Maximum number of retries checking for vLLM server startup.
1341
- Only used when openai_api_server=True.
1342
- retry_alert_frequency: int
1343
- Frequency of alert logs for vLLM server startup retries.
1344
- Only used when openai_api_server=True.
1345
- engine_args : dict
1346
- Additional keyword arguments to pass to the vLLM engine.
1347
- For example, `tensor_parallel_size=2`.
1248
+ gpu : int
1249
+ Number of GPUs to use.
1250
+ gpu_type : str
1251
+ Type of Nvidia GPU to use.
1252
+ queue_timeout : int
1253
+ Time to keep the job in NVCF's queue.
1348
1254
  """
1349
1255
  ...
1350
1256
 
1351
- @typing.overload
1352
- 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]]:
1257
+ 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]]]:
1353
1258
  """
1354
- Specifies the event(s) that this flow depends on.
1355
-
1356
- ```
1357
- @trigger(event='foo')
1358
- ```
1359
- or
1360
- ```
1361
- @trigger(events=['foo', 'bar'])
1362
- ```
1363
-
1364
- Additionally, you can specify the parameter mappings
1365
- to map event payload to Metaflow parameters for the flow.
1366
- ```
1367
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1368
- ```
1369
- or
1370
- ```
1371
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1372
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1373
- ```
1374
-
1375
- 'parameters' can also be a list of strings and tuples like so:
1376
- ```
1377
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1378
- ```
1379
- This is equivalent to:
1380
- ```
1381
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1382
- ```
1259
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
1383
1260
 
1261
+ > Examples
1384
1262
 
1385
- Parameters
1386
- ----------
1387
- event : Union[str, Dict[str, Any]], optional, default None
1388
- Event dependency for this flow.
1389
- events : List[Union[str, Dict[str, Any]]], default []
1390
- Events dependency for this flow.
1391
- options : Dict[str, Any], default {}
1392
- Backend-specific configuration for tuning eventing behavior.
1393
- """
1394
- ...
1395
-
1396
- @typing.overload
1397
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1398
- ...
1399
-
1400
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, 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] = {}):
1401
- """
1402
- Specifies the event(s) that this flow depends on.
1263
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
1264
+ ```python
1265
+ @huggingface_hub
1266
+ @step
1267
+ def pull_model_from_huggingface(self):
1268
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
1269
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
1270
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
1271
+ # value of the function is a reference to the model in the backend storage.
1272
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
1403
1273
 
1404
- ```
1405
- @trigger(event='foo')
1406
- ```
1407
- or
1408
- ```
1409
- @trigger(events=['foo', 'bar'])
1274
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
1275
+ self.llama_model = current.huggingface_hub.snapshot_download(
1276
+ repo_id=self.model_id,
1277
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
1278
+ )
1279
+ self.next(self.train)
1410
1280
  ```
1411
1281
 
1412
- Additionally, you can specify the parameter mappings
1413
- to map event payload to Metaflow parameters for the flow.
1414
- ```
1415
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1416
- ```
1417
- or
1418
- ```
1419
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1420
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1282
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
1283
+ ```python
1284
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
1285
+ @step
1286
+ def pull_model_from_huggingface(self):
1287
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1421
1288
  ```
1422
1289
 
1423
- 'parameters' can also be a list of strings and tuples like so:
1424
- ```
1425
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1426
- ```
1427
- This is equivalent to:
1428
- ```
1429
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1290
+ ```python
1291
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
1292
+ @step
1293
+ def finetune_model(self):
1294
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1295
+ # path_to_model will be /my-directory
1430
1296
  ```
1431
1297
 
1432
-
1433
- Parameters
1434
- ----------
1435
- event : Union[str, Dict[str, Any]], optional, default None
1436
- Event dependency for this flow.
1437
- events : List[Union[str, Dict[str, Any]]], default []
1438
- Events dependency for this flow.
1439
- options : Dict[str, Any], default {}
1440
- Backend-specific configuration for tuning eventing behavior.
1441
- """
1442
- ...
1443
-
1444
- @typing.overload
1445
- 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]]:
1446
- """
1447
- Specifies the Conda environment for all steps of the flow.
1448
-
1449
- Use `@conda_base` to set common libraries required by all
1450
- steps and use `@conda` to specify step-specific additions.
1298
+ ```python
1299
+ # Takes all the arguments passed to `snapshot_download`
1300
+ # except for `local_dir`
1301
+ @huggingface_hub(load=[
1302
+ {
1303
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
1304
+ },
1305
+ {
1306
+ "repo_id": "myorg/mistral-lora",
1307
+ "repo_type": "model",
1308
+ },
1309
+ ])
1310
+ @step
1311
+ def finetune_model(self):
1312
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1313
+ # path_to_model will be /my-directory
1314
+ ```
1451
1315
 
1452
1316
 
1453
1317
  Parameters
1454
1318
  ----------
1455
- packages : Dict[str, str], default {}
1456
- Packages to use for this flow. The key is the name of the package
1457
- and the value is the version to use.
1458
- libraries : Dict[str, str], default {}
1459
- Supported for backward compatibility. When used with packages, packages will take precedence.
1460
- python : str, optional, default None
1461
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1462
- that the version used will correspond to the version of the Python interpreter used to start the run.
1463
- disabled : bool, default False
1464
- If set to True, disables Conda.
1465
- """
1466
- ...
1467
-
1468
- @typing.overload
1469
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1470
- ...
1471
-
1472
- 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):
1473
- """
1474
- Specifies the Conda environment for all steps of the flow.
1319
+ temp_dir_root : str, optional
1320
+ The root directory that will hold the temporary directory where objects will be downloaded.
1475
1321
 
1476
- Use `@conda_base` to set common libraries required by all
1477
- steps and use `@conda` to specify step-specific additions.
1322
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
1323
+ The list of repos (models/datasets) to load.
1478
1324
 
1325
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
1479
1326
 
1480
- Parameters
1481
- ----------
1482
- packages : Dict[str, str], default {}
1483
- Packages to use for this flow. The key is the name of the package
1484
- and the value is the version to use.
1485
- libraries : Dict[str, str], default {}
1486
- Supported for backward compatibility. When used with packages, packages will take precedence.
1487
- python : str, optional, default None
1488
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1489
- that the version used will correspond to the version of the Python interpreter used to start the run.
1490
- disabled : bool, default False
1491
- If set to True, disables Conda.
1327
+ - If repo (model/dataset) is not found in the datastore:
1328
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
1329
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
1330
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
1331
+
1332
+ - If repo is found in the datastore:
1333
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
1492
1334
  """
1493
1335
  ...
1494
1336
 
1495
- @typing.overload
1496
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1337
+ 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]]]:
1497
1338
  """
1498
- Specifies the PyPI packages for all steps of the flow.
1339
+ Specifies that this step should execute on DGX cloud.
1499
1340
 
1500
- Use `@pypi_base` to set common packages required by all
1501
- steps and use `@pypi` to specify step-specific overrides.
1502
1341
 
1503
1342
  Parameters
1504
1343
  ----------
1505
- packages : Dict[str, str], default: {}
1506
- Packages to use for this flow. The key is the name of the package
1507
- and the value is the version to use.
1508
- python : str, optional, default: None
1509
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1510
- that the version used will correspond to the version of the Python interpreter used to start the run.
1344
+ gpu : int
1345
+ Number of GPUs to use.
1346
+ gpu_type : str
1347
+ Type of Nvidia GPU to use.
1511
1348
  """
1512
1349
  ...
1513
1350
 
1514
- @typing.overload
1515
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1516
- ...
1517
-
1518
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1351
+ 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]]:
1519
1352
  """
1520
- Specifies the PyPI packages for all steps of the flow.
1353
+ 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)
1354
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1355
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1356
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1357
+ starts only after all sensors finish.
1521
1358
 
1522
- Use `@pypi_base` to set common packages required by all
1523
- steps and use `@pypi` to specify step-specific overrides.
1524
1359
 
1525
1360
  Parameters
1526
1361
  ----------
1527
- packages : Dict[str, str], default: {}
1528
- Packages to use for this flow. The key is the name of the package
1529
- and the value is the version to use.
1530
- python : str, optional, default: None
1531
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1532
- that the version used will correspond to the version of the Python interpreter used to start the run.
1362
+ timeout : int
1363
+ Time, in seconds before the task times out and fails. (Default: 3600)
1364
+ poke_interval : int
1365
+ Time in seconds that the job should wait in between each try. (Default: 60)
1366
+ mode : str
1367
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1368
+ exponential_backoff : bool
1369
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1370
+ pool : str
1371
+ the slot pool this task should run in,
1372
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1373
+ soft_fail : bool
1374
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1375
+ name : str
1376
+ Name of the sensor on Airflow
1377
+ description : str
1378
+ Description of sensor in the Airflow UI
1379
+ bucket_key : Union[str, List[str]]
1380
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1381
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1382
+ bucket_name : str
1383
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1384
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1385
+ wildcard_match : bool
1386
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1387
+ aws_conn_id : str
1388
+ a reference to the s3 connection on Airflow. (Default: None)
1389
+ verify : bool
1390
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1533
1391
  """
1534
1392
  ...
1535
1393
 
@@ -1584,6 +1442,90 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1584
1442
  """
1585
1443
  ...
1586
1444
 
1445
+ @typing.overload
1446
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1447
+ """
1448
+ Specifies the PyPI packages for all steps of the flow.
1449
+
1450
+ Use `@pypi_base` to set common packages required by all
1451
+ steps and use `@pypi` to specify step-specific overrides.
1452
+
1453
+ Parameters
1454
+ ----------
1455
+ packages : Dict[str, str], default: {}
1456
+ Packages to use for this flow. The key is the name of the package
1457
+ and the value is the version to use.
1458
+ python : str, optional, default: None
1459
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1460
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1461
+ """
1462
+ ...
1463
+
1464
+ @typing.overload
1465
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1466
+ ...
1467
+
1468
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1469
+ """
1470
+ Specifies the PyPI packages for all steps of the flow.
1471
+
1472
+ Use `@pypi_base` to set common packages required by all
1473
+ steps and use `@pypi` to specify step-specific overrides.
1474
+
1475
+ Parameters
1476
+ ----------
1477
+ packages : Dict[str, str], default: {}
1478
+ Packages to use for this flow. The key is the name of the package
1479
+ and the value is the version to use.
1480
+ python : str, optional, default: None
1481
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1482
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1483
+ """
1484
+ ...
1485
+
1486
+ 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]]:
1487
+ """
1488
+ 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.
1489
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1490
+
1491
+
1492
+ Parameters
1493
+ ----------
1494
+ timeout : int
1495
+ Time, in seconds before the task times out and fails. (Default: 3600)
1496
+ poke_interval : int
1497
+ Time in seconds that the job should wait in between each try. (Default: 60)
1498
+ mode : str
1499
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1500
+ exponential_backoff : bool
1501
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1502
+ pool : str
1503
+ the slot pool this task should run in,
1504
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1505
+ soft_fail : bool
1506
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1507
+ name : str
1508
+ Name of the sensor on Airflow
1509
+ description : str
1510
+ Description of sensor in the Airflow UI
1511
+ external_dag_id : str
1512
+ The dag_id that contains the task you want to wait for.
1513
+ external_task_ids : List[str]
1514
+ The list of task_ids that you want to wait for.
1515
+ If None (default value) the sensor waits for the DAG. (Default: None)
1516
+ allowed_states : List[str]
1517
+ Iterable of allowed states, (Default: ['success'])
1518
+ failed_states : List[str]
1519
+ Iterable of failed or dis-allowed states. (Default: None)
1520
+ execution_delta : datetime.timedelta
1521
+ time difference with the previous execution to look at,
1522
+ the default is the same logical date as the current task or DAG. (Default: None)
1523
+ check_existence: bool
1524
+ Set to True to check if the external task exists or check if
1525
+ the DAG to wait for exists. (Default: True)
1526
+ """
1527
+ ...
1528
+
1587
1529
  def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1588
1530
  """
1589
1531
  Specifies what flows belong to the same project.
@@ -1619,46 +1561,104 @@ def project(*, name: str, branch: typing.Optional[str] = None, production: bool
1619
1561
  """
1620
1562
  ...
1621
1563
 
1622
- 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]]:
1564
+ @typing.overload
1565
+ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1566
+ """
1567
+ Specifies the flow(s) that this flow depends on.
1568
+
1569
+ ```
1570
+ @trigger_on_finish(flow='FooFlow')
1571
+ ```
1572
+ or
1573
+ ```
1574
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1575
+ ```
1576
+ This decorator respects the @project decorator and triggers the flow
1577
+ when upstream runs within the same namespace complete successfully
1578
+
1579
+ Additionally, you can specify project aware upstream flow dependencies
1580
+ by specifying the fully qualified project_flow_name.
1581
+ ```
1582
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1583
+ ```
1584
+ or
1585
+ ```
1586
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1587
+ ```
1588
+
1589
+ You can also specify just the project or project branch (other values will be
1590
+ inferred from the current project or project branch):
1591
+ ```
1592
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1593
+ ```
1594
+
1595
+ Note that `branch` is typically one of:
1596
+ - `prod`
1597
+ - `user.bob`
1598
+ - `test.my_experiment`
1599
+ - `prod.staging`
1600
+
1601
+
1602
+ Parameters
1603
+ ----------
1604
+ flow : Union[str, Dict[str, str]], optional, default None
1605
+ Upstream flow dependency for this flow.
1606
+ flows : List[Union[str, Dict[str, str]]], default []
1607
+ Upstream flow dependencies for this flow.
1608
+ options : Dict[str, Any], default {}
1609
+ Backend-specific configuration for tuning eventing behavior.
1610
+ """
1611
+ ...
1612
+
1613
+ @typing.overload
1614
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1615
+ ...
1616
+
1617
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1623
1618
  """
1624
- 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)
1625
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1626
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1627
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1628
- starts only after all sensors finish.
1619
+ Specifies the flow(s) that this flow depends on.
1620
+
1621
+ ```
1622
+ @trigger_on_finish(flow='FooFlow')
1623
+ ```
1624
+ or
1625
+ ```
1626
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1627
+ ```
1628
+ This decorator respects the @project decorator and triggers the flow
1629
+ when upstream runs within the same namespace complete successfully
1630
+
1631
+ Additionally, you can specify project aware upstream flow dependencies
1632
+ by specifying the fully qualified project_flow_name.
1633
+ ```
1634
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1635
+ ```
1636
+ or
1637
+ ```
1638
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1639
+ ```
1640
+
1641
+ You can also specify just the project or project branch (other values will be
1642
+ inferred from the current project or project branch):
1643
+ ```
1644
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1645
+ ```
1646
+
1647
+ Note that `branch` is typically one of:
1648
+ - `prod`
1649
+ - `user.bob`
1650
+ - `test.my_experiment`
1651
+ - `prod.staging`
1629
1652
 
1630
1653
 
1631
1654
  Parameters
1632
1655
  ----------
1633
- timeout : int
1634
- Time, in seconds before the task times out and fails. (Default: 3600)
1635
- poke_interval : int
1636
- Time in seconds that the job should wait in between each try. (Default: 60)
1637
- mode : str
1638
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1639
- exponential_backoff : bool
1640
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1641
- pool : str
1642
- the slot pool this task should run in,
1643
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1644
- soft_fail : bool
1645
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1646
- name : str
1647
- Name of the sensor on Airflow
1648
- description : str
1649
- Description of sensor in the Airflow UI
1650
- bucket_key : Union[str, List[str]]
1651
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1652
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1653
- bucket_name : str
1654
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1655
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1656
- wildcard_match : bool
1657
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1658
- aws_conn_id : str
1659
- a reference to the s3 connection on Airflow. (Default: None)
1660
- verify : bool
1661
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1656
+ flow : Union[str, Dict[str, str]], optional, default None
1657
+ Upstream flow dependency for this flow.
1658
+ flows : List[Union[str, Dict[str, str]]], default []
1659
+ Upstream flow dependencies for this flow.
1660
+ options : Dict[str, Any], default {}
1661
+ Backend-specific configuration for tuning eventing behavior.
1662
1662
  """
1663
1663
  ...
1664
1664
 
@@ -1777,148 +1777,148 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1777
1777
  ...
1778
1778
 
1779
1779
  @typing.overload
1780
- def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1780
+ 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]]:
1781
1781
  """
1782
- Specifies the flow(s) that this flow depends on.
1782
+ Specifies the Conda environment for all steps of the flow.
1783
+
1784
+ Use `@conda_base` to set common libraries required by all
1785
+ steps and use `@conda` to specify step-specific additions.
1786
+
1787
+
1788
+ Parameters
1789
+ ----------
1790
+ packages : Dict[str, str], default {}
1791
+ Packages to use for this flow. The key is the name of the package
1792
+ and the value is the version to use.
1793
+ libraries : Dict[str, str], default {}
1794
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1795
+ python : str, optional, default None
1796
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1797
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1798
+ disabled : bool, default False
1799
+ If set to True, disables Conda.
1800
+ """
1801
+ ...
1802
+
1803
+ @typing.overload
1804
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1805
+ ...
1806
+
1807
+ 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):
1808
+ """
1809
+ Specifies the Conda environment for all steps of the flow.
1810
+
1811
+ Use `@conda_base` to set common libraries required by all
1812
+ steps and use `@conda` to specify step-specific additions.
1813
+
1814
+
1815
+ Parameters
1816
+ ----------
1817
+ packages : Dict[str, str], default {}
1818
+ Packages to use for this flow. The key is the name of the package
1819
+ and the value is the version to use.
1820
+ libraries : Dict[str, str], default {}
1821
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1822
+ python : str, optional, default None
1823
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1824
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1825
+ disabled : bool, default False
1826
+ If set to True, disables Conda.
1827
+ """
1828
+ ...
1829
+
1830
+ @typing.overload
1831
+ 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]]:
1832
+ """
1833
+ Specifies the event(s) that this flow depends on.
1783
1834
 
1784
1835
  ```
1785
- @trigger_on_finish(flow='FooFlow')
1836
+ @trigger(event='foo')
1786
1837
  ```
1787
1838
  or
1788
1839
  ```
1789
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1840
+ @trigger(events=['foo', 'bar'])
1790
1841
  ```
1791
- This decorator respects the @project decorator and triggers the flow
1792
- when upstream runs within the same namespace complete successfully
1793
1842
 
1794
- Additionally, you can specify project aware upstream flow dependencies
1795
- by specifying the fully qualified project_flow_name.
1843
+ Additionally, you can specify the parameter mappings
1844
+ to map event payload to Metaflow parameters for the flow.
1796
1845
  ```
1797
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1846
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1798
1847
  ```
1799
1848
  or
1800
1849
  ```
1801
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1850
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1851
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1802
1852
  ```
1803
1853
 
1804
- You can also specify just the project or project branch (other values will be
1805
- inferred from the current project or project branch):
1854
+ 'parameters' can also be a list of strings and tuples like so:
1806
1855
  ```
1807
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1856
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1857
+ ```
1858
+ This is equivalent to:
1859
+ ```
1860
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1808
1861
  ```
1809
-
1810
- Note that `branch` is typically one of:
1811
- - `prod`
1812
- - `user.bob`
1813
- - `test.my_experiment`
1814
- - `prod.staging`
1815
1862
 
1816
1863
 
1817
1864
  Parameters
1818
1865
  ----------
1819
- flow : Union[str, Dict[str, str]], optional, default None
1820
- Upstream flow dependency for this flow.
1821
- flows : List[Union[str, Dict[str, str]]], default []
1822
- Upstream flow dependencies for this flow.
1866
+ event : Union[str, Dict[str, Any]], optional, default None
1867
+ Event dependency for this flow.
1868
+ events : List[Union[str, Dict[str, Any]]], default []
1869
+ Events dependency for this flow.
1823
1870
  options : Dict[str, Any], default {}
1824
1871
  Backend-specific configuration for tuning eventing behavior.
1825
1872
  """
1826
1873
  ...
1827
1874
 
1828
1875
  @typing.overload
1829
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1876
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1830
1877
  ...
1831
1878
 
1832
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1879
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, 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] = {}):
1833
1880
  """
1834
- Specifies the flow(s) that this flow depends on.
1881
+ Specifies the event(s) that this flow depends on.
1835
1882
 
1836
1883
  ```
1837
- @trigger_on_finish(flow='FooFlow')
1884
+ @trigger(event='foo')
1838
1885
  ```
1839
1886
  or
1840
1887
  ```
1841
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1888
+ @trigger(events=['foo', 'bar'])
1842
1889
  ```
1843
- This decorator respects the @project decorator and triggers the flow
1844
- when upstream runs within the same namespace complete successfully
1845
1890
 
1846
- Additionally, you can specify project aware upstream flow dependencies
1847
- by specifying the fully qualified project_flow_name.
1891
+ Additionally, you can specify the parameter mappings
1892
+ to map event payload to Metaflow parameters for the flow.
1848
1893
  ```
1849
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1894
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1850
1895
  ```
1851
1896
  or
1852
1897
  ```
1853
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1898
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1899
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1854
1900
  ```
1855
1901
 
1856
- You can also specify just the project or project branch (other values will be
1857
- inferred from the current project or project branch):
1902
+ 'parameters' can also be a list of strings and tuples like so:
1858
1903
  ```
1859
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1904
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1905
+ ```
1906
+ This is equivalent to:
1907
+ ```
1908
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1860
1909
  ```
1861
-
1862
- Note that `branch` is typically one of:
1863
- - `prod`
1864
- - `user.bob`
1865
- - `test.my_experiment`
1866
- - `prod.staging`
1867
1910
 
1868
1911
 
1869
1912
  Parameters
1870
1913
  ----------
1871
- flow : Union[str, Dict[str, str]], optional, default None
1872
- Upstream flow dependency for this flow.
1873
- flows : List[Union[str, Dict[str, str]]], default []
1874
- Upstream flow dependencies for this flow.
1914
+ event : Union[str, Dict[str, Any]], optional, default None
1915
+ Event dependency for this flow.
1916
+ events : List[Union[str, Dict[str, Any]]], default []
1917
+ Events dependency for this flow.
1875
1918
  options : Dict[str, Any], default {}
1876
1919
  Backend-specific configuration for tuning eventing behavior.
1877
1920
  """
1878
1921
  ...
1879
1922
 
1880
- 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]]:
1881
- """
1882
- 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.
1883
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1884
-
1885
-
1886
- Parameters
1887
- ----------
1888
- timeout : int
1889
- Time, in seconds before the task times out and fails. (Default: 3600)
1890
- poke_interval : int
1891
- Time in seconds that the job should wait in between each try. (Default: 60)
1892
- mode : str
1893
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1894
- exponential_backoff : bool
1895
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1896
- pool : str
1897
- the slot pool this task should run in,
1898
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1899
- soft_fail : bool
1900
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1901
- name : str
1902
- Name of the sensor on Airflow
1903
- description : str
1904
- Description of sensor in the Airflow UI
1905
- external_dag_id : str
1906
- The dag_id that contains the task you want to wait for.
1907
- external_task_ids : List[str]
1908
- The list of task_ids that you want to wait for.
1909
- If None (default value) the sensor waits for the DAG. (Default: None)
1910
- allowed_states : List[str]
1911
- Iterable of allowed states, (Default: ['success'])
1912
- failed_states : List[str]
1913
- Iterable of failed or dis-allowed states. (Default: None)
1914
- execution_delta : datetime.timedelta
1915
- time difference with the previous execution to look at,
1916
- the default is the same logical date as the current task or DAG. (Default: None)
1917
- check_existence: bool
1918
- Set to True to check if the external task exists or check if
1919
- the DAG to wait for exists. (Default: True)
1920
- """
1921
- ...
1922
-
1923
1923
  pkg_name: str
1924
1924