ob-metaflow-stubs 6.0.10.1__py2.py3-none-any.whl → 6.0.10.2rc0__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 +897 -896
  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 +4 -4
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +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 +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +61 -61
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +1 -1
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +2 -2
  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 +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +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 +3 -3
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +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 +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/__init__.pyi +2 -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 +3 -3
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/click_importer.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/__init__.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/code_packager.pyi +2 -2
  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 +3 -3
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +3 -3
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +3 -3
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +9 -3
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/experimental/__init__.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/perimeters.pyi +1 -1
  88. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/utils.pyi +2 -2
  89. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +1 -1
  90. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +1 -1
  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 +3 -3
  96. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +3 -3
  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 +1 -1
  100. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +1 -1
  101. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  102. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  103. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  104. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  105. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  106. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  107. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  108. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  109. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  110. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  111. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  112. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  113. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  114. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -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 +2 -1
  118. metaflow-stubs/packaging_sys/__init__.pyi +5 -5
  119. metaflow-stubs/packaging_sys/backend.pyi +1 -1
  120. metaflow-stubs/packaging_sys/distribution_support.pyi +4 -4
  121. metaflow-stubs/packaging_sys/tar_backend.pyi +4 -4
  122. metaflow-stubs/packaging_sys/utils.pyi +1 -1
  123. metaflow-stubs/packaging_sys/v1.pyi +3 -3
  124. metaflow-stubs/parameters.pyi +3 -3
  125. metaflow-stubs/plugins/__init__.pyi +11 -11
  126. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  127. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  128. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  129. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  130. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  131. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  132. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  133. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  134. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  135. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  136. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  137. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  138. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  139. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  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 +3 -3
  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 +2 -2
  157. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  158. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  159. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  160. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  161. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  162. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  163. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  164. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  165. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  166. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  167. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  168. metaflow-stubs/plugins/cards/card_decorator.pyi +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 +3 -3
  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 +3 -3
  196. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  197. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  198. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  199. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  200. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  201. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  202. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  203. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  204. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  205. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  206. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  207. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  208. metaflow-stubs/plugins/optuna/__init__.pyi +1 -1
  209. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  210. metaflow-stubs/plugins/perimeters.pyi +1 -1
  211. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  212. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  213. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  214. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  215. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  216. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  217. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  218. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  219. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  220. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  221. metaflow-stubs/plugins/secrets/__init__.pyi +1 -1
  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 +2 -2
  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 +1 -1
  234. metaflow-stubs/profilers/__init__.pyi +1 -1
  235. metaflow-stubs/pylint_wrapper.pyi +1 -1
  236. metaflow-stubs/runner/__init__.pyi +1 -1
  237. metaflow-stubs/runner/deployer.pyi +34 -34
  238. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  239. metaflow-stubs/runner/metaflow_runner.pyi +2 -2
  240. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  241. metaflow-stubs/runner/nbrun.pyi +1 -1
  242. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  243. metaflow-stubs/runner/utils.pyi +2 -2
  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 +2 -2
  251. metaflow-stubs/user_configs/config_parameters.pyi +3 -3
  252. metaflow-stubs/user_decorators/__init__.pyi +1 -1
  253. metaflow-stubs/user_decorators/common.pyi +1 -1
  254. metaflow-stubs/user_decorators/mutable_flow.pyi +3 -3
  255. metaflow-stubs/user_decorators/mutable_step.pyi +4 -4
  256. metaflow-stubs/user_decorators/user_flow_decorator.pyi +4 -4
  257. metaflow-stubs/user_decorators/user_step_decorator.pyi +5 -5
  258. {ob_metaflow_stubs-6.0.10.1.dist-info → ob_metaflow_stubs-6.0.10.2rc0.dist-info}/METADATA +1 -1
  259. ob_metaflow_stubs-6.0.10.2rc0.dist-info/RECORD +262 -0
  260. ob_metaflow_stubs-6.0.10.1.dist-info/RECORD +0 -262
  261. {ob_metaflow_stubs-6.0.10.1.dist-info → ob_metaflow_stubs-6.0.10.2rc0.dist-info}/WHEEL +0 -0
  262. {ob_metaflow_stubs-6.0.10.1.dist-info → ob_metaflow_stubs-6.0.10.2rc0.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.18.3.2+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-09-09T09:20:35.730188 #
4
+ # Generated on 2025-09-09T23:55:12.839647 #
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,8 +39,8 @@ 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 cards as cards
43
42
  from . import tuple_util as tuple_util
43
+ from . import cards as cards
44
44
  from . import metaflow_git as metaflow_git
45
45
  from . import events as events
46
46
  from . import runner as runner
@@ -48,9 +48,9 @@ from . import plugins as plugins
48
48
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
49
49
  from . import includefile as includefile
50
50
  from .includefile import IncludeFile as IncludeFile
51
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
51
52
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
52
53
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
53
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
54
54
  from . import client as client
55
55
  from .client.core import namespace as namespace
56
56
  from .client.core import get_namespace as get_namespace
@@ -83,6 +83,7 @@ from .mf_extensions.outerbounds.plugins.checkpoint_datastores.nebius import nebi
83
83
  from .mf_extensions.outerbounds.plugins.checkpoint_datastores.coreweave import coreweave_checkpoints as coreweave_checkpoints
84
84
  from .mf_extensions.outerbounds.plugins.aws.assume_role_decorator import assume_role as assume_role
85
85
  from .mf_extensions.outerbounds.plugins.apps.core.deployer import AppDeployer as AppDeployer
86
+ from .mf_extensions.outerbounds.plugins.apps.core.deployer import DeployedApp as DeployedApp
86
87
  from . import system as system
87
88
  from . import cli_components as cli_components
88
89
  from . import pylint_wrapper as pylint_wrapper
@@ -167,104 +168,36 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
167
168
  """
168
169
  ...
169
170
 
170
- 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]]]:
171
- """
172
- This decorator is used to run vllm APIs as Metaflow task sidecars.
173
-
174
- User code call
175
- --------------
176
- @vllm(
177
- model="...",
178
- ...
179
- )
180
-
181
- Valid backend options
182
- ---------------------
183
- - 'local': Run as a separate process on the local task machine.
184
-
185
- Valid model options
186
- -------------------
187
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
188
-
189
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
190
- If you need multiple models, you must create multiple @vllm decorators.
191
-
192
-
193
- Parameters
194
- ----------
195
- model: str
196
- HuggingFace model identifier to be served by vLLM.
197
- backend: str
198
- Determines where and how to run the vLLM process.
199
- openai_api_server: bool
200
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
201
- Default is False (uses native engine).
202
- Set to True for backward compatibility with existing code.
203
- debug: bool
204
- Whether to turn on verbose debugging logs.
205
- card_refresh_interval: int
206
- Interval in seconds for refreshing the vLLM status card.
207
- Only used when openai_api_server=True.
208
- max_retries: int
209
- Maximum number of retries checking for vLLM server startup.
210
- Only used when openai_api_server=True.
211
- retry_alert_frequency: int
212
- Frequency of alert logs for vLLM server startup retries.
213
- Only used when openai_api_server=True.
214
- engine_args : dict
215
- Additional keyword arguments to pass to the vLLM engine.
216
- For example, `tensor_parallel_size=2`.
217
- """
218
- ...
219
-
220
171
  @typing.overload
221
- 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]]]:
172
+ 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]]]:
222
173
  """
223
- Specifies that the step will success under all circumstances.
224
-
225
- The decorator will create an optional artifact, specified by `var`, which
226
- contains the exception raised. You can use it to detect the presence
227
- of errors, indicating that all happy-path artifacts produced by the step
228
- are missing.
174
+ Specifies environment variables to be set prior to the execution of a step.
229
175
 
230
176
 
231
177
  Parameters
232
178
  ----------
233
- var : str, optional, default None
234
- Name of the artifact in which to store the caught exception.
235
- If not specified, the exception is not stored.
236
- print_exception : bool, default True
237
- Determines whether or not the exception is printed to
238
- stdout when caught.
179
+ vars : Dict[str, str], default {}
180
+ Dictionary of environment variables to set.
239
181
  """
240
182
  ...
241
183
 
242
184
  @typing.overload
243
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
185
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
244
186
  ...
245
187
 
246
188
  @typing.overload
247
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
189
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
248
190
  ...
249
191
 
250
- 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):
192
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
251
193
  """
252
- Specifies that the step will success under all circumstances.
253
-
254
- The decorator will create an optional artifact, specified by `var`, which
255
- contains the exception raised. You can use it to detect the presence
256
- of errors, indicating that all happy-path artifacts produced by the step
257
- are missing.
194
+ Specifies environment variables to be set prior to the execution of a step.
258
195
 
259
196
 
260
197
  Parameters
261
198
  ----------
262
- var : str, optional, default None
263
- Name of the artifact in which to store the caught exception.
264
- If not specified, the exception is not stored.
265
- print_exception : bool, default True
266
- Determines whether or not the exception is printed to
267
- stdout when caught.
199
+ vars : Dict[str, str], default {}
200
+ Dictionary of environment variables to set.
268
201
  """
269
202
  ...
270
203
 
@@ -284,306 +217,40 @@ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[t
284
217
  """
285
218
  ...
286
219
 
287
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
220
+ 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]]]:
288
221
  """
289
- Specifies that this step should execute on Kubernetes.
222
+ Specifies that this step should execute on DGX cloud.
290
223
 
291
224
 
292
225
  Parameters
293
226
  ----------
294
- cpu : int, default 1
295
- Number of CPUs required for this step. If `@resources` is
296
- also present, the maximum value from all decorators is used.
297
- memory : int, default 4096
298
- Memory size (in MB) required for this step. If
299
- `@resources` is also present, the maximum value from all decorators is
300
- used.
301
- disk : int, default 10240
302
- Disk size (in MB) required for this step. If
303
- `@resources` is also present, the maximum value from all decorators is
304
- used.
305
- image : str, optional, default None
306
- Docker image to use when launching on Kubernetes. If not specified, and
307
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
308
- not, a default Docker image mapping to the current version of Python is used.
309
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
310
- If given, the imagePullPolicy to be applied to the Docker image of the step.
311
- image_pull_secrets: List[str], default []
312
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
313
- Kubernetes image pull secrets to use when pulling container images
314
- in Kubernetes.
315
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
316
- Kubernetes service account to use when launching pod in Kubernetes.
317
- secrets : List[str], optional, default None
318
- Kubernetes secrets to use when launching pod in Kubernetes. These
319
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
320
- in Metaflow configuration.
321
- node_selector: Union[Dict[str,str], str], optional, default None
322
- Kubernetes node selector(s) to apply to the pod running the task.
323
- Can be passed in as a comma separated string of values e.g.
324
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
325
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
326
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
327
- Kubernetes namespace to use when launching pod in Kubernetes.
328
- gpu : int, optional, default None
329
- Number of GPUs required for this step. A value of zero implies that
330
- the scheduled node should not have GPUs.
331
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
332
- The vendor of the GPUs to be used for this step.
333
- tolerations : List[Dict[str,str]], default []
334
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
335
- Kubernetes tolerations to use when launching pod in Kubernetes.
336
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
337
- Kubernetes labels to use when launching pod in Kubernetes.
338
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
339
- Kubernetes annotations to use when launching pod in Kubernetes.
340
- use_tmpfs : bool, default False
341
- This enables an explicit tmpfs mount for this step.
342
- tmpfs_tempdir : bool, default True
343
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
344
- tmpfs_size : int, optional, default: None
345
- The value for the size (in MiB) of the tmpfs mount for this step.
346
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
347
- memory allocated for this step.
348
- tmpfs_path : str, optional, default /metaflow_temp
349
- Path to tmpfs mount for this step.
350
- persistent_volume_claims : Dict[str, str], optional, default None
351
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
352
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
353
- shared_memory: int, optional
354
- Shared memory size (in MiB) required for this step
355
- port: int, optional
356
- Port number to specify in the Kubernetes job object
357
- compute_pool : str, optional, default None
358
- Compute pool to be used for for this step.
359
- If not specified, any accessible compute pool within the perimeter is used.
360
- hostname_resolution_timeout: int, default 10 * 60
361
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
362
- Only applicable when @parallel is used.
363
- qos: str, default: Burstable
364
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
365
-
366
- security_context: Dict[str, Any], optional, default None
367
- Container security context. Applies to the task container. Allows the following keys:
368
- - privileged: bool, optional, default None
369
- - allow_privilege_escalation: bool, optional, default None
370
- - run_as_user: int, optional, default None
371
- - run_as_group: int, optional, default None
372
- - run_as_non_root: bool, optional, default None
227
+ gpu : int
228
+ Number of GPUs to use.
229
+ gpu_type : str
230
+ Type of Nvidia GPU to use.
373
231
  """
374
232
  ...
375
233
 
376
234
  @typing.overload
377
- def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
235
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
378
236
  """
379
- Enables loading / saving of models within a step.
380
-
381
- > Examples
382
- - Saving Models
383
- ```python
384
- @model
385
- @step
386
- def train(self):
387
- # current.model.save returns a dictionary reference to the model saved
388
- self.my_model = current.model.save(
389
- path_to_my_model,
390
- label="my_model",
391
- metadata={
392
- "epochs": 10,
393
- "batch-size": 32,
394
- "learning-rate": 0.001,
395
- }
396
- )
397
- self.next(self.test)
398
-
399
- @model(load="my_model")
400
- @step
401
- def test(self):
402
- # `current.model.loaded` returns a dictionary of the loaded models
403
- # where the key is the name of the artifact and the value is the path to the model
404
- print(os.listdir(current.model.loaded["my_model"]))
405
- self.next(self.end)
406
- ```
407
-
408
- - Loading models
409
- ```python
410
- @step
411
- def train(self):
412
- # current.model.load returns the path to the model loaded
413
- checkpoint_path = current.model.load(
414
- self.checkpoint_key,
415
- )
416
- model_path = current.model.load(
417
- self.model,
418
- )
419
- self.next(self.test)
420
- ```
421
-
422
-
423
- Parameters
424
- ----------
425
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
426
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
427
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
428
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
429
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
430
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
431
-
432
- temp_dir_root : str, default: None
433
- The root directory under which `current.model.loaded` will store loaded models
237
+ Internal decorator to support Fast bakery
434
238
  """
435
239
  ...
436
240
 
437
241
  @typing.overload
438
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
242
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
439
243
  ...
440
244
 
441
- @typing.overload
442
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
245
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
246
+ """
247
+ Internal decorator to support Fast bakery
248
+ """
443
249
  ...
444
250
 
445
- def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
251
+ 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]]]:
446
252
  """
447
- Enables loading / saving of models within a step.
448
-
449
- > Examples
450
- - Saving Models
451
- ```python
452
- @model
453
- @step
454
- def train(self):
455
- # current.model.save returns a dictionary reference to the model saved
456
- self.my_model = current.model.save(
457
- path_to_my_model,
458
- label="my_model",
459
- metadata={
460
- "epochs": 10,
461
- "batch-size": 32,
462
- "learning-rate": 0.001,
463
- }
464
- )
465
- self.next(self.test)
466
-
467
- @model(load="my_model")
468
- @step
469
- def test(self):
470
- # `current.model.loaded` returns a dictionary of the loaded models
471
- # where the key is the name of the artifact and the value is the path to the model
472
- print(os.listdir(current.model.loaded["my_model"]))
473
- self.next(self.end)
474
- ```
475
-
476
- - Loading models
477
- ```python
478
- @step
479
- def train(self):
480
- # current.model.load returns the path to the model loaded
481
- checkpoint_path = current.model.load(
482
- self.checkpoint_key,
483
- )
484
- model_path = current.model.load(
485
- self.model,
486
- )
487
- self.next(self.test)
488
- ```
489
-
490
-
491
- Parameters
492
- ----------
493
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
494
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
495
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
496
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
497
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
498
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
499
-
500
- temp_dir_root : str, default: None
501
- The root directory under which `current.model.loaded` will store loaded models
502
- """
503
- ...
504
-
505
- @typing.overload
506
- 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]]]:
507
- """
508
- Specifies the resources needed when executing this step.
509
-
510
- Use `@resources` to specify the resource requirements
511
- independently of the specific compute layer (`@batch`, `@kubernetes`).
512
-
513
- You can choose the compute layer on the command line by executing e.g.
514
- ```
515
- python myflow.py run --with batch
516
- ```
517
- or
518
- ```
519
- python myflow.py run --with kubernetes
520
- ```
521
- which executes the flow on the desired system using the
522
- requirements specified in `@resources`.
523
-
524
-
525
- Parameters
526
- ----------
527
- cpu : int, default 1
528
- Number of CPUs required for this step.
529
- gpu : int, optional, default None
530
- Number of GPUs required for this step.
531
- disk : int, optional, default None
532
- Disk size (in MB) required for this step. Only applies on Kubernetes.
533
- memory : int, default 4096
534
- Memory size (in MB) required for this step.
535
- shared_memory : int, optional, default None
536
- The value for the size (in MiB) of the /dev/shm volume for this step.
537
- This parameter maps to the `--shm-size` option in Docker.
538
- """
539
- ...
540
-
541
- @typing.overload
542
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
543
- ...
544
-
545
- @typing.overload
546
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
547
- ...
548
-
549
- 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):
550
- """
551
- Specifies the resources needed when executing this step.
552
-
553
- Use `@resources` to specify the resource requirements
554
- independently of the specific compute layer (`@batch`, `@kubernetes`).
555
-
556
- You can choose the compute layer on the command line by executing e.g.
557
- ```
558
- python myflow.py run --with batch
559
- ```
560
- or
561
- ```
562
- python myflow.py run --with kubernetes
563
- ```
564
- which executes the flow on the desired system using the
565
- requirements specified in `@resources`.
566
-
567
-
568
- Parameters
569
- ----------
570
- cpu : int, default 1
571
- Number of CPUs required for this step.
572
- gpu : int, optional, default None
573
- Number of GPUs required for this step.
574
- disk : int, optional, default None
575
- Disk size (in MB) required for this step. Only applies on Kubernetes.
576
- memory : int, default 4096
577
- Memory size (in MB) required for this step.
578
- shared_memory : int, optional, default None
579
- The value for the size (in MiB) of the /dev/shm volume for this step.
580
- This parameter maps to the `--shm-size` option in Docker.
581
- """
582
- ...
583
-
584
- 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]]]:
585
- """
586
- S3 Proxy decorator for routing S3 requests through a local proxy service.
253
+ S3 Proxy decorator for routing S3 requests through a local proxy service.
587
254
 
588
255
 
589
256
  Parameters
@@ -663,86 +330,41 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
663
330
  ...
664
331
 
665
332
  @typing.overload
666
- 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]]]:
667
- """
668
- Specifies environment variables to be set prior to the execution of a step.
669
-
670
-
671
- Parameters
672
- ----------
673
- vars : Dict[str, str], default {}
674
- Dictionary of environment variables to set.
675
- """
676
- ...
677
-
678
- @typing.overload
679
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
680
- ...
681
-
682
- @typing.overload
683
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
684
- ...
685
-
686
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
687
- """
688
- Specifies environment variables to be set prior to the execution of a step.
689
-
690
-
691
- Parameters
692
- ----------
693
- vars : Dict[str, str], default {}
694
- Dictionary of environment variables to set.
695
- """
696
- ...
697
-
698
- @typing.overload
699
- 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]]]:
333
+ 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]]]:
700
334
  """
701
- Specifies the PyPI packages for the step.
702
-
703
- Information in this decorator will augment any
704
- attributes set in the `@pyi_base` flow-level decorator. Hence,
705
- you can use `@pypi_base` to set packages required by all
706
- steps and use `@pypi` to specify step-specific overrides.
335
+ Specifies secrets to be retrieved and injected as environment variables prior to
336
+ the execution of a step.
707
337
 
708
338
 
709
339
  Parameters
710
340
  ----------
711
- packages : Dict[str, str], default: {}
712
- Packages to use for this step. The key is the name of the package
713
- and the value is the version to use.
714
- python : str, optional, default: None
715
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
716
- that the version used will correspond to the version of the Python interpreter used to start the run.
341
+ sources : List[Union[str, Dict[str, Any]]], default: []
342
+ List of secret specs, defining how the secrets are to be retrieved
343
+ role : str, optional, default: None
344
+ Role to use for fetching secrets
717
345
  """
718
346
  ...
719
347
 
720
348
  @typing.overload
721
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
349
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
722
350
  ...
723
351
 
724
352
  @typing.overload
725
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
353
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
726
354
  ...
727
355
 
728
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
356
+ 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):
729
357
  """
730
- Specifies the PyPI packages for the step.
731
-
732
- Information in this decorator will augment any
733
- attributes set in the `@pyi_base` flow-level decorator. Hence,
734
- you can use `@pypi_base` to set packages required by all
735
- steps and use `@pypi` to specify step-specific overrides.
358
+ Specifies secrets to be retrieved and injected as environment variables prior to
359
+ the execution of a step.
736
360
 
737
361
 
738
362
  Parameters
739
363
  ----------
740
- packages : Dict[str, str], default: {}
741
- Packages to use for this step. The key is the name of the package
742
- and the value is the version to use.
743
- python : str, optional, default: None
744
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
745
- that the version used will correspond to the version of the Python interpreter used to start the run.
364
+ sources : List[Union[str, Dict[str, Any]]], default: []
365
+ List of secret specs, defining how the secrets are to be retrieved
366
+ role : str, optional, default: None
367
+ Role to use for fetching secrets
746
368
  """
747
369
  ...
748
370
 
@@ -827,168 +449,251 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.
827
449
  ...
828
450
 
829
451
  @typing.overload
830
- 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]]]:
452
+ def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
831
453
  """
832
- Enables checkpointing for a step.
454
+ CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
455
+ It exists to make it easier for users to know that this decorator should only be used with
456
+ a Neo Cloud like CoreWeave.
457
+ """
458
+ ...
459
+
460
+ @typing.overload
461
+ def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
462
+ ...
463
+
464
+ def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
465
+ """
466
+ CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
467
+ It exists to make it easier for users to know that this decorator should only be used with
468
+ a Neo Cloud like CoreWeave.
469
+ """
470
+ ...
471
+
472
+ @typing.overload
473
+ 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]]]:
474
+ """
475
+ Creates a human-readable report, a Metaflow Card, after this step completes.
833
476
 
834
- > Examples
477
+ Note that you may add multiple `@card` decorators in a step with different parameters.
835
478
 
836
- - Saving Checkpoints
837
479
 
838
- ```python
839
- @checkpoint
840
- @step
841
- def train(self):
842
- model = create_model(self.parameters, checkpoint_path = None)
843
- for i in range(self.epochs):
844
- # some training logic
845
- loss = model.train(self.dataset)
846
- if i % 10 == 0:
847
- model.save(
848
- current.checkpoint.directory,
849
- )
850
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
851
- # and returns a reference dictionary to the checkpoint saved in the datastore
852
- self.latest_checkpoint = current.checkpoint.save(
853
- name="epoch_checkpoint",
854
- metadata={
855
- "epoch": i,
856
- "loss": loss,
857
- }
858
- )
859
- ```
480
+ Parameters
481
+ ----------
482
+ type : str, default 'default'
483
+ Card type.
484
+ id : str, optional, default None
485
+ If multiple cards are present, use this id to identify this card.
486
+ options : Dict[str, Any], default {}
487
+ Options passed to the card. The contents depend on the card type.
488
+ timeout : int, default 45
489
+ Interrupt reporting if it takes more than this many seconds.
490
+ """
491
+ ...
492
+
493
+ @typing.overload
494
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
495
+ ...
496
+
497
+ @typing.overload
498
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
499
+ ...
500
+
501
+ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
502
+ """
503
+ Creates a human-readable report, a Metaflow Card, after this step completes.
504
+
505
+ Note that you may add multiple `@card` decorators in a step with different parameters.
860
506
 
861
- - Using Loaded Checkpoints
862
507
 
508
+ Parameters
509
+ ----------
510
+ type : str, default 'default'
511
+ Card type.
512
+ id : str, optional, default None
513
+ If multiple cards are present, use this id to identify this card.
514
+ options : Dict[str, Any], default {}
515
+ Options passed to the card. The contents depend on the card type.
516
+ timeout : int, default 45
517
+ Interrupt reporting if it takes more than this many seconds.
518
+ """
519
+ ...
520
+
521
+ @typing.overload
522
+ def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
523
+ """
524
+ Enables loading / saving of models within a step.
525
+
526
+ > Examples
527
+ - Saving Models
863
528
  ```python
864
- @retry(times=3)
865
- @checkpoint
529
+ @model
866
530
  @step
867
531
  def train(self):
868
- # Assume that the task has restarted and the previous attempt of the task
869
- # saved a checkpoint
870
- checkpoint_path = None
871
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
872
- print("Loaded checkpoint from the previous attempt")
873
- checkpoint_path = current.checkpoint.directory
532
+ # current.model.save returns a dictionary reference to the model saved
533
+ self.my_model = current.model.save(
534
+ path_to_my_model,
535
+ label="my_model",
536
+ metadata={
537
+ "epochs": 10,
538
+ "batch-size": 32,
539
+ "learning-rate": 0.001,
540
+ }
541
+ )
542
+ self.next(self.test)
874
543
 
875
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
876
- for i in range(self.epochs):
877
- ...
544
+ @model(load="my_model")
545
+ @step
546
+ def test(self):
547
+ # `current.model.loaded` returns a dictionary of the loaded models
548
+ # where the key is the name of the artifact and the value is the path to the model
549
+ print(os.listdir(current.model.loaded["my_model"]))
550
+ self.next(self.end)
551
+ ```
552
+
553
+ - Loading models
554
+ ```python
555
+ @step
556
+ def train(self):
557
+ # current.model.load returns the path to the model loaded
558
+ checkpoint_path = current.model.load(
559
+ self.checkpoint_key,
560
+ )
561
+ model_path = current.model.load(
562
+ self.model,
563
+ )
564
+ self.next(self.test)
878
565
  ```
879
566
 
880
567
 
881
568
  Parameters
882
569
  ----------
883
- load_policy : str, default: "fresh"
884
- The policy for loading the checkpoint. The following policies are supported:
885
- - "eager": Loads the the latest available checkpoint within the namespace.
886
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
887
- will be loaded at the start of the task.
888
- - "none": Do not load any checkpoint
889
- - "fresh": Loads the lastest checkpoint created within the running Task.
890
- This mode helps loading checkpoints across various retry attempts of the same task.
891
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
892
- created within the task will be loaded when the task is retries execution on failure.
570
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
571
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
572
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
573
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
574
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
575
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
893
576
 
894
577
  temp_dir_root : str, default: None
895
- The root directory under which `current.checkpoint.directory` will be created.
578
+ The root directory under which `current.model.loaded` will store loaded models
896
579
  """
897
580
  ...
898
581
 
899
582
  @typing.overload
900
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
583
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
901
584
  ...
902
585
 
903
586
  @typing.overload
904
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
587
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
905
588
  ...
906
589
 
907
- 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):
590
+ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
908
591
  """
909
- Enables checkpointing for a step.
592
+ Enables loading / saving of models within a step.
910
593
 
911
594
  > Examples
912
-
913
- - Saving Checkpoints
914
-
595
+ - Saving Models
915
596
  ```python
916
- @checkpoint
597
+ @model
917
598
  @step
918
599
  def train(self):
919
- model = create_model(self.parameters, checkpoint_path = None)
920
- for i in range(self.epochs):
921
- # some training logic
922
- loss = model.train(self.dataset)
923
- if i % 10 == 0:
924
- model.save(
925
- current.checkpoint.directory,
926
- )
927
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
928
- # and returns a reference dictionary to the checkpoint saved in the datastore
929
- self.latest_checkpoint = current.checkpoint.save(
930
- name="epoch_checkpoint",
931
- metadata={
932
- "epoch": i,
933
- "loss": loss,
934
- }
935
- )
936
- ```
600
+ # current.model.save returns a dictionary reference to the model saved
601
+ self.my_model = current.model.save(
602
+ path_to_my_model,
603
+ label="my_model",
604
+ metadata={
605
+ "epochs": 10,
606
+ "batch-size": 32,
607
+ "learning-rate": 0.001,
608
+ }
609
+ )
610
+ self.next(self.test)
937
611
 
938
- - Using Loaded Checkpoints
612
+ @model(load="my_model")
613
+ @step
614
+ def test(self):
615
+ # `current.model.loaded` returns a dictionary of the loaded models
616
+ # where the key is the name of the artifact and the value is the path to the model
617
+ print(os.listdir(current.model.loaded["my_model"]))
618
+ self.next(self.end)
619
+ ```
939
620
 
621
+ - Loading models
940
622
  ```python
941
- @retry(times=3)
942
- @checkpoint
943
623
  @step
944
624
  def train(self):
945
- # Assume that the task has restarted and the previous attempt of the task
946
- # saved a checkpoint
947
- checkpoint_path = None
948
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
949
- print("Loaded checkpoint from the previous attempt")
950
- checkpoint_path = current.checkpoint.directory
951
-
952
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
953
- for i in range(self.epochs):
954
- ...
625
+ # current.model.load returns the path to the model loaded
626
+ checkpoint_path = current.model.load(
627
+ self.checkpoint_key,
628
+ )
629
+ model_path = current.model.load(
630
+ self.model,
631
+ )
632
+ self.next(self.test)
955
633
  ```
956
634
 
957
635
 
958
636
  Parameters
959
637
  ----------
960
- load_policy : str, default: "fresh"
961
- The policy for loading the checkpoint. The following policies are supported:
962
- - "eager": Loads the the latest available checkpoint within the namespace.
963
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
964
- will be loaded at the start of the task.
965
- - "none": Do not load any checkpoint
966
- - "fresh": Loads the lastest checkpoint created within the running Task.
967
- This mode helps loading checkpoints across various retry attempts of the same task.
968
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
969
- created within the task will be loaded when the task is retries execution on failure.
638
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
639
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
640
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
641
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
642
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
643
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
970
644
 
971
645
  temp_dir_root : str, default: None
972
- The root directory under which `current.checkpoint.directory` will be created.
973
- """
974
- ...
975
-
976
- @typing.overload
977
- def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
978
- """
979
- A simple decorator that demonstrates using CardDecoratorInjector
980
- to inject a card and render simple markdown content.
646
+ The root directory under which `current.model.loaded` will store loaded models
981
647
  """
982
648
  ...
983
649
 
984
- @typing.overload
985
- def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
986
- ...
987
-
988
- def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
650
+ 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]]]:
989
651
  """
990
- A simple decorator that demonstrates using CardDecoratorInjector
991
- to inject a card and render simple markdown content.
652
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
653
+
654
+ User code call
655
+ --------------
656
+ @vllm(
657
+ model="...",
658
+ ...
659
+ )
660
+
661
+ Valid backend options
662
+ ---------------------
663
+ - 'local': Run as a separate process on the local task machine.
664
+
665
+ Valid model options
666
+ -------------------
667
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
668
+
669
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
670
+ If you need multiple models, you must create multiple @vllm decorators.
671
+
672
+
673
+ Parameters
674
+ ----------
675
+ model: str
676
+ HuggingFace model identifier to be served by vLLM.
677
+ backend: str
678
+ Determines where and how to run the vLLM process.
679
+ openai_api_server: bool
680
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
681
+ Default is False (uses native engine).
682
+ Set to True for backward compatibility with existing code.
683
+ debug: bool
684
+ Whether to turn on verbose debugging logs.
685
+ card_refresh_interval: int
686
+ Interval in seconds for refreshing the vLLM status card.
687
+ Only used when openai_api_server=True.
688
+ max_retries: int
689
+ Maximum number of retries checking for vLLM server startup.
690
+ Only used when openai_api_server=True.
691
+ retry_alert_frequency: int
692
+ Frequency of alert logs for vLLM server startup retries.
693
+ Only used when openai_api_server=True.
694
+ engine_args : dict
695
+ Additional keyword arguments to pass to the vLLM engine.
696
+ For example, `tensor_parallel_size=2`.
992
697
  """
993
698
  ...
994
699
 
@@ -1011,66 +716,202 @@ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
1011
716
  """
1012
717
  ...
1013
718
 
719
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
720
+ """
721
+ Specifies that this step should execute on Kubernetes.
722
+
723
+
724
+ Parameters
725
+ ----------
726
+ cpu : int, default 1
727
+ Number of CPUs required for this step. If `@resources` is
728
+ also present, the maximum value from all decorators is used.
729
+ memory : int, default 4096
730
+ Memory size (in MB) required for this step. If
731
+ `@resources` is also present, the maximum value from all decorators is
732
+ used.
733
+ disk : int, default 10240
734
+ Disk size (in MB) required for this step. If
735
+ `@resources` is also present, the maximum value from all decorators is
736
+ used.
737
+ image : str, optional, default None
738
+ Docker image to use when launching on Kubernetes. If not specified, and
739
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
740
+ not, a default Docker image mapping to the current version of Python is used.
741
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
742
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
743
+ image_pull_secrets: List[str], default []
744
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
745
+ Kubernetes image pull secrets to use when pulling container images
746
+ in Kubernetes.
747
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
748
+ Kubernetes service account to use when launching pod in Kubernetes.
749
+ secrets : List[str], optional, default None
750
+ Kubernetes secrets to use when launching pod in Kubernetes. These
751
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
752
+ in Metaflow configuration.
753
+ node_selector: Union[Dict[str,str], str], optional, default None
754
+ Kubernetes node selector(s) to apply to the pod running the task.
755
+ Can be passed in as a comma separated string of values e.g.
756
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
757
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
758
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
759
+ Kubernetes namespace to use when launching pod in Kubernetes.
760
+ gpu : int, optional, default None
761
+ Number of GPUs required for this step. A value of zero implies that
762
+ the scheduled node should not have GPUs.
763
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
764
+ The vendor of the GPUs to be used for this step.
765
+ tolerations : List[Dict[str,str]], default []
766
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
767
+ Kubernetes tolerations to use when launching pod in Kubernetes.
768
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
769
+ Kubernetes labels to use when launching pod in Kubernetes.
770
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
771
+ Kubernetes annotations to use when launching pod in Kubernetes.
772
+ use_tmpfs : bool, default False
773
+ This enables an explicit tmpfs mount for this step.
774
+ tmpfs_tempdir : bool, default True
775
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
776
+ tmpfs_size : int, optional, default: None
777
+ The value for the size (in MiB) of the tmpfs mount for this step.
778
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
779
+ memory allocated for this step.
780
+ tmpfs_path : str, optional, default /metaflow_temp
781
+ Path to tmpfs mount for this step.
782
+ persistent_volume_claims : Dict[str, str], optional, default None
783
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
784
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
785
+ shared_memory: int, optional
786
+ Shared memory size (in MiB) required for this step
787
+ port: int, optional
788
+ Port number to specify in the Kubernetes job object
789
+ compute_pool : str, optional, default None
790
+ Compute pool to be used for for this step.
791
+ If not specified, any accessible compute pool within the perimeter is used.
792
+ hostname_resolution_timeout: int, default 10 * 60
793
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
794
+ Only applicable when @parallel is used.
795
+ qos: str, default: Burstable
796
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
797
+
798
+ security_context: Dict[str, Any], optional, default None
799
+ Container security context. Applies to the task container. Allows the following keys:
800
+ - privileged: bool, optional, default None
801
+ - allow_privilege_escalation: bool, optional, default None
802
+ - run_as_user: int, optional, default None
803
+ - run_as_group: int, optional, default None
804
+ - run_as_non_root: bool, optional, default None
805
+ """
806
+ ...
807
+
1014
808
  @typing.overload
1015
- 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]]]:
809
+ 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]]]:
1016
810
  """
1017
- Creates a human-readable report, a Metaflow Card, after this step completes.
811
+ Specifies the Conda environment for the step.
1018
812
 
1019
- Note that you may add multiple `@card` decorators in a step with different parameters.
813
+ Information in this decorator will augment any
814
+ attributes set in the `@conda_base` flow-level decorator. Hence,
815
+ you can use `@conda_base` to set packages required by all
816
+ steps and use `@conda` to specify step-specific overrides.
1020
817
 
1021
818
 
1022
819
  Parameters
1023
820
  ----------
1024
- type : str, default 'default'
1025
- Card type.
1026
- id : str, optional, default None
1027
- If multiple cards are present, use this id to identify this card.
1028
- options : Dict[str, Any], default {}
1029
- Options passed to the card. The contents depend on the card type.
1030
- timeout : int, default 45
1031
- Interrupt reporting if it takes more than this many seconds.
821
+ packages : Dict[str, str], default {}
822
+ Packages to use for this step. The key is the name of the package
823
+ and the value is the version to use.
824
+ libraries : Dict[str, str], default {}
825
+ Supported for backward compatibility. When used with packages, packages will take precedence.
826
+ python : str, optional, default None
827
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
828
+ that the version used will correspond to the version of the Python interpreter used to start the run.
829
+ disabled : bool, default False
830
+ If set to True, disables @conda.
1032
831
  """
1033
832
  ...
1034
833
 
1035
834
  @typing.overload
1036
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
835
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1037
836
  ...
1038
837
 
1039
838
  @typing.overload
1040
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
839
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1041
840
  ...
1042
841
 
1043
- def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
842
+ 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):
1044
843
  """
1045
- Creates a human-readable report, a Metaflow Card, after this step completes.
844
+ Specifies the Conda environment for the step.
1046
845
 
1047
- Note that you may add multiple `@card` decorators in a step with different parameters.
846
+ Information in this decorator will augment any
847
+ attributes set in the `@conda_base` flow-level decorator. Hence,
848
+ you can use `@conda_base` to set packages required by all
849
+ steps and use `@conda` to specify step-specific overrides.
1048
850
 
1049
851
 
1050
852
  Parameters
1051
853
  ----------
1052
- type : str, default 'default'
1053
- Card type.
1054
- id : str, optional, default None
1055
- If multiple cards are present, use this id to identify this card.
1056
- options : Dict[str, Any], default {}
1057
- Options passed to the card. The contents depend on the card type.
1058
- timeout : int, default 45
1059
- Interrupt reporting if it takes more than this many seconds.
854
+ packages : Dict[str, str], default {}
855
+ Packages to use for this step. The key is the name of the package
856
+ and the value is the version to use.
857
+ libraries : Dict[str, str], default {}
858
+ Supported for backward compatibility. When used with packages, packages will take precedence.
859
+ python : str, optional, default None
860
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
861
+ that the version used will correspond to the version of the Python interpreter used to start the run.
862
+ disabled : bool, default False
863
+ If set to True, disables @conda.
1060
864
  """
1061
865
  ...
1062
866
 
1063
- 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]]]:
867
+ @typing.overload
868
+ 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]]]:
1064
869
  """
1065
- Specifies that this step should execute on DGX cloud.
870
+ Specifies the PyPI packages for the step.
871
+
872
+ Information in this decorator will augment any
873
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
874
+ you can use `@pypi_base` to set packages required by all
875
+ steps and use `@pypi` to specify step-specific overrides.
1066
876
 
1067
877
 
1068
878
  Parameters
1069
879
  ----------
1070
- gpu : int
1071
- Number of GPUs to use.
1072
- gpu_type : str
1073
- Type of Nvidia GPU to use.
880
+ packages : Dict[str, str], default: {}
881
+ Packages to use for this step. The key is the name of the package
882
+ and the value is the version to use.
883
+ python : str, optional, default: None
884
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
885
+ that the version used will correspond to the version of the Python interpreter used to start the run.
886
+ """
887
+ ...
888
+
889
+ @typing.overload
890
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
891
+ ...
892
+
893
+ @typing.overload
894
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
895
+ ...
896
+
897
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
898
+ """
899
+ Specifies the PyPI packages for the step.
900
+
901
+ Information in this decorator will augment any
902
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
903
+ you can use `@pypi_base` to set packages required by all
904
+ steps and use `@pypi` to specify step-specific overrides.
905
+
906
+
907
+ Parameters
908
+ ----------
909
+ packages : Dict[str, str], default: {}
910
+ Packages to use for this step. The key is the name of the package
911
+ and the value is the version to use.
912
+ python : str, optional, default: None
913
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
914
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1074
915
  """
1075
916
  ...
1076
917
 
@@ -1094,23 +935,149 @@ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
1094
935
  ...
1095
936
 
1096
937
  @typing.overload
1097
- def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1098
- """
1099
- Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1100
- It exists to make it easier for users to know that this decorator should only be used with
1101
- a Neo Cloud like Nebius.
1102
- """
938
+ 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]]]:
939
+ """
940
+ Enables checkpointing for a step.
941
+
942
+ > Examples
943
+
944
+ - Saving Checkpoints
945
+
946
+ ```python
947
+ @checkpoint
948
+ @step
949
+ def train(self):
950
+ model = create_model(self.parameters, checkpoint_path = None)
951
+ for i in range(self.epochs):
952
+ # some training logic
953
+ loss = model.train(self.dataset)
954
+ if i % 10 == 0:
955
+ model.save(
956
+ current.checkpoint.directory,
957
+ )
958
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
959
+ # and returns a reference dictionary to the checkpoint saved in the datastore
960
+ self.latest_checkpoint = current.checkpoint.save(
961
+ name="epoch_checkpoint",
962
+ metadata={
963
+ "epoch": i,
964
+ "loss": loss,
965
+ }
966
+ )
967
+ ```
968
+
969
+ - Using Loaded Checkpoints
970
+
971
+ ```python
972
+ @retry(times=3)
973
+ @checkpoint
974
+ @step
975
+ def train(self):
976
+ # Assume that the task has restarted and the previous attempt of the task
977
+ # saved a checkpoint
978
+ checkpoint_path = None
979
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
980
+ print("Loaded checkpoint from the previous attempt")
981
+ checkpoint_path = current.checkpoint.directory
982
+
983
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
984
+ for i in range(self.epochs):
985
+ ...
986
+ ```
987
+
988
+
989
+ Parameters
990
+ ----------
991
+ load_policy : str, default: "fresh"
992
+ The policy for loading the checkpoint. The following policies are supported:
993
+ - "eager": Loads the the latest available checkpoint within the namespace.
994
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
995
+ will be loaded at the start of the task.
996
+ - "none": Do not load any checkpoint
997
+ - "fresh": Loads the lastest checkpoint created within the running Task.
998
+ This mode helps loading checkpoints across various retry attempts of the same task.
999
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1000
+ created within the task will be loaded when the task is retries execution on failure.
1001
+
1002
+ temp_dir_root : str, default: None
1003
+ The root directory under which `current.checkpoint.directory` will be created.
1004
+ """
1005
+ ...
1006
+
1007
+ @typing.overload
1008
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1103
1009
  ...
1104
1010
 
1105
1011
  @typing.overload
1106
- def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1012
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1107
1013
  ...
1108
1014
 
1109
- def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1015
+ 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):
1110
1016
  """
1111
- Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1112
- It exists to make it easier for users to know that this decorator should only be used with
1113
- a Neo Cloud like Nebius.
1017
+ Enables checkpointing for a step.
1018
+
1019
+ > Examples
1020
+
1021
+ - Saving Checkpoints
1022
+
1023
+ ```python
1024
+ @checkpoint
1025
+ @step
1026
+ def train(self):
1027
+ model = create_model(self.parameters, checkpoint_path = None)
1028
+ for i in range(self.epochs):
1029
+ # some training logic
1030
+ loss = model.train(self.dataset)
1031
+ if i % 10 == 0:
1032
+ model.save(
1033
+ current.checkpoint.directory,
1034
+ )
1035
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
1036
+ # and returns a reference dictionary to the checkpoint saved in the datastore
1037
+ self.latest_checkpoint = current.checkpoint.save(
1038
+ name="epoch_checkpoint",
1039
+ metadata={
1040
+ "epoch": i,
1041
+ "loss": loss,
1042
+ }
1043
+ )
1044
+ ```
1045
+
1046
+ - Using Loaded Checkpoints
1047
+
1048
+ ```python
1049
+ @retry(times=3)
1050
+ @checkpoint
1051
+ @step
1052
+ def train(self):
1053
+ # Assume that the task has restarted and the previous attempt of the task
1054
+ # saved a checkpoint
1055
+ checkpoint_path = None
1056
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1057
+ print("Loaded checkpoint from the previous attempt")
1058
+ checkpoint_path = current.checkpoint.directory
1059
+
1060
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1061
+ for i in range(self.epochs):
1062
+ ...
1063
+ ```
1064
+
1065
+
1066
+ Parameters
1067
+ ----------
1068
+ load_policy : str, default: "fresh"
1069
+ The policy for loading the checkpoint. The following policies are supported:
1070
+ - "eager": Loads the the latest available checkpoint within the namespace.
1071
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1072
+ will be loaded at the start of the task.
1073
+ - "none": Do not load any checkpoint
1074
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1075
+ This mode helps loading checkpoints across various retry attempts of the same task.
1076
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1077
+ created within the task will be loaded when the task is retries execution on failure.
1078
+
1079
+ temp_dir_root : str, default: None
1080
+ The root directory under which `current.checkpoint.directory` will be created.
1114
1081
  """
1115
1082
  ...
1116
1083
 
@@ -1158,138 +1125,172 @@ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy:
1158
1125
  ...
1159
1126
 
1160
1127
  @typing.overload
1161
- 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]]]:
1128
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1162
1129
  """
1163
- Specifies the Conda environment for the step.
1164
-
1165
- Information in this decorator will augment any
1166
- attributes set in the `@conda_base` flow-level decorator. Hence,
1167
- you can use `@conda_base` to set packages required by all
1168
- steps and use `@conda` to specify step-specific overrides.
1169
-
1170
-
1171
- Parameters
1172
- ----------
1173
- packages : Dict[str, str], default {}
1174
- Packages to use for this step. The key is the name of the package
1175
- and the value is the version to use.
1176
- libraries : Dict[str, str], default {}
1177
- Supported for backward compatibility. When used with packages, packages will take precedence.
1178
- python : str, optional, default None
1179
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1180
- that the version used will correspond to the version of the Python interpreter used to start the run.
1181
- disabled : bool, default False
1182
- If set to True, disables @conda.
1130
+ A simple decorator that demonstrates using CardDecoratorInjector
1131
+ to inject a card and render simple markdown content.
1183
1132
  """
1184
1133
  ...
1185
1134
 
1186
1135
  @typing.overload
1187
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1136
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1137
+ ...
1138
+
1139
+ def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1140
+ """
1141
+ A simple decorator that demonstrates using CardDecoratorInjector
1142
+ to inject a card and render simple markdown content.
1143
+ """
1188
1144
  ...
1189
1145
 
1190
1146
  @typing.overload
1191
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1147
+ def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1148
+ """
1149
+ Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1150
+ It exists to make it easier for users to know that this decorator should only be used with
1151
+ a Neo Cloud like Nebius.
1152
+ """
1192
1153
  ...
1193
1154
 
1194
- 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):
1155
+ @typing.overload
1156
+ def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1157
+ ...
1158
+
1159
+ def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1195
1160
  """
1196
- Specifies the Conda environment for the step.
1197
-
1198
- Information in this decorator will augment any
1199
- attributes set in the `@conda_base` flow-level decorator. Hence,
1200
- you can use `@conda_base` to set packages required by all
1201
- steps and use `@conda` to specify step-specific overrides.
1202
-
1203
-
1204
- Parameters
1205
- ----------
1206
- packages : Dict[str, str], default {}
1207
- Packages to use for this step. The key is the name of the package
1208
- and the value is the version to use.
1209
- libraries : Dict[str, str], default {}
1210
- Supported for backward compatibility. When used with packages, packages will take precedence.
1211
- python : str, optional, default None
1212
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1213
- that the version used will correspond to the version of the Python interpreter used to start the run.
1214
- disabled : bool, default False
1215
- If set to True, disables @conda.
1161
+ Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1162
+ It exists to make it easier for users to know that this decorator should only be used with
1163
+ a Neo Cloud like Nebius.
1216
1164
  """
1217
1165
  ...
1218
1166
 
1219
1167
  @typing.overload
1220
- 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]]]:
1168
+ 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]]]:
1221
1169
  """
1222
- Specifies secrets to be retrieved and injected as environment variables prior to
1223
- the execution of a step.
1170
+ Specifies the resources needed when executing this step.
1171
+
1172
+ Use `@resources` to specify the resource requirements
1173
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1174
+
1175
+ You can choose the compute layer on the command line by executing e.g.
1176
+ ```
1177
+ python myflow.py run --with batch
1178
+ ```
1179
+ or
1180
+ ```
1181
+ python myflow.py run --with kubernetes
1182
+ ```
1183
+ which executes the flow on the desired system using the
1184
+ requirements specified in `@resources`.
1224
1185
 
1225
1186
 
1226
1187
  Parameters
1227
1188
  ----------
1228
- sources : List[Union[str, Dict[str, Any]]], default: []
1229
- List of secret specs, defining how the secrets are to be retrieved
1230
- role : str, optional, default: None
1231
- Role to use for fetching secrets
1189
+ cpu : int, default 1
1190
+ Number of CPUs required for this step.
1191
+ gpu : int, optional, default None
1192
+ Number of GPUs required for this step.
1193
+ disk : int, optional, default None
1194
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1195
+ memory : int, default 4096
1196
+ Memory size (in MB) required for this step.
1197
+ shared_memory : int, optional, default None
1198
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1199
+ This parameter maps to the `--shm-size` option in Docker.
1232
1200
  """
1233
1201
  ...
1234
1202
 
1235
1203
  @typing.overload
1236
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1204
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1237
1205
  ...
1238
1206
 
1239
1207
  @typing.overload
1240
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1208
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1241
1209
  ...
1242
1210
 
1243
- 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):
1211
+ 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):
1244
1212
  """
1245
- Specifies secrets to be retrieved and injected as environment variables prior to
1246
- the execution of a step.
1213
+ Specifies the resources needed when executing this step.
1214
+
1215
+ Use `@resources` to specify the resource requirements
1216
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1217
+
1218
+ You can choose the compute layer on the command line by executing e.g.
1219
+ ```
1220
+ python myflow.py run --with batch
1221
+ ```
1222
+ or
1223
+ ```
1224
+ python myflow.py run --with kubernetes
1225
+ ```
1226
+ which executes the flow on the desired system using the
1227
+ requirements specified in `@resources`.
1247
1228
 
1248
1229
 
1249
1230
  Parameters
1250
1231
  ----------
1251
- sources : List[Union[str, Dict[str, Any]]], default: []
1252
- List of secret specs, defining how the secrets are to be retrieved
1253
- role : str, optional, default: None
1254
- Role to use for fetching secrets
1255
- """
1256
- ...
1257
-
1258
- @typing.overload
1259
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1260
- """
1261
- Internal decorator to support Fast bakery
1232
+ cpu : int, default 1
1233
+ Number of CPUs required for this step.
1234
+ gpu : int, optional, default None
1235
+ Number of GPUs required for this step.
1236
+ disk : int, optional, default None
1237
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1238
+ memory : int, default 4096
1239
+ Memory size (in MB) required for this step.
1240
+ shared_memory : int, optional, default None
1241
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1242
+ This parameter maps to the `--shm-size` option in Docker.
1262
1243
  """
1263
1244
  ...
1264
1245
 
1265
1246
  @typing.overload
1266
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1267
- ...
1268
-
1269
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1247
+ 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]]]:
1270
1248
  """
1271
- Internal decorator to support Fast bakery
1249
+ Specifies that the step will success under all circumstances.
1250
+
1251
+ The decorator will create an optional artifact, specified by `var`, which
1252
+ contains the exception raised. You can use it to detect the presence
1253
+ of errors, indicating that all happy-path artifacts produced by the step
1254
+ are missing.
1255
+
1256
+
1257
+ Parameters
1258
+ ----------
1259
+ var : str, optional, default None
1260
+ Name of the artifact in which to store the caught exception.
1261
+ If not specified, the exception is not stored.
1262
+ print_exception : bool, default True
1263
+ Determines whether or not the exception is printed to
1264
+ stdout when caught.
1272
1265
  """
1273
1266
  ...
1274
1267
 
1275
1268
  @typing.overload
1276
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1277
- """
1278
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1279
- It exists to make it easier for users to know that this decorator should only be used with
1280
- a Neo Cloud like CoreWeave.
1281
- """
1269
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1282
1270
  ...
1283
1271
 
1284
1272
  @typing.overload
1285
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1273
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1286
1274
  ...
1287
1275
 
1288
- def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1276
+ 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):
1289
1277
  """
1290
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1291
- It exists to make it easier for users to know that this decorator should only be used with
1292
- a Neo Cloud like CoreWeave.
1278
+ Specifies that the step will success under all circumstances.
1279
+
1280
+ The decorator will create an optional artifact, specified by `var`, which
1281
+ contains the exception raised. You can use it to detect the presence
1282
+ of errors, indicating that all happy-path artifacts produced by the step
1283
+ are missing.
1284
+
1285
+
1286
+ Parameters
1287
+ ----------
1288
+ var : str, optional, default None
1289
+ Name of the artifact in which to store the caught exception.
1290
+ If not specified, the exception is not stored.
1291
+ print_exception : bool, default True
1292
+ Determines whether or not the exception is printed to
1293
+ stdout when caught.
1293
1294
  """
1294
1295
  ...
1295
1296
 
@@ -1390,136 +1391,106 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1390
1391
  ...
1391
1392
 
1392
1393
  @typing.overload
1393
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1394
+ 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]]:
1394
1395
  """
1395
- Specifies the times when the flow should be run when running on a
1396
- production scheduler.
1396
+ Specifies the Conda environment for all steps of the flow.
1397
+
1398
+ Use `@conda_base` to set common libraries required by all
1399
+ steps and use `@conda` to specify step-specific additions.
1397
1400
 
1398
1401
 
1399
1402
  Parameters
1400
1403
  ----------
1401
- hourly : bool, default False
1402
- Run the workflow hourly.
1403
- daily : bool, default True
1404
- Run the workflow daily.
1405
- weekly : bool, default False
1406
- Run the workflow weekly.
1407
- cron : str, optional, default None
1408
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1409
- specified by this expression.
1410
- timezone : str, optional, default None
1411
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1412
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1404
+ packages : Dict[str, str], default {}
1405
+ Packages to use for this flow. The key is the name of the package
1406
+ and the value is the version to use.
1407
+ libraries : Dict[str, str], default {}
1408
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1409
+ python : str, optional, default None
1410
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1411
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1412
+ disabled : bool, default False
1413
+ If set to True, disables Conda.
1413
1414
  """
1414
1415
  ...
1415
1416
 
1416
1417
  @typing.overload
1417
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1418
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1418
1419
  ...
1419
1420
 
1420
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1421
+ 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):
1421
1422
  """
1422
- Specifies the times when the flow should be run when running on a
1423
- production scheduler.
1423
+ Specifies the Conda environment for all steps of the flow.
1424
+
1425
+ Use `@conda_base` to set common libraries required by all
1426
+ steps and use `@conda` to specify step-specific additions.
1424
1427
 
1425
1428
 
1426
1429
  Parameters
1427
1430
  ----------
1428
- hourly : bool, default False
1429
- Run the workflow hourly.
1430
- daily : bool, default True
1431
- Run the workflow daily.
1432
- weekly : bool, default False
1433
- Run the workflow weekly.
1434
- cron : str, optional, default None
1435
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1436
- specified by this expression.
1437
- timezone : str, optional, default None
1438
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1439
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1431
+ packages : Dict[str, str], default {}
1432
+ Packages to use for this flow. The key is the name of the package
1433
+ and the value is the version to use.
1434
+ libraries : Dict[str, str], default {}
1435
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1436
+ python : str, optional, default None
1437
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1438
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1439
+ disabled : bool, default False
1440
+ If set to True, disables Conda.
1440
1441
  """
1441
1442
  ...
1442
1443
 
1443
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1444
+ @typing.overload
1445
+ 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]]:
1444
1446
  """
1445
- Specifies what flows belong to the same project.
1446
-
1447
- A project-specific namespace is created for all flows that
1448
- use the same `@project(name)`.
1449
-
1447
+ Specifies the event(s) that this flow depends on.
1450
1448
 
1451
- Parameters
1452
- ----------
1453
- name : str
1454
- Project name. Make sure that the name is unique amongst all
1455
- projects that use the same production scheduler. The name may
1456
- contain only lowercase alphanumeric characters and underscores.
1449
+ ```
1450
+ @trigger(event='foo')
1451
+ ```
1452
+ or
1453
+ ```
1454
+ @trigger(events=['foo', 'bar'])
1455
+ ```
1457
1456
 
1458
- branch : Optional[str], default None
1459
- The branch to use. If not specified, the branch is set to
1460
- `user.<username>` unless `production` is set to `True`. This can
1461
- also be set on the command line using `--branch` as a top-level option.
1462
- It is an error to specify `branch` in the decorator and on the command line.
1457
+ Additionally, you can specify the parameter mappings
1458
+ to map event payload to Metaflow parameters for the flow.
1459
+ ```
1460
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1461
+ ```
1462
+ or
1463
+ ```
1464
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1465
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1466
+ ```
1463
1467
 
1464
- production : bool, default False
1465
- Whether or not the branch is the production branch. This can also be set on the
1466
- command line using `--production` as a top-level option. It is an error to specify
1467
- `production` in the decorator and on the command line.
1468
- The project branch name will be:
1469
- - if `branch` is specified:
1470
- - if `production` is True: `prod.<branch>`
1471
- - if `production` is False: `test.<branch>`
1472
- - if `branch` is not specified:
1473
- - if `production` is True: `prod`
1474
- - if `production` is False: `user.<username>`
1475
- """
1476
- ...
1477
-
1478
- def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1479
- """
1480
- The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1481
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1482
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1483
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1484
- starts only after all sensors finish.
1468
+ 'parameters' can also be a list of strings and tuples like so:
1469
+ ```
1470
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1471
+ ```
1472
+ This is equivalent to:
1473
+ ```
1474
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1475
+ ```
1485
1476
 
1486
1477
 
1487
1478
  Parameters
1488
1479
  ----------
1489
- timeout : int
1490
- Time, in seconds before the task times out and fails. (Default: 3600)
1491
- poke_interval : int
1492
- Time in seconds that the job should wait in between each try. (Default: 60)
1493
- mode : str
1494
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1495
- exponential_backoff : bool
1496
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1497
- pool : str
1498
- the slot pool this task should run in,
1499
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1500
- soft_fail : bool
1501
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1502
- name : str
1503
- Name of the sensor on Airflow
1504
- description : str
1505
- Description of sensor in the Airflow UI
1506
- bucket_key : Union[str, List[str]]
1507
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1508
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1509
- bucket_name : str
1510
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1511
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1512
- wildcard_match : bool
1513
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1514
- aws_conn_id : str
1515
- a reference to the s3 connection on Airflow. (Default: None)
1516
- verify : bool
1517
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1480
+ event : Union[str, Dict[str, Any]], optional, default None
1481
+ Event dependency for this flow.
1482
+ events : List[Union[str, Dict[str, Any]]], default []
1483
+ Events dependency for this flow.
1484
+ options : Dict[str, Any], default {}
1485
+ Backend-specific configuration for tuning eventing behavior.
1518
1486
  """
1519
1487
  ...
1520
1488
 
1521
1489
  @typing.overload
1522
- 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]]:
1490
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1491
+ ...
1492
+
1493
+ 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] = {}):
1523
1494
  """
1524
1495
  Specifies the event(s) that this flow depends on.
1525
1496
 
@@ -1564,48 +1535,101 @@ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = No
1564
1535
  ...
1565
1536
 
1566
1537
  @typing.overload
1567
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1538
+ 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]]:
1539
+ """
1540
+ Specifies the flow(s) that this flow depends on.
1541
+
1542
+ ```
1543
+ @trigger_on_finish(flow='FooFlow')
1544
+ ```
1545
+ or
1546
+ ```
1547
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1548
+ ```
1549
+ This decorator respects the @project decorator and triggers the flow
1550
+ when upstream runs within the same namespace complete successfully
1551
+
1552
+ Additionally, you can specify project aware upstream flow dependencies
1553
+ by specifying the fully qualified project_flow_name.
1554
+ ```
1555
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1556
+ ```
1557
+ or
1558
+ ```
1559
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1560
+ ```
1561
+
1562
+ You can also specify just the project or project branch (other values will be
1563
+ inferred from the current project or project branch):
1564
+ ```
1565
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1566
+ ```
1567
+
1568
+ Note that `branch` is typically one of:
1569
+ - `prod`
1570
+ - `user.bob`
1571
+ - `test.my_experiment`
1572
+ - `prod.staging`
1573
+
1574
+
1575
+ Parameters
1576
+ ----------
1577
+ flow : Union[str, Dict[str, str]], optional, default None
1578
+ Upstream flow dependency for this flow.
1579
+ flows : List[Union[str, Dict[str, str]]], default []
1580
+ Upstream flow dependencies for this flow.
1581
+ options : Dict[str, Any], default {}
1582
+ Backend-specific configuration for tuning eventing behavior.
1583
+ """
1584
+ ...
1585
+
1586
+ @typing.overload
1587
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1568
1588
  ...
1569
1589
 
1570
- 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] = {}):
1590
+ 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] = {}):
1571
1591
  """
1572
- Specifies the event(s) that this flow depends on.
1592
+ Specifies the flow(s) that this flow depends on.
1573
1593
 
1574
1594
  ```
1575
- @trigger(event='foo')
1595
+ @trigger_on_finish(flow='FooFlow')
1576
1596
  ```
1577
1597
  or
1578
1598
  ```
1579
- @trigger(events=['foo', 'bar'])
1599
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1580
1600
  ```
1601
+ This decorator respects the @project decorator and triggers the flow
1602
+ when upstream runs within the same namespace complete successfully
1581
1603
 
1582
- Additionally, you can specify the parameter mappings
1583
- to map event payload to Metaflow parameters for the flow.
1604
+ Additionally, you can specify project aware upstream flow dependencies
1605
+ by specifying the fully qualified project_flow_name.
1584
1606
  ```
1585
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1607
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1586
1608
  ```
1587
1609
  or
1588
1610
  ```
1589
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1590
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1611
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1591
1612
  ```
1592
1613
 
1593
- 'parameters' can also be a list of strings and tuples like so:
1594
- ```
1595
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1596
- ```
1597
- This is equivalent to:
1614
+ You can also specify just the project or project branch (other values will be
1615
+ inferred from the current project or project branch):
1598
1616
  ```
1599
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1617
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1600
1618
  ```
1601
1619
 
1620
+ Note that `branch` is typically one of:
1621
+ - `prod`
1622
+ - `user.bob`
1623
+ - `test.my_experiment`
1624
+ - `prod.staging`
1625
+
1602
1626
 
1603
1627
  Parameters
1604
1628
  ----------
1605
- event : Union[str, Dict[str, Any]], optional, default None
1606
- Event dependency for this flow.
1607
- events : List[Union[str, Dict[str, Any]]], default []
1608
- Events dependency for this flow.
1629
+ flow : Union[str, Dict[str, str]], optional, default None
1630
+ Upstream flow dependency for this flow.
1631
+ flows : List[Union[str, Dict[str, str]]], default []
1632
+ Upstream flow dependencies for this flow.
1609
1633
  options : Dict[str, Any], default {}
1610
1634
  Backend-specific configuration for tuning eventing behavior.
1611
1635
  """
@@ -1655,53 +1679,53 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1655
1679
  ...
1656
1680
 
1657
1681
  @typing.overload
1658
- 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]]:
1682
+ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1659
1683
  """
1660
- Specifies the Conda environment for all steps of the flow.
1661
-
1662
- Use `@conda_base` to set common libraries required by all
1663
- steps and use `@conda` to specify step-specific additions.
1684
+ Specifies the times when the flow should be run when running on a
1685
+ production scheduler.
1664
1686
 
1665
1687
 
1666
1688
  Parameters
1667
1689
  ----------
1668
- packages : Dict[str, str], default {}
1669
- Packages to use for this flow. The key is the name of the package
1670
- and the value is the version to use.
1671
- libraries : Dict[str, str], default {}
1672
- Supported for backward compatibility. When used with packages, packages will take precedence.
1673
- python : str, optional, default None
1674
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1675
- that the version used will correspond to the version of the Python interpreter used to start the run.
1676
- disabled : bool, default False
1677
- If set to True, disables Conda.
1690
+ hourly : bool, default False
1691
+ Run the workflow hourly.
1692
+ daily : bool, default True
1693
+ Run the workflow daily.
1694
+ weekly : bool, default False
1695
+ Run the workflow weekly.
1696
+ cron : str, optional, default None
1697
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1698
+ specified by this expression.
1699
+ timezone : str, optional, default None
1700
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1701
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1678
1702
  """
1679
1703
  ...
1680
1704
 
1681
1705
  @typing.overload
1682
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1706
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1683
1707
  ...
1684
1708
 
1685
- 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):
1709
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1686
1710
  """
1687
- Specifies the Conda environment for all steps of the flow.
1688
-
1689
- Use `@conda_base` to set common libraries required by all
1690
- steps and use `@conda` to specify step-specific additions.
1711
+ Specifies the times when the flow should be run when running on a
1712
+ production scheduler.
1691
1713
 
1692
1714
 
1693
1715
  Parameters
1694
1716
  ----------
1695
- packages : Dict[str, str], default {}
1696
- Packages to use for this flow. The key is the name of the package
1697
- and the value is the version to use.
1698
- libraries : Dict[str, str], default {}
1699
- Supported for backward compatibility. When used with packages, packages will take precedence.
1700
- python : str, optional, default None
1701
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1702
- that the version used will correspond to the version of the Python interpreter used to start the run.
1703
- disabled : bool, default False
1704
- If set to True, disables Conda.
1717
+ hourly : bool, default False
1718
+ Run the workflow hourly.
1719
+ daily : bool, default True
1720
+ Run the workflow daily.
1721
+ weekly : bool, default False
1722
+ Run the workflow weekly.
1723
+ cron : str, optional, default None
1724
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1725
+ specified by this expression.
1726
+ timezone : str, optional, default None
1727
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1728
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1705
1729
  """
1706
1730
  ...
1707
1731
 
@@ -1819,104 +1843,81 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1819
1843
  """
1820
1844
  ...
1821
1845
 
1822
- @typing.overload
1823
- 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]]:
1846
+ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1824
1847
  """
1825
- Specifies the flow(s) that this flow depends on.
1826
-
1827
- ```
1828
- @trigger_on_finish(flow='FooFlow')
1829
- ```
1830
- or
1831
- ```
1832
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1833
- ```
1834
- This decorator respects the @project decorator and triggers the flow
1835
- when upstream runs within the same namespace complete successfully
1836
-
1837
- Additionally, you can specify project aware upstream flow dependencies
1838
- by specifying the fully qualified project_flow_name.
1839
- ```
1840
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1841
- ```
1842
- or
1843
- ```
1844
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1845
- ```
1846
-
1847
- You can also specify just the project or project branch (other values will be
1848
- inferred from the current project or project branch):
1849
- ```
1850
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1851
- ```
1852
-
1853
- Note that `branch` is typically one of:
1854
- - `prod`
1855
- - `user.bob`
1856
- - `test.my_experiment`
1857
- - `prod.staging`
1848
+ The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1849
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1850
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1851
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1852
+ starts only after all sensors finish.
1858
1853
 
1859
1854
 
1860
1855
  Parameters
1861
1856
  ----------
1862
- flow : Union[str, Dict[str, str]], optional, default None
1863
- Upstream flow dependency for this flow.
1864
- flows : List[Union[str, Dict[str, str]]], default []
1865
- Upstream flow dependencies for this flow.
1866
- options : Dict[str, Any], default {}
1867
- Backend-specific configuration for tuning eventing behavior.
1857
+ timeout : int
1858
+ Time, in seconds before the task times out and fails. (Default: 3600)
1859
+ poke_interval : int
1860
+ Time in seconds that the job should wait in between each try. (Default: 60)
1861
+ mode : str
1862
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1863
+ exponential_backoff : bool
1864
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1865
+ pool : str
1866
+ the slot pool this task should run in,
1867
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1868
+ soft_fail : bool
1869
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1870
+ name : str
1871
+ Name of the sensor on Airflow
1872
+ description : str
1873
+ Description of sensor in the Airflow UI
1874
+ bucket_key : Union[str, List[str]]
1875
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1876
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1877
+ bucket_name : str
1878
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1879
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1880
+ wildcard_match : bool
1881
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1882
+ aws_conn_id : str
1883
+ a reference to the s3 connection on Airflow. (Default: None)
1884
+ verify : bool
1885
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1868
1886
  """
1869
1887
  ...
1870
1888
 
1871
- @typing.overload
1872
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1873
- ...
1874
-
1875
- 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] = {}):
1889
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1876
1890
  """
1877
- Specifies the flow(s) that this flow depends on.
1878
-
1879
- ```
1880
- @trigger_on_finish(flow='FooFlow')
1881
- ```
1882
- or
1883
- ```
1884
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1885
- ```
1886
- This decorator respects the @project decorator and triggers the flow
1887
- when upstream runs within the same namespace complete successfully
1888
-
1889
- Additionally, you can specify project aware upstream flow dependencies
1890
- by specifying the fully qualified project_flow_name.
1891
- ```
1892
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1893
- ```
1894
- or
1895
- ```
1896
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1897
- ```
1898
-
1899
- You can also specify just the project or project branch (other values will be
1900
- inferred from the current project or project branch):
1901
- ```
1902
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1903
- ```
1891
+ Specifies what flows belong to the same project.
1904
1892
 
1905
- Note that `branch` is typically one of:
1906
- - `prod`
1907
- - `user.bob`
1908
- - `test.my_experiment`
1909
- - `prod.staging`
1893
+ A project-specific namespace is created for all flows that
1894
+ use the same `@project(name)`.
1910
1895
 
1911
1896
 
1912
1897
  Parameters
1913
1898
  ----------
1914
- flow : Union[str, Dict[str, str]], optional, default None
1915
- Upstream flow dependency for this flow.
1916
- flows : List[Union[str, Dict[str, str]]], default []
1917
- Upstream flow dependencies for this flow.
1918
- options : Dict[str, Any], default {}
1919
- Backend-specific configuration for tuning eventing behavior.
1899
+ name : str
1900
+ Project name. Make sure that the name is unique amongst all
1901
+ projects that use the same production scheduler. The name may
1902
+ contain only lowercase alphanumeric characters and underscores.
1903
+
1904
+ branch : Optional[str], default None
1905
+ The branch to use. If not specified, the branch is set to
1906
+ `user.<username>` unless `production` is set to `True`. This can
1907
+ also be set on the command line using `--branch` as a top-level option.
1908
+ It is an error to specify `branch` in the decorator and on the command line.
1909
+
1910
+ production : bool, default False
1911
+ Whether or not the branch is the production branch. This can also be set on the
1912
+ command line using `--production` as a top-level option. It is an error to specify
1913
+ `production` in the decorator and on the command line.
1914
+ The project branch name will be:
1915
+ - if `branch` is specified:
1916
+ - if `production` is True: `prod.<branch>`
1917
+ - if `production` is False: `test.<branch>`
1918
+ - if `branch` is not specified:
1919
+ - if `production` is True: `prod`
1920
+ - if `production` is False: `user.<username>`
1920
1921
  """
1921
1922
  ...
1922
1923