Added lightning studio check (#2910)

* Added lightning studio check * Fixed ruff errors * Fixed PR comments * Auto-update of NLP template --------- Co-authored-by: GitHub Actions <[email protected]>
zenml-io · Aug 7, 2024 · a981e50 · a981e50
1 parent 919b50c
commit a981e50
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/src/zenml/enums.py b/src/zenml/enums.py
@@ -302,6 +302,7 @@ class EnvironmentType(StrEnum):
     NOTEBOOK = "notebook"
     PAPERSPACE = "paperspace"
     WSL = "wsl"
+    LIGHTNING_AI_STUDIO = "lightning_ai_studio"
 
 
 class ModelStages(StrEnum):

diff --git a/src/zenml/environment.py b/src/zenml/environment.py
@@ -70,6 +70,8 @@ def get_environment() -> str:
         return EnvironmentType.BITBUCKET_CI
     elif Environment.in_ci():
         return EnvironmentType.GENERIC_CI
+    elif Environment.in_lightning_ai_studio():
+        return EnvironmentType.LIGHTNING_AI_STUDIO
     elif Environment.in_docker():
         return EnvironmentType.DOCKER
     elif Environment.in_container():
@@ -339,6 +341,19 @@ def in_wsl() -> bool:
         """
         return "microsoft-standard" in platform.uname().release
 
+    @staticmethod
+    def in_lightning_ai_studio() -> bool:
+        """If the current Python process is running in Lightning.ai studios.
+
+        Returns:
+            `True` if the current Python process is running in Lightning.ai studios,
+            `False` otherwise.
+        """
+        return (
+            "LIGHTNING_CLOUD_URL" in os.environ
+            and "LIGHTNING_CLOUDSPACE_HOST" in os.environ
+        )
+
     def register_component(
         self, component: "BaseEnvironmentComponent"
     ) -> "BaseEnvironmentComponent":