ob-metaflow-extensions 1.2.10__py2.py3-none-any.whl → 1.3.0__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-extensions might be problematic. Click here for more details.

@@ -1,12 +1,10 @@
1
- from metaflow.user_configs.config_decorators import (
2
- MutableFlow,
3
- MutableStep,
4
- CustomFlowDecorator,
5
- )
1
+ from metaflow.user_decorators.mutable_flow import MutableFlow
2
+ from metaflow.user_decorators.mutable_step import MutableStep
3
+ from metaflow.user_decorators.user_flow_decorator import FlowMutator
6
4
  from .assume_role import OBP_ASSUME_ROLE_ARN_ENV_VAR
7
5
 
8
6
 
9
- class assume_role(CustomFlowDecorator):
7
+ class assume_role(FlowMutator):
10
8
  """
11
9
  Flow-level decorator for assuming AWS IAM roles.
12
10
 
@@ -42,7 +40,7 @@ class assume_role(CustomFlowDecorator):
42
40
  "`role_arn` must be a valid AWS IAM role ARN starting with 'arn:aws:iam::'"
43
41
  )
44
42
 
45
- def evaluate(self, mutable_flow: MutableFlow) -> None:
43
+ def pre_mutate(self, mutable_flow: MutableFlow) -> None:
46
44
  """
47
45
  This method is called by Metaflow to apply the decorator to the flow.
48
46
  It sets up environment variables that will be used by the AWS client
@@ -51,14 +49,29 @@ class assume_role(CustomFlowDecorator):
51
49
  # Import environment decorator at runtime to avoid circular imports
52
50
  from metaflow import environment
53
51
 
52
+ def _swap_environment_variables(step: MutableStep, role_arn: str) -> None:
53
+ _step_has_env_set = True
54
+ _env_kwargs = {OBP_ASSUME_ROLE_ARN_ENV_VAR: role_arn}
55
+ for d in step.decorator_specs:
56
+ name, _, _, deco_kwargs = d
57
+ if name == "environment":
58
+ _env_kwargs.update(deco_kwargs["vars"])
59
+ _step_has_env_set = True
60
+
61
+ if _step_has_env_set:
62
+ # remove the environment decorator
63
+ step.remove_decorator("environment")
64
+
65
+ # add the environment decorator
66
+ step.add_decorator(
67
+ environment,
68
+ deco_kwargs=dict(vars=_env_kwargs),
69
+ )
70
+
54
71
  # Set the role ARN as an environment variable that will be picked up
55
72
  # by the get_aws_client function
56
73
  def _setup_role_assumption(step: MutableStep) -> None:
57
- # We'll inject the role assumption by adding an environment decorator
58
- # The role will be available through an environment variable
59
- step.add_decorator(
60
- environment, vars={OBP_ASSUME_ROLE_ARN_ENV_VAR: self.role_arn}
61
- )
74
+ _swap_environment_variables(step, self.role_arn)
62
75
 
63
76
  # Apply the role assumption setup to all steps in the flow
64
77
  for _, step in mutable_flow.steps:
@@ -1,12 +1,11 @@
1
- from metaflow.user_configs.config_decorators import (
2
- MutableFlow,
3
- MutableStep,
4
- CustomFlowDecorator,
5
- )
1
+ from metaflow.user_decorators.user_flow_decorator import FlowMutator
2
+ from metaflow.user_decorators.mutable_flow import MutableFlow
3
+ from metaflow.user_decorators.mutable_step import MutableStep
4
+ from .external_chckpt import _ExternalCheckpointFlowDeco
6
5
  import os
7
6
 
8
7
 
9
- class coreweave_checkpoints(CustomFlowDecorator):
8
+ class coreweave_checkpoints(_ExternalCheckpointFlowDeco):
10
9
 
11
10
  """
12
11
 
@@ -46,78 +45,14 @@ class coreweave_checkpoints(CustomFlowDecorator):
46
45
  super().__init__(*args, **kwargs)
47
46
 
48
47
  def init(self, *args, **kwargs):
49
- self.bucket_path = kwargs.get("bucket_path", None)
50
-
51
- self.secrets = kwargs.get("secrets", [])
52
- if self.bucket_path is None:
53
- raise ValueError(
54
- "`bucket_path` keyword argument is required for the coreweave_datastore"
55
- )
56
- if not self.bucket_path.startswith("s3://"):
57
- raise ValueError(
58
- "`bucket_path` must start with `s3://` for the coreweave_datastore"
59
- )
60
-
48
+ super().init(*args, **kwargs)
61
49
  self.coreweave_endpoint_url = f"https://cwobject.com"
62
- if self.secrets is None:
63
- raise ValueError(
64
- "`secrets` keyword argument is required for the coreweave_datastore"
65
- )
66
50
 
67
- def evaluate(self, mutable_flow: MutableFlow) -> None:
51
+ def pre_mutate(self, mutable_flow: MutableFlow) -> None:
68
52
  from metaflow import (
69
- checkpoint,
70
- model,
71
- huggingface_hub,
72
- secrets,
73
53
  with_artifact_store,
74
54
  )
75
55
 
76
- def _add_secrets(step: MutableStep) -> None:
77
- decos_to_add = []
78
- swapping_decos = {
79
- "huggingface_hub": huggingface_hub,
80
- "model": model,
81
- "checkpoint": checkpoint,
82
- }
83
- already_has_secrets = False
84
- secrets_present_in_deco = []
85
- for d in step.decorators:
86
- if d.name in swapping_decos:
87
- decos_to_add.append((d.name, d.attributes))
88
- elif d.name == "secrets":
89
- already_has_secrets = True
90
- secrets_present_in_deco.extend(d.attributes["sources"])
91
-
92
- # If the step aleady has secrets then take all the sources in
93
- # the secrets and add the addtional secrets to the existing secrets
94
- secrets_to_add = self.secrets
95
- if already_has_secrets:
96
- secrets_to_add.extend(secrets_present_in_deco)
97
-
98
- secrets_to_add = list(set(secrets_to_add))
99
-
100
- if len(decos_to_add) == 0:
101
- if already_has_secrets:
102
- step.remove_decorator("secrets")
103
-
104
- step.add_decorator(
105
- secrets,
106
- sources=secrets_to_add,
107
- )
108
- return
109
-
110
- for d, _ in decos_to_add:
111
- step.remove_decorator(d)
112
-
113
- step.add_decorator(
114
- secrets,
115
- sources=secrets_to_add,
116
- )
117
- for d, attrs in decos_to_add:
118
- _deco_to_add = swapping_decos[d]
119
- step.add_decorator(_deco_to_add, **attrs)
120
-
121
56
  def _coreweave_config():
122
57
  return {
123
58
  "root": self.bucket_path,
@@ -131,9 +66,6 @@ class coreweave_checkpoints(CustomFlowDecorator):
131
66
 
132
67
  mutable_flow.add_decorator(
133
68
  with_artifact_store,
134
- type="coreweave",
135
- config=_coreweave_config,
69
+ deco_kwargs=dict(type="coreweave", config=_coreweave_config),
136
70
  )
137
-
138
- for step_name, step in mutable_flow.steps:
139
- _add_secrets(step)
71
+ self._swap_secrets(mutable_flow)
@@ -0,0 +1,85 @@
1
+ from metaflow.user_decorators.user_flow_decorator import FlowMutator
2
+ from metaflow.user_decorators.mutable_flow import MutableFlow
3
+ from metaflow.user_decorators.mutable_step import MutableStep
4
+ import os
5
+
6
+
7
+ class _ExternalCheckpointFlowDeco(FlowMutator):
8
+ def init(self, *args, **kwargs):
9
+ self.bucket_path = kwargs.get("bucket_path", None)
10
+
11
+ self.secrets = kwargs.get("secrets", [])
12
+ if self.bucket_path is None:
13
+ raise ValueError(
14
+ "`bucket_path` keyword argument is required for the coreweave_datastore"
15
+ )
16
+ if not self.bucket_path.startswith("s3://"):
17
+ raise ValueError(
18
+ "`bucket_path` must start with `s3://` for the coreweave_datastore"
19
+ )
20
+ if self.secrets is None:
21
+ raise ValueError(
22
+ "`secrets` keyword argument is required for the coreweave_datastore"
23
+ )
24
+
25
+ def _swap_secrets(self, mutable_flow: MutableFlow) -> None:
26
+ from metaflow import (
27
+ checkpoint,
28
+ model,
29
+ huggingface_hub,
30
+ secrets,
31
+ with_artifact_store,
32
+ )
33
+
34
+ def _add_secrets(step: MutableStep) -> None:
35
+ decos_to_add = []
36
+ swapping_decos = {
37
+ "huggingface_hub": huggingface_hub,
38
+ "model": model,
39
+ "checkpoint": checkpoint,
40
+ }
41
+ already_has_secrets = False
42
+ secrets_present_in_deco = []
43
+ for d in step.decorator_specs:
44
+ name, _, _, deco_kwargs = d
45
+ if name in swapping_decos:
46
+ decos_to_add.append((name, deco_kwargs))
47
+ elif name == "secrets":
48
+ already_has_secrets = True
49
+ secrets_present_in_deco.extend(deco_kwargs["sources"])
50
+
51
+ # If the step aleady has secrets then take all the sources in
52
+ # the secrets and add the addtional secrets to the existing secrets
53
+ secrets_to_add = self.secrets
54
+ if already_has_secrets:
55
+ secrets_to_add.extend(secrets_present_in_deco)
56
+
57
+ secrets_to_add = list(set(secrets_to_add))
58
+
59
+ if len(decos_to_add) == 0:
60
+ if already_has_secrets:
61
+ step.remove_decorator("secrets")
62
+
63
+ step.add_decorator(
64
+ secrets,
65
+ deco_kwargs=dict(
66
+ sources=secrets_to_add,
67
+ ),
68
+ )
69
+ return
70
+
71
+ for d, _ in decos_to_add:
72
+ step.remove_decorator(d)
73
+
74
+ step.add_decorator(
75
+ secrets,
76
+ deco_kwargs=dict(
77
+ sources=secrets_to_add,
78
+ ),
79
+ )
80
+ for d, attrs in decos_to_add:
81
+ _deco_to_add = swapping_decos[d]
82
+ step.add_decorator(_deco_to_add, deco_kwargs=attrs)
83
+
84
+ for step_name, step in mutable_flow.steps:
85
+ _add_secrets(step)
@@ -1,14 +1,11 @@
1
- from metaflow.user_configs.config_decorators import (
2
- MutableFlow,
3
- MutableStep,
4
- CustomFlowDecorator,
5
- )
1
+ from metaflow.user_decorators.mutable_flow import MutableFlow
2
+ from .external_chckpt import _ExternalCheckpointFlowDeco
6
3
  import os
7
4
 
8
5
  NEBIUS_ENDPOINT_URL = "https://storage.eu-north1.nebius.cloud:443"
9
6
 
10
7
 
11
- class nebius_checkpoints(CustomFlowDecorator):
8
+ class nebius_checkpoints(_ExternalCheckpointFlowDeco):
12
9
 
13
10
  """
14
11
 
@@ -52,78 +49,14 @@ class nebius_checkpoints(CustomFlowDecorator):
52
49
  super().__init__(*args, **kwargs)
53
50
 
54
51
  def init(self, *args, **kwargs):
55
- self.bucket_path = kwargs.get("bucket_path", None)
56
-
57
- self.secrets = kwargs.get("secrets", [])
58
- if self.bucket_path is None:
59
- raise ValueError(
60
- "`bucket_path` keyword argument is required for the coreweave_datastore"
61
- )
62
- if not self.bucket_path.startswith("s3://"):
63
- raise ValueError(
64
- "`bucket_path` must start with `s3://` for the coreweave_datastore"
65
- )
66
-
52
+ super().init(*args, **kwargs)
67
53
  self.nebius_endpoint_url = kwargs.get("endpoint_url", NEBIUS_ENDPOINT_URL)
68
- if self.secrets is None:
69
- raise ValueError(
70
- "`secrets` keyword argument is required for the coreweave_datastore"
71
- )
72
54
 
73
- def evaluate(self, mutable_flow: MutableFlow) -> None:
55
+ def pre_mutate(self, mutable_flow: MutableFlow) -> None:
74
56
  from metaflow import (
75
- checkpoint,
76
- model,
77
- huggingface_hub,
78
- secrets,
79
57
  with_artifact_store,
80
58
  )
81
59
 
82
- def _add_secrets(step: MutableStep) -> None:
83
- decos_to_add = []
84
- swapping_decos = {
85
- "huggingface_hub": huggingface_hub,
86
- "model": model,
87
- "checkpoint": checkpoint,
88
- }
89
- already_has_secrets = False
90
- secrets_present_in_deco = []
91
- for d in step.decorators:
92
- if d.name in swapping_decos:
93
- decos_to_add.append((d.name, d.attributes))
94
- elif d.name == "secrets":
95
- already_has_secrets = True
96
- secrets_present_in_deco.extend(d.attributes["sources"])
97
-
98
- # If the step aleady has secrets then take all the sources in
99
- # the secrets and add the addtional secrets to the existing secrets
100
- secrets_to_add = self.secrets
101
- if already_has_secrets:
102
- secrets_to_add.extend(secrets_present_in_deco)
103
-
104
- secrets_to_add = list(set(secrets_to_add))
105
-
106
- if len(decos_to_add) == 0:
107
- if already_has_secrets:
108
- step.remove_decorator("secrets")
109
-
110
- step.add_decorator(
111
- secrets,
112
- sources=secrets_to_add,
113
- )
114
- return
115
-
116
- for d, _ in decos_to_add:
117
- step.remove_decorator(d)
118
-
119
- step.add_decorator(
120
- secrets,
121
- sources=secrets_to_add,
122
- )
123
- for d, attrs in decos_to_add:
124
- _deco_to_add = swapping_decos[d]
125
- step.add_decorator(_deco_to_add, **attrs)
126
-
127
60
  def _nebius_config():
128
61
  return {
129
62
  "root": self.bucket_path,
@@ -135,10 +68,6 @@ class nebius_checkpoints(CustomFlowDecorator):
135
68
  }
136
69
 
137
70
  mutable_flow.add_decorator(
138
- with_artifact_store,
139
- type="s3",
140
- config=_nebius_config,
71
+ with_artifact_store, deco_kwargs=dict(type="s3", config=_nebius_config)
141
72
  )
142
-
143
- for step_name, step in mutable_flow.steps:
144
- _add_secrets(step)
73
+ self._swap_secrets(mutable_flow)
@@ -351,12 +351,16 @@ class DockerEnvironment(MetaflowEnvironment):
351
351
  config.append("--disable=F0401")
352
352
  return config
353
353
 
354
- def get_package_commands(self, codepackage_url, datastore_type):
354
+ def get_package_commands(
355
+ self, codepackage_url, datastore_type, code_package_metadata=None
356
+ ):
355
357
  # we must set the skip install flag at this stage in order to skip package downloads,
356
358
  # doing so in bootstrap_commands is too late in the lifecycle.
357
359
  return [
358
360
  "export METAFLOW_SKIP_INSTALL_DEPENDENCIES=$FASTBAKERY_IMAGE",
359
- ] + super().get_package_commands(codepackage_url, datastore_type)
361
+ ] + super().get_package_commands(
362
+ codepackage_url, datastore_type, code_package_metadata=code_package_metadata
363
+ )
360
364
 
361
365
  def bootstrap_commands(self, step_name, datastore_type):
362
366
  if step_name in self.skipped_steps:
@@ -1,13 +1,13 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: ob-metaflow-extensions
3
- Version: 1.2.10
3
+ Version: 1.3.0
4
4
  Summary: Outerbounds Platform Extensions for Metaflow
5
5
  Author: Outerbounds, Inc.
6
6
  License: Commercial
7
7
  Description-Content-Type: text/markdown
8
8
  Requires-Dist: boto3
9
9
  Requires-Dist: kubernetes
10
- Requires-Dist: ob-metaflow (==2.15.21.5)
10
+ Requires-Dist: ob-metaflow (==2.16.8.1)
11
11
 
12
12
  # Outerbounds platform package
13
13
 
@@ -41,17 +41,18 @@ metaflow_extensions/outerbounds/plugins/apps/core/config/unified_config.py,sha25
41
41
  metaflow_extensions/outerbounds/plugins/apps/core/experimental/__init__.py,sha256=rd4qGTkHndKYfJmoAKZWiY0KK4j5BK6RBrtle-it1Mg,2746
42
42
  metaflow_extensions/outerbounds/plugins/aws/__init__.py,sha256=VBGdjNKeFLXGZuqh4jVk8cFtO1AWof73a6k_cnbAOYA,145
43
43
  metaflow_extensions/outerbounds/plugins/aws/assume_role.py,sha256=mBewNlnSYsR2rFXFkX-DUH6ku01h2yOcMcLHoCL7eyI,161
44
- metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py,sha256=b3te8-RJxsj5cqXbQTFxxeIfnqTERAFInzLgXLVLyYA,2391
44
+ metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py,sha256=MrExIdwnQ4ZVKbrWqEQAIuZO4lWEy1rE0ihGUx7GMJA,2976
45
45
  metaflow_extensions/outerbounds/plugins/card_utilities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
46
46
  metaflow_extensions/outerbounds/plugins/card_utilities/async_cards.py,sha256=6rQhtZXK5DenXPfCRS1ul0jvLJYd48jrJAlnodID21w,4347
47
47
  metaflow_extensions/outerbounds/plugins/card_utilities/extra_components.py,sha256=D8rgCaSc7PuLD0MHJjqsjN0g0PQMN1H-ySOJqi5uIOE,19111
48
48
  metaflow_extensions/outerbounds/plugins/card_utilities/injector.py,sha256=UQtHrviFs7Z5LBh0nbaIv8aEnHnyXFjs9TiGCjtcIWc,2452
49
49
  metaflow_extensions/outerbounds/plugins/checkpoint_datastores/__init__.py,sha256=H96Bdv2XKTewKIU587sjNbU538SKMZOYWQ2qCp1JaNU,84
50
- metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py,sha256=_WzoOROFjoFa8TzsMNFp-r_1Zz7NUp-5ljn_kKlczXA,4534
51
- metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py,sha256=zgqDLFewCeF5jqh-hUNKmC_OAjld09ln0bb8Lkeqapc,4659
50
+ metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py,sha256=_SKSs0qfRt864rmz3gcm87zEPEh16_ai-MJ-D47v15Y,2340
51
+ metaflow_extensions/outerbounds/plugins/checkpoint_datastores/external_chckpt.py,sha256=BRigyuyYYv5H1CxfifFJPPvgMdQeCNIT8Hm0j1xo-YA,3025
52
+ metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py,sha256=xFGdSxlHcwZwgTSKj8efEv6ydCWYqRGrreGqXhf02ko,2321
52
53
  metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
53
54
  metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py,sha256=ShE5omFBr83wkvEhL_ptRFvDNMs6wefg4BjaafQjTcM,3602
54
- metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=nmp_INGIAiWyrhyJ71BH38eRLu1xCIEEKejmXNQ6RlA,15378
55
+ metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=vKN3Waz0SgbOALVTRcf_8MCGhxdRueRFZJ7JmWbRJRM,15487
55
56
  metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=9iqIBowKk95vUYvnMrFNWmmSuK-whKnqJwcCmEcRjTU,5727
56
57
  metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_cli.py,sha256=kqFyu2bJSnc9_9aYfBpz5xK6L6luWFZK_NMuh8f1eVk,1494
57
58
  metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_decorator.py,sha256=MXSIp05-jvt8Q2uGaLKjtuM_ToLeRLxhtMbfHc9Kcko,1515
@@ -116,7 +117,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py,sha256=GRSz2
116
117
  metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
117
118
  metaflow_extensions/outerbounds/toplevel/plugins/torchtune/__init__.py,sha256=uTVkdSk3xZ7hEKYfdlyVteWj5KeDwaM1hU9WT-_YKfI,50
118
119
  metaflow_extensions/outerbounds/toplevel/plugins/vllm/__init__.py,sha256=ekcgD3KVydf-a0xMI60P4uy6ePkSEoFHiGnDq1JM940,45
119
- ob_metaflow_extensions-1.2.10.dist-info/METADATA,sha256=5BoeRaxkfZYm-hBikHcUdYdH4dF-PWkG5MHnJ4fosT4,520
120
- ob_metaflow_extensions-1.2.10.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
121
- ob_metaflow_extensions-1.2.10.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
122
- ob_metaflow_extensions-1.2.10.dist-info/RECORD,,
120
+ ob_metaflow_extensions-1.3.0.dist-info/METADATA,sha256=AqSHKfzK42NBfzMZAHKk3eqN2WWVkxbQQNp_EQA_EnA,518
121
+ ob_metaflow_extensions-1.3.0.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
122
+ ob_metaflow_extensions-1.3.0.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
123
+ ob_metaflow_extensions-1.3.0.dist-info/RECORD,,