fix lightning orchestrator

src/zenml/integrations/lightning/__init__.py

@@ -28,7 +28,7 @@ class LightningIntegration(Integration):
     """Definition of Lightning Integration for ZenML."""
 
     NAME = LIGHTNING
-    REQUIREMENTS = ["lightning-sdk"]
+    REQUIREMENTS = ["lightning-sdk>=0.1.17"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

src/zenml/integrations/lightning/flavors/lightning_orchestrator_flavor.py

@@ -85,6 +85,15 @@ def is_synchronous(self) -> bool:
         """
         return self.synchronous
 
+    @property
+    def is_schedulable(self) -> bool:
+        """Whether the orchestrator is schedulable or not.
+
+        Returns:
+            Whether the orchestrator is schedulable or not.
+        """
+        return False
+
 
 class LightningOrchestratorFlavor(BaseOrchestratorFlavor):
     """Lightning orchestrator flavor."""

src/zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

@@ -103,20 +103,29 @@ def _set_lightning_env_vars(
 
         Args:
             deployment: The pipeline deployment to prepare or run.
+
+        Raises:
+            ValueError: If the user id and api key or username and organization
         """
         settings = cast(
             LightningOrchestratorSettings, self.get_settings(deployment)
         )
-        if settings.user_id:
-            os.environ["LIGHTNING_USER_ID"] = settings.user_id
-        if settings.api_key:
-            os.environ["LIGHTNING_API_KEY"] = settings.api_key
+        if not settings.user_id or not settings.api_key:
+            raise ValueError(
+                "Lightning orchestrator requires `user_id` and `api_key` both to be set in the settings."
+            )
+        os.environ["LIGHTNING_USER_ID"] = settings.user_id
+        os.environ["LIGHTNING_API_KEY"] = settings.api_key
         if settings.username:
             os.environ["LIGHTNING_USERNAME"] = settings.username
+        elif settings.organization:
+            os.environ["LIGHTNING_ORG"] = settings.organization
+        else:
+            raise ValueError(
+                "Lightning orchestrator requires either `username` or `organization` to be set in the settings."
+            )
         if settings.teamspace:
             os.environ["LIGHTNING_TEAMSPACE"] = settings.teamspace
-        if settings.organization:
-            os.environ["LIGHTNING_ORG"] = settings.organization
 
     @property
     def config(self) -> LightningOrchestratorConfig:
@@ -267,9 +276,7 @@ def prepare_or_run_pipeline(
         ) as code_file:
             code_archive.write_archive(code_file)
             code_path = code_file.name
-
         filename = f"{orchestrator_run_name}.tar.gz"
-
         # Construct the env variables for the pipeline
         env_vars = environment.copy()
         orchestrator_run_id = str(uuid4())
@@ -392,9 +399,7 @@ def _construct_lightning_steps(
                 f"Installing requirements: {pipeline_requirements_to_string}"
             )
             studio.run(f"uv pip install {pipeline_requirements_to_string}")
-            studio.run(
-                "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-            )
+            studio.run("pip install zenml -y")
 
             for custom_command in settings.custom_commands or []:
                 studio.run(
@@ -488,9 +493,7 @@ def _upload_and_run_pipeline(
             )
             studio.run("pip install uv")
             studio.run(f"uv pip install {requirements}")
-            studio.run(
-                "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-            )
+            studio.run("pip install zenml -y")
             # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
             for command in settings.custom_commands or []:
                 output = studio.run(
@@ -563,9 +566,7 @@ def _run_step_in_new_studio(
         )
         studio.run("pip install uv")
         studio.run(f"uv pip install {details['requirements']}")
-        studio.run(
-            "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-        )
+        studio.run("pip install zenml -y")
         # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
         for command in custom_commands or []:
             output = studio.run(

src/zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint.py

@@ -166,9 +166,7 @@ def main() -> None:
         f"uv pip install {pipeline_requirements_to_string}"
     )
     logger.info(output)
-    output = main_studio.run(
-        "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-    )
+    output = main_studio.run("pip install zenml -y")
     logger.info(output)
 
     for command in pipeline_settings.custom_commands or []:
@@ -250,9 +248,7 @@ def run_step_on_lightning_studio(step_name: str) -> None:
                     f"uv pip install {step_requirements_to_string}"
                 )
                 logger.info(output)
-                output = studio.run(
-                    "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
-                )
+                output = studio.run("pip install zenml -y")
                 logger.info(output)
                 for command in step_settings.custom_commands or []:
                     output = studio.run(