ob-metaflow 2.15.18.1__py2.py3-none-any.whl → 2.16.0.1__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


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

Files changed (93) hide show
  1. metaflow/__init__.py +7 -1
  2. metaflow/_vendor/imghdr/__init__.py +180 -0
  3. metaflow/cli.py +16 -1
  4. metaflow/cli_components/init_cmd.py +1 -0
  5. metaflow/cli_components/run_cmds.py +6 -2
  6. metaflow/client/core.py +22 -30
  7. metaflow/cmd/develop/stub_generator.py +19 -2
  8. metaflow/datastore/task_datastore.py +0 -1
  9. metaflow/debug.py +5 -0
  10. metaflow/decorators.py +230 -70
  11. metaflow/extension_support/__init__.py +15 -8
  12. metaflow/extension_support/_empty_file.py +2 -2
  13. metaflow/flowspec.py +80 -53
  14. metaflow/graph.py +24 -2
  15. metaflow/meta_files.py +13 -0
  16. metaflow/metadata_provider/metadata.py +7 -1
  17. metaflow/metaflow_config.py +5 -0
  18. metaflow/metaflow_environment.py +82 -25
  19. metaflow/metaflow_version.py +1 -1
  20. metaflow/package/__init__.py +664 -0
  21. metaflow/packaging_sys/__init__.py +870 -0
  22. metaflow/packaging_sys/backend.py +113 -0
  23. metaflow/packaging_sys/distribution_support.py +153 -0
  24. metaflow/packaging_sys/tar_backend.py +86 -0
  25. metaflow/packaging_sys/utils.py +91 -0
  26. metaflow/packaging_sys/v1.py +476 -0
  27. metaflow/plugins/__init__.py +3 -0
  28. metaflow/plugins/airflow/airflow.py +11 -1
  29. metaflow/plugins/airflow/airflow_cli.py +15 -4
  30. metaflow/plugins/argo/argo_workflows.py +346 -301
  31. metaflow/plugins/argo/argo_workflows_cli.py +16 -4
  32. metaflow/plugins/argo/exit_hooks.py +209 -0
  33. metaflow/plugins/aws/aws_utils.py +1 -1
  34. metaflow/plugins/aws/batch/batch.py +22 -3
  35. metaflow/plugins/aws/batch/batch_cli.py +3 -0
  36. metaflow/plugins/aws/batch/batch_decorator.py +13 -5
  37. metaflow/plugins/aws/step_functions/step_functions.py +10 -1
  38. metaflow/plugins/aws/step_functions/step_functions_cli.py +15 -4
  39. metaflow/plugins/cards/card_cli.py +20 -1
  40. metaflow/plugins/cards/card_creator.py +24 -1
  41. metaflow/plugins/cards/card_decorator.py +57 -6
  42. metaflow/plugins/cards/card_modules/convert_to_native_type.py +5 -2
  43. metaflow/plugins/cards/card_modules/test_cards.py +16 -0
  44. metaflow/plugins/cards/metadata.py +22 -0
  45. metaflow/plugins/exit_hook/__init__.py +0 -0
  46. metaflow/plugins/exit_hook/exit_hook_decorator.py +46 -0
  47. metaflow/plugins/exit_hook/exit_hook_script.py +52 -0
  48. metaflow/plugins/kubernetes/kubernetes.py +8 -1
  49. metaflow/plugins/kubernetes/kubernetes_cli.py +3 -0
  50. metaflow/plugins/kubernetes/kubernetes_decorator.py +13 -5
  51. metaflow/plugins/package_cli.py +25 -23
  52. metaflow/plugins/parallel_decorator.py +4 -2
  53. metaflow/plugins/pypi/bootstrap.py +8 -2
  54. metaflow/plugins/pypi/conda_decorator.py +39 -82
  55. metaflow/plugins/pypi/conda_environment.py +6 -2
  56. metaflow/plugins/pypi/pypi_decorator.py +4 -4
  57. metaflow/plugins/secrets/__init__.py +3 -0
  58. metaflow/plugins/secrets/secrets_decorator.py +9 -173
  59. metaflow/plugins/secrets/secrets_func.py +49 -0
  60. metaflow/plugins/secrets/secrets_spec.py +101 -0
  61. metaflow/plugins/secrets/utils.py +74 -0
  62. metaflow/plugins/test_unbounded_foreach_decorator.py +2 -2
  63. metaflow/plugins/timeout_decorator.py +0 -1
  64. metaflow/plugins/uv/bootstrap.py +11 -0
  65. metaflow/plugins/uv/uv_environment.py +4 -2
  66. metaflow/pylint_wrapper.py +5 -1
  67. metaflow/runner/click_api.py +5 -4
  68. metaflow/runner/metaflow_runner.py +16 -1
  69. metaflow/runner/subprocess_manager.py +14 -2
  70. metaflow/runtime.py +82 -11
  71. metaflow/task.py +91 -7
  72. metaflow/user_configs/config_options.py +13 -8
  73. metaflow/user_configs/config_parameters.py +0 -4
  74. metaflow/user_decorators/__init__.py +0 -0
  75. metaflow/user_decorators/common.py +144 -0
  76. metaflow/user_decorators/mutable_flow.py +499 -0
  77. metaflow/user_decorators/mutable_step.py +424 -0
  78. metaflow/user_decorators/user_flow_decorator.py +263 -0
  79. metaflow/user_decorators/user_step_decorator.py +712 -0
  80. metaflow/util.py +4 -1
  81. metaflow/version.py +1 -1
  82. {ob_metaflow-2.15.18.1.data → ob_metaflow-2.16.0.1.data}/data/share/metaflow/devtools/Tiltfile +27 -2
  83. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.16.0.1.dist-info}/METADATA +2 -2
  84. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.16.0.1.dist-info}/RECORD +90 -70
  85. metaflow/info_file.py +0 -25
  86. metaflow/package.py +0 -203
  87. metaflow/user_configs/config_decorators.py +0 -568
  88. {ob_metaflow-2.15.18.1.data → ob_metaflow-2.16.0.1.data}/data/share/metaflow/devtools/Makefile +0 -0
  89. {ob_metaflow-2.15.18.1.data → ob_metaflow-2.16.0.1.data}/data/share/metaflow/devtools/pick_services.sh +0 -0
  90. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.16.0.1.dist-info}/WHEEL +0 -0
  91. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.16.0.1.dist-info}/entry_points.txt +0 -0
  92. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.16.0.1.dist-info}/licenses/LICENSE +0 -0
  93. {ob_metaflow-2.15.18.1.dist-info → ob_metaflow-2.16.0.1.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,74 @@
1
+ import os
2
+ import re
3
+ from metaflow.exception import MetaflowException
4
+
5
+ DISALLOWED_SECRETS_ENV_VAR_PREFIXES = ["METAFLOW_"]
6
+
7
+
8
+ def get_default_secrets_backend_type():
9
+ from metaflow.metaflow_config import DEFAULT_SECRETS_BACKEND_TYPE
10
+
11
+ if DEFAULT_SECRETS_BACKEND_TYPE is None:
12
+ raise MetaflowException(
13
+ "No default secrets backend type configured, but needed by @secrets. "
14
+ "Set METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE."
15
+ )
16
+ return DEFAULT_SECRETS_BACKEND_TYPE
17
+
18
+
19
+ def validate_env_vars_across_secrets(all_secrets_env_vars):
20
+ vars_injected_by = {}
21
+ for secret_spec, env_vars in all_secrets_env_vars:
22
+ for k in env_vars:
23
+ if k in vars_injected_by:
24
+ raise MetaflowException(
25
+ "Secret '%s' will inject '%s' as env var, and it is also added by '%s'"
26
+ % (secret_spec, k, vars_injected_by[k])
27
+ )
28
+ vars_injected_by[k] = secret_spec
29
+
30
+
31
+ def validate_env_vars_vs_existing_env(all_secrets_env_vars):
32
+ for secret_spec, env_vars in all_secrets_env_vars:
33
+ for k in env_vars:
34
+ if k in os.environ:
35
+ raise MetaflowException(
36
+ "Secret '%s' will inject '%s' as env var, but it already exists in env"
37
+ % (secret_spec, k)
38
+ )
39
+
40
+
41
+ def validate_env_vars(env_vars):
42
+ for k, v in env_vars.items():
43
+ if not isinstance(k, str):
44
+ raise MetaflowException("Found non string key %s (%s)" % (str(k), type(k)))
45
+ if not isinstance(v, str):
46
+ raise MetaflowException(
47
+ "Found non string value %s (%s)" % (str(v), type(v))
48
+ )
49
+ if not re.fullmatch("[a-zA-Z_][a-zA-Z0-9_]*", k):
50
+ raise MetaflowException("Found invalid env var name '%s'." % k)
51
+ for disallowed_prefix in DISALLOWED_SECRETS_ENV_VAR_PREFIXES:
52
+ if k.startswith(disallowed_prefix):
53
+ raise MetaflowException(
54
+ "Found disallowed env var name '%s' (starts with '%s')."
55
+ % (k, disallowed_prefix)
56
+ )
57
+
58
+
59
+ def get_secrets_backend_provider(secrets_backend_type):
60
+ from metaflow.plugins import SECRETS_PROVIDERS
61
+
62
+ try:
63
+ provider_cls = [
64
+ pc for pc in SECRETS_PROVIDERS if pc.TYPE == secrets_backend_type
65
+ ][0]
66
+ return provider_cls()
67
+ except IndexError:
68
+ raise MetaflowException(
69
+ "Unknown secrets backend type %s (available types: %s)"
70
+ % (
71
+ secrets_backend_type,
72
+ ", ".join(pc.TYPE for pc in SECRETS_PROVIDERS if pc.TYPE != "inline"),
73
+ )
74
+ )
@@ -56,9 +56,9 @@ class InternalTestUnboundedForeachDecorator(StepDecorator):
56
56
  name = "unbounded_test_foreach_internal"
57
57
  results_dict = {}
58
58
 
59
- def __init__(self, attributes=None, statically_defined=False):
59
+ def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
60
60
  super(InternalTestUnboundedForeachDecorator, self).__init__(
61
- attributes, statically_defined
61
+ attributes, statically_defined, inserted_by
62
62
  )
63
63
 
64
64
  def step_init(
@@ -38,7 +38,6 @@ class TimeoutDecorator(StepDecorator):
38
38
  defaults = {"seconds": 0, "minutes": 0, "hours": 0}
39
39
 
40
40
  def init(self):
41
- super().init()
42
41
  # Initialize secs in __init__ so other decorators could safely use this
43
42
  # value without worrying about decorator order.
44
43
  # Convert values in attributes to type:int since they can be type:str
@@ -1,4 +1,5 @@
1
1
  import os
2
+ import shutil
2
3
  import subprocess
3
4
  import sys
4
5
  import time
@@ -6,6 +7,7 @@ import time
6
7
  from metaflow.util import which
7
8
  from metaflow.info_file import read_info_file
8
9
  from metaflow.metaflow_config import get_pinned_conda_libs
10
+ from metaflow.packaging_sys import MetaflowCodeContent, ContentType
9
11
  from urllib.request import Request, urlopen
10
12
  from urllib.error import URLError
11
13
 
@@ -93,6 +95,15 @@ if __name__ == "__main__":
93
95
  return skip_pkgs
94
96
 
95
97
  def sync_uv_project(datastore_type):
98
+ # Move the files to the current directory so uv can find them.
99
+ for filename in ["uv.lock", "pyproject.toml"]:
100
+ path_to_file = MetaflowCodeContent.get_filename(
101
+ filename, ContentType.OTHER_CONTENT
102
+ )
103
+ if path_to_file is None:
104
+ raise RuntimeError(f"Could not find {filename} in the package.")
105
+ shutil.move(path_to_file, os.path.join(os.getcwd(), filename))
106
+
96
107
  print("Syncing uv project...")
97
108
  dependencies = " ".join(get_dependencies(datastore_type))
98
109
  skip_pkgs = " ".join(
@@ -2,6 +2,7 @@ import os
2
2
 
3
3
  from metaflow.exception import MetaflowException
4
4
  from metaflow.metaflow_environment import MetaflowEnvironment
5
+ from metaflow.packaging_sys import ContentType
5
6
 
6
7
 
7
8
  class UVException(MetaflowException):
@@ -12,6 +13,7 @@ class UVEnvironment(MetaflowEnvironment):
12
13
  TYPE = "uv"
13
14
 
14
15
  def __init__(self, flow):
16
+ super().__init__(flow)
15
17
  self.flow = flow
16
18
 
17
19
  def validate_environment(self, logger, datastore_type):
@@ -43,8 +45,8 @@ class UVEnvironment(MetaflowEnvironment):
43
45
  pyproject_path = _find("pyproject.toml")
44
46
  uv_lock_path = _find("uv.lock")
45
47
  files = [
46
- (uv_lock_path, "uv.lock"),
47
- (pyproject_path, "pyproject.toml"),
48
+ (uv_lock_path, "uv.lock", ContentType.OTHER_CONTENT),
49
+ (pyproject_path, "pyproject.toml", ContentType.OTHER_CONTENT),
48
50
  ]
49
51
  return files
50
52
 
@@ -28,7 +28,11 @@ class PyLint(object):
28
28
  return self._run is not None
29
29
 
30
30
  def run(self, logger=None, warnings=False, pylint_config=[]):
31
- args = [self._fname]
31
+ args = [
32
+ self._fname,
33
+ "--signature-mutators",
34
+ "metaflow.user_decorators.user_step_decorator.user_step_decorator",
35
+ ]
32
36
  if not warnings:
33
37
  args.append("--errors-only")
34
38
  if pylint_config:
@@ -46,7 +46,6 @@ from metaflow.exception import MetaflowException
46
46
  from metaflow.includefile import FilePathClass
47
47
  from metaflow.metaflow_config import CLICK_API_PROCESS_CONFIG
48
48
  from metaflow.parameters import JSONTypeClass, flow_context
49
- from metaflow.user_configs.config_decorators import CustomFlowDecorator
50
49
  from metaflow.user_configs.config_options import (
51
50
  ConfigValue,
52
51
  ConvertDictOrStr,
@@ -55,6 +54,7 @@ from metaflow.user_configs.config_options import (
55
54
  MultipleTuple,
56
55
  config_options_with_config_input,
57
56
  )
57
+ from metaflow.user_decorators.user_flow_decorator import FlowMutator
58
58
 
59
59
  # Define a recursive type alias for JSON
60
60
  JSON = Union[Dict[str, "JSON"], List["JSON"], str, int, float, bool, None]
@@ -264,12 +264,12 @@ def extract_flow_class_from_file(flow_file: str) -> FlowSpec:
264
264
  loaded_modules[flow_file] = module
265
265
 
266
266
  classes = inspect.getmembers(
267
- module, lambda x: inspect.isclass(x) or isinstance(x, CustomFlowDecorator)
267
+ module, lambda x: inspect.isclass(x) or isinstance(x, FlowMutator)
268
268
  )
269
269
  flow_cls = None
270
270
 
271
271
  for _, kls in classes:
272
- if isinstance(kls, CustomFlowDecorator):
272
+ if isinstance(kls, FlowMutator):
273
273
  kls = kls._flow_cls
274
274
  if (
275
275
  kls is not FlowSpec
@@ -512,7 +512,7 @@ class MetaflowAPI(object):
512
512
 
513
513
  # At this point, we are like in start() in cli.py -- we obtained the
514
514
  # properly processed config_options which we can now use to process
515
- # the config decorators (including CustomStep/FlowDecorators)
515
+ # the config decorators (including StepMutator/FlowMutator)
516
516
  # Note that if CLICK_API_PROCESS_CONFIG is False, we still do this because
517
517
  # it will init all parameters (config_options will be None)
518
518
  # We ignore any errors if we don't check the configs in the click API.
@@ -658,6 +658,7 @@ if __name__ == "__main__":
658
658
  .kubernetes()
659
659
  .step(
660
660
  step_name="process",
661
+ code_package_metadata="some_version",
661
662
  code_package_sha="some_sha",
662
663
  code_package_url="some_url",
663
664
  )
@@ -1,4 +1,5 @@
1
1
  import importlib
2
+ import inspect
2
3
  import os
3
4
  import sys
4
5
  import json
@@ -200,8 +201,22 @@ class RunnerMeta(type):
200
201
  def f(self, *args, **kwargs):
201
202
  return runner_subcommand(self, *args, **kwargs)
202
203
 
203
- f.__doc__ = runner_subcommand.__doc__ or ""
204
+ f.__doc__ = runner_subcommand.__init__.__doc__ or ""
204
205
  f.__name__ = subcommand_name
206
+ sig = inspect.signature(runner_subcommand)
207
+ # We take all the same parameters except replace the first with
208
+ # simple "self"
209
+ new_parameters = {}
210
+ for name, param in sig.parameters.items():
211
+ if new_parameters:
212
+ new_parameters[name] = param
213
+ else:
214
+ new_parameters["self"] = inspect.Parameter(
215
+ "self", inspect.Parameter.POSITIONAL_OR_KEYWORD
216
+ )
217
+ f.__signature__ = inspect.Signature(
218
+ list(new_parameters.values()), return_annotation=runner_subcommand
219
+ )
205
220
 
206
221
  return f
207
222
 
@@ -9,6 +9,8 @@ import tempfile
9
9
  import threading
10
10
  from typing import Callable, Dict, Iterator, List, Optional, Tuple
11
11
 
12
+ from metaflow.packaging_sys import MetaflowCodeContent
13
+ from metaflow.util import get_metaflow_root
12
14
  from .utils import check_process_exited
13
15
 
14
16
 
@@ -150,7 +152,12 @@ class SubprocessManager(object):
150
152
  int
151
153
  The process ID of the subprocess.
152
154
  """
153
-
155
+ updated_env = MetaflowCodeContent.get_env_vars_for_packaged_metaflow(
156
+ get_metaflow_root()
157
+ )
158
+ if updated_env:
159
+ env = env or {}
160
+ env.update(updated_env)
154
161
  command_obj = CommandManager(command, env, cwd)
155
162
  pid = command_obj.run(show_output=show_output)
156
163
  self.commands[pid] = command_obj
@@ -181,7 +188,12 @@ class SubprocessManager(object):
181
188
  int
182
189
  The process ID of the subprocess.
183
190
  """
184
-
191
+ updated_env = MetaflowCodeContent.get_env_vars_for_packaged_metaflow(
192
+ get_metaflow_root()
193
+ )
194
+ if updated_env:
195
+ env = env or {}
196
+ env.update(updated_env)
185
197
  command_obj = CommandManager(command, env, cwd)
186
198
  pid = await command_obj.async_run()
187
199
  self.commands[pid] = command_obj
metaflow/runtime.py CHANGED
@@ -16,6 +16,7 @@ import time
16
16
  import subprocess
17
17
  from datetime import datetime
18
18
  from io import BytesIO
19
+ from itertools import chain
19
20
  from functools import partial
20
21
  from concurrent import futures
21
22
 
@@ -24,7 +25,7 @@ from contextlib import contextmanager
24
25
 
25
26
  from . import get_namespace
26
27
  from .metadata_provider import MetaDatum
27
- from .metaflow_config import MAX_ATTEMPTS, UI_URL
28
+ from .metaflow_config import FEAT_ALWAYS_UPLOAD_CODE_PACKAGE, MAX_ATTEMPTS, UI_URL
28
29
  from .exception import (
29
30
  MetaflowException,
30
31
  MetaflowInternalError,
@@ -95,6 +96,7 @@ class NativeRuntime(object):
95
96
  max_num_splits=MAX_NUM_SPLITS,
96
97
  max_log_size=MAX_LOG_SIZE,
97
98
  resume_identifier=None,
99
+ skip_decorator_hooks=False,
98
100
  ):
99
101
  if run_id is None:
100
102
  self._run_id = metadata.new_run_id()
@@ -107,6 +109,7 @@ class NativeRuntime(object):
107
109
  self._flow_datastore = flow_datastore
108
110
  self._metadata = metadata
109
111
  self._environment = environment
112
+ self._package = package
110
113
  self._logger = logger
111
114
  self._max_workers = max_workers
112
115
  self._active_tasks = dict() # Key: step name;
@@ -128,6 +131,7 @@ class NativeRuntime(object):
128
131
  self._ran_or_scheduled_task_index = set()
129
132
  self._reentrant = reentrant
130
133
  self._run_url = None
134
+ self._skip_decorator_hooks = skip_decorator_hooks
131
135
 
132
136
  # If steps_to_rerun is specified, we will not clone them in resume mode.
133
137
  self._steps_to_rerun = steps_to_rerun or {}
@@ -179,9 +183,10 @@ class NativeRuntime(object):
179
183
  # finished.
180
184
  self._control_num_splits = {} # control_task -> num_splits mapping
181
185
 
182
- for step in flow:
183
- for deco in step.decorators:
184
- deco.runtime_init(flow, graph, package, self._run_id)
186
+ if not self._skip_decorator_hooks:
187
+ for step in flow:
188
+ for deco in step.decorators:
189
+ deco.runtime_init(flow, graph, package, self._run_id)
185
190
 
186
191
  def _new_task(self, step, input_paths=None, **kwargs):
187
192
  if input_paths is None:
@@ -192,7 +197,7 @@ class NativeRuntime(object):
192
197
  if step in self._steps_to_rerun:
193
198
  may_clone = False
194
199
 
195
- if step == "_parameters":
200
+ if step == "_parameters" or self._skip_decorator_hooks:
196
201
  decos = []
197
202
  else:
198
203
  decos = getattr(self._flow, step).decorators
@@ -566,9 +571,11 @@ class NativeRuntime(object):
566
571
  raise
567
572
  finally:
568
573
  # on finish clean tasks
569
- for step in self._flow:
570
- for deco in step.decorators:
571
- deco.runtime_finished(exception)
574
+ if not self._skip_decorator_hooks:
575
+ for step in self._flow:
576
+ for deco in step.decorators:
577
+ deco.runtime_finished(exception)
578
+ self._run_exit_hooks()
572
579
 
573
580
  # assert that end was executed and it was successful
574
581
  if ("end", ()) in self._finished:
@@ -591,6 +598,50 @@ class NativeRuntime(object):
591
598
  "The *end* step was not successful by the end of flow."
592
599
  )
593
600
 
601
+ def _run_exit_hooks(self):
602
+ try:
603
+ exit_hook_decos = self._flow._flow_decorators.get("exit_hook", [])
604
+ if not exit_hook_decos:
605
+ return
606
+
607
+ successful = ("end", ()) in self._finished or self._clone_only
608
+ pathspec = f"{self._graph.name}/{self._run_id}"
609
+ flow_file = self._environment.get_environment_info()["script"]
610
+
611
+ def _call(fn_name):
612
+ try:
613
+ result = (
614
+ subprocess.check_output(
615
+ args=[
616
+ sys.executable,
617
+ "-m",
618
+ "metaflow.plugins.exit_hook.exit_hook_script",
619
+ flow_file,
620
+ fn_name,
621
+ pathspec,
622
+ ],
623
+ env=os.environ,
624
+ )
625
+ .decode()
626
+ .strip()
627
+ )
628
+ print(result)
629
+ except subprocess.CalledProcessError as e:
630
+ print(f"[exit_hook] Hook '{fn_name}' failed with error: {e}")
631
+ except Exception as e:
632
+ print(f"[exit_hook] Unexpected error in hook '{fn_name}': {e}")
633
+
634
+ # Call all exit hook functions regardless of individual failures
635
+ for fn_name in [
636
+ name
637
+ for deco in exit_hook_decos
638
+ for name in (deco.success_hooks if successful else deco.error_hooks)
639
+ ]:
640
+ _call(fn_name)
641
+
642
+ except Exception as ex:
643
+ pass # do not fail due to exit hooks for whatever reason.
644
+
594
645
  def _killall(self):
595
646
  # If we are here, all children have received a signal and are shutting down.
596
647
  # We want to give them an opportunity to do so and then kill
@@ -622,7 +673,6 @@ class NativeRuntime(object):
622
673
  # Given the current task information (task_index), the type of transition,
623
674
  # and the split index, return the new task index.
624
675
  def _translate_index(self, task, next_step, type, split_index=None):
625
-
626
676
  match = re.match(r"^(.+)\[(.*)\]$", task.task_index)
627
677
  if match:
628
678
  _, foreach_index = match.groups()
@@ -960,6 +1010,22 @@ class NativeRuntime(object):
960
1010
  # Initialize the task (which can be expensive using remote datastores)
961
1011
  # before launching the worker so that cost is amortized over time, instead
962
1012
  # of doing it during _queue_push.
1013
+ if (
1014
+ FEAT_ALWAYS_UPLOAD_CODE_PACKAGE
1015
+ and "METAFLOW_CODE_SHA" not in os.environ
1016
+ ):
1017
+ # We check if the code package is uploaded and, if so, we set the
1018
+ # environment variables that will cause the metadata service to
1019
+ # register the code package with the task created in _new_task below
1020
+ code_sha = self._package.package_sha(timeout=0.01)
1021
+ if code_sha:
1022
+ os.environ["METAFLOW_CODE_SHA"] = code_sha
1023
+ os.environ["METAFLOW_CODE_URL"] = self._package.package_url()
1024
+ os.environ["METAFLOW_CODE_DS"] = self._flow_datastore.TYPE
1025
+ os.environ["METAFLOW_CODE_METADATA"] = (
1026
+ self._package.package_metadata
1027
+ )
1028
+
963
1029
  task = self._new_task(step, **task_kwargs)
964
1030
  self._launch_worker(task)
965
1031
 
@@ -1511,6 +1577,7 @@ class CLIArgs(object):
1511
1577
  def __init__(self, task):
1512
1578
  self.task = task
1513
1579
  self.entrypoint = list(task.entrypoint)
1580
+ step_obj = getattr(self.task.flow, self.task.step)
1514
1581
  self.top_level_options = {
1515
1582
  "quiet": True,
1516
1583
  "metadata": self.task.metadata_type,
@@ -1522,8 +1589,12 @@ class CLIArgs(object):
1522
1589
  "datastore-root": self.task.datastore_sysroot,
1523
1590
  "with": [
1524
1591
  deco.make_decorator_spec()
1525
- for deco in self.task.decos
1526
- if not deco.statically_defined
1592
+ for deco in chain(
1593
+ self.task.decos,
1594
+ step_obj.wrappers,
1595
+ step_obj.config_decorators,
1596
+ )
1597
+ if not deco.statically_defined and deco.inserted_by is None
1527
1598
  ],
1528
1599
  }
1529
1600
 
metaflow/task.py CHANGED
@@ -6,6 +6,7 @@ import os
6
6
  import time
7
7
  import traceback
8
8
 
9
+
9
10
  from types import MethodType, FunctionType
10
11
 
11
12
  from metaflow.sidecar import Message, MessageTypes
@@ -24,6 +25,7 @@ from .unbounded_foreach import UBF_CONTROL
24
25
  from .util import all_equal, get_username, resolve_identity, unicode_type
25
26
  from .clone_util import clone_task_helper
26
27
  from .metaflow_current import current
28
+ from metaflow.user_configs.config_parameters import ConfigValue
27
29
  from metaflow.system import _system_logger, _system_monitor
28
30
  from metaflow.tracing import get_trace_id
29
31
  from metaflow.tuple_util import ForeachFrame
@@ -57,11 +59,89 @@ class MetaflowTask(object):
57
59
  self.monitor = monitor
58
60
  self.ubf_context = ubf_context
59
61
 
60
- def _exec_step_function(self, step_function, input_obj=None):
61
- if input_obj is None:
62
- step_function()
63
- else:
64
- step_function(input_obj)
62
+ def _exec_step_function(self, step_function, orig_step_func, input_obj=None):
63
+ wrappers_stack = []
64
+ wrapped_func = None
65
+ do_next = False
66
+ raised_exception = None
67
+ # If we have wrappers w1, w2 and w3, we need to execute
68
+ # - w3_pre
69
+ # - w2_pre
70
+ # - w1_pre
71
+ # - step_function
72
+ # - w1_post
73
+ # - w2_post
74
+ # - w3_post
75
+ # in that order. We do this by maintaining a stack of generators.
76
+ # Note that if any of the pre functions returns a function, we execute that
77
+ # instead of the rest of the inside part. This is useful if you want to create
78
+ # no-op function for example.
79
+ for w in reversed(orig_step_func.wrappers):
80
+ wrapped_func = w.pre_step(orig_step_func.name, self.flow, input_obj)
81
+ wrappers_stack.append(w)
82
+ if w.skip_step:
83
+ # We have nothing to run
84
+ do_next = w.skip_step
85
+ break
86
+ if wrapped_func:
87
+ break # We have nothing left to do since we now execute the
88
+ # wrapped function
89
+ # Else, we continue down the list of wrappers
90
+ try:
91
+ if not do_next:
92
+ if input_obj is None:
93
+ if wrapped_func:
94
+ do_next = wrapped_func(self.flow)
95
+ if not do_next:
96
+ do_next = True
97
+ else:
98
+ step_function()
99
+ else:
100
+ if wrapped_func:
101
+ do_next = wrapped_func(self.flow, input_obj)
102
+ if not do_next:
103
+ do_next = True
104
+ else:
105
+ step_function(input_obj)
106
+ except Exception as ex:
107
+ raised_exception = ex
108
+
109
+ if do_next:
110
+ # If we are skipping the step, or executed a wrapped function,
111
+ # we need to set the transition variables
112
+ # properly. We call the next function as needed
113
+ graph_node = self.flow._graph[step_function.name]
114
+ out_funcs = [getattr(self.flow, f) for f in graph_node.out_funcs]
115
+ if out_funcs:
116
+ if isinstance(do_next, bool):
117
+ # We need to extract things from the self.next. This is not possible
118
+ # in the case where there was a num_parallel.
119
+ if graph_node.parallel_foreach:
120
+ raise RuntimeError(
121
+ "Skipping a parallel foreach step without providing "
122
+ "the arguments to the self.next call is not supported. "
123
+ )
124
+ if graph_node.foreach_param:
125
+ self.flow.next(*out_funcs, foreach=graph_node.foreach_param)
126
+ else:
127
+ self.flow.next(*out_funcs)
128
+ elif isinstance(do_next, dict):
129
+ # Here it is a dictionary so we just call the next method with
130
+ # those arguments
131
+ self.flow.next(*out_funcs, **do_next)
132
+ else:
133
+ raise RuntimeError(
134
+ "Invalid value passed to self.next; expected "
135
+ " bool of a dictionary; got: %s" % do_next
136
+ )
137
+ # We back out of the stack of generators
138
+ for w in reversed(wrappers_stack):
139
+ raised_exception = w.post_step(
140
+ step_function.name, self.flow, raised_exception
141
+ )
142
+ if raised_exception:
143
+ # We have an exception that we need to propagate
144
+ raise raised_exception
65
145
 
66
146
  def _init_parameters(self, parameter_ds, passdown=True):
67
147
  cls = self.flow.__class__
@@ -538,6 +618,9 @@ class MetaflowTask(object):
538
618
  output.save_metadata(
539
619
  {
540
620
  "task_begin": {
621
+ "code_package_metadata": os.environ.get(
622
+ "METAFLOW_CODE_METADATA", ""
623
+ ),
541
624
  "code_package_sha": os.environ.get("METAFLOW_CODE_SHA"),
542
625
  "code_package_ds": os.environ.get("METAFLOW_CODE_DS"),
543
626
  "code_package_url": os.environ.get("METAFLOW_CODE_URL"),
@@ -651,6 +734,7 @@ class MetaflowTask(object):
651
734
  inputs,
652
735
  )
653
736
 
737
+ orig_step_func = step_func
654
738
  for deco in decorators:
655
739
  # decorators can actually decorate the step function,
656
740
  # or they can replace it altogether. This functionality
@@ -667,9 +751,9 @@ class MetaflowTask(object):
667
751
  )
668
752
 
669
753
  if join_type:
670
- self._exec_step_function(step_func, input_obj)
754
+ self._exec_step_function(step_func, orig_step_func, input_obj)
671
755
  else:
672
- self._exec_step_function(step_func)
756
+ self._exec_step_function(step_func, orig_step_func)
673
757
 
674
758
  for deco in decorators:
675
759
  deco.task_post_step(
@@ -7,8 +7,9 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
7
7
  from metaflow._vendor import click
8
8
  from metaflow.debug import debug
9
9
 
10
- from .config_parameters import CONFIG_FILE, ConfigValue
10
+ from .config_parameters import ConfigValue
11
11
  from ..exception import MetaflowException, MetaflowInternalError
12
+ from ..packaging_sys import MetaflowCodeContent
12
13
  from ..parameters import DeployTimeField, ParameterContext, current_flow
13
14
  from ..util import get_username
14
15
 
@@ -24,12 +25,16 @@ _CONVERTED_DEFAULT_NO_FILE = _CONVERTED_DEFAULT + _NO_FILE
24
25
 
25
26
  def _load_config_values(info_file: Optional[str] = None) -> Optional[Dict[Any, Any]]:
26
27
  if info_file is None:
27
- info_file = os.path.basename(CONFIG_FILE)
28
- try:
29
- with open(info_file, encoding="utf-8") as contents:
30
- return json.load(contents).get("user_configs", {})
31
- except IOError:
32
- return None
28
+ config_content = MetaflowCodeContent.get_config()
29
+ else:
30
+ try:
31
+ with open(info_file, encoding="utf-8") as f:
32
+ config_content = json.load(f)
33
+ except IOError:
34
+ return None
35
+ if config_content:
36
+ return config_content.get("user_configs", {})
37
+ return None
33
38
 
34
39
 
35
40
  class ConvertPath(click.Path):
@@ -437,7 +442,7 @@ class LocalFileInput(click.Path):
437
442
  # Small wrapper around click.Path to set the value from which to read configuration
438
443
  # values. This is set immediately upon processing the --local-config-file
439
444
  # option and will therefore then be available when processing any of the other
440
- # --config options (which will call ConfigInput.process_configs
445
+ # --config options (which will call ConfigInput.process_configs)
441
446
  name = "LocalFileInput"
442
447
 
443
448
  def convert(self, value, param, ctx):
@@ -49,10 +49,6 @@ if TYPE_CHECKING:
49
49
 
50
50
  # return tracefunc_closure
51
51
 
52
- CONFIG_FILE = os.path.join(
53
- os.path.dirname(os.path.abspath(__file__)), "CONFIG_PARAMETERS"
54
- )
55
-
56
52
  ID_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
57
53
 
58
54
  UNPACK_KEY = "_unpacked_delayed_"
File without changes