metaflow 2.15.21__py2.py3-none-any.whl → 2.16.1__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (74) hide show
  1. metaflow/__init__.py +7 -1
  2. metaflow/cli.py +19 -1
  3. metaflow/cli_components/init_cmd.py +1 -0
  4. metaflow/cli_components/run_cmds.py +8 -2
  5. metaflow/client/core.py +22 -30
  6. metaflow/datastore/task_datastore.py +0 -1
  7. metaflow/debug.py +5 -0
  8. metaflow/decorators.py +236 -70
  9. metaflow/extension_support/__init__.py +15 -8
  10. metaflow/extension_support/_empty_file.py +2 -2
  11. metaflow/flowspec.py +92 -60
  12. metaflow/graph.py +24 -2
  13. metaflow/meta_files.py +13 -0
  14. metaflow/metadata_provider/metadata.py +7 -1
  15. metaflow/metaflow_config.py +5 -0
  16. metaflow/metaflow_environment.py +82 -25
  17. metaflow/metaflow_version.py +1 -1
  18. metaflow/package/__init__.py +664 -0
  19. metaflow/packaging_sys/__init__.py +870 -0
  20. metaflow/packaging_sys/backend.py +113 -0
  21. metaflow/packaging_sys/distribution_support.py +153 -0
  22. metaflow/packaging_sys/tar_backend.py +86 -0
  23. metaflow/packaging_sys/utils.py +91 -0
  24. metaflow/packaging_sys/v1.py +480 -0
  25. metaflow/plugins/airflow/airflow.py +5 -1
  26. metaflow/plugins/airflow/airflow_cli.py +16 -5
  27. metaflow/plugins/argo/argo_workflows.py +15 -4
  28. metaflow/plugins/argo/argo_workflows_cli.py +17 -4
  29. metaflow/plugins/aws/batch/batch.py +22 -3
  30. metaflow/plugins/aws/batch/batch_cli.py +3 -0
  31. metaflow/plugins/aws/batch/batch_decorator.py +13 -5
  32. metaflow/plugins/aws/step_functions/step_functions.py +4 -1
  33. metaflow/plugins/aws/step_functions/step_functions_cli.py +16 -4
  34. metaflow/plugins/cards/card_decorator.py +0 -5
  35. metaflow/plugins/kubernetes/kubernetes.py +8 -1
  36. metaflow/plugins/kubernetes/kubernetes_cli.py +3 -0
  37. metaflow/plugins/kubernetes/kubernetes_decorator.py +13 -5
  38. metaflow/plugins/package_cli.py +25 -23
  39. metaflow/plugins/parallel_decorator.py +4 -2
  40. metaflow/plugins/pypi/bootstrap.py +8 -2
  41. metaflow/plugins/pypi/conda_decorator.py +39 -82
  42. metaflow/plugins/pypi/conda_environment.py +6 -2
  43. metaflow/plugins/pypi/pypi_decorator.py +4 -4
  44. metaflow/plugins/test_unbounded_foreach_decorator.py +2 -2
  45. metaflow/plugins/timeout_decorator.py +0 -1
  46. metaflow/plugins/uv/bootstrap.py +12 -1
  47. metaflow/plugins/uv/uv_environment.py +4 -2
  48. metaflow/pylint_wrapper.py +5 -1
  49. metaflow/runner/click_api.py +5 -4
  50. metaflow/runner/subprocess_manager.py +14 -2
  51. metaflow/runtime.py +37 -11
  52. metaflow/task.py +92 -7
  53. metaflow/user_configs/config_options.py +13 -8
  54. metaflow/user_configs/config_parameters.py +0 -4
  55. metaflow/user_decorators/__init__.py +0 -0
  56. metaflow/user_decorators/common.py +144 -0
  57. metaflow/user_decorators/mutable_flow.py +499 -0
  58. metaflow/user_decorators/mutable_step.py +424 -0
  59. metaflow/user_decorators/user_flow_decorator.py +264 -0
  60. metaflow/user_decorators/user_step_decorator.py +712 -0
  61. metaflow/util.py +4 -1
  62. metaflow/version.py +1 -1
  63. {metaflow-2.15.21.dist-info → metaflow-2.16.1.dist-info}/METADATA +2 -2
  64. {metaflow-2.15.21.dist-info → metaflow-2.16.1.dist-info}/RECORD +71 -60
  65. metaflow/info_file.py +0 -25
  66. metaflow/package.py +0 -203
  67. metaflow/user_configs/config_decorators.py +0 -568
  68. {metaflow-2.15.21.data → metaflow-2.16.1.data}/data/share/metaflow/devtools/Makefile +0 -0
  69. {metaflow-2.15.21.data → metaflow-2.16.1.data}/data/share/metaflow/devtools/Tiltfile +0 -0
  70. {metaflow-2.15.21.data → metaflow-2.16.1.data}/data/share/metaflow/devtools/pick_services.sh +0 -0
  71. {metaflow-2.15.21.dist-info → metaflow-2.16.1.dist-info}/WHEEL +0 -0
  72. {metaflow-2.15.21.dist-info → metaflow-2.16.1.dist-info}/entry_points.txt +0 -0
  73. {metaflow-2.15.21.dist-info → metaflow-2.16.1.dist-info}/licenses/LICENSE +0 -0
  74. {metaflow-2.15.21.dist-info → metaflow-2.16.1.dist-info}/top_level.txt +0 -0
metaflow/__init__.py CHANGED
@@ -104,7 +104,13 @@ from .flowspec import FlowSpec
104
104
  from .parameters import Parameter, JSONTypeClass, JSONType
105
105
 
106
106
  from .user_configs.config_parameters import Config, ConfigValue, config_expr
107
- from .user_configs.config_decorators import CustomFlowDecorator, CustomStepDecorator
107
+ from .user_decorators.user_step_decorator import (
108
+ UserStepDecorator,
109
+ StepMutator,
110
+ user_step_decorator,
111
+ USER_SKIP_STEP,
112
+ )
113
+ from .user_decorators.user_flow_decorator import FlowMutator
108
114
 
109
115
  # data layer
110
116
  # For historical reasons, we make metaflow.plugins.datatools accessible as
metaflow/cli.py CHANGED
@@ -28,6 +28,7 @@ from .metaflow_config import (
28
28
  from .metaflow_current import current
29
29
  from metaflow.system import _system_monitor, _system_logger
30
30
  from .metaflow_environment import MetaflowEnvironment
31
+ from .packaging_sys import MetaflowCodeContent
31
32
  from .plugins import (
32
33
  DATASTORES,
33
34
  ENVIRONMENTS,
@@ -152,8 +153,13 @@ def check(obj, warnings=False):
152
153
  def show(obj):
153
154
  echo_always("\n%s" % obj.graph.doc)
154
155
  for node_name in obj.graph.sorted_nodes:
156
+ echo_always("")
155
157
  node = obj.graph[node_name]
156
- echo_always("\nStep *%s*" % node.name, err=False)
158
+ for deco in node.decorators:
159
+ echo_always("@%s" % deco.name, err=False)
160
+ for deco in node.wrappers:
161
+ echo_always("@%s" % deco.decorator_name, err=False)
162
+ echo_always("Step *%s*" % node.name, err=False)
157
163
  echo_always(node.doc if node.doc else "?", indent=True, err=False)
158
164
  if node.type != "end":
159
165
  echo_always(
@@ -336,6 +342,11 @@ def start(
336
342
  echo(" executing *%s*" % ctx.obj.flow.name, fg="magenta", nl=False)
337
343
  echo(" for *%s*" % resolve_identity(), fg="magenta")
338
344
 
345
+ # Check if we need to setup the distribution finder (if running )
346
+ dist_info = MetaflowCodeContent.get_distribution_finder()
347
+ if dist_info:
348
+ sys.meta_path.append(dist_info)
349
+
339
350
  # Setup the context
340
351
  cli_args._set_top_kwargs(ctx.params)
341
352
  ctx.obj.echo = echo
@@ -436,6 +447,10 @@ def start(
436
447
  # be raised. For resume, since we ignore those options, we ignore the error.
437
448
  raise ctx.obj.delayed_config_exception
438
449
 
450
+ # Init all values in the flow mutators and then process them
451
+ for decorator in ctx.obj.flow._flow_state.get(_FlowState.FLOW_MUTATORS, []):
452
+ decorator.external_init()
453
+
439
454
  new_cls = ctx.obj.flow._process_config_decorators(config_options)
440
455
  if new_cls:
441
456
  ctx.obj.flow = new_cls(use_cli=False)
@@ -547,6 +562,9 @@ def start(
547
562
  ctx.obj.logger,
548
563
  )
549
564
 
565
+ # Check the graph again (mutators may have changed it)
566
+ ctx.obj.graph = ctx.obj.flow._graph
567
+
550
568
  # TODO (savin): Enable lazy instantiation of package
551
569
  ctx.obj.package = None
552
570
 
@@ -46,6 +46,7 @@ def init(obj, run_id=None, task_id=None, tags=None, **kwargs):
46
46
  obj.event_logger,
47
47
  obj.monitor,
48
48
  run_id=run_id,
49
+ skip_decorator_hooks=True,
49
50
  )
50
51
  obj.flow._set_constants(obj.graph, kwargs, obj.config_options)
51
52
  runtime.persist_constants(task_id=task_id)
@@ -8,7 +8,7 @@ from .. import decorators, namespace, parameters, tracing
8
8
  from ..exception import CommandException
9
9
  from ..graph import FlowGraph
10
10
  from ..metaflow_current import current
11
- from ..metaflow_config import DEFAULT_DECOSPECS
11
+ from ..metaflow_config import DEFAULT_DECOSPECS, FEAT_ALWAYS_UPLOAD_CODE_PACKAGE
12
12
  from ..package import MetaflowPackage
13
13
  from ..runtime import NativeRuntime
14
14
  from ..system import _system_logger
@@ -54,6 +54,8 @@ def before_run(obj, tags, decospecs):
54
54
  decorators._init_step_decorators(
55
55
  obj.flow, obj.graph, obj.environment, obj.flow_datastore, obj.logger
56
56
  )
57
+ # Re-read graph since it may have been modified by mutators
58
+ obj.graph = obj.flow._graph
57
59
 
58
60
  obj.metadata.add_sticky_tags(tags=tags)
59
61
 
@@ -61,7 +63,11 @@ def before_run(obj, tags, decospecs):
61
63
  # We explicitly avoid doing this in `start` since it is invoked for every
62
64
  # step in the run.
63
65
  obj.package = MetaflowPackage(
64
- obj.flow, obj.environment, obj.echo, obj.package_suffixes
66
+ obj.flow,
67
+ obj.environment,
68
+ obj.echo,
69
+ suffixes=obj.package_suffixes,
70
+ flow_datastore=obj.flow_datastore if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE else None,
65
71
  )
66
72
 
67
73
 
metaflow/client/core.py CHANGED
@@ -32,11 +32,12 @@ from metaflow.exception import (
32
32
  from metaflow.includefile import IncludedFile
33
33
  from metaflow.metaflow_config import DEFAULT_METADATA, MAX_ATTEMPTS
34
34
  from metaflow.metaflow_environment import MetaflowEnvironment
35
+ from metaflow.package import MetaflowPackage
36
+ from metaflow.packaging_sys import ContentType
35
37
  from metaflow.plugins import ENVIRONMENTS, METADATA_PROVIDERS
36
38
  from metaflow.unbounded_foreach import CONTROL_TASK_TAG
37
39
  from metaflow.util import cached_property, is_stringish, resolve_identity, to_unicode
38
40
 
39
- from ..info_file import INFO_FILE
40
41
  from .filecache import FileCache
41
42
 
42
43
  if TYPE_CHECKING:
@@ -816,20 +817,26 @@ class MetaflowCode(object):
816
817
  self._path = info["location"]
817
818
  self._ds_type = info["ds_type"]
818
819
  self._sha = info["sha"]
820
+ self._code_metadata = info.get(
821
+ "metadata",
822
+ '{"version": 0, "archive_format": "tgz", "mfcontent_version": 0}',
823
+ )
824
+
825
+ self._backend = MetaflowPackage.get_backend(self._code_metadata)
819
826
 
820
827
  if filecache is None:
821
828
  filecache = FileCache()
822
829
  _, blobdata = filecache.get_data(
823
830
  self._ds_type, self._flow_name, self._path, self._sha
824
831
  )
825
- code_obj = BytesIO(blobdata)
826
- self._tar = tarfile.open(fileobj=code_obj, mode="r:gz")
827
- # The JSON module in Python3 deals with Unicode. Tar gives bytes.
828
- info_str = (
829
- self._tar.extractfile(os.path.basename(INFO_FILE)).read().decode("utf-8")
830
- )
831
- self._info = json.loads(info_str)
832
- self._flowspec = self._tar.extractfile(self._info["script"]).read()
832
+ self._code_obj = BytesIO(blobdata)
833
+ self._info = MetaflowPackage.cls_get_info(self._code_metadata, self._code_obj)
834
+ if self._info:
835
+ self._flowspec = MetaflowPackage.cls_get_content(
836
+ self._code_metadata, self._code_obj, self._info["script"]
837
+ )
838
+ else:
839
+ raise MetaflowInternalError("Code package metadata is invalid.")
833
840
 
834
841
  @property
835
842
  def path(self) -> str:
@@ -877,7 +884,9 @@ class MetaflowCode(object):
877
884
  TarFile
878
885
  TarFile for everything in this code package
879
886
  """
880
- return self._tar
887
+ if self._backend.type == "tgz":
888
+ return self._backend.cls_open(self._code_obj)
889
+ raise RuntimeError("Archive is not a tarball")
881
890
 
882
891
  def extract(self) -> TemporaryDirectory:
883
892
  """
@@ -908,27 +917,10 @@ class MetaflowCode(object):
908
917
  The directory and its contents are automatically deleted when
909
918
  this object is garbage collected.
910
919
  """
911
- exclusions = [
912
- "metaflow/",
913
- "metaflow_extensions/",
914
- "INFO",
915
- "CONFIG_PARAMETERS",
916
- "conda.manifest",
917
- # This file is created when using the conda/pypi features available in
918
- # nflx-metaflow-extensions: https://github.com/Netflix/metaflow-nflx-extensions
919
- "condav2-1.cnd",
920
- ]
921
- members = [
922
- m
923
- for m in self.tarball.getmembers()
924
- if not any(
925
- (x.endswith("/") and m.name.startswith(x)) or (m.name == x)
926
- for x in exclusions
927
- )
928
- ]
929
-
930
920
  tmp = TemporaryDirectory()
931
- self.tarball.extractall(tmp.name, members)
921
+ MetaflowPackage.cls_extract_into(
922
+ self._code_metadata, self._code_obj, tmp.name, ContentType.USER_CONTENT
923
+ )
932
924
  return tmp
933
925
 
934
926
  @property
@@ -99,7 +99,6 @@ class TaskDataStore(object):
99
99
  mode="r",
100
100
  allow_not_done=False,
101
101
  ):
102
-
103
102
  self._storage_impl = flow_datastore._storage_impl
104
103
  self.TYPE = self._storage_impl.TYPE
105
104
  self._ca_store = flow_datastore.ca_store
metaflow/debug.py CHANGED
@@ -42,6 +42,11 @@ class Debug(object):
42
42
  filename = inspect.stack()[1][1]
43
43
  print("debug[%s %s:%s]: %s" % (typ, filename, lineno, s), file=sys.stderr)
44
44
 
45
+ def __getattr__(self, name):
46
+ # Small piece of code to get pyright and other linters to recognize that there
47
+ # are dynamic attributes.
48
+ return getattr(self, name)
49
+
45
50
  def noop(self, args):
46
51
  pass
47
52
 
metaflow/decorators.py CHANGED
@@ -1,37 +1,35 @@
1
- from functools import partial
1
+ import importlib
2
2
  import json
3
3
  import re
4
4
 
5
- from typing import Any, Callable, NewType, TypeVar, Union, overload
5
+ from functools import partial
6
+ from typing import Any, Callable, Dict, List, NewType, Tuple, TypeVar, Union, overload
6
7
 
7
- from .flowspec import FlowSpec
8
+ from .flowspec import FlowSpec, _FlowState
8
9
  from .exception import (
9
10
  MetaflowInternalError,
10
11
  MetaflowException,
11
12
  InvalidDecoratorAttribute,
12
13
  )
13
14
 
15
+ from .debug import debug
14
16
  from .parameters import current_flow
15
- from .user_configs.config_decorators import CustomStepDecorator
16
17
  from .user_configs.config_parameters import (
17
18
  UNPACK_KEY,
18
19
  resolve_delayed_evaluator,
19
20
  unpack_delayed_evaluator,
20
21
  )
22
+ from .user_decorators.mutable_flow import MutableFlow
23
+ from .user_decorators.mutable_step import MutableStep
24
+ from .user_decorators.user_flow_decorator import FlowMutator, FlowMutatorMeta
25
+ from .user_decorators.user_step_decorator import (
26
+ StepMutator,
27
+ UserStepDecoratorBase,
28
+ UserStepDecoratorMeta,
29
+ )
21
30
 
22
31
  from metaflow._vendor import click
23
32
 
24
- try:
25
- unicode
26
- except NameError:
27
- unicode = str
28
- basestring = str
29
-
30
- # Contains the decorators on which _init was called. We want to ensure it is called
31
- # only once on each decorator and, as the _init() function below can be called in
32
- # several places, we need to track which decorator had their init function called
33
- _inited_decorators = set()
34
-
35
33
 
36
34
  class BadStepDecoratorException(MetaflowException):
37
35
  headline = "Syntax error"
@@ -63,11 +61,14 @@ class UnknownStepDecoratorException(MetaflowException):
63
61
  headline = "Unknown step decorator"
64
62
 
65
63
  def __init__(self, deconame):
66
- from .plugins import STEP_DECORATORS
67
-
68
64
  decos = ", ".join(
69
- t.name for t in STEP_DECORATORS if not t.name.endswith("_internal")
65
+ [
66
+ x
67
+ for x in UserStepDecoratorMeta.all_decorators().keys()
68
+ if not x.endswith("_internal")
69
+ ]
70
70
  )
71
+
71
72
  msg = (
72
73
  "Unknown step decorator *{deconame}*. The following decorators are "
73
74
  "supported: *{decos}*".format(deconame=deconame, decos=decos)
@@ -92,9 +93,7 @@ class UnknownFlowDecoratorException(MetaflowException):
92
93
  headline = "Unknown flow decorator"
93
94
 
94
95
  def __init__(self, deconame):
95
- from .plugins import FLOW_DECORATORS
96
-
97
- decos = ", ".join(t.name for t in FLOW_DECORATORS)
96
+ decos = ", ".join(FlowMutatorMeta.all_decorators().keys())
98
97
  msg = (
99
98
  "Unknown flow decorator *{deconame}*. The following decorators are "
100
99
  "supported: *{decos}*".format(deconame=deconame, decos=decos)
@@ -123,9 +122,10 @@ class Decorator(object):
123
122
  # `allow_multiple` allows setting many decorators of the same type to a step/flow.
124
123
  allow_multiple = False
125
124
 
126
- def __init__(self, attributes=None, statically_defined=False):
125
+ def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
127
126
  self.attributes = self.defaults.copy()
128
127
  self.statically_defined = statically_defined
128
+ self.inserted_by = inserted_by
129
129
  self._user_defined_attributes = set()
130
130
  self._ran_init = False
131
131
 
@@ -143,7 +143,9 @@ class Decorator(object):
143
143
  Initializes the decorator. In general, any operation you would do in __init__
144
144
  should be done here.
145
145
  """
146
+ pass
146
147
 
148
+ def external_init(self):
147
149
  # In some cases (specifically when using remove_decorator), we may need to call
148
150
  # init multiple times. Short-circuit re-evaluating.
149
151
  if self._ran_init:
@@ -154,12 +156,14 @@ class Decorator(object):
154
156
  self._user_defined_attributes.update(new_user_attributes)
155
157
  self.attributes = resolve_delayed_evaluator(self.attributes, to_dict=True)
156
158
 
159
+ if "init" in self.__class__.__dict__:
160
+ self.init()
157
161
  self._ran_init = True
158
162
 
159
163
  @classmethod
160
- def _parse_decorator_spec(cls, deco_spec):
164
+ def extract_args_kwargs_from_decorator_spec(cls, deco_spec):
161
165
  if len(deco_spec) == 0:
162
- return cls()
166
+ return [], {}
163
167
 
164
168
  attrs = {}
165
169
  # TODO: Do we really want to allow spaces in the names of attributes?!?
@@ -179,9 +183,20 @@ class Decorator(object):
179
183
  val_parsed = val.strip()
180
184
 
181
185
  attrs[name.strip()] = val_parsed
182
- return cls(attributes=attrs)
186
+
187
+ return [], attrs
188
+
189
+ @classmethod
190
+ def parse_decorator_spec(cls, deco_spec):
191
+ if len(deco_spec) == 0:
192
+ return cls()
193
+
194
+ _, kwargs = cls.extract_args_kwargs_from_decorator_spec(deco_spec)
195
+ return cls(attributes=kwargs)
183
196
 
184
197
  def make_decorator_spec(self):
198
+ # Make sure all attributes are evaluated
199
+ self.external_init()
185
200
  attrs = {k: v for k, v in self.attributes.items() if v is not None}
186
201
  if attrs:
187
202
  attr_list = []
@@ -189,7 +204,7 @@ class Decorator(object):
189
204
  # escaping but for more complex types (typically dictionaries or lists),
190
205
  # we dump using JSON.
191
206
  for k, v in attrs.items():
192
- if isinstance(v, (int, float, unicode, basestring)):
207
+ if isinstance(v, (int, float, str)):
193
208
  attr_list.append("%s=%s" % (k, str(v)))
194
209
  else:
195
210
  attr_list.append("%s=%s" % (k, json.dumps(v).replace('"', '\\"')))
@@ -199,8 +214,21 @@ class Decorator(object):
199
214
  else:
200
215
  return self.name
201
216
 
217
+ def get_args_kwargs(self) -> Tuple[List[Any], Dict[str, Any]]:
218
+ """
219
+ Get the arguments and keyword arguments of the decorator.
220
+
221
+ Returns
222
+ -------
223
+ Tuple[List[Any], Dict[str, Any]]
224
+ A tuple containing a list of arguments and a dictionary of keyword arguments.
225
+ """
226
+ return [], dict(self.attributes)
227
+
202
228
  def __str__(self):
203
229
  mode = "static" if self.statically_defined else "dynamic"
230
+ if self.inserted_by:
231
+ mode += " (inserted by %s)" % " from ".join(self.inserted_by)
204
232
  attrs = " ".join("%s=%s" % x for x in self.attributes.items())
205
233
  if attrs:
206
234
  attrs = " " + attrs
@@ -315,15 +343,36 @@ class StepDecorator(Decorator):
315
343
 
316
344
  def add_to_package(self):
317
345
  """
318
- Called to add custom packages needed for a decorator. This hook will be
346
+ Called to add custom files needed for this environment. This hook will be
319
347
  called in the `MetaflowPackage` class where metaflow compiles the code package
320
- tarball. This hook is invoked in the `MetaflowPackage`'s `path_tuples`
321
- function. The `path_tuples` function is a generator that yields a tuple of
322
- `(file_path, arcname)`.`file_path` is the path of the file in the local file system;
323
- the `arcname` is the path of the file in the constructed tarball or the path of the file
324
- after decompressing the tarball.
325
-
326
- Returns a list of tuples where each tuple represents (file_path, arcname)
348
+ tarball. This hook can return one of two things (the first is for backwards
349
+ compatibility -- move to the second):
350
+ - a generator yielding a tuple of `(file_path, arcname)` to add files to
351
+ the code package. `file_path` is the path to the file on the local filesystem
352
+ and `arcname` is the path relative to the packaged code.
353
+ - a generator yielding a tuple of `(content, arcname, type)` where:
354
+ - type is one of
355
+ ContentType.{USER_CONTENT, CODE_CONTENT, MODULE_CONTENT, OTHER_CONTENT}
356
+ - for USER_CONTENT:
357
+ - the file will be included relative to the directory containing the
358
+ user's flow file.
359
+ - content: path to the file to include
360
+ - arcname: path relative to the directory containing the user's flow file
361
+ - for CODE_CONTENT:
362
+ - the file will be included relative to the code directory in the package.
363
+ This will be the directory containing `metaflow`.
364
+ - content: path to the file to include
365
+ - arcname: path relative to the code directory in the package
366
+ - for MODULE_CONTENT:
367
+ - the module will be added to the code package as a python module. It will
368
+ be accessible as usual (import <module_name>)
369
+ - content: name of the module
370
+ - arcname: None (ignored)
371
+ - for OTHER_CONTENT:
372
+ - the file will be included relative to any other configuration/metadata
373
+ files for the flow
374
+ - content: path to the file to include
375
+ - arcname: path relative to the config directory in the package
327
376
  """
328
377
  return []
329
378
 
@@ -472,7 +521,7 @@ def _base_step_decorator(decotype, *args, **kwargs):
472
521
  # No keyword arguments specified for the decorator, e.g. @foobar.
473
522
  # The first argument is the function to be decorated.
474
523
  func = args[0]
475
- if isinstance(func, CustomStepDecorator):
524
+ if isinstance(func, StepMutator):
476
525
  func = func._my_step
477
526
  if not hasattr(func, "is_step"):
478
527
  raise BadStepDecoratorException(decotype.name, func)
@@ -498,9 +547,10 @@ def _base_step_decorator(decotype, *args, **kwargs):
498
547
 
499
548
 
500
549
  _all_step_decos = None
550
+ _all_flow_decos = None
501
551
 
502
552
 
503
- def _get_all_step_decos():
553
+ def get_all_step_decos():
504
554
  global _all_step_decos
505
555
  if _all_step_decos is None:
506
556
  from .plugins import STEP_DECORATORS
@@ -509,6 +559,67 @@ def _get_all_step_decos():
509
559
  return _all_step_decos
510
560
 
511
561
 
562
+ def get_all_flow_decos():
563
+ global _all_flow_decos
564
+ if _all_flow_decos is None:
565
+ from .plugins import FLOW_DECORATORS
566
+
567
+ _all_flow_decos = {decotype.name: decotype for decotype in FLOW_DECORATORS}
568
+ return _all_flow_decos
569
+
570
+
571
+ def extract_step_decorator_from_decospec(decospec: str):
572
+ splits = decospec.split(":", 1)
573
+ deconame = splits[0]
574
+
575
+ # Check if it is a user-defined decorator or metaflow decorator
576
+ deco_cls = UserStepDecoratorMeta.get_decorator_by_name(deconame)
577
+ if deco_cls is not None:
578
+ return (
579
+ deco_cls.parse_decorator_spec(splits[1] if len(splits) > 1 else ""),
580
+ len(splits) > 1,
581
+ )
582
+
583
+ # Check if this is a decorator we can import
584
+ if "." in deconame:
585
+ # We consider this to be a import path to a user decorator so
586
+ # something like "my_package.my_decorator"
587
+ module_name, class_name = deconame.rsplit(".", 1)
588
+ try:
589
+ module = importlib.import_module(module_name)
590
+ except ImportError as e:
591
+ raise MetaflowException(
592
+ "Could not import user decorator %s" % deconame
593
+ ) from e
594
+ deco_cls = getattr(module, class_name, None)
595
+ if (
596
+ deco_cls is None
597
+ or not isinstance(deco_cls, type)
598
+ or not issubclass(deco_cls, UserStepDecoratorBase)
599
+ ):
600
+ raise UnknownStepDecoratorException(deconame)
601
+ return (
602
+ deco_cls.parse_decorator_spec(splits[1] if len(splits) > 1 else ""),
603
+ len(splits) > 1,
604
+ )
605
+
606
+ raise UnknownStepDecoratorException(deconame)
607
+
608
+
609
+ def extract_flow_decorator_from_decospec(decospec: str):
610
+ splits = decospec.split(":", 1)
611
+ deconame = splits[0]
612
+ # Check if it is a user-defined decorator or metaflow decorator
613
+ deco_cls = FlowMutatorMeta.get_decorator_by_name(deconame)
614
+ if deco_cls is not None:
615
+ return (
616
+ deco_cls.parse_decorator_spec(splits[1] if len(splits) > 1 else ""),
617
+ len(splits) > 1,
618
+ )
619
+ else:
620
+ raise UnknownFlowDecoratorException(deconame)
621
+
622
+
512
623
  def _attach_decorators(flow, decospecs):
513
624
  """
514
625
  Attach decorators to all steps during runtime. This has the same
@@ -532,42 +643,33 @@ def _attach_decorators_to_step(step, decospecs):
532
643
  effect as if you defined the decorators statically in the source for
533
644
  the step.
534
645
  """
535
-
536
- decos = _get_all_step_decos()
537
-
538
646
  for decospec in decospecs:
539
- splits = decospec.split(":", 1)
540
- deconame = splits[0]
541
- if deconame not in decos:
542
- raise UnknownStepDecoratorException(deconame)
543
- # Attach the decorator to step if it doesn't have the decorator
544
- # already. This means that statically defined decorators are always
545
- # preferred over runtime decorators.
546
- if (
547
- deconame not in [deco.name for deco in step.decorators]
548
- or decos[deconame].allow_multiple
549
- ):
550
- # if the decorator is present in a step and is of type allow_multiple
551
- # then add the decorator to the step
552
- deco = decos[deconame]._parse_decorator_spec(
553
- splits[1] if len(splits) > 1 else ""
554
- )
555
- step.decorators.append(deco)
647
+ step_deco, _ = extract_step_decorator_from_decospec(decospec)
648
+ if isinstance(step_deco, StepDecorator):
649
+ # Check multiple
650
+ if (
651
+ step_deco.name not in [deco.name for deco in step.decorators]
652
+ or step_deco.allow_multiple
653
+ ):
654
+ step.decorators.append(step_deco)
655
+ # Else it is ignored -- this is a non-static decorator
656
+
657
+ else:
658
+ step_deco.add_or_raise(step, False, 1, None)
556
659
 
557
660
 
558
661
  def _init(flow, only_non_static=False):
559
662
  for decorators in flow._flow_decorators.values():
560
663
  for deco in decorators:
561
- if deco in _inited_decorators:
562
- continue
563
- deco.init()
564
- _inited_decorators.add(deco)
664
+ deco.external_init()
665
+
565
666
  for flowstep in flow:
566
667
  for deco in flowstep.decorators:
567
- if deco in _inited_decorators:
568
- continue
569
- deco.init()
570
- _inited_decorators.add(deco)
668
+ deco.external_init()
669
+ for deco in flowstep.config_decorators or []:
670
+ deco.external_init()
671
+ for deco in flowstep.wrappers or []:
672
+ deco.external_init()
571
673
 
572
674
 
573
675
  def _init_flow_decorators(
@@ -613,6 +715,74 @@ def _init_flow_decorators(
613
715
 
614
716
 
615
717
  def _init_step_decorators(flow, graph, environment, flow_datastore, logger):
718
+ # NOTE: We don't need graph but keeping it for backwards compatibility with
719
+ # extensions that use it directly. We will remove it at some point.
720
+
721
+ # We call the mutate method for both the flow and step mutators.
722
+ cls = flow.__class__
723
+ # Run all the decorators. We first run the flow-level decorators
724
+ # and then the step level ones to maintain a consistent order with how
725
+ # other decorators are run.
726
+
727
+ for deco in cls._flow_state.get(_FlowState.FLOW_MUTATORS, []):
728
+ if isinstance(deco, FlowMutator):
729
+ inserted_by_value = [deco.decorator_name] + (deco.inserted_by or [])
730
+ mutable_flow = MutableFlow(
731
+ cls,
732
+ pre_mutate=False,
733
+ statically_defined=deco.statically_defined,
734
+ inserted_by=inserted_by_value,
735
+ )
736
+ # Sanity check to make sure we are applying the decorator to the right
737
+ # class
738
+ if not deco._flow_cls == cls and not issubclass(cls, deco._flow_cls):
739
+ raise MetaflowInternalError(
740
+ "FlowMutator registered on the wrong flow -- "
741
+ "expected %s but got %s" % (deco._flow_cls.__name__, cls.__name__)
742
+ )
743
+ debug.userconf_exec(
744
+ "Evaluating flow level decorator %s (post)" % deco.__class__.__name__
745
+ )
746
+ deco.mutate(mutable_flow)
747
+ # We reset cached_parameters on the very off chance that the user added
748
+ # more configurations based on the configuration
749
+ if _FlowState.CACHED_PARAMETERS in cls._flow_state:
750
+ del cls._flow_state[_FlowState.CACHED_PARAMETERS]
751
+ else:
752
+ raise MetaflowInternalError(
753
+ "A non FlowMutator found in flow custom decorators"
754
+ )
755
+
756
+ for step in cls._steps:
757
+ for deco in step.config_decorators:
758
+ inserted_by_value = [deco.decorator_name] + (deco.inserted_by or [])
759
+
760
+ if isinstance(deco, StepMutator):
761
+ debug.userconf_exec(
762
+ "Evaluating step level decorator %s (post) for %s"
763
+ % (deco.__class__.__name__, step.name)
764
+ )
765
+ deco.mutate(
766
+ MutableStep(
767
+ cls,
768
+ step,
769
+ pre_mutate=False,
770
+ statically_defined=deco.statically_defined,
771
+ inserted_by=inserted_by_value,
772
+ )
773
+ )
774
+ else:
775
+ raise MetaflowInternalError(
776
+ "A non StepMutator found in step custom decorators"
777
+ )
778
+
779
+ if step.config_decorators:
780
+ # We remove all mention of the custom step decorator
781
+ setattr(cls, step.name, step)
782
+
783
+ cls._init_graph()
784
+ graph = flow._graph
785
+
616
786
  for step in flow:
617
787
  for deco in step.decorators:
618
788
  deco.step_init(
@@ -685,12 +855,8 @@ def step(
685
855
  f.is_step = True
686
856
  f.decorators = []
687
857
  f.config_decorators = []
688
- try:
689
- # python 3
690
- f.name = f.__name__
691
- except:
692
- # python 2
693
- f.name = f.__func__.func_name
858
+ f.wrappers = []
859
+ f.name = f.__name__
694
860
  return f
695
861
 
696
862