Skip to content

Commit

Permalink
fix lightning orchestrator (#3010)
Browse files Browse the repository at this point in the history
* fix lightning orchestrator

* set review

* fix zenml install
safoinme authored Sep 24, 2024
1 parent 522c9ce commit b05cfeb
Showing 4 changed files with 30 additions and 24 deletions.
2 changes: 1 addition & 1 deletion src/zenml/integrations/lightning/__init__.py
Original file line number Diff line number Diff line change
@@ -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]]:
Original file line number Diff line number Diff line change
@@ -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."""
Original file line number Diff line number Diff line change
@@ -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(
Original file line number Diff line number Diff line change
@@ -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(

0 comments on commit b05cfeb

Please sign in to comment.