ob-metaflow-stubs 6.0.7.3__py2.py3-none-any.whl → 6.0.8.1__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.

Potentially problematic release.


This version of ob-metaflow-stubs might be problematic. Click here for more details.

Files changed (262) hide show
  1. metaflow-stubs/__init__.pyi +912 -912
  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 +1 -1
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  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 +1 -1
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +50 -50
  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 +2 -2
  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 +1 -1
  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 +4 -4
  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 +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +4 -4
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +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 +2 -2
  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 +2 -2
  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 +2 -2
  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 +1 -1
  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 +3 -3
  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 +1 -1
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +1 -1
  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 +1 -1
  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 +2 -2
  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 +3 -3
  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 +2 -2
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +3 -3
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +3 -3
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +2 -2
  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 +4 -4
  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 +2 -2
  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 +2 -2
  99. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +2 -2
  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 +1 -1
  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 +6 -6
  119. metaflow-stubs/packaging_sys/backend.pyi +3 -3
  120. metaflow-stubs/packaging_sys/distribution_support.pyi +3 -3
  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 +2 -2
  124. metaflow-stubs/parameters.pyi +3 -3
  125. metaflow-stubs/plugins/__init__.pyi +10 -10
  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 +2 -2
  137. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +1 -1
  138. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  139. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +1 -1
  140. metaflow-stubs/plugins/argo/exit_hooks.pyi +2 -2
  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 +2 -2
  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 +3 -3
  156. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +1 -1
  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 +2 -2
  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 +1 -1
  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 +2 -2
  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 +1 -1
  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 +1 -1
  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 +3 -3
  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 +2 -2
  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 +1 -1
  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 +1 -1
  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 +4 -4
  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 +2 -2
  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 +1 -1
  230. metaflow-stubs/plugins/timeout_decorator.pyi +1 -1
  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 +2 -2
  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 +4 -4
  238. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  239. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  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 +1 -1
  244. metaflow-stubs/system/__init__.pyi +1 -1
  245. metaflow-stubs/system/system_logger.pyi +2 -2
  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 +1 -1
  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 +4 -4
  255. metaflow-stubs/user_decorators/mutable_step.pyi +4 -4
  256. metaflow-stubs/user_decorators/user_flow_decorator.pyi +4 -4
  257. metaflow-stubs/user_decorators/user_step_decorator.pyi +4 -4
  258. {ob_metaflow_stubs-6.0.7.3.dist-info → ob_metaflow_stubs-6.0.8.1.dist-info}/METADATA +1 -1
  259. ob_metaflow_stubs-6.0.8.1.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.8.1.dist-info}/WHEEL +0 -0
  262. {ob_metaflow_stubs-6.0.7.3.dist-info → ob_metaflow_stubs-6.0.8.1.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-21T23:31:59.916830 #
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
 
@@ -39,10 +39,10 @@ from .user_decorators.user_step_decorator import UserStepDecorator as UserStepDe
39
39
  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
- from . import events as events
43
- from . import tuple_util as tuple_util
44
42
  from . import cards as cards
43
+ from . import tuple_util as tuple_util
45
44
  from . import metaflow_git as metaflow_git
45
+ from . import events as events
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
@@ -167,91 +167,152 @@ 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 ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
171
171
  """
172
- Specifies that this step should execute on DGX cloud.
172
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
173
+
174
+ User code call
175
+ --------------
176
+ @ollama(
177
+ models=[...],
178
+ ...
179
+ )
180
+
181
+ Valid backend options
182
+ ---------------------
183
+ - 'local': Run as a separate process on the local task machine.
184
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
185
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
186
+
187
+ Valid model options
188
+ -------------------
189
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
173
190
 
174
191
 
175
192
  Parameters
176
193
  ----------
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.
194
+ models: list[str]
195
+ List of Ollama containers running models in sidecars.
196
+ backend: str
197
+ Determines where and how to run the Ollama process.
198
+ force_pull: bool
199
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
200
+ cache_update_policy: str
201
+ Cache update policy: "auto", "force", or "never".
202
+ force_cache_update: bool
203
+ Simple override for "force" cache update policy.
204
+ debug: bool
205
+ Whether to turn on verbose debugging logs.
206
+ circuit_breaker_config: dict
207
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
208
+ timeout_config: dict
209
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
183
210
  """
184
211
  ...
185
212
 
186
213
  @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]]]:
214
+ 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]]]:
188
215
  """
189
- Specifies environment variables to be set prior to the execution of a step.
216
+ Specifies the number of times the task corresponding
217
+ to a step needs to be retried.
218
+
219
+ This decorator is useful for handling transient errors, such as networking issues.
220
+ If your task contains operations that can't be retried safely, e.g. database updates,
221
+ it is advisable to annotate it with `@retry(times=0)`.
222
+
223
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
224
+ decorator will execute a no-op task after all retries have been exhausted,
225
+ ensuring that the flow execution can continue.
190
226
 
191
227
 
192
228
  Parameters
193
229
  ----------
194
- vars : Dict[str, str], default {}
195
- Dictionary of environment variables to set.
230
+ times : int, default 3
231
+ Number of times to retry this task.
232
+ minutes_between_retries : int, default 2
233
+ Number of minutes between retries.
196
234
  """
197
235
  ...
198
236
 
199
237
  @typing.overload
200
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
238
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
201
239
  ...
202
240
 
203
241
  @typing.overload
204
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
242
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
205
243
  ...
206
244
 
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] = {}):
245
+ 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):
208
246
  """
209
- Specifies environment variables to be set prior to the execution of a step.
247
+ Specifies the number of times the task corresponding
248
+ to a step needs to be retried.
249
+
250
+ This decorator is useful for handling transient errors, such as networking issues.
251
+ If your task contains operations that can't be retried safely, e.g. database updates,
252
+ it is advisable to annotate it with `@retry(times=0)`.
253
+
254
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
255
+ decorator will execute a no-op task after all retries have been exhausted,
256
+ ensuring that the flow execution can continue.
210
257
 
211
258
 
212
259
  Parameters
213
260
  ----------
214
- vars : Dict[str, str], default {}
215
- Dictionary of environment variables to set.
261
+ times : int, default 3
262
+ Number of times to retry this task.
263
+ minutes_between_retries : int, default 2
264
+ Number of minutes between retries.
216
265
  """
217
266
  ...
218
267
 
219
268
  @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]]]:
269
+ 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
270
  """
222
- Specifies secrets to be retrieved and injected as environment variables prior to
223
- the execution of a step.
271
+ Specifies that the step will success under all circumstances.
272
+
273
+ The decorator will create an optional artifact, specified by `var`, which
274
+ contains the exception raised. You can use it to detect the presence
275
+ of errors, indicating that all happy-path artifacts produced by the step
276
+ are missing.
224
277
 
225
278
 
226
279
  Parameters
227
280
  ----------
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
281
+ var : str, optional, default None
282
+ Name of the artifact in which to store the caught exception.
283
+ If not specified, the exception is not stored.
284
+ print_exception : bool, default True
285
+ Determines whether or not the exception is printed to
286
+ stdout when caught.
232
287
  """
233
288
  ...
234
289
 
235
290
  @typing.overload
236
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
291
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
237
292
  ...
238
293
 
239
294
  @typing.overload
240
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
295
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
241
296
  ...
242
297
 
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):
298
+ 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
299
  """
245
- Specifies secrets to be retrieved and injected as environment variables prior to
246
- the execution of a step.
300
+ Specifies that the step will success under all circumstances.
301
+
302
+ The decorator will create an optional artifact, specified by `var`, which
303
+ contains the exception raised. You can use it to detect the presence
304
+ of errors, indicating that all happy-path artifacts produced by the step
305
+ are missing.
247
306
 
248
307
 
249
308
  Parameters
250
309
  ----------
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
310
+ var : str, optional, default None
311
+ Name of the artifact in which to store the caught exception.
312
+ If not specified, the exception is not stored.
313
+ print_exception : bool, default True
314
+ Determines whether or not the exception is printed to
315
+ stdout when caught.
255
316
  """
256
317
  ...
257
318
 
@@ -384,97 +445,53 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
384
445
  """
385
446
  ...
386
447
 
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]]]:
388
- """
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.
398
- """
399
- ...
400
-
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]]]:
448
+ 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]]]:
402
449
  """
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"])`
450
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
417
451
 
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
- ```
452
+ User code call
453
+ --------------
454
+ @vllm(
455
+ model="...",
456
+ ...
457
+ )
425
458
 
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
- ```
459
+ Valid backend options
460
+ ---------------------
461
+ - 'local': Run as a separate process on the local task machine.
433
462
 
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
- ```
463
+ Valid model options
464
+ -------------------
465
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
441
466
 
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
- ```
467
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
468
+ If you need multiple models, you must create multiple @vllm decorators.
459
469
 
460
470
 
461
471
  Parameters
462
472
  ----------
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)
473
+ model: str
474
+ HuggingFace model identifier to be served by vLLM.
475
+ backend: str
476
+ Determines where and how to run the vLLM process.
477
+ openai_api_server: bool
478
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
479
+ Default is False (uses native engine).
480
+ Set to True for backward compatibility with existing code.
481
+ debug: bool
482
+ Whether to turn on verbose debugging logs.
483
+ card_refresh_interval: int
484
+ Interval in seconds for refreshing the vLLM status card.
485
+ Only used when openai_api_server=True.
486
+ max_retries: int
487
+ Maximum number of retries checking for vLLM server startup.
488
+ Only used when openai_api_server=True.
489
+ retry_alert_frequency: int
490
+ Frequency of alert logs for vLLM server startup retries.
491
+ Only used when openai_api_server=True.
492
+ engine_args : dict
493
+ Additional keyword arguments to pass to the vLLM engine.
494
+ For example, `tensor_parallel_size=2`.
478
495
  """
479
496
  ...
480
497
 
@@ -537,46 +554,230 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
537
554
  """
538
555
  ...
539
556
 
540
- def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
557
+ @typing.overload
558
+ 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]]]:
541
559
  """
542
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
560
+ Enables checkpointing for a step.
543
561
 
544
- User code call
545
- --------------
546
- @ollama(
547
- models=[...],
548
- ...
549
- )
562
+ > Examples
550
563
 
551
- Valid backend options
552
- ---------------------
553
- - 'local': Run as a separate process on the local task machine.
554
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
555
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
564
+ - Saving Checkpoints
556
565
 
557
- Valid model options
558
- -------------------
559
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
566
+ ```python
567
+ @checkpoint
568
+ @step
569
+ def train(self):
570
+ model = create_model(self.parameters, checkpoint_path = None)
571
+ for i in range(self.epochs):
572
+ # some training logic
573
+ loss = model.train(self.dataset)
574
+ if i % 10 == 0:
575
+ model.save(
576
+ current.checkpoint.directory,
577
+ )
578
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
579
+ # and returns a reference dictionary to the checkpoint saved in the datastore
580
+ self.latest_checkpoint = current.checkpoint.save(
581
+ name="epoch_checkpoint",
582
+ metadata={
583
+ "epoch": i,
584
+ "loss": loss,
585
+ }
586
+ )
587
+ ```
588
+
589
+ - Using Loaded Checkpoints
590
+
591
+ ```python
592
+ @retry(times=3)
593
+ @checkpoint
594
+ @step
595
+ def train(self):
596
+ # Assume that the task has restarted and the previous attempt of the task
597
+ # saved a checkpoint
598
+ checkpoint_path = None
599
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
600
+ print("Loaded checkpoint from the previous attempt")
601
+ checkpoint_path = current.checkpoint.directory
602
+
603
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
604
+ for i in range(self.epochs):
605
+ ...
606
+ ```
560
607
 
561
608
 
562
609
  Parameters
563
610
  ----------
564
- models: list[str]
565
- List of Ollama containers running models in sidecars.
566
- backend: str
567
- Determines where and how to run the Ollama process.
568
- force_pull: bool
569
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
570
- cache_update_policy: str
571
- Cache update policy: "auto", "force", or "never".
572
- force_cache_update: bool
573
- Simple override for "force" cache update policy.
574
- debug: bool
575
- Whether to turn on verbose debugging logs.
576
- circuit_breaker_config: dict
577
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
578
- timeout_config: dict
579
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
611
+ load_policy : str, default: "fresh"
612
+ The policy for loading the checkpoint. The following policies are supported:
613
+ - "eager": Loads the the latest available checkpoint within the namespace.
614
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
615
+ will be loaded at the start of the task.
616
+ - "none": Do not load any checkpoint
617
+ - "fresh": Loads the lastest checkpoint created within the running Task.
618
+ This mode helps loading checkpoints across various retry attempts of the same task.
619
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
620
+ created within the task will be loaded when the task is retries execution on failure.
621
+
622
+ temp_dir_root : str, default: None
623
+ The root directory under which `current.checkpoint.directory` will be created.
624
+ """
625
+ ...
626
+
627
+ @typing.overload
628
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
629
+ ...
630
+
631
+ @typing.overload
632
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
633
+ ...
634
+
635
+ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
636
+ """
637
+ Enables checkpointing for a step.
638
+
639
+ > Examples
640
+
641
+ - Saving Checkpoints
642
+
643
+ ```python
644
+ @checkpoint
645
+ @step
646
+ def train(self):
647
+ model = create_model(self.parameters, checkpoint_path = None)
648
+ for i in range(self.epochs):
649
+ # some training logic
650
+ loss = model.train(self.dataset)
651
+ if i % 10 == 0:
652
+ model.save(
653
+ current.checkpoint.directory,
654
+ )
655
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
656
+ # and returns a reference dictionary to the checkpoint saved in the datastore
657
+ self.latest_checkpoint = current.checkpoint.save(
658
+ name="epoch_checkpoint",
659
+ metadata={
660
+ "epoch": i,
661
+ "loss": loss,
662
+ }
663
+ )
664
+ ```
665
+
666
+ - Using Loaded Checkpoints
667
+
668
+ ```python
669
+ @retry(times=3)
670
+ @checkpoint
671
+ @step
672
+ def train(self):
673
+ # Assume that the task has restarted and the previous attempt of the task
674
+ # saved a checkpoint
675
+ checkpoint_path = None
676
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
677
+ print("Loaded checkpoint from the previous attempt")
678
+ checkpoint_path = current.checkpoint.directory
679
+
680
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
681
+ for i in range(self.epochs):
682
+ ...
683
+ ```
684
+
685
+
686
+ Parameters
687
+ ----------
688
+ load_policy : str, default: "fresh"
689
+ The policy for loading the checkpoint. The following policies are supported:
690
+ - "eager": Loads the the latest available checkpoint within the namespace.
691
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
692
+ will be loaded at the start of the task.
693
+ - "none": Do not load any checkpoint
694
+ - "fresh": Loads the lastest checkpoint created within the running Task.
695
+ This mode helps loading checkpoints across various retry attempts of the same task.
696
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
697
+ created within the task will be loaded when the task is retries execution on failure.
698
+
699
+ temp_dir_root : str, default: None
700
+ The root directory under which `current.checkpoint.directory` will be created.
701
+ """
702
+ ...
703
+
704
+ 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]]]:
705
+ """
706
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
707
+
708
+ > Examples
709
+
710
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
711
+ ```python
712
+ @huggingface_hub
713
+ @step
714
+ def pull_model_from_huggingface(self):
715
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
716
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
717
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
718
+ # value of the function is a reference to the model in the backend storage.
719
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
720
+
721
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
722
+ self.llama_model = current.huggingface_hub.snapshot_download(
723
+ repo_id=self.model_id,
724
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
725
+ )
726
+ self.next(self.train)
727
+ ```
728
+
729
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
730
+ ```python
731
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
732
+ @step
733
+ def pull_model_from_huggingface(self):
734
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
735
+ ```
736
+
737
+ ```python
738
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
739
+ @step
740
+ def finetune_model(self):
741
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
742
+ # path_to_model will be /my-directory
743
+ ```
744
+
745
+ ```python
746
+ # Takes all the arguments passed to `snapshot_download`
747
+ # except for `local_dir`
748
+ @huggingface_hub(load=[
749
+ {
750
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
751
+ },
752
+ {
753
+ "repo_id": "myorg/mistral-lora",
754
+ "repo_type": "model",
755
+ },
756
+ ])
757
+ @step
758
+ def finetune_model(self):
759
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
760
+ # path_to_model will be /my-directory
761
+ ```
762
+
763
+
764
+ Parameters
765
+ ----------
766
+ temp_dir_root : str, optional
767
+ The root directory that will hold the temporary directory where objects will be downloaded.
768
+
769
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
770
+ The list of repos (models/datasets) to load.
771
+
772
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
773
+
774
+ - If repo (model/dataset) is not found in the datastore:
775
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
776
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
777
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
778
+
779
+ - If repo is found in the datastore:
780
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
580
781
  """
581
782
  ...
582
783
 
@@ -670,14 +871,47 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
670
871
  ...
671
872
 
672
873
  @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]]]:
874
+ 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]]]:
674
875
  """
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.
876
+ Specifies environment variables to be set prior to the execution of a step.
877
+
878
+
879
+ Parameters
880
+ ----------
881
+ vars : Dict[str, str], default {}
882
+ Dictionary of environment variables to set.
883
+ """
884
+ ...
885
+
886
+ @typing.overload
887
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
888
+ ...
889
+
890
+ @typing.overload
891
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
892
+ ...
893
+
894
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
895
+ """
896
+ Specifies environment variables to be set prior to the execution of a step.
897
+
898
+
899
+ Parameters
900
+ ----------
901
+ vars : Dict[str, str], default {}
902
+ Dictionary of environment variables to set.
903
+ """
904
+ ...
905
+
906
+ @typing.overload
907
+ 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]]]:
908
+ """
909
+ Specifies the PyPI packages for the step.
910
+
911
+ Information in this decorator will augment any
912
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
913
+ you can use `@pypi_base` to set packages required by all
914
+ steps and use `@pypi` to specify step-specific overrides.
681
915
 
682
916
 
683
917
  Parameters
@@ -720,22 +954,19 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
720
954
  """
721
955
  ...
722
956
 
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):
957
+ 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]]]:
736
958
  """
737
- Decorator prototype for all step decorators. This function gets specialized
738
- and imported for all decorators types by _import_plugin_decorators().
959
+ Specifies that this step should execute on DGX cloud.
960
+
961
+
962
+ Parameters
963
+ ----------
964
+ gpu : int
965
+ Number of GPUs to use.
966
+ gpu_type : str
967
+ Type of Nvidia GPU to use.
968
+ queue_timeout : int
969
+ Time to keep the job in NVCF's queue.
739
970
  """
740
971
  ...
741
972
 
@@ -779,58 +1010,135 @@ def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag],
779
1010
  """
780
1011
  ...
781
1012
 
1013
+ 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]]]:
1014
+ """
1015
+ Specifies that this step should execute on DGX cloud.
1016
+
1017
+
1018
+ Parameters
1019
+ ----------
1020
+ gpu : int
1021
+ Number of GPUs to use.
1022
+ gpu_type : str
1023
+ Type of Nvidia GPU to use.
1024
+ """
1025
+ ...
1026
+
782
1027
  @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]]]:
1028
+ 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]]]:
784
1029
  """
785
- Specifies the number of times the task corresponding
786
- to a step needs to be retried.
1030
+ Specifies secrets to be retrieved and injected as environment variables prior to
1031
+ the execution of a step.
787
1032
 
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)`.
791
1033
 
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.
1034
+ Parameters
1035
+ ----------
1036
+ sources : List[Union[str, Dict[str, Any]]], default: []
1037
+ List of secret specs, defining how the secrets are to be retrieved
1038
+ role : str, optional, default: None
1039
+ Role to use for fetching secrets
1040
+ """
1041
+ ...
1042
+
1043
+ @typing.overload
1044
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1045
+ ...
1046
+
1047
+ @typing.overload
1048
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1049
+ ...
1050
+
1051
+ 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):
1052
+ """
1053
+ Specifies secrets to be retrieved and injected as environment variables prior to
1054
+ the execution of a step.
795
1055
 
796
1056
 
797
1057
  Parameters
798
1058
  ----------
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.
1059
+ sources : List[Union[str, Dict[str, Any]]], default: []
1060
+ List of secret specs, defining how the secrets are to be retrieved
1061
+ role : str, optional, default: None
1062
+ Role to use for fetching secrets
803
1063
  """
804
1064
  ...
805
1065
 
806
1066
  @typing.overload
807
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1067
+ 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]]]:
1068
+ """
1069
+ Specifies the resources needed when executing this step.
1070
+
1071
+ Use `@resources` to specify the resource requirements
1072
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1073
+
1074
+ You can choose the compute layer on the command line by executing e.g.
1075
+ ```
1076
+ python myflow.py run --with batch
1077
+ ```
1078
+ or
1079
+ ```
1080
+ python myflow.py run --with kubernetes
1081
+ ```
1082
+ which executes the flow on the desired system using the
1083
+ requirements specified in `@resources`.
1084
+
1085
+
1086
+ Parameters
1087
+ ----------
1088
+ cpu : int, default 1
1089
+ Number of CPUs required for this step.
1090
+ gpu : int, optional, default None
1091
+ Number of GPUs required for this step.
1092
+ disk : int, optional, default None
1093
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1094
+ memory : int, default 4096
1095
+ Memory size (in MB) required for this step.
1096
+ shared_memory : int, optional, default None
1097
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1098
+ This parameter maps to the `--shm-size` option in Docker.
1099
+ """
808
1100
  ...
809
1101
 
810
1102
  @typing.overload
811
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1103
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
812
1104
  ...
813
1105
 
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):
1106
+ @typing.overload
1107
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1108
+ ...
1109
+
1110
+ 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
1111
  """
816
- Specifies the number of times the task corresponding
817
- to a step needs to be retried.
1112
+ Specifies the resources needed when executing this step.
818
1113
 
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)`.
1114
+ Use `@resources` to specify the resource requirements
1115
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
822
1116
 
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.
1117
+ You can choose the compute layer on the command line by executing e.g.
1118
+ ```
1119
+ python myflow.py run --with batch
1120
+ ```
1121
+ or
1122
+ ```
1123
+ python myflow.py run --with kubernetes
1124
+ ```
1125
+ which executes the flow on the desired system using the
1126
+ requirements specified in `@resources`.
826
1127
 
827
1128
 
828
1129
  Parameters
829
1130
  ----------
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.
1131
+ cpu : int, default 1
1132
+ Number of CPUs required for this step.
1133
+ gpu : int, optional, default None
1134
+ Number of GPUs required for this step.
1135
+ disk : int, optional, default None
1136
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1137
+ memory : int, default 4096
1138
+ Memory size (in MB) required for this step.
1139
+ shared_memory : int, optional, default None
1140
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1141
+ This parameter maps to the `--shm-size` option in Docker.
834
1142
  """
835
1143
  ...
836
1144
 
@@ -853,387 +1161,51 @@ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
853
1161
  """
854
1162
  ...
855
1163
 
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]]]:
1164
+ 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]]]:
858
1165
  """
859
- Specifies the Conda environment for the step.
860
-
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.
1166
+ S3 Proxy decorator for routing S3 requests through a local proxy service.
865
1167
 
866
1168
 
867
1169
  Parameters
868
1170
  ----------
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.
1171
+ integration_name : str, optional
1172
+ Name of the S3 proxy integration. If not specified, will use the only
1173
+ available S3 proxy integration in the namespace (fails if multiple exist).
1174
+ write_mode : str, optional
1175
+ The desired behavior during write operations to target (origin) S3 bucket.
1176
+ allowed options are:
1177
+ "origin-and-cache" -> write to both the target S3 bucket and local object
1178
+ storage
1179
+ "origin" -> only write to the target S3 bucket
1180
+ "cache" -> only write to the object storage service used for caching
1181
+ debug : bool, optional
1182
+ Enable debug logging for proxy operations.
879
1183
  """
880
1184
  ...
881
1185
 
882
1186
  @typing.overload
883
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1187
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1188
+ """
1189
+ Internal decorator to support Fast bakery
1190
+ """
884
1191
  ...
885
1192
 
886
1193
  @typing.overload
887
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1194
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
888
1195
  ...
889
1196
 
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):
1197
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
891
1198
  """
892
- Specifies the Conda environment for the step.
1199
+ Internal decorator to support Fast bakery
1200
+ """
1201
+ ...
1202
+
1203
+ @typing.overload
1204
+ 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]]]:
1205
+ """
1206
+ Creates a human-readable report, a Metaflow Card, after this step completes.
893
1207
 
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.
898
-
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
923
-
924
- ```python
925
- @checkpoint
926
- @step
927
- def train(self):
928
- model = create_model(self.parameters, checkpoint_path = None)
929
- for i in range(self.epochs):
930
- # some training logic
931
- loss = model.train(self.dataset)
932
- if i % 10 == 0:
933
- model.save(
934
- current.checkpoint.directory,
935
- )
936
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
937
- # and returns a reference dictionary to the checkpoint saved in the datastore
938
- self.latest_checkpoint = current.checkpoint.save(
939
- name="epoch_checkpoint",
940
- metadata={
941
- "epoch": i,
942
- "loss": loss,
943
- }
944
- )
945
- ```
946
-
947
- - Using Loaded Checkpoints
948
-
949
- ```python
950
- @retry(times=3)
951
- @checkpoint
952
- @step
953
- def train(self):
954
- # Assume that the task has restarted and the previous attempt of the task
955
- # saved a checkpoint
956
- checkpoint_path = None
957
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
958
- print("Loaded checkpoint from the previous attempt")
959
- checkpoint_path = current.checkpoint.directory
960
-
961
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
962
- for i in range(self.epochs):
963
- ...
964
- ```
965
-
966
-
967
- Parameters
968
- ----------
969
- load_policy : str, default: "fresh"
970
- The policy for loading the checkpoint. The following policies are supported:
971
- - "eager": Loads the the latest available checkpoint within the namespace.
972
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
973
- will be loaded at the start of the task.
974
- - "none": Do not load any checkpoint
975
- - "fresh": Loads the lastest checkpoint created within the running Task.
976
- This mode helps loading checkpoints across various retry attempts of the same task.
977
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
978
- created within the task will be loaded when the task is retries execution on failure.
979
-
980
- temp_dir_root : str, default: None
981
- The root directory under which `current.checkpoint.directory` will be created.
982
- """
983
- ...
984
-
985
- @typing.overload
986
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
987
- ...
988
-
989
- @typing.overload
990
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
991
- ...
992
-
993
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
994
- """
995
- Enables checkpointing for a step.
996
-
997
- > Examples
998
-
999
- - Saving Checkpoints
1000
-
1001
- ```python
1002
- @checkpoint
1003
- @step
1004
- def train(self):
1005
- model = create_model(self.parameters, checkpoint_path = None)
1006
- for i in range(self.epochs):
1007
- # some training logic
1008
- loss = model.train(self.dataset)
1009
- if i % 10 == 0:
1010
- model.save(
1011
- current.checkpoint.directory,
1012
- )
1013
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
1014
- # and returns a reference dictionary to the checkpoint saved in the datastore
1015
- self.latest_checkpoint = current.checkpoint.save(
1016
- name="epoch_checkpoint",
1017
- metadata={
1018
- "epoch": i,
1019
- "loss": loss,
1020
- }
1021
- )
1022
- ```
1023
-
1024
- - Using Loaded Checkpoints
1025
-
1026
- ```python
1027
- @retry(times=3)
1028
- @checkpoint
1029
- @step
1030
- def train(self):
1031
- # Assume that the task has restarted and the previous attempt of the task
1032
- # saved a checkpoint
1033
- checkpoint_path = None
1034
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1035
- print("Loaded checkpoint from the previous attempt")
1036
- checkpoint_path = current.checkpoint.directory
1037
-
1038
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1039
- for i in range(self.epochs):
1040
- ...
1041
- ```
1042
-
1043
-
1044
- Parameters
1045
- ----------
1046
- load_policy : str, default: "fresh"
1047
- The policy for loading the checkpoint. The following policies are supported:
1048
- - "eager": Loads the the latest available checkpoint within the namespace.
1049
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1050
- will be loaded at the start of the task.
1051
- - "none": Do not load any checkpoint
1052
- - "fresh": Loads the lastest checkpoint created within the running Task.
1053
- This mode helps loading checkpoints across various retry attempts of the same task.
1054
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1055
- created within the task will be loaded when the task is retries execution on failure.
1056
-
1057
- temp_dir_root : str, default: None
1058
- The root directory under which `current.checkpoint.directory` will be created.
1059
- """
1060
- ...
1061
-
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
- @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]]]:
1233
- """
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.
1208
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1237
1209
 
1238
1210
 
1239
1211
  Parameters
@@ -1298,197 +1270,331 @@ def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFla
1298
1270
  """
1299
1271
  ...
1300
1272
 
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]]]:
1273
+ @typing.overload
1274
+ 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]]]:
1302
1275
  """
1303
- This decorator is used to run vllm APIs as Metaflow task sidecars.
1276
+ Specifies the Conda environment for the step.
1304
1277
 
1305
- User code call
1306
- --------------
1307
- @vllm(
1308
- model="...",
1309
- ...
1310
- )
1278
+ Information in this decorator will augment any
1279
+ attributes set in the `@conda_base` flow-level decorator. Hence,
1280
+ you can use `@conda_base` to set packages required by all
1281
+ steps and use `@conda` to specify step-specific overrides.
1311
1282
 
1312
- Valid backend options
1313
- ---------------------
1314
- - 'local': Run as a separate process on the local task machine.
1315
1283
 
1316
- Valid model options
1317
- -------------------
1318
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
1284
+ Parameters
1285
+ ----------
1286
+ packages : Dict[str, str], default {}
1287
+ Packages to use for this step. The key is the name of the package
1288
+ and the value is the version to use.
1289
+ libraries : Dict[str, str], default {}
1290
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1291
+ python : str, optional, default None
1292
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1293
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1294
+ disabled : bool, default False
1295
+ If set to True, disables @conda.
1296
+ """
1297
+ ...
1298
+
1299
+ @typing.overload
1300
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1301
+ ...
1302
+
1303
+ @typing.overload
1304
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1305
+ ...
1306
+
1307
+ 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):
1308
+ """
1309
+ Specifies the Conda environment for the step.
1319
1310
 
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.
1311
+ Information in this decorator will augment any
1312
+ attributes set in the `@conda_base` flow-level decorator. Hence,
1313
+ you can use `@conda_base` to set packages required by all
1314
+ steps and use `@conda` to specify step-specific overrides.
1322
1315
 
1323
1316
 
1324
1317
  Parameters
1325
1318
  ----------
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`.
1319
+ packages : Dict[str, str], default {}
1320
+ Packages to use for this step. The key is the name of the package
1321
+ and the value is the version to use.
1322
+ libraries : Dict[str, str], default {}
1323
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1324
+ python : str, optional, default None
1325
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1326
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1327
+ disabled : bool, default False
1328
+ If set to True, disables @conda.
1329
+ """
1330
+ ...
1331
+
1332
+ @typing.overload
1333
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1334
+ """
1335
+ Decorator prototype for all step decorators. This function gets specialized
1336
+ and imported for all decorators types by _import_plugin_decorators().
1337
+ """
1338
+ ...
1339
+
1340
+ @typing.overload
1341
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1342
+ ...
1343
+
1344
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1345
+ """
1346
+ Decorator prototype for all step decorators. This function gets specialized
1347
+ and imported for all decorators types by _import_plugin_decorators().
1348
+ """
1349
+ ...
1350
+
1351
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1352
+ """
1353
+ Allows setting external datastores to save data for the
1354
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1355
+
1356
+ This decorator is useful when users wish to save data to a different datastore
1357
+ than what is configured in Metaflow. This can be for variety of reasons:
1358
+
1359
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1360
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1361
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1362
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1363
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1364
+
1365
+ Usage:
1366
+ ----------
1367
+
1368
+ - Using a custom IAM role to access the datastore.
1369
+
1370
+ ```python
1371
+ @with_artifact_store(
1372
+ type="s3",
1373
+ config=lambda: {
1374
+ "root": "s3://my-bucket-foo/path/to/root",
1375
+ "role_arn": ROLE,
1376
+ },
1377
+ )
1378
+ class MyFlow(FlowSpec):
1379
+
1380
+ @checkpoint
1381
+ @step
1382
+ def start(self):
1383
+ with open("my_file.txt", "w") as f:
1384
+ f.write("Hello, World!")
1385
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1386
+ self.next(self.end)
1387
+
1388
+ ```
1389
+
1390
+ - Using credentials to access the s3-compatible datastore.
1391
+
1392
+ ```python
1393
+ @with_artifact_store(
1394
+ type="s3",
1395
+ config=lambda: {
1396
+ "root": "s3://my-bucket-foo/path/to/root",
1397
+ "client_params": {
1398
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1399
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1400
+ },
1401
+ },
1402
+ )
1403
+ class MyFlow(FlowSpec):
1404
+
1405
+ @checkpoint
1406
+ @step
1407
+ def start(self):
1408
+ with open("my_file.txt", "w") as f:
1409
+ f.write("Hello, World!")
1410
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1411
+ self.next(self.end)
1412
+
1413
+ ```
1414
+
1415
+ - Accessing objects stored in external datastores after task execution.
1416
+
1417
+ ```python
1418
+ run = Run("CheckpointsTestsFlow/8992")
1419
+ with artifact_store_from(run=run, config={
1420
+ "client_params": {
1421
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1422
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1423
+ },
1424
+ }):
1425
+ with Checkpoint() as cp:
1426
+ latest = cp.list(
1427
+ task=run["start"].task
1428
+ )[0]
1429
+ print(latest)
1430
+ cp.load(
1431
+ latest,
1432
+ "test-checkpoints"
1433
+ )
1434
+
1435
+ task = Task("TorchTuneFlow/8484/train/53673")
1436
+ with artifact_store_from(run=run, config={
1437
+ "client_params": {
1438
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1439
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1440
+ },
1441
+ }):
1442
+ load_model(
1443
+ task.data.model_ref,
1444
+ "test-models"
1445
+ )
1446
+ ```
1447
+ Parameters:
1448
+ ----------
1449
+
1450
+ type: str
1451
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1452
+
1453
+ config: dict or Callable
1454
+ Dictionary of configuration options for the datastore. The following keys are required:
1455
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1456
+ - example: 's3://bucket-name/path/to/root'
1457
+ - example: 'gs://bucket-name/path/to/root'
1458
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1459
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1460
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1461
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1348
1462
  """
1349
1463
  ...
1350
1464
 
1351
1465
  @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]]:
1466
+ 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]]:
1353
1467
  """
1354
- Specifies the event(s) that this flow depends on.
1468
+ Specifies the flow(s) that this flow depends on.
1355
1469
 
1356
1470
  ```
1357
- @trigger(event='foo')
1471
+ @trigger_on_finish(flow='FooFlow')
1358
1472
  ```
1359
1473
  or
1360
1474
  ```
1361
- @trigger(events=['foo', 'bar'])
1475
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1362
1476
  ```
1477
+ This decorator respects the @project decorator and triggers the flow
1478
+ when upstream runs within the same namespace complete successfully
1363
1479
 
1364
- Additionally, you can specify the parameter mappings
1365
- to map event payload to Metaflow parameters for the flow.
1480
+ Additionally, you can specify project aware upstream flow dependencies
1481
+ by specifying the fully qualified project_flow_name.
1366
1482
  ```
1367
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1483
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1368
1484
  ```
1369
1485
  or
1370
1486
  ```
1371
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1372
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1487
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1373
1488
  ```
1374
1489
 
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:
1490
+ You can also specify just the project or project branch (other values will be
1491
+ inferred from the current project or project branch):
1380
1492
  ```
1381
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1493
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1382
1494
  ```
1383
1495
 
1496
+ Note that `branch` is typically one of:
1497
+ - `prod`
1498
+ - `user.bob`
1499
+ - `test.my_experiment`
1500
+ - `prod.staging`
1501
+
1384
1502
 
1385
1503
  Parameters
1386
1504
  ----------
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.
1505
+ flow : Union[str, Dict[str, str]], optional, default None
1506
+ Upstream flow dependency for this flow.
1507
+ flows : List[Union[str, Dict[str, str]]], default []
1508
+ Upstream flow dependencies for this flow.
1391
1509
  options : Dict[str, Any], default {}
1392
1510
  Backend-specific configuration for tuning eventing behavior.
1393
1511
  """
1394
1512
  ...
1395
1513
 
1396
1514
  @typing.overload
1397
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1515
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1398
1516
  ...
1399
1517
 
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] = {}):
1518
+ 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] = {}):
1401
1519
  """
1402
- Specifies the event(s) that this flow depends on.
1520
+ Specifies the flow(s) that this flow depends on.
1403
1521
 
1404
1522
  ```
1405
- @trigger(event='foo')
1523
+ @trigger_on_finish(flow='FooFlow')
1406
1524
  ```
1407
1525
  or
1408
1526
  ```
1409
- @trigger(events=['foo', 'bar'])
1527
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1410
1528
  ```
1529
+ This decorator respects the @project decorator and triggers the flow
1530
+ when upstream runs within the same namespace complete successfully
1411
1531
 
1412
- Additionally, you can specify the parameter mappings
1413
- to map event payload to Metaflow parameters for the flow.
1532
+ Additionally, you can specify project aware upstream flow dependencies
1533
+ by specifying the fully qualified project_flow_name.
1414
1534
  ```
1415
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1535
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1416
1536
  ```
1417
1537
  or
1418
1538
  ```
1419
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1420
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1539
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1421
1540
  ```
1422
1541
 
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:
1542
+ You can also specify just the project or project branch (other values will be
1543
+ inferred from the current project or project branch):
1428
1544
  ```
1429
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1545
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1430
1546
  ```
1431
1547
 
1548
+ Note that `branch` is typically one of:
1549
+ - `prod`
1550
+ - `user.bob`
1551
+ - `test.my_experiment`
1552
+ - `prod.staging`
1553
+
1432
1554
 
1433
1555
  Parameters
1434
1556
  ----------
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.
1557
+ flow : Union[str, Dict[str, str]], optional, default None
1558
+ Upstream flow dependency for this flow.
1559
+ flows : List[Union[str, Dict[str, str]]], default []
1560
+ Upstream flow dependencies for this flow.
1439
1561
  options : Dict[str, Any], default {}
1440
1562
  Backend-specific configuration for tuning eventing behavior.
1441
1563
  """
1442
1564
  ...
1443
1565
 
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]]:
1566
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1446
1567
  """
1447
- Specifies the Conda environment for all steps of the flow.
1568
+ Specifies what flows belong to the same project.
1448
1569
 
1449
- Use `@conda_base` to set common libraries required by all
1450
- steps and use `@conda` to specify step-specific additions.
1570
+ A project-specific namespace is created for all flows that
1571
+ use the same `@project(name)`.
1451
1572
 
1452
1573
 
1453
1574
  Parameters
1454
1575
  ----------
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.
1475
-
1476
- Use `@conda_base` to set common libraries required by all
1477
- steps and use `@conda` to specify step-specific additions.
1576
+ name : str
1577
+ Project name. Make sure that the name is unique amongst all
1578
+ projects that use the same production scheduler. The name may
1579
+ contain only lowercase alphanumeric characters and underscores.
1478
1580
 
1581
+ branch : Optional[str], default None
1582
+ The branch to use. If not specified, the branch is set to
1583
+ `user.<username>` unless `production` is set to `True`. This can
1584
+ also be set on the command line using `--branch` as a top-level option.
1585
+ It is an error to specify `branch` in the decorator and on the command line.
1479
1586
 
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.
1587
+ production : bool, default False
1588
+ Whether or not the branch is the production branch. This can also be set on the
1589
+ command line using `--production` as a top-level option. It is an error to specify
1590
+ `production` in the decorator and on the command line.
1591
+ The project branch name will be:
1592
+ - if `branch` is specified:
1593
+ - if `production` is True: `prod.<branch>`
1594
+ - if `production` is False: `test.<branch>`
1595
+ - if `branch` is not specified:
1596
+ - if `production` is True: `prod`
1597
+ - if `production` is False: `user.<username>`
1492
1598
  """
1493
1599
  ...
1494
1600
 
@@ -1584,38 +1690,97 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1584
1690
  """
1585
1691
  ...
1586
1692
 
1587
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1693
+ 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]]:
1588
1694
  """
1589
- Specifies what flows belong to the same project.
1590
-
1591
- A project-specific namespace is created for all flows that
1592
- use the same `@project(name)`.
1695
+ 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.
1696
+ 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.
1593
1697
 
1594
1698
 
1595
1699
  Parameters
1596
1700
  ----------
1701
+ timeout : int
1702
+ Time, in seconds before the task times out and fails. (Default: 3600)
1703
+ poke_interval : int
1704
+ Time in seconds that the job should wait in between each try. (Default: 60)
1705
+ mode : str
1706
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1707
+ exponential_backoff : bool
1708
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1709
+ pool : str
1710
+ the slot pool this task should run in,
1711
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1712
+ soft_fail : bool
1713
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1597
1714
  name : str
1598
- Project name. Make sure that the name is unique amongst all
1599
- projects that use the same production scheduler. The name may
1600
- contain only lowercase alphanumeric characters and underscores.
1715
+ Name of the sensor on Airflow
1716
+ description : str
1717
+ Description of sensor in the Airflow UI
1718
+ external_dag_id : str
1719
+ The dag_id that contains the task you want to wait for.
1720
+ external_task_ids : List[str]
1721
+ The list of task_ids that you want to wait for.
1722
+ If None (default value) the sensor waits for the DAG. (Default: None)
1723
+ allowed_states : List[str]
1724
+ Iterable of allowed states, (Default: ['success'])
1725
+ failed_states : List[str]
1726
+ Iterable of failed or dis-allowed states. (Default: None)
1727
+ execution_delta : datetime.timedelta
1728
+ time difference with the previous execution to look at,
1729
+ the default is the same logical date as the current task or DAG. (Default: None)
1730
+ check_existence: bool
1731
+ Set to True to check if the external task exists or check if
1732
+ the DAG to wait for exists. (Default: True)
1733
+ """
1734
+ ...
1735
+
1736
+ @typing.overload
1737
+ 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]]:
1738
+ """
1739
+ Specifies the Conda environment for all steps of the flow.
1601
1740
 
1602
- branch : Optional[str], default None
1603
- The branch to use. If not specified, the branch is set to
1604
- `user.<username>` unless `production` is set to `True`. This can
1605
- also be set on the command line using `--branch` as a top-level option.
1606
- It is an error to specify `branch` in the decorator and on the command line.
1741
+ Use `@conda_base` to set common libraries required by all
1742
+ steps and use `@conda` to specify step-specific additions.
1607
1743
 
1608
- production : bool, default False
1609
- Whether or not the branch is the production branch. This can also be set on the
1610
- command line using `--production` as a top-level option. It is an error to specify
1611
- `production` in the decorator and on the command line.
1612
- The project branch name will be:
1613
- - if `branch` is specified:
1614
- - if `production` is True: `prod.<branch>`
1615
- - if `production` is False: `test.<branch>`
1616
- - if `branch` is not specified:
1617
- - if `production` is True: `prod`
1618
- - if `production` is False: `user.<username>`
1744
+
1745
+ Parameters
1746
+ ----------
1747
+ packages : Dict[str, str], default {}
1748
+ Packages to use for this flow. The key is the name of the package
1749
+ and the value is the version to use.
1750
+ libraries : Dict[str, str], default {}
1751
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1752
+ python : str, optional, default None
1753
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1754
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1755
+ disabled : bool, default False
1756
+ If set to True, disables Conda.
1757
+ """
1758
+ ...
1759
+
1760
+ @typing.overload
1761
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1762
+ ...
1763
+
1764
+ 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):
1765
+ """
1766
+ Specifies the Conda environment for all steps of the flow.
1767
+
1768
+ Use `@conda_base` to set common libraries required by all
1769
+ steps and use `@conda` to specify step-specific additions.
1770
+
1771
+
1772
+ Parameters
1773
+ ----------
1774
+ packages : Dict[str, str], default {}
1775
+ Packages to use for this flow. The key is the name of the package
1776
+ and the value is the version to use.
1777
+ libraries : Dict[str, str], default {}
1778
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1779
+ python : str, optional, default None
1780
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1781
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1782
+ disabled : bool, default False
1783
+ If set to True, disables Conda.
1619
1784
  """
1620
1785
  ...
1621
1786
 
@@ -1662,263 +1827,98 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1662
1827
  """
1663
1828
  ...
1664
1829
 
1665
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1666
- """
1667
- Allows setting external datastores to save data for the
1668
- `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1669
-
1670
- This decorator is useful when users wish to save data to a different datastore
1671
- than what is configured in Metaflow. This can be for variety of reasons:
1672
-
1673
- 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1674
- 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1675
- - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1676
- 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1677
- - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1678
-
1679
- Usage:
1680
- ----------
1681
-
1682
- - Using a custom IAM role to access the datastore.
1683
-
1684
- ```python
1685
- @with_artifact_store(
1686
- type="s3",
1687
- config=lambda: {
1688
- "root": "s3://my-bucket-foo/path/to/root",
1689
- "role_arn": ROLE,
1690
- },
1691
- )
1692
- class MyFlow(FlowSpec):
1693
-
1694
- @checkpoint
1695
- @step
1696
- def start(self):
1697
- with open("my_file.txt", "w") as f:
1698
- f.write("Hello, World!")
1699
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1700
- self.next(self.end)
1701
-
1702
- ```
1703
-
1704
- - Using credentials to access the s3-compatible datastore.
1705
-
1706
- ```python
1707
- @with_artifact_store(
1708
- type="s3",
1709
- config=lambda: {
1710
- "root": "s3://my-bucket-foo/path/to/root",
1711
- "client_params": {
1712
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1713
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1714
- },
1715
- },
1716
- )
1717
- class MyFlow(FlowSpec):
1718
-
1719
- @checkpoint
1720
- @step
1721
- def start(self):
1722
- with open("my_file.txt", "w") as f:
1723
- f.write("Hello, World!")
1724
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1725
- self.next(self.end)
1726
-
1727
- ```
1728
-
1729
- - Accessing objects stored in external datastores after task execution.
1730
-
1731
- ```python
1732
- run = Run("CheckpointsTestsFlow/8992")
1733
- with artifact_store_from(run=run, config={
1734
- "client_params": {
1735
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1736
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1737
- },
1738
- }):
1739
- with Checkpoint() as cp:
1740
- latest = cp.list(
1741
- task=run["start"].task
1742
- )[0]
1743
- print(latest)
1744
- cp.load(
1745
- latest,
1746
- "test-checkpoints"
1747
- )
1748
-
1749
- task = Task("TorchTuneFlow/8484/train/53673")
1750
- with artifact_store_from(run=run, config={
1751
- "client_params": {
1752
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1753
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1754
- },
1755
- }):
1756
- load_model(
1757
- task.data.model_ref,
1758
- "test-models"
1759
- )
1760
- ```
1761
- Parameters:
1762
- ----------
1763
-
1764
- type: str
1765
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1766
-
1767
- config: dict or Callable
1768
- Dictionary of configuration options for the datastore. The following keys are required:
1769
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1770
- - example: 's3://bucket-name/path/to/root'
1771
- - example: 'gs://bucket-name/path/to/root'
1772
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1773
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1774
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1775
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1776
- """
1777
- ...
1778
-
1779
1830
  @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]]:
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]]:
1781
1832
  """
1782
- Specifies the flow(s) that this flow depends on.
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