Environment
zenml.environment
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."""
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."""
self.deactivate()
active: bool
property
readonly
Check if the environment component is currently active.
__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.
Source code in zenml/environment.py
def __exit__(self, *args: Any) -> None:
"""Environment component 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."""
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."""
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."""
return platform.python_version()
@staticmethod
def in_docker() -> bool:
"""If the current python process is running in a docker container."""
# 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."""
return "COLAB_GPU" in os.environ
@staticmethod
def in_notebook() -> bool:
"""If the current python process is running in a notebook."""
try:
from IPython import get_ipython # type: ignore
if get_ipython() is None:
# IPython is installed but not running from a notebook
return False
else:
return True
except ImportError:
# We do not even have IPython installed
return False
@staticmethod
def in_paperspace_gradient() -> bool:
"""If the current python process is running in Paperspace Gradient."""
return "PAPERSPACE_NOTEBOOK_REPO_ID" in os.environ
def register_component(
self, component: "BaseEnvironmentComponent"
) -> "BaseEnvironmentComponent":
"""Registers an environment component.
Args:
name: the environment component name.
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."""
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. Make sure you're calling this in the context of a "
f"step function that has all relevant integrations enabled."
)
@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.
__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 |
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. Make sure you're calling this in the context of a "
f"step function that has all relevant integrations enabled."
)
__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.
Source code in zenml/environment.py
def get_components(
self,
) -> Dict[str, "BaseEnvironmentComponent"]:
"""Get all registered environment components."""
return self._components.copy()
get_system_info()
staticmethod
Information about the operating system.
Source code in zenml/environment.py
@staticmethod
def get_system_info() -> Dict[str, Any]:
"""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.
Source code in zenml/environment.py
@staticmethod
def in_docker() -> bool:
"""If the current python process is running in a docker container."""
# 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.
Source code in zenml/environment.py
@staticmethod
def in_google_colab() -> bool:
"""If the current python process is running in a Google Colab."""
return "COLAB_GPU" in os.environ
in_notebook()
staticmethod
If the current python process is running in a notebook.
Source code in zenml/environment.py
@staticmethod
def in_notebook() -> bool:
"""If the current python process is running in a notebook."""
try:
from IPython import get_ipython # type: ignore
if get_ipython() is None:
# IPython is installed but not running from a notebook
return False
else:
return True
except ImportError:
# We do not even have IPython installed
return False
in_paperspace_gradient()
staticmethod
If the current python process is running in Paperspace Gradient.
Source code in zenml/environment.py
@staticmethod
def in_paperspace_gradient() -> bool:
"""If the current python process is running in Paperspace Gradient."""
return "PAPERSPACE_NOTEBOOK_REPO_ID" in os.environ
python_version()
staticmethod
Returns the python version of the running interpreter.
Source code in zenml/environment.py
@staticmethod
def python_version() -> str:
"""Returns the python version of the running interpreter."""
return platform.python_version()
register_component(self, component)
Registers an environment component.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
name |
the environment component name. |
required | |
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:
name: the environment component name.
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."""
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.
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."""
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