Misc bugfixes (#3234)

* Set stack before importing pipeline * Improve logic when fetching neptune run * Fix error message * Auto-update of NLP template * Missing return * Auto-update of LLM Finetuning template * Auto-update of Starter template * Auto-update of E2E template --------- Co-authored-by: GitHub Actions <[email protected]>
zenml-io · Dec 2, 2024 · 2b843f4 · 2b843f4
1 parent 634d034
commit 2b843f4
Show file tree

Hide file tree

Showing 8 changed files with 130 additions and 143 deletions.
diff --git a/examples/e2e/.copier-answers.yml b/examples/e2e/.copier-answers.yml
@@ -1,5 +1,5 @@
 # Changes here will be overwritten by Copier
-_commit: 2024.11.20-2-g760142f
+_commit: 2024.11.28
 _src_path: gh:zenml-io/template-e2e-batch
 data_quality_checks: true
 email: [email protected]

diff --git a/examples/e2e_nlp/.copier-answers.yml b/examples/e2e_nlp/.copier-answers.yml
@@ -1,5 +1,5 @@
 # Changes here will be overwritten by Copier
-_commit: 2024.10.30-2-g1ae14e3
+_commit: 2024.11.28
 _src_path: gh:zenml-io/template-nlp
 accelerator: cpu
 cloud_of_choice: aws

diff --git a/examples/llm_finetuning/.copier-answers.yml b/examples/llm_finetuning/.copier-answers.yml
@@ -1,5 +1,5 @@
 # Changes here will be overwritten by Copier
-_commit: 2024.11.08-2-gece1d46
+_commit: 2024.11.28
 _src_path: gh:zenml-io/template-llm-finetuning
 bf16: true
 cuda_version: cuda11.8

diff --git a/examples/mlops_starter/.copier-answers.yml b/examples/mlops_starter/.copier-answers.yml
@@ -1,5 +1,5 @@
 # Changes here will be overwritten by Copier
-_commit: 2024.10.30-7-gb60e441
+_commit: 2024.11.28
 _src_path: gh:zenml-io/template-starter
 email: [email protected]
 full_name: ZenML GmbH

diff --git a/src/zenml/cli/pipeline.py b/src/zenml/cli/pipeline.py
@@ -40,6 +40,35 @@
 logger = get_logger(__name__)
 
 
+def _import_pipeline(source: str) -> Pipeline:
+    """Import a pipeline.
+
+    Args:
+        source: The pipeline source.
+
+    Returns:
+        The pipeline.
+    """
+    try:
+        pipeline_instance = source_utils.load(source)
+    except ModuleNotFoundError as e:
+        source_root = source_utils.get_source_root()
+        cli_utils.error(
+            f"Unable to import module `{e.name}`. Make sure the source path is "
+            f"relative to your source root `{source_root}`."
+        )
+    except AttributeError as e:
+        cli_utils.error("Unable to load attribute from module: " + str(e))
+
+    if not isinstance(pipeline_instance, Pipeline):
+        cli_utils.error(
+            f"The given source path `{source}` does not resolve to a pipeline "
+            "object."
+        )
+
+    return pipeline_instance
+
+
 @cli.group(cls=TagGroup, tag=CliCategories.MANAGEMENT_TOOLS)
 def pipeline() -> None:
     """Interact with pipelines, runs and schedules."""
@@ -85,22 +114,7 @@ def register_pipeline(
             "source code root."
         )
 
-    try:
-        pipeline_instance = source_utils.load(source)
-    except ModuleNotFoundError as e:
-        source_root = source_utils.get_source_root()
-        cli_utils.error(
-            f"Unable to import module `{e.name}`. Make sure the source path is "
-            f"relative to your source root `{source_root}`."
-        )
-    except AttributeError as e:
-        cli_utils.error("Unable to load attribute from module: " + str(e))
-
-    if not isinstance(pipeline_instance, Pipeline):
-        cli_utils.error(
-            f"The given source path `{source}` does not resolve to a pipeline "
-            "object."
-        )
+    pipeline_instance = _import_pipeline(source=source)
 
     parameters: Dict[str, Any] = {}
     if parameters_path:
@@ -176,24 +190,9 @@ def build_pipeline(
             "your source code root."
         )
 
-    try:
-        pipeline_instance = source_utils.load(source)
-    except ModuleNotFoundError as e:
-        source_root = source_utils.get_source_root()
-        cli_utils.error(
-            f"Unable to import module `{e.name}`. Make sure the source path is "
-            f"relative to your source root `{source_root}`."
-        )
-    except AttributeError as e:
-        cli_utils.error("Unable to load attribute from module: " + str(e))
-
-    if not isinstance(pipeline_instance, Pipeline):
-        cli_utils.error(
-            f"The given source path `{source}` does not resolve to a pipeline "
-            "object."
-        )
-
     with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
+        pipeline_instance = _import_pipeline(source=source)
+
         pipeline_instance = pipeline_instance.with_options(
             config_path=config_path
         )
@@ -277,36 +276,21 @@ def run_pipeline(
             "your source code root."
         )
 
-    try:
-        pipeline_instance = source_utils.load(source)
-    except ModuleNotFoundError as e:
-        source_root = source_utils.get_source_root()
-        cli_utils.error(
-            f"Unable to import module `{e.name}`. Make sure the source path is "
-            f"relative to your source root `{source_root}`."
-        )
-    except AttributeError as e:
-        cli_utils.error("Unable to load attribute from module: " + str(e))
-
-    if not isinstance(pipeline_instance, Pipeline):
-        cli_utils.error(
-            f"The given source path `{source}` does not resolve to a pipeline "
-            "object."
-        )
-
-    build: Union[str, PipelineBuildBase, None] = None
-    if build_path_or_id:
-        if uuid_utils.is_valid_uuid(build_path_or_id):
-            build = build_path_or_id
-        elif os.path.exists(build_path_or_id):
-            build = PipelineBuildBase.from_yaml(build_path_or_id)
-        else:
-            cli_utils.error(
-                f"The specified build {build_path_or_id} is not a valid UUID "
-                "or file path."
-            )
-
     with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
+        pipeline_instance = _import_pipeline(source=source)
+
+        build: Union[str, PipelineBuildBase, None] = None
+        if build_path_or_id:
+            if uuid_utils.is_valid_uuid(build_path_or_id):
+                build = build_path_or_id
+            elif os.path.exists(build_path_or_id):
+                build = PipelineBuildBase.from_yaml(build_path_or_id)
+            else:
+                cli_utils.error(
+                    f"The specified build {build_path_or_id} is not a valid UUID "
+                    "or file path."
+                )
+
         pipeline_instance = pipeline_instance.with_options(
             config_path=config_path,
             build=build,
@@ -369,24 +353,9 @@ def create_run_template(
             "init` at your source code root."
         )
 
-    try:
-        pipeline_instance = source_utils.load(source)
-    except ModuleNotFoundError as e:
-        source_root = source_utils.get_source_root()
-        cli_utils.error(
-            f"Unable to import module `{e.name}`. Make sure the source path is "
-            f"relative to your source root `{source_root}`."
-        )
-    except AttributeError as e:
-        cli_utils.error("Unable to load attribute from module: " + str(e))
-
-    if not isinstance(pipeline_instance, Pipeline):
-        cli_utils.error(
-            f"The given source path `{source}` does not resolve to a pipeline "
-            "object."
-        )
-
     with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
+        pipeline_instance = _import_pipeline(source=source)
+
         pipeline_instance = pipeline_instance.with_options(
             config_path=config_path
         )

diff --git a/src/zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py b/src/zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py
@@ -77,10 +77,12 @@ def prepare_step_run(self, info: "StepRunInfo") -> None:
             NeptuneExperimentTrackerSettings, self.get_settings(info)
         )
 
-        self.run_state.token = self.config.api_token
-        self.run_state.project = self.config.project
-        self.run_state.run_name = info.run_name
-        self.run_state.tags = list(settings.tags)
+        self.run_state.initialize(
+            project=self.config.project,
+            token=self.config.api_token,
+            run_name=info.run_name,
+            tags=list(settings.tags),
+        )
 
     def get_step_run_metadata(
         self, info: "StepRunInfo"
@@ -107,4 +109,4 @@ def cleanup_step_run(self, info: "StepRunInfo", step_failed: bool) -> None:
         """
         self.run_state.active_run.sync()
         self.run_state.active_run.stop()
-        self.run_state.reset_active_run()
+        self.run_state.reset()