From 48a98409aaefdad576cc5edd19c576d6f1982648 Mon Sep 17 00:00:00 2001
From: Hamza Tahir <hamza@zenml.io>
Date: Tue, 6 Aug 2024 22:01:34 +0200
Subject: [PATCH] Added lightning studio check

---
 src/zenml/enums.py       |  1 +
 src/zenml/environment.py | 15 +++++++++++++++
 2 files changed, 16 insertions(+)

diff --git a/src/zenml/enums.py b/src/zenml/enums.py
index 2df9e6af793..e5fba4e640e 100644
--- a/src/zenml/enums.py
+++ b/src/zenml/enums.py
@@ -303,6 +303,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
index ac80bbf9be9..0d870c84954 100644
--- 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":