d3-tb-agent-comms 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- d3_tb_agent_comms-0.1.0/PKG-INFO +11 -0
- d3_tb_agent_comms-0.1.0/README.md +1 -0
- d3_tb_agent_comms-0.1.0/pyproject.toml +21 -0
- d3_tb_agent_comms-0.1.0/src/d3_tb_agent_comms/__init__.py +4 -0
- d3_tb_agent_comms-0.1.0/src/d3_tb_agent_comms/custom_types.py +27 -0
- d3_tb_agent_comms-0.1.0/src/d3_tb_agent_comms/main.py +184 -0
- d3_tb_agent_comms-0.1.0/src/d3_tb_agent_comms/py.typed +0 -0
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: d3-tb-agent-comms
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Lib that holds a Class for easily interacting with TB-Agent Endpoints
|
|
5
|
+
Author: Oscar Mitchell
|
|
6
|
+
Author-email: Oscar Mitchell <oscar.mitchell@disguise.one>
|
|
7
|
+
Requires-Dist: requests>=2.32.5
|
|
8
|
+
Requires-Python: >=3.12
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
|
|
11
|
+
This is a readme
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
This is a readme
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "d3_tb_agent_comms"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Lib that holds a Class for easily interacting with TB-Agent Endpoints"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
authors = [
|
|
7
|
+
{ name = "Oscar Mitchell", email = "oscar.mitchell@disguise.one" }
|
|
8
|
+
]
|
|
9
|
+
requires-python = ">=3.12"
|
|
10
|
+
dependencies = [
|
|
11
|
+
"requests>=2.32.5",
|
|
12
|
+
]
|
|
13
|
+
|
|
14
|
+
[build-system]
|
|
15
|
+
requires = ["uv_build>=0.8.15,<0.9.0"]
|
|
16
|
+
build-backend = "uv_build"
|
|
17
|
+
|
|
18
|
+
[dependency-groups]
|
|
19
|
+
dev = [
|
|
20
|
+
"ruff>=0.12.12",
|
|
21
|
+
]
|
|
@@ -0,0 +1,4 @@
|
|
|
1
|
+
from main import AgentHandler as AgentHandler
|
|
2
|
+
from custom_types import D3SystemInfo as D3SystemInfo, MachineHealthInfo as MachineHealthInfo, AgentHandlerException as AgentHandlerException
|
|
3
|
+
|
|
4
|
+
__all__ = ["AgentHandler", "D3SystemInfo", "MachineHealthInfo", "AgentHandlerException"]
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
from dataclasses import dataclass
|
|
2
|
+
from typing import Final
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
@dataclass(slots=True)
|
|
6
|
+
class D3SystemInfo(Final):
|
|
7
|
+
major: int
|
|
8
|
+
minor: int
|
|
9
|
+
patch: int
|
|
10
|
+
build: int
|
|
11
|
+
branch: str
|
|
12
|
+
api_port: int
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
@dataclass(slots=True)
|
|
16
|
+
class MachineHealthInfo(Final):
|
|
17
|
+
revision_number: int
|
|
18
|
+
project_folder_check: bool
|
|
19
|
+
project_folder_shared: bool
|
|
20
|
+
is_service_running: bool
|
|
21
|
+
is_buddy_running: bool
|
|
22
|
+
can_see_storage_server: bool
|
|
23
|
+
can_see_teamcity: bool
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class AgentHandlerException(Exception):
|
|
27
|
+
pass
|
|
@@ -0,0 +1,184 @@
|
|
|
1
|
+
# ruff: noqa: N802 - caps in function names
|
|
2
|
+
# In this case the caps serve a purpose of decerning verb vs noun
|
|
3
|
+
# As in we are not, "getting a ping", we are using a "GET" method called "ping"
|
|
4
|
+
from __future__ import annotations
|
|
5
|
+
|
|
6
|
+
import inspect
|
|
7
|
+
import json
|
|
8
|
+
from typing import TYPE_CHECKING, Final
|
|
9
|
+
|
|
10
|
+
import requests
|
|
11
|
+
|
|
12
|
+
from src.tb_agent_comms.types import AgentHandlerException, D3SystemInfo, MachineHealthInfo
|
|
13
|
+
|
|
14
|
+
if TYPE_CHECKING:
|
|
15
|
+
from logging import Logger
|
|
16
|
+
|
|
17
|
+
SUCCESS: Final[int] = 200
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class AgentHandler:
|
|
21
|
+
"""Handles Agent API requests. Automatic retry and JSON encode/decode included"""
|
|
22
|
+
|
|
23
|
+
def __init__(self, ip: str, port: int, *, logger: Logger) -> None:
|
|
24
|
+
"""CTOR for AgentHandler
|
|
25
|
+
|
|
26
|
+
:param ip: the ip of the target TB-Agent
|
|
27
|
+
:param port: the port of the target TB-Agent
|
|
28
|
+
:param logger: a logger instance to use for logging
|
|
29
|
+
"""
|
|
30
|
+
self.ip = ip
|
|
31
|
+
self.port = port
|
|
32
|
+
self.logger = logger
|
|
33
|
+
|
|
34
|
+
self.version = -1
|
|
35
|
+
|
|
36
|
+
def _request_handler(self, data: dict | None = None, timeout: int = 5) -> dict:
|
|
37
|
+
"""Helper function for handling all request
|
|
38
|
+
|
|
39
|
+
:param data: the data to send to the endpoint
|
|
40
|
+
:param timeout: the timeout for the endpoint
|
|
41
|
+
|
|
42
|
+
:return: the decoded JSON returned from the endpoint
|
|
43
|
+
|
|
44
|
+
:raises:
|
|
45
|
+
"""
|
|
46
|
+
calling_func_name = inspect.stack()[1].function
|
|
47
|
+
split_name = calling_func_name.split("_", 1)
|
|
48
|
+
method, endpoint = split_name[0], split_name[1].replace("_", "-")
|
|
49
|
+
url = f"http://{self.ip}:{self.port}/{endpoint}"
|
|
50
|
+
headers = {"accept": "application/json"}
|
|
51
|
+
|
|
52
|
+
attempts = 0
|
|
53
|
+
for _ in range(5):
|
|
54
|
+
attempts += 1
|
|
55
|
+
self.logger.info(f"Attempting #{attempts} call to {endpoint}")
|
|
56
|
+
|
|
57
|
+
try:
|
|
58
|
+
if method == "GET":
|
|
59
|
+
response = requests.get(url=url, headers=headers, timeout=timeout)
|
|
60
|
+
elif method == "POST":
|
|
61
|
+
response = requests.post(url=url, data=data, headers=headers, timeout=timeout)
|
|
62
|
+
else:
|
|
63
|
+
raise NameError(
|
|
64
|
+
f"AgentHandler: Name of calling function did not fit expected format: {calling_func_name}\nThis is a bug with the library" # noqa: E501
|
|
65
|
+
)
|
|
66
|
+
except requests.exceptions.RequestException as e:
|
|
67
|
+
self.logger.error(f"!!!!! Failed to get a response from TB-Agent due to: {e}")
|
|
68
|
+
continue
|
|
69
|
+
|
|
70
|
+
if not response.ok:
|
|
71
|
+
self.logger.warning(
|
|
72
|
+
f"!!! The response from the TB-Agent API was not ok\nCode: {response.status_code}\nResponse text: {response.text}" # noqa: E501
|
|
73
|
+
)
|
|
74
|
+
continue
|
|
75
|
+
|
|
76
|
+
try:
|
|
77
|
+
return json.loads(response.text)
|
|
78
|
+
except ValueError as e:
|
|
79
|
+
self.logger.error(f"!!!!! Failed to decode response from TB-Agent: {e}\nRaw text: {response.text}")
|
|
80
|
+
continue
|
|
81
|
+
|
|
82
|
+
raise AgentHandlerException(f"!!!!! Maximum number of attempts to get a response from /{endpoint} reached")
|
|
83
|
+
|
|
84
|
+
def RAW_ping(self, timeout: int = 2) -> bool:
|
|
85
|
+
"""Raw-er version of `GET_ping` for Agent updater checks
|
|
86
|
+
|
|
87
|
+
:return: whether a code 200 came back from the request
|
|
88
|
+
"""
|
|
89
|
+
url = f"http://{self.ip}:{self.port}/ping"
|
|
90
|
+
headers = {"accept": "application/json"}
|
|
91
|
+
try:
|
|
92
|
+
response = requests.get(url, headers=headers, timeout=timeout)
|
|
93
|
+
except requests.exceptions.RequestException:
|
|
94
|
+
return False
|
|
95
|
+
else:
|
|
96
|
+
return response.status_code == SUCCESS
|
|
97
|
+
|
|
98
|
+
def GET_ping(self) -> None:
|
|
99
|
+
return self._request_handler()["return"]
|
|
100
|
+
|
|
101
|
+
def GET_d3installed(self) -> bool:
|
|
102
|
+
return self._request_handler()["d3installed"]
|
|
103
|
+
|
|
104
|
+
def GET_d3system(self) -> D3SystemInfo:
|
|
105
|
+
body = self._request_handler()
|
|
106
|
+
return D3SystemInfo(
|
|
107
|
+
major=body["major"],
|
|
108
|
+
minor=body["minor"],
|
|
109
|
+
patch=body["patch"],
|
|
110
|
+
build=body["build"],
|
|
111
|
+
branch=body["branch"],
|
|
112
|
+
api_port=body["apiPort"],
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
def GET_agent_version(self) -> int:
|
|
116
|
+
return self._request_handler()["version"]
|
|
117
|
+
|
|
118
|
+
def GET_running_project(self) -> str:
|
|
119
|
+
return self._request_handler()["projectName"]
|
|
120
|
+
|
|
121
|
+
def GET_summarise_machine_health(self) -> MachineHealthInfo:
|
|
122
|
+
body = self._request_handler()
|
|
123
|
+
return MachineHealthInfo(
|
|
124
|
+
revision_number=body["revision number"],
|
|
125
|
+
project_folder_check=body["project folder check"],
|
|
126
|
+
project_folder_shared=body["shared project folder check"],
|
|
127
|
+
is_service_running=body["is d3 service running"],
|
|
128
|
+
is_buddy_running=body["is d3 buddy running"],
|
|
129
|
+
can_see_storage_server=body["can machine see storage server"],
|
|
130
|
+
can_see_teamcity=body["can machine see teamcity"],
|
|
131
|
+
)
|
|
132
|
+
|
|
133
|
+
def GET_gpu_vendor(self) -> str:
|
|
134
|
+
return self._request_handler()["gpu_output"]
|
|
135
|
+
|
|
136
|
+
def GET_current_os(self) -> str:
|
|
137
|
+
return self._request_handler()["OS"]
|
|
138
|
+
|
|
139
|
+
def GET_licence_codes(self) -> str:
|
|
140
|
+
return self._request_handler()["licenceCodes"]
|
|
141
|
+
|
|
142
|
+
def POST_restart_local_machine(self, timeout: int = 1) -> str:
|
|
143
|
+
data = {"timeout_duration": timeout}
|
|
144
|
+
return self._request_handler(data)["msg"]
|
|
145
|
+
|
|
146
|
+
def POST_download_d3(
|
|
147
|
+
self, revision_number: int, destination: str, timeout: int = 0, download_time: int = 360, host: str = ""
|
|
148
|
+
) -> str:
|
|
149
|
+
data = {"rev": revision_number, "destination": destination}
|
|
150
|
+
if timeout:
|
|
151
|
+
data["timeout"] = timeout
|
|
152
|
+
if host:
|
|
153
|
+
data["host"] = host
|
|
154
|
+
return self._request_handler(data, timeout=download_time)["msg"]
|
|
155
|
+
|
|
156
|
+
def POST_install_d3(self, path: str) -> str:
|
|
157
|
+
data = {"path": path}
|
|
158
|
+
return self._request_handler(data, timeout=15)["taskName"]
|
|
159
|
+
|
|
160
|
+
def POST_check_pid(self, pid: int) -> bool:
|
|
161
|
+
data = {"pid": pid}
|
|
162
|
+
return self._request_handler(data)["runnning"]
|
|
163
|
+
|
|
164
|
+
def POST_update_agent(self, github_key: str) -> None:
|
|
165
|
+
data = {"githubKey": github_key}
|
|
166
|
+
self._request_handler(data, timeout=25)
|
|
167
|
+
|
|
168
|
+
def POST_start_d3_project(self, project_folder: str, project_name: str) -> int:
|
|
169
|
+
data = {"projectFolder": project_folder, "projectName": project_name}
|
|
170
|
+
return self._request_handler(data)["msg"]
|
|
171
|
+
|
|
172
|
+
def POST_terminate_d3(self) -> str:
|
|
173
|
+
return self._request_handler({})["msg"]
|
|
174
|
+
|
|
175
|
+
def POST_delete_d3_installer(self, installer_path: str) -> str:
|
|
176
|
+
data = {"path": installer_path}
|
|
177
|
+
return self._request_handler(data)["msg"]
|
|
178
|
+
|
|
179
|
+
def POST_check_task_running(self, task_name: str) -> bool:
|
|
180
|
+
data = {"taskName": task_name}
|
|
181
|
+
return self._request_handler(data, timeout=15)["isRunning"]
|
|
182
|
+
|
|
183
|
+
def POST_restart_d3service(self) -> str:
|
|
184
|
+
return self._request_handler()["msg"]
|
|
File without changes
|