Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 49 additions & 55 deletions ankaios_sdk/_components/complete_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,9 @@
Classes
-------

- CompleteState:
- :class:`CompleteState`:
Represents the complete state of the Ankaios cluster.
- AgentAttributes:
- :class:`AgentAttributes`:
Represents the attributes of an agent in the Ankaios cluster.

Usage
Expand Down Expand Up @@ -62,13 +62,6 @@
.. code-block:: python

workload_states = complete_state.get_workload_states()

- Get the tags and status of a specific agent:
.. code-block:: python

agent_attributes = AgentAttributes()
tags = agent_attributes.tags
status = agent_attributes.status
"""

__all__ = ["CompleteState", "AgentAttributes"]
Expand Down Expand Up @@ -99,13 +92,14 @@ def __init__(
"""
Initializes a CompleteState instance with the provided data.

Args:
manifest (Manifest): The manifest to initialize the complete state.
configs (dict): The configurations to set in the complete state.
workloads (list[Workload]): The workloads to set
in the complete state.
_proto (_ank_base.CompleteState): The proto message to initialize
the complete state.
:param manifest: The manifest to initialize the complete state.
:type manifest: Manifest
:param configs: The configurations to set in the complete state.
:type configs: dict
:param workloads: The workloads to set in the complete state.
:type workloads: list[Workload]
:param _proto: The proto message to initialize the complete state.
:type _proto: _ank_base.CompleteState
"""
self._complete_state = _ank_base.CompleteState()
self._set_api_version(SUPPORTED_API_VERSION)
Expand Down Expand Up @@ -136,39 +130,38 @@ def __str__(self) -> str:
"""
Returns the string representation of the complete state.

Returns:
str: The string representation of the complete state.
:returns: The string representation of the complete state.
:rtype: str
"""
return str(self._to_proto())

def _set_api_version(self, version: str) -> None:
"""
Sets the API version for the complete state.

Args:
version (str): The API version to set.
:param version: The API version to set.
:type version: str
"""
self._complete_state.desiredState.apiVersion = version

def get_api_version(self) -> str:
"""
Gets the API version of the complete state.

Returns:
str: The API version of the complete state.
:returns: The API version of the complete state.
:rtype: str
"""
return str(self._complete_state.desiredState.apiVersion)

def get_workload(self, workload_name: str) -> Workload:
"""
Gets a workload from the complete state by its name.

Args:
workload_name (str): The name of the workload to retrieve.
:param workload_name: The name of the workload to retrieve.
:type workload_name: str

Returns:
Workload: The workload with the specified name,
or None if not found.
:returns: The workload with the specified name, or None if not found.
:rtype: Workload
"""
if (
workload_name
Expand All @@ -188,8 +181,8 @@ def get_workloads(self) -> list[Workload]:
"""
Gets a list of workloads from the complete state.

Returns:
list[Workload]: A list of workloads in the complete state.
:returns: A list of workloads in the complete state.
:rtype: list[Workload]
"""
workloads = []
for (
Expand All @@ -205,8 +198,8 @@ def get_workload_states(self) -> WorkloadStateCollection:
"""
Gets the workload states.

Returns:
WorkloadStateCollection: The collection of workload states.
:returns: The collection of workload states.
:rtype: WorkloadStateCollection
"""
workload_state_collection = WorkloadStateCollection()
workload_state_collection._from_proto(
Expand All @@ -218,9 +211,8 @@ def get_agents(self) -> dict[str, "AgentAttributes"]:
"""
Gets the connected agents and their attributes.

Returns:
dict[str, AgentAttributes]: A dict with the agents and
their attributes.
:returns: A dict with the agents and their attributes.
:rtype: dict[str, AgentAttributes]
"""
agents = {}
for name, attributes in self._complete_state.agents.agents.items():
Expand All @@ -231,9 +223,10 @@ def set_agent_tags(self, agent_name: str, tags: dict[str, str]) -> None:
"""
Sets the tags for a specific agent.

Args:
agent_name (str): The name of the agent.
tags (dict[str, str]): The tags to set for the agent.
:param agent_name: The name of the agent.
:type agent_name: str
:param tags: The tags to set for the agent.
:type tags: dict[str, str]
"""
agent_attributes = self._complete_state.agents.agents[agent_name]
agent_attributes.tags.tags.clear()
Expand All @@ -244,8 +237,8 @@ def set_configs(self, configs: dict) -> None:
"""
Sets the configurations in the complete state.

Args:
configs (dict): The configurations to set in the complete state.
:param configs: The configurations to set in the complete state.
:type configs: dict
"""
self._complete_state.desiredState.configs.configs.clear()
for key, value in configs.items():
Expand All @@ -257,8 +250,8 @@ def get_configs(self) -> dict:
"""
Gets the configurations from the complete state.

Returns:
dict: The configurations from the complete state
:returns: The configurations from the complete state.
:rtype: dict
"""

def _from_config_item(
Expand Down Expand Up @@ -287,10 +280,10 @@ def _from_config_item(

def to_dict(self) -> dict:
"""
Returns the CompleteState as a dictionary
Returns the CompleteState as a dictionary.

Returns:
dict: The CompleteState as a dictionary.
:returns: The CompleteState as a dictionary.
:rtype: dict
"""
data = {
"desired_state": {
Expand Down Expand Up @@ -323,16 +316,18 @@ def _to_proto(self) -> _ank_base.CompleteState:
"""
Returns the CompleteState as a proto message.

Returns:
_ank_base.CompleteState: The protobuf message representing
the complete state.
:returns: The protobuf message representing the complete state.
:rtype: _ank_base.CompleteState
"""
return self._complete_state


class AgentAttributes:
"""
A class to represent the attributes of an agent.

:var dict[str, str] tags:
The agent's tags
"""

def __init__(self) -> None:
Expand All @@ -347,21 +342,21 @@ def status(self) -> dict[str, str]:
"""
Gets the status of the agent.

Returns:
dict[str, str]: The status of the agent.
:returns: The status of the agent.
:rtype: dict[str, str]
"""
status = {
"cpu_usage": int(self._status.cpu_usage.cpu_usage),
"free_memory": self._status.free_memory.free_memory,
}
return status

def to_dict(self) -> dict:
def to_dict(self) -> dict[str, Union[dict, str]]:
"""
Returns the AgentAttributes as a dictionary.

Returns:
dict: The AgentAttributes as a dictionary.
:returns: The AgentAttributes as a dictionary.
:rtype: dict[str, Union[dict, str]]
"""
data = {
"tags": self.tags,
Expand All @@ -374,9 +369,8 @@ def _from_proto(proto: _ank_base.AgentAttributes) -> "AgentAttributes":
"""
Initializes the AgentAttributes instance from a proto message.

Args:
proto (_ank_base.AgentAttributes): The proto message to
initialize the agent attributes.
:param proto: The proto message to initialize the agent attributes.
:type proto: _ank_base.AgentAttributes
"""
obj = AgentAttributes()
obj.tags = dict(proto.tags.tags)
Expand Down
Loading