确定代码是否在 Databricks 或 IDE (PyCharm) 上运行

Question

我正在构建一个 Python 包，数据科学家可以使用该包来管理他们的 MLOps 生命周期。现在，这个包可以在本地（通常在 PyCharm 上）或 Databricks 上使用。

我希望包的某些功能取决于它的运行位置，即，如果它被 Databricks 笔记本调用，我希望它做一些不同的事情，如果它在本地运行，我希望它做一些完全不同的事情。

有什么方法可以确定它是从哪里调用的吗？

我有点怀疑我们是否可以使用类似下面的内容来检查您的代码是否在笔记本上运行，或者因为这将是一个将存储在您的 Databricks 环境中的包，
我如何检查代码是否在笔记本上运行在IPython笔记本中执行？

Answer 1

我发现有效的解决方法是检查数据块特定的环境变量。

import os
def is_running_in_databricks() -> bool:
    return "DATABRICKS_RUNTIME_VERSION" in os.environ