Add support for overriding artifacts

talsperre · talsperre · commit db1cc4a695c5 · 2025-03-06T00:55:23.000-08:00
diff --git a/metaflow/cli.py b/metaflow/cli.py
@@ -153,7 +153,7 @@ def show(obj):
     echo_always("\n%s" % obj.graph.doc)
     for node_name in obj.graph.sorted_nodes:
         node = obj.graph[node_name]
-        echo_always("\nStep *%s* and type: *%s*" % (node.name, node.type), err=False)
+        echo_always("\nStep *%s* " % node.name, err=False)
         echo_always(node.doc if node.doc else "?", indent=True, err=False)
         if node.type != "end":
             echo_always(
@@ -474,8 +474,8 @@ def start(
         # ctx.obj.monitor = MONITOR_SIDECARS["nullSidecarMonitor"](
         #     flow=ctx.obj.flow, env=ctx.obj.environment
         # )
-        # ctx.obj.spin_datastore_impl = [d for d in DATASTORES if d.TYPE == "local"][0]
-        ctx.obj.spin_datastore_impl = [d for d in DATASTORES if d.TYPE == "s3"][0]
+        ctx.obj.spin_datastore_impl = [d for d in DATASTORES if d.TYPE == "local"][0]
+        # ctx.obj.spin_datastore_impl = [d for d in DATASTORES if d.TYPE == "s3"][0]
         if datastore_root is None:
             datastore_root = ctx.obj.spin_datastore_impl.get_datastore_root_from_config(
                 ctx.obj.echo
diff --git a/metaflow/cli_components/run_cmds.py b/metaflow/cli_components/run_cmds.py
@@ -398,10 +398,11 @@ def run(
 
 
 @click.command(help="Spins up a task for a given step from a previous run locally.")
-@click.argument(
-    "step-name",
-    required=True,
-    type=str,
+@click.option(
+    "--step-name",
+    default=None,
+    show_default=True,
+    help="Step name to spin up. Must provide either step-name or spin-pathspec.",
 )
 @click.option(
     "--spin-pathspec",
@@ -418,6 +419,15 @@ def run(
     show_default=True,
     help="Skip decorators attached to the step.",
 )
+@click.option(
+    "--artifacts-module",
+    default=None,
+    show_default=True,
+    help="Path to a module that contains artifacts to be used in the spun step. The artifacts should "
+    "be defined as a dictionary called ARTIFACTS with keys as the artifact names and values as the "
+    "artifact values. The artifact values will overwrite the default values of the artifacts used in "
+    "the spun step.",
+)
 @click.option(
     "--max-log-size",
     default=10,
@@ -430,20 +440,31 @@ def run(
 @click.pass_obj
 def spin(
     obj,
-    step_name,
+    step_name=None,
     spin_pathspec=None,
+    artifacts_module=None,
     skip_decorators=False,
     max_log_size=None,
     run_id_file=None,
     runner_attribute_file=None,
     **kwargs
 ):
     before_run(obj, [], [])
-    if spin_pathspec is None:
+    if step_name and spin_pathspec:
+        raise CommandException(
+            "Cannot specify both step-name and spin-pathspec. Please specify only one."
+        )
+    if not (step_name or spin_pathspec):
+        raise CommandException(
+            "Please specify either step-name or spin-pathspec to spin a task."
+        )
+    if step_name is not None:
         spin_pathspec = get_latest_task_pathspec(obj.flow.name, step_name)
+    else:
+        step_name = spin_pathspec.split("/")[2]
 
     obj.echo(
-        f"Spinning up step *{step_name}* locally with task pathspec *{spin_pathspec}*"
+        f"Spinning up step *{step_name}* locally using previous task pathspec *{spin_pathspec}*"
     )
     obj.flow._set_constants(obj.graph, kwargs, obj.config_options)
     step_func = getattr(obj.flow, step_name)
@@ -464,6 +485,7 @@ def spin(
         step_func,
         spin_pathspec,
         skip_decorators,
+        artifacts_module,
         max_log_size * 1024 * 1024,
     )
     _system_logger.log_event(
diff --git a/metaflow/cli_components/step_cmd.py b/metaflow/cli_components/step_cmd.py
@@ -6,7 +6,7 @@
 from ..exception import CommandException
 from ..task import MetaflowTask
 from ..unbounded_foreach import UBF_CONTROL, UBF_TASK
-from ..util import decompress_list
+from ..util import decompress_list, read_artifacts_module
 import metaflow.tracing as tracing
 
 
@@ -232,6 +232,15 @@ def step(
     show_default=True,
     help="Skip decorators attached to the step.",
 )
+@click.option(
+    "--artifacts-module",
+    default=None,
+    show_default=True,
+    help="Path to a module that contains artifacts to be used in the spun step. The artifacts should "
+    "be defined as a dictionary called ARTIFACTS with keys as the artifact names and values as the "
+    "artifact values. The artifact values will overwrite the default values of the artifacts used in "
+    "the spun step.",
+)
 @click.pass_context
 def spin_internal(
     ctx,
@@ -245,6 +254,7 @@ def spin_internal(
     max_user_code_retries=None,
     namespace=None,
     skip_decorators=False,
+    artifacts_module=None,
 ):
     import time
 
@@ -259,16 +269,18 @@ def spin_internal(
         echo = echo_always
 
     input_paths = decompress_list(input_paths) if input_paths else []
-    # print(f"Input paths: {input_paths} and type: {type(input_paths)}")
-    # print(f"Split index: {split_index} and type: {type(split_index)}")
-    # print(f"Retry count: {retry_count} and type: {type(retry_count)}")
-    # print(f"Max user code retries: {max_user_code_retries} and type: {type(max_user_code_retries)}")
-
     echo(
         f"Spinning a task, *{step_name}* with previous task pathspec: {spin_pathspec}",
         fg="magenta",
         bold=False,
     )
+    # if namespace is not None:
+    #     namespace(namespace or None)
+
+    spin_artifacts = read_artifacts_module(artifacts_module) if artifacts_module else {}
+    spin_artifacts = spin_artifacts.get("ARTIFACTS", {})
+
+    print(f"spin_artifacts: {spin_artifacts}")
 
     task = MetaflowTask(
         ctx.obj.flow,
@@ -298,19 +310,8 @@ def spin_internal(
         retry_count,
         max_user_code_retries,
         spin_pathspec,
+        skip_decorators,
+        spin_artifacts,
     )
 
-    #
-    # task.run_spin_step(
-    #     step_name,
-    #     task_pathspec,
-    #     run_id,
-    #     task_id,
-    #     input_paths,
-    #     split_index,
-    #     retry_count,
-    #     max_user_code_retries,
-    #     namespace,
-    #     skip_decorators,
-    # )
-    # print("Time taken for the whole thing: ", time.time() - start)
+    echo(f"Time taken for the whole thing: {time.time() - start}")
diff --git a/metaflow/runtime.py b/metaflow/runtime.py
@@ -91,6 +91,7 @@ def __init__(
         step_func,
         task_pathspec,
         skip_decorators=False,
+        artifacts_module=None,
         max_log_size=MAX_LOG_SIZE,
     ):
         from metaflow import Task
@@ -109,7 +110,6 @@ def __init__(
         # Spin specific metadata, event_logger, monitor, and flow_datastore
         self._spin_metadata = spin_metadata
         self._spin_flow_datastore = spin_flow_datastore
-        print(f"Spin Flow Datastore in SpinRuntime: {self._spin_flow_datastore}")
 
         self._step_func = step_func
         self._task_pathspec = task_pathspec
@@ -119,6 +119,7 @@ def __init__(
         self._whitelist_decorators = None
         self._config_file_name = None
         self._skip_decorators = skip_decorators
+        self._artifacts_module = artifacts_module
         self._max_log_size = max_log_size
         self._encoding = sys.stdout.encoding or "UTF-8"
 
@@ -139,7 +140,7 @@ def split_index(self):
 
     @property
     def input_paths(self):
-        start_time = time.time()
+        st = time.time()
 
         def _format_input_paths(task_pathspec):
             _, run_id, step_name, task_id = task_pathspec.split("/")
@@ -158,13 +159,12 @@ def _format_input_paths(task_pathspec):
             ).task
             self._input_paths = [_format_input_paths(task.pathspec)]
         else:
-            # print("I am in else of input_paths")
             self._input_paths = [
                 _format_input_paths(task_pathspec)
                 for task_pathspec in self._task.parent_task_pathspecs
             ]
-        end_time = time.time()
-        # print(f"Time taken to get input paths: {end_time - start_time}")
+        et = time.time()
+        print(f"Time taken to get input paths: {et - st}")
         return self._input_paths
 
     @property
@@ -232,6 +232,8 @@ def _launch_and_monitor_task(self):
             self._max_log_size,
             self._config_file_name,
             spin_pathspec=self._task_pathspec,
+            skip_decorators=self._skip_decorators,
+            artifacts_module=self._artifacts_module,
         )
 
         # print("Worker created")
@@ -1710,9 +1712,13 @@ class CLIArgs(object):
     for step execution in StepDecorator.runtime_step_cli().
     """
 
-    def __init__(self, task, spin_pathspec=None):
+    def __init__(
+        self, task, spin_pathspec=None, skip_decorators=False, artifacts_module=None
+    ):
         self.task = task
         self.spin_pathspec = spin_pathspec
+        self.skip_decorators = skip_decorators
+        self.artifacts_module = artifacts_module
         self.entrypoint = list(task.entrypoint)
         self.top_level_options = {
             "quiet": False,
@@ -1779,6 +1785,8 @@ def spin_args(self):
             "max-user-code-retries": self.task.user_code_retries,
             "namespace": get_namespace() or "",
             "spin-pathspec": self.spin_pathspec,
+            "skip-decorators": self.skip_decorators,
+            "artifacts-module": self.artifacts_module,
         }
         self.env = {}
 
@@ -1820,10 +1828,20 @@ def __str__(self):
 
 
 class Worker(object):
-    def __init__(self, task, max_logs_size, config_file_name, spin_pathspec=None):
+    def __init__(
+        self,
+        task,
+        max_logs_size,
+        config_file_name,
+        spin_pathspec=None,
+        skip_decorators=False,
+        artifacts_module=None,
+    ):
         self.task = task
         self._config_file_name = config_file_name
         self.spin_pathspec = spin_pathspec
+        self.skip_decorators = skip_decorators
+        self.artifacts_module = artifacts_module
         self._proc = self._launch()
 
         if task.retries > task.user_code_retries:
@@ -1855,7 +1873,12 @@ def __init__(self, task, max_logs_size, config_file_name, spin_pathspec=None):
         # not it is properly shut down)
 
     def _launch(self):
-        args = CLIArgs(self.task, spin_pathspec=self.spin_pathspec)
+        args = CLIArgs(
+            self.task,
+            spin_pathspec=self.spin_pathspec,
+            skip_decorators=self.skip_decorators,
+            artifacts_module=self.artifacts_module,
+        )
         env = dict(os.environ)
 
         if self.task.clone_run_id:
diff --git a/metaflow/task.py b/metaflow/task.py
diff --git a/metaflow/util.py b/metaflow/util.py
