def _get_agent_metadata(self, node_ip: str) -> AgentMetadata:
node = self._nodes_by_ip.get(node_ip)
if not node:
return AgentMetadata(state=AgentState.ORPHANED)
return AgentMetadata(
agent_id=node.metadata.name,
allocated_resources=allocated_node_resources(
self._pods_by_ip[node_ip]),
batch_task_count=self._count_batch_tasks(node_ip),
state=(AgentState.RUNNING
if self._pods_by_ip[node_ip] else AgentState.IDLE),
task_count=len(self._pods_by_ip[node_ip]),
total_resources=total_node_resources(node),
)
def _make_metadata(
rg_id,
instance_id,
agent_state=AgentState.RUNNING,
is_stale=False,
weight=1,
tasks=5,
batch_tasks=0,
):
return ClusterNodeMetadata(
AgentMetadata(
agent_id='foo',
batch_task_count=batch_tasks,
state=agent_state,
task_count=tasks,
),
InstanceMetadata(
group_id=rg_id,
hostname='host1',
instance_id=instance_id,
ip_address='1.2.3.4',
is_stale=is_stale,
market='market-1',
state='running',
uptime=1000,
weight=weight,
),
)
def _get_agent_metadata(self, instance_ip: str) -> AgentMetadata:
agent_dict = self._agents_by_ip.get(instance_ip)
if not agent_dict:
return AgentMetadata(state=AgentState.ORPHANED)
allocated_resources = allocated_agent_resources(agent_dict)
return AgentMetadata(
agent_id=agent_dict['id'],
allocated_resources=allocated_agent_resources(agent_dict),
batch_task_count=self._task_count_per_agent[
agent_dict['id']]['batch_tasks'],
state=(AgentState.RUNNING
if any(allocated_resources) else AgentState.IDLE),
task_count=self._task_count_per_agent[agent_dict['id']]
['all_tasks'],
total_resources=total_agent_resources(agent_dict),
)
def test_get_agent_metadata(mock_cluster_connector):
instance = list(mock_cluster_connector.simulator.aws_clusters[0].instances.values())[0]
mesos_resources = ClustermanResources(
get_market_resources(TEST_MARKET).cpus,
get_market_resources(TEST_MARKET).mem * 1000,
get_market_resources(TEST_MARKET).disk * 1000,
)
assert mock_cluster_connector.get_agent_metadata(instance.ip_address) == AgentMetadata(
agent_id=mock.ANY,
state=AgentState.IDLE,
total_resources=mesos_resources,
)
def _get_agent_metadata(self, instance_ip: str) -> AgentMetadata:
for c in self.simulator.aws_clusters:
for i in c.instances.values():
if instance_ip == i.ip_address:
return AgentMetadata(
agent_id=str(uuid.uuid4()),
state=(
AgentState.ORPHANED
if self.simulator.current_time < i.join_time
else AgentState.IDLE
),
total_resources=ClustermanResources(
cpus=i.resources.cpus,
mem=i.resources.mem * 1000,
disk=(i.resources.disk or staticconf.read_int('ebs_volume_size', 0)) * 1000,
gpus=(i.resources.gpus),
)
)
# if we don't know the given IP then it's orphaned
return AgentMetadata(state=AgentState.ORPHANED)
def test_get_agent_metadata_unknown(mock_cluster_connector):
assert mock_cluster_connector.get_agent_metadata('1.2.3.4') == AgentMetadata(
state=AgentState.ORPHANED,
)