Environment
zenml.environment
Environment implementation.
BaseEnvironmentComponent
Base Environment component class.
All Environment components must inherit from this class and provide a unique
value for the NAME
attribute.
Different code components can independently contribute with information to the global Environment by extending and instantiating this class:
from zenml.environment import BaseEnvironmentComponent
MY_ENV_NAME = "my_env"
class MyEnvironmentComponent(BaseEnvironmentComponent):
NAME = MY_ENV_NAME
def __init__(self, my_env_attr: str) -> None:
super().__init__()
self._my_env_attr = my_env_attr
@property
def my_env_attr(self) -> str:
return self._my_env_attr
my_env = MyEnvironmentComponent()
There are two ways to register and deregister a BaseEnvironmentComponent
instance with the global Environment:
- by explicitly calling its
activate
anddeactivate
methods:
my_env.activate()
# ... environment component is active
# and registered in the global Environment
my_env.deactivate()
# ... environment component is not active
- by using the instance as a context:
with my_env:
# ... environment component is active
# and registered in the global Environment
# ... environment component is not active
While active, environment components can be discovered and accessed from the global environment:
from foo.bar.my_env import MY_ENV_NAME
from zenml.environment import Environment
my_env = Environment.get_component(MY_ENV_NAME)
# this works too, but throws an error if the component is not active:
my_env = Environment[MY_ENV_NAME]
Attributes:
Name | Type | Description |
---|---|---|
NAME |
str |
a unique name for this component. This name will be used to
register this component in the global Environment and to
subsequently retrieve it by calling |
Source code in zenml/environment.py
class BaseEnvironmentComponent(metaclass=EnvironmentComponentMeta):
"""Base Environment component class.
All Environment components must inherit from this class and provide a unique
value for the `NAME` attribute.
Different code components can independently contribute with information to
the global Environment by extending and instantiating this class:
```python
from zenml.environment import BaseEnvironmentComponent
MY_ENV_NAME = "my_env"
class MyEnvironmentComponent(BaseEnvironmentComponent):
NAME = MY_ENV_NAME
def __init__(self, my_env_attr: str) -> None:
super().__init__()
self._my_env_attr = my_env_attr
@property
def my_env_attr(self) -> str:
return self._my_env_attr
my_env = MyEnvironmentComponent()
```
There are two ways to register and deregister a `BaseEnvironmentComponent`
instance with the global Environment:
1. by explicitly calling its `activate` and `deactivate` methods:
```python
my_env.activate()
# ... environment component is active
# and registered in the global Environment
my_env.deactivate()
# ... environment component is not active
```
2. by using the instance as a context:
```python
with my_env:
# ... environment component is active
# and registered in the global Environment
# ... environment component is not active
```
While active, environment components can be discovered and accessed from
the global environment:
```python
from foo.bar.my_env import MY_ENV_NAME
from zenml.environment import Environment
my_env = Environment.get_component(MY_ENV_NAME)
# this works too, but throws an error if the component is not active:
my_env = Environment[MY_ENV_NAME]
```
Attributes:
NAME: a unique name for this component. This name will be used to
register this component in the global Environment and to
subsequently retrieve it by calling `Environment().get_component`.
"""
NAME: str = _BASE_ENVIRONMENT_COMPONENT_NAME
def __init__(self) -> None:
"""Initialize an environment component."""
self._active = False
def activate(self) -> None:
"""Activate the environment component and register it in the global Environment.
Raises:
RuntimeError: if the component is already active.
"""
if self._active:
raise RuntimeError(
f"Environment component {self.NAME} is already active."
)
Environment().register_component(self)
self._active = True
def deactivate(self) -> None:
"""Deactivate the environment component and deregister it from the global Environment.
Raises:
RuntimeError: if the component is not active.
"""
if not self._active:
raise RuntimeError(
f"Environment component {self.NAME} is not active."
)
Environment().deregister_component(self)
self._active = False
@property
def active(self) -> bool:
"""Check if the environment component is currently active.
Returns:
`True` if the environment component is currently active, `False`
otherwise.
"""
return self._active
def __enter__(self) -> "BaseEnvironmentComponent":
"""Environment component context entry point.
Returns:
The BaseEnvironmentComponent instance.
"""
self.activate()
return self
def __exit__(self, *args: Any) -> None:
"""Environment component context exit point.
Args:
*args: the arguments passed to the context exit point.
"""
self.deactivate()
active: bool
property
readonly
Check if the environment component is currently active.
Returns:
Type | Description |
---|---|
bool |
|
__enter__(self)
special
Environment component context entry point.
Returns:
Type | Description |
---|---|
BaseEnvironmentComponent |
The BaseEnvironmentComponent instance. |
Source code in zenml/environment.py
def __enter__(self) -> "BaseEnvironmentComponent":
"""Environment component context entry point.
Returns:
The BaseEnvironmentComponent instance.
"""
self.activate()
return self
__exit__(self, *args)
special
Environment component context exit point.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
*args |
Any |
the arguments passed to the context exit point. |
() |
Source code in zenml/environment.py
def __exit__(self, *args: Any) -> None:
"""Environment component context exit point.
Args:
*args: the arguments passed to the context exit point.
"""
self.deactivate()
__init__(self)
special
Initialize an environment component.
Source code in zenml/environment.py
def __init__(self) -> None:
"""Initialize an environment component."""
self._active = False
activate(self)
Activate the environment component and register it in the global Environment.
Exceptions:
Type | Description |
---|---|
RuntimeError |
if the component is already active. |
Source code in zenml/environment.py
def activate(self) -> None:
"""Activate the environment component and register it in the global Environment.
Raises:
RuntimeError: if the component is already active.
"""
if self._active:
raise RuntimeError(
f"Environment component {self.NAME} is already active."
)
Environment().register_component(self)
self._active = True
deactivate(self)
Deactivate the environment component and deregister it from the global Environment.
Exceptions:
Type | Description |
---|---|
RuntimeError |
if the component is not active. |
Source code in zenml/environment.py
def deactivate(self) -> None:
"""Deactivate the environment component and deregister it from the global Environment.
Raises:
RuntimeError: if the component is not active.
"""
if not self._active:
raise RuntimeError(
f"Environment component {self.NAME} is not active."
)
Environment().deregister_component(self)
self._active = False
Environment
Provides environment information.
Individual environment components can be registered separately to extend
the global Environment object with additional information (see
BaseEnvironmentComponent
).
Source code in zenml/environment.py
class Environment(metaclass=SingletonMetaClass):
"""Provides environment information.
Individual environment components can be registered separately to extend
the global Environment object with additional information (see
`BaseEnvironmentComponent`).
"""
def __init__(self) -> None:
"""Initializes an Environment instance.
Note: Environment is a singleton class, which means this method will
only get called once. All following `Environment()` calls will return
the previously initialized instance.
"""
self._components: Dict[str, "BaseEnvironmentComponent"] = {}
@property
def step_is_running(self) -> bool:
"""Returns if a step is currently running.
Returns:
`True` if a step is currently running, `False` otherwise.
"""
from zenml.steps import STEP_ENVIRONMENT_NAME
# A step is considered to be running if there is an active step
# environment
return self.has_component(STEP_ENVIRONMENT_NAME)
@staticmethod
def get_system_info() -> Dict[str, Any]:
"""Information about the operating system.
Returns:
A dictionary containing information about the operating system.
"""
system = platform.system()
if system == "Windows":
release, version, csd, ptype = platform.win32_ver()
return {
"os": "windows",
"windows_version_release": release,
"windows_version": version,
"windows_version_service_pack": csd,
"windows_version_os_type": ptype,
}
if system == "Darwin":
return {"os": "mac", "mac_version": platform.mac_ver()[0]}
if system == "Linux":
return {
"os": "linux",
"linux_distro": distro.id(),
"linux_distro_like": distro.like(),
"linux_distro_version": distro.version(),
}
# We don't collect data for any other system.
return {"os": "unknown"}
@staticmethod
def python_version() -> str:
"""Returns the python version of the running interpreter.
Returns:
str: the python version
"""
return platform.python_version()
@staticmethod
def in_docker() -> bool:
"""If the current python process is running in a docker container.
Returns:
`True` if the current python process is running in a docker
container, `False` otherwise.
"""
# TODO [ENG-167]: Make this more reliable and add test.
try:
with open("/proc/1/cgroup", "rt") as ifh:
info = ifh.read()
return "docker" in info or "kubepod" in info
except (FileNotFoundError, Exception):
return False
@staticmethod
def in_google_colab() -> bool:
"""If the current Python process is running in a Google Colab.
Returns:
`True` if the current Python process is running in a Google Colab,
`False` otherwise.
"""
try:
import google.colab # noqa
return True
except ModuleNotFoundError:
return False
@staticmethod
def in_notebook() -> bool:
"""If the current Python process is running in a notebook.
Returns:
`True` if the current Python process is running in a notebook,
`False` otherwise.
"""
if find_spec("IPython") is not None:
from IPython import get_ipython # type: ignore
if get_ipython().__class__.__name__ in [
"TerminalInteractiveShell",
"ZMQInteractiveShell",
]:
return True
return False
@staticmethod
def in_paperspace_gradient() -> bool:
"""If the current Python process is running in Paperspace Gradient.
Returns:
`True` if the current Python process is running in Paperspace
Gradient, `False` otherwise.
"""
return "PAPERSPACE_NOTEBOOK_REPO_ID" in os.environ
def register_component(
self, component: "BaseEnvironmentComponent"
) -> "BaseEnvironmentComponent":
"""Registers an environment component.
Args:
component: a BaseEnvironmentComponent instance.
Returns:
The newly registered environment component, or the environment
component that was already registered under the given name.
"""
if component.NAME not in self._components:
self._components[component.NAME] = component
logger.debug(f"Registered environment component {component.NAME}")
return component
else:
logger.warning(
f"Ignoring attempt to overwrite an existing Environment "
f"component registered under the name {component.NAME}."
)
return self._components[component.NAME]
def deregister_component(
self, component: "BaseEnvironmentComponent"
) -> None:
"""Deregisters an environment component.
Args:
component: a BaseEnvironmentComponent instance.
"""
if self._components.get(component.NAME) is component:
del self._components[component.NAME]
logger.debug(f"Deregistered environment component {component.NAME}")
else:
logger.warning(
f"Ignoring attempt to deregister an inexistent Environment "
f"component with the name {component.NAME}."
)
def get_component(self, name: str) -> Optional["BaseEnvironmentComponent"]:
"""Get the environment component with a known name.
Args:
name: the environment component name.
Returns:
The environment component that is registered under the given name,
or None if no such component is registered.
"""
return self._components.get(name)
def get_components(
self,
) -> Dict[str, "BaseEnvironmentComponent"]:
"""Get all registered environment components.
Returns:
A dictionary containing all registered environment components.
"""
return self._components.copy()
def has_component(self, name: str) -> bool:
"""Check if the environment component with a known name is currently available.
Args:
name: the environment component name.
Returns:
`True` if an environment component with the given name is
currently registered for the given name, `False` otherwise.
"""
return name in self._components
def __getitem__(self, name: str) -> "BaseEnvironmentComponent":
"""Get the environment component with the given name.
Args:
name: the environment component name.
Returns:
`BaseEnvironmentComponent` instance that was registered for the
given name.
Raises:
KeyError: if no environment component is registered for the given name.
"""
if name in self._components:
return self._components[name]
else:
raise KeyError(
f"No environment component with name {name} is currently "
f"registered. This could happen for example if you're trying "
f"to access an environment component that is only available "
f"in the context of a step function, or, in the case of "
f"globally available environment components, if a relevant "
f"integration has not been activated yet."
)
@property
def step_environment(self) -> "StepEnvironment":
"""Get the current step environment component, if one is available.
This should only be called in the context of a step function.
Returns:
The `StepEnvironment` that describes the current step.
"""
from zenml.steps import STEP_ENVIRONMENT_NAME, StepEnvironment
return cast(StepEnvironment, self[STEP_ENVIRONMENT_NAME])
step_environment: StepEnvironment
property
readonly
Get the current step environment component, if one is available.
This should only be called in the context of a step function.
Returns:
Type | Description |
---|---|
StepEnvironment |
The |
step_is_running: bool
property
readonly
Returns if a step is currently running.
Returns:
Type | Description |
---|---|
bool |
|
__getitem__(self, name)
special
Get the environment component with the given name.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
name |
str |
the environment component name. |
required |
Returns:
Type | Description |
---|---|
BaseEnvironmentComponent |
|
Exceptions:
Type | Description |
---|---|
KeyError |
if no environment component is registered for the given name. |
Source code in zenml/environment.py
def __getitem__(self, name: str) -> "BaseEnvironmentComponent":
"""Get the environment component with the given name.
Args:
name: the environment component name.
Returns:
`BaseEnvironmentComponent` instance that was registered for the
given name.
Raises:
KeyError: if no environment component is registered for the given name.
"""
if name in self._components:
return self._components[name]
else:
raise KeyError(
f"No environment component with name {name} is currently "
f"registered. This could happen for example if you're trying "
f"to access an environment component that is only available "
f"in the context of a step function, or, in the case of "
f"globally available environment components, if a relevant "
f"integration has not been activated yet."
)
__init__(self)
special
Initializes an Environment instance.
Note: Environment is a singleton class, which means this method will
only get called once. All following Environment()
calls will return
the previously initialized instance.
Source code in zenml/environment.py
def __init__(self) -> None:
"""Initializes an Environment instance.
Note: Environment is a singleton class, which means this method will
only get called once. All following `Environment()` calls will return
the previously initialized instance.
"""
self._components: Dict[str, "BaseEnvironmentComponent"] = {}
deregister_component(self, component)
Deregisters an environment component.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
component |
BaseEnvironmentComponent |
a BaseEnvironmentComponent instance. |
required |
Source code in zenml/environment.py
def deregister_component(
self, component: "BaseEnvironmentComponent"
) -> None:
"""Deregisters an environment component.
Args:
component: a BaseEnvironmentComponent instance.
"""
if self._components.get(component.NAME) is component:
del self._components[component.NAME]
logger.debug(f"Deregistered environment component {component.NAME}")
else:
logger.warning(
f"Ignoring attempt to deregister an inexistent Environment "
f"component with the name {component.NAME}."
)
get_component(self, name)
Get the environment component with a known name.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
name |
str |
the environment component name. |
required |
Returns:
Type | Description |
---|---|
Optional[BaseEnvironmentComponent] |
The environment component that is registered under the given name, or None if no such component is registered. |
Source code in zenml/environment.py
def get_component(self, name: str) -> Optional["BaseEnvironmentComponent"]:
"""Get the environment component with a known name.
Args:
name: the environment component name.
Returns:
The environment component that is registered under the given name,
or None if no such component is registered.
"""
return self._components.get(name)
get_components(self)
Get all registered environment components.
Returns:
Type | Description |
---|---|
Dict[str, BaseEnvironmentComponent] |
A dictionary containing all registered environment components. |
Source code in zenml/environment.py
def get_components(
self,
) -> Dict[str, "BaseEnvironmentComponent"]:
"""Get all registered environment components.
Returns:
A dictionary containing all registered environment components.
"""
return self._components.copy()
get_system_info()
staticmethod
Information about the operating system.
Returns:
Type | Description |
---|---|
Dict[str, Any] |
A dictionary containing information about the operating system. |
Source code in zenml/environment.py
@staticmethod
def get_system_info() -> Dict[str, Any]:
"""Information about the operating system.
Returns:
A dictionary containing information about the operating system.
"""
system = platform.system()
if system == "Windows":
release, version, csd, ptype = platform.win32_ver()
return {
"os": "windows",
"windows_version_release": release,
"windows_version": version,
"windows_version_service_pack": csd,
"windows_version_os_type": ptype,
}
if system == "Darwin":
return {"os": "mac", "mac_version": platform.mac_ver()[0]}
if system == "Linux":
return {
"os": "linux",
"linux_distro": distro.id(),
"linux_distro_like": distro.like(),
"linux_distro_version": distro.version(),
}
# We don't collect data for any other system.
return {"os": "unknown"}
has_component(self, name)
Check if the environment component with a known name is currently available.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
name |
str |
the environment component name. |
required |
Returns:
Type | Description |
---|---|
bool |
|
Source code in zenml/environment.py
def has_component(self, name: str) -> bool:
"""Check if the environment component with a known name is currently available.
Args:
name: the environment component name.
Returns:
`True` if an environment component with the given name is
currently registered for the given name, `False` otherwise.
"""
return name in self._components
in_docker()
staticmethod
If the current python process is running in a docker container.
Returns:
Type | Description |
---|---|
bool |
|
Source code in zenml/environment.py
@staticmethod
def in_docker() -> bool:
"""If the current python process is running in a docker container.
Returns:
`True` if the current python process is running in a docker
container, `False` otherwise.
"""
# TODO [ENG-167]: Make this more reliable and add test.
try:
with open("/proc/1/cgroup", "rt") as ifh:
info = ifh.read()
return "docker" in info or "kubepod" in info
except (FileNotFoundError, Exception):
return False
in_google_colab()
staticmethod
If the current Python process is running in a Google Colab.
Returns:
Type | Description |
---|---|
bool |
|
Source code in zenml/environment.py
@staticmethod
def in_google_colab() -> bool:
"""If the current Python process is running in a Google Colab.
Returns:
`True` if the current Python process is running in a Google Colab,
`False` otherwise.
"""
try:
import google.colab # noqa
return True
except ModuleNotFoundError:
return False
in_notebook()
staticmethod
If the current Python process is running in a notebook.
Returns:
Type | Description |
---|---|
bool |
|
Source code in zenml/environment.py
@staticmethod
def in_notebook() -> bool:
"""If the current Python process is running in a notebook.
Returns:
`True` if the current Python process is running in a notebook,
`False` otherwise.
"""
if find_spec("IPython") is not None:
from IPython import get_ipython # type: ignore
if get_ipython().__class__.__name__ in [
"TerminalInteractiveShell",
"ZMQInteractiveShell",
]:
return True
return False
in_paperspace_gradient()
staticmethod
If the current Python process is running in Paperspace Gradient.
Returns:
Type | Description |
---|---|
bool |
|
Source code in zenml/environment.py
@staticmethod
def in_paperspace_gradient() -> bool:
"""If the current Python process is running in Paperspace Gradient.
Returns:
`True` if the current Python process is running in Paperspace
Gradient, `False` otherwise.
"""
return "PAPERSPACE_NOTEBOOK_REPO_ID" in os.environ
python_version()
staticmethod
Returns the python version of the running interpreter.
Returns:
Type | Description |
---|---|
str |
the python version |
Source code in zenml/environment.py
@staticmethod
def python_version() -> str:
"""Returns the python version of the running interpreter.
Returns:
str: the python version
"""
return platform.python_version()
register_component(self, component)
Registers an environment component.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
component |
BaseEnvironmentComponent |
a BaseEnvironmentComponent instance. |
required |
Returns:
Type | Description |
---|---|
BaseEnvironmentComponent |
The newly registered environment component, or the environment component that was already registered under the given name. |
Source code in zenml/environment.py
def register_component(
self, component: "BaseEnvironmentComponent"
) -> "BaseEnvironmentComponent":
"""Registers an environment component.
Args:
component: a BaseEnvironmentComponent instance.
Returns:
The newly registered environment component, or the environment
component that was already registered under the given name.
"""
if component.NAME not in self._components:
self._components[component.NAME] = component
logger.debug(f"Registered environment component {component.NAME}")
return component
else:
logger.warning(
f"Ignoring attempt to overwrite an existing Environment "
f"component registered under the name {component.NAME}."
)
return self._components[component.NAME]
EnvironmentComponentMeta (type)
Metaclass responsible for registering different EnvironmentComponent instances in the global Environment.
Source code in zenml/environment.py
class EnvironmentComponentMeta(type):
"""Metaclass responsible for registering different EnvironmentComponent instances in the global Environment."""
def __new__(
mcs, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]
) -> "EnvironmentComponentMeta":
"""Hook into creation of an BaseEnvironmentComponent class.
Args:
name: the name of the class being created.
bases: the base classes of the class being created.
dct: the dictionary of attributes of the class being created.
Returns:
The newly created class.
"""
cls = cast(
Type["BaseEnvironmentComponent"],
super().__new__(mcs, name, bases, dct),
)
if name != "BaseEnvironmentComponent":
assert (
cls.NAME and cls.NAME != _BASE_ENVIRONMENT_COMPONENT_NAME
), "You should specify a unique NAME when creating an EnvironmentComponent !"
return cls
__new__(mcs, name, bases, dct)
special
staticmethod
Hook into creation of an BaseEnvironmentComponent class.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
name |
str |
the name of the class being created. |
required |
bases |
Tuple[Type[Any], ...] |
the base classes of the class being created. |
required |
dct |
Dict[str, Any] |
the dictionary of attributes of the class being created. |
required |
Returns:
Type | Description |
---|---|
EnvironmentComponentMeta |
The newly created class. |
Source code in zenml/environment.py
def __new__(
mcs, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]
) -> "EnvironmentComponentMeta":
"""Hook into creation of an BaseEnvironmentComponent class.
Args:
name: the name of the class being created.
bases: the base classes of the class being created.
dct: the dictionary of attributes of the class being created.
Returns:
The newly created class.
"""
cls = cast(
Type["BaseEnvironmentComponent"],
super().__new__(mcs, name, bases, dct),
)
if name != "BaseEnvironmentComponent":
assert (
cls.NAME and cls.NAME != _BASE_ENVIRONMENT_COMPONENT_NAME
), "You should specify a unique NAME when creating an EnvironmentComponent !"
return cls
get_environment()
Returns a string representing the execution environment of the pipeline.
Currently, one of docker
, paperspace
, 'colab', or native
.
Returns:
Type | Description |
---|---|
str |
the execution environment |
Source code in zenml/environment.py
def get_environment() -> str:
"""Returns a string representing the execution environment of the pipeline.
Currently, one of `docker`, `paperspace`, 'colab', or `native`.
Returns:
str: the execution environment
"""
if Environment.in_docker():
return "docker"
elif Environment.in_google_colab():
return "colab"
elif Environment.in_paperspace_gradient():
return "paperspace"
elif Environment.in_notebook():
return "notebook"
else:
return "native"
get_system_details()
Returns OS, python and ZenML information.
Returns:
Type | Description |
---|---|
str |
OS, python and ZenML information |
Source code in zenml/environment.py
def get_system_details() -> str:
"""Returns OS, python and ZenML information.
Returns:
str: OS, python and ZenML information
"""
from zenml.integrations.registry import integration_registry
info = {
"ZenML version": __version__,
"Install path": Path(__file__).resolve().parent,
"Python version": Environment.python_version(),
"Platform information": Environment.get_system_info(),
"Environment": get_environment(),
"Integrations": integration_registry.get_installed_integrations(),
}
return "\n".join(
"{:>10} {}".format(k + ":", str(v).replace("\n", " "))
for k, v in info.items()
)