xpk 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl
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.
- xpk/__init__.py +15 -0
- xpk/commands/__init__.py +15 -0
- xpk/commands/batch.py +109 -0
- xpk/commands/cluster.py +784 -0
- xpk/commands/cluster_gcluster.py +185 -0
- xpk/commands/info.py +245 -0
- xpk/commands/inspector.py +363 -0
- xpk/commands/job.py +197 -0
- xpk/commands/kind.py +253 -0
- xpk/commands/shell.py +120 -0
- xpk/commands/version.py +39 -0
- xpk/commands/workload.py +692 -0
- xpk/core/__init__.py +15 -0
- xpk/core/blueprint/__init__.py +15 -0
- xpk/core/blueprint/blueprint_definitions.py +61 -0
- xpk/core/blueprint/blueprint_generator.py +652 -0
- xpk/core/cluster_private.py +197 -0
- xpk/core/commands.py +352 -0
- xpk/core/core.py +2824 -0
- xpk/core/docker_manager.py +308 -0
- xpk/core/gcluster_manager.py +158 -0
- xpk/core/kjob.py +205 -0
- xpk/core/kueue.py +352 -0
- xpk/core/nap.py +349 -0
- xpk/core/pathways.py +298 -0
- xpk/core/ray.py +222 -0
- xpk/core/system_characteristics.py +1395 -0
- xpk/core/workload.py +133 -0
- xpk/core/workload_decorators/__init__.py +15 -0
- xpk/core/workload_decorators/rdma_decorator.py +109 -0
- xpk/core/workload_decorators/tcpxo_decorator.py +157 -0
- xpk/main.py +73 -0
- xpk/parser/__init__.py +15 -0
- xpk/parser/batch.py +184 -0
- xpk/parser/cluster.py +621 -0
- xpk/parser/common.py +71 -0
- xpk/parser/core.py +109 -0
- xpk/parser/info.py +63 -0
- xpk/parser/inspector.py +65 -0
- xpk/parser/job.py +126 -0
- xpk/parser/kind.py +94 -0
- xpk/parser/shell.py +50 -0
- xpk/parser/validators.py +39 -0
- xpk/parser/version.py +23 -0
- xpk/parser/workload.py +684 -0
- xpk/utils/__init__.py +15 -0
- xpk/utils/console.py +55 -0
- xpk/utils/file.py +82 -0
- xpk/utils/network.py +168 -0
- xpk/utils/objects.py +85 -0
- xpk/utils/yaml.py +30 -0
- {xpk-0.5.0.dist-info → xpk-0.6.0.dist-info}/METADATA +301 -28
- xpk-0.6.0.dist-info/RECORD +57 -0
- {xpk-0.5.0.dist-info → xpk-0.6.0.dist-info}/WHEEL +1 -1
- xpk-0.6.0.dist-info/entry_points.txt +2 -0
- xpk-0.5.0.dist-info/RECORD +0 -7
- xpk-0.5.0.dist-info/entry_points.txt +0 -2
- xpk.py +0 -7282
- {xpk-0.5.0.dist-info → xpk-0.6.0.dist-info}/LICENSE +0 -0
- {xpk-0.5.0.dist-info → xpk-0.6.0.dist-info}/top_level.txt +0 -0
xpk/utils/console.py
ADDED
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright 2024 Google LLC
|
|
3
|
+
|
|
4
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
you may not use this file except in compliance with the License.
|
|
6
|
+
You may obtain a copy of the License at
|
|
7
|
+
|
|
8
|
+
https://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
|
|
10
|
+
Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
See the License for the specific language governing permissions and
|
|
14
|
+
limitations under the License.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
import sys
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def xpk_print(*args, **kwargs):
|
|
21
|
+
"""Helper function to print a prefix before function provided args.
|
|
22
|
+
|
|
23
|
+
Args:
|
|
24
|
+
*args: user provided print args.
|
|
25
|
+
**kwargs: user provided print args.
|
|
26
|
+
"""
|
|
27
|
+
sys.stdout.write('[XPK] ')
|
|
28
|
+
print(*args, **kwargs)
|
|
29
|
+
sys.stdout.flush()
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def xpk_exit(error_code):
|
|
33
|
+
"""Helper function to exit xpk with an associated error code.
|
|
34
|
+
|
|
35
|
+
Args:
|
|
36
|
+
error_code: If the code provided is zero, then no issues occurred.
|
|
37
|
+
"""
|
|
38
|
+
if error_code == 0:
|
|
39
|
+
xpk_print('Exiting XPK cleanly')
|
|
40
|
+
sys.exit(0)
|
|
41
|
+
else:
|
|
42
|
+
xpk_print(f'XPK failed, error code {error_code}')
|
|
43
|
+
sys.exit(error_code)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def get_user_input(input_msg):
|
|
47
|
+
"""Function to get the user input for a prompt.
|
|
48
|
+
|
|
49
|
+
Args:
|
|
50
|
+
input_msg: message to be displayed by the prompt.
|
|
51
|
+
Returns:
|
|
52
|
+
True if user enter y or yes at the prompt, False otherwise.
|
|
53
|
+
"""
|
|
54
|
+
user_input = input(input_msg)
|
|
55
|
+
return user_input in ('y', 'yes')
|
xpk/utils/file.py
ADDED
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright 2024 Google LLC
|
|
3
|
+
|
|
4
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
you may not use this file except in compliance with the License.
|
|
6
|
+
You may obtain a copy of the License at
|
|
7
|
+
|
|
8
|
+
https://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
|
|
10
|
+
Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
See the License for the specific language governing permissions and
|
|
14
|
+
limitations under the License.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
import tempfile
|
|
18
|
+
import os
|
|
19
|
+
from .console import xpk_print
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def make_tmp_files(per_command_name):
|
|
23
|
+
"""Make temporary files for each command.
|
|
24
|
+
|
|
25
|
+
Args:
|
|
26
|
+
per_command_name: list of command names.
|
|
27
|
+
|
|
28
|
+
Returns:
|
|
29
|
+
A list of temporary files for each command.
|
|
30
|
+
"""
|
|
31
|
+
# Supports removal of spaces from command names before converting to file name.
|
|
32
|
+
return [
|
|
33
|
+
tempfile.NamedTemporaryFile(
|
|
34
|
+
delete=False, prefix=command.replace(' ', '-') + '-'
|
|
35
|
+
)
|
|
36
|
+
for command in per_command_name
|
|
37
|
+
]
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def write_tmp_file(payload):
|
|
41
|
+
"""Writes `payload` to a temporary file.
|
|
42
|
+
|
|
43
|
+
Args:
|
|
44
|
+
payload: The string to be written to the file.
|
|
45
|
+
|
|
46
|
+
Returns:
|
|
47
|
+
A file object that was written to.
|
|
48
|
+
"""
|
|
49
|
+
with tempfile.NamedTemporaryFile(delete=False) as tmp:
|
|
50
|
+
with open(file=tmp.name, mode='w', encoding='utf=8') as f:
|
|
51
|
+
f.write(payload)
|
|
52
|
+
f.flush()
|
|
53
|
+
return tmp
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def append_tmp_file(payload, file):
|
|
57
|
+
"""Appends `payload` to an already created file.
|
|
58
|
+
|
|
59
|
+
Use `write_temporary_file` to create a file.
|
|
60
|
+
|
|
61
|
+
Args:
|
|
62
|
+
payload: The string to be written to the file.
|
|
63
|
+
file: The file to append to.
|
|
64
|
+
|
|
65
|
+
Returns:
|
|
66
|
+
A file object that was written to.
|
|
67
|
+
"""
|
|
68
|
+
with open(file=file.name, mode='a', encoding='utf=8') as f:
|
|
69
|
+
f.write(payload)
|
|
70
|
+
f.flush()
|
|
71
|
+
return file
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def ensure_directory_exists(directory_path):
|
|
75
|
+
"""Checks if a directory exists and creates it if it doesn't.
|
|
76
|
+
|
|
77
|
+
Args:
|
|
78
|
+
directory_path: The path to the directory.
|
|
79
|
+
"""
|
|
80
|
+
if not os.path.exists(directory_path):
|
|
81
|
+
os.makedirs(directory_path)
|
|
82
|
+
xpk_print(f"Directory '{directory_path}' created successfully.")
|
xpk/utils/network.py
ADDED
|
@@ -0,0 +1,168 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright 2024 Google LLC
|
|
3
|
+
|
|
4
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
you may not use this file except in compliance with the License.
|
|
6
|
+
You may obtain a copy of the License at
|
|
7
|
+
|
|
8
|
+
https://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
|
|
10
|
+
Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
See the License for the specific language governing permissions and
|
|
14
|
+
limitations under the License.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
import ipaddress
|
|
18
|
+
import socket
|
|
19
|
+
import requests
|
|
20
|
+
from .console import xpk_print
|
|
21
|
+
|
|
22
|
+
# Retrives machine's external IP address
|
|
23
|
+
ip_resolver_url = "http://api.ipify.org"
|
|
24
|
+
all_IPs_cidr = "0.0.0.0/0"
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def get_current_machine_ip(external_ip=True):
|
|
28
|
+
"""
|
|
29
|
+
Gets the IP address of the current machine.
|
|
30
|
+
|
|
31
|
+
Args:
|
|
32
|
+
external: If True (default), retrieves the external IP address.
|
|
33
|
+
If False, retrieves the internal IP address.
|
|
34
|
+
|
|
35
|
+
Returns:
|
|
36
|
+
The IP address as a string.
|
|
37
|
+
"""
|
|
38
|
+
|
|
39
|
+
try:
|
|
40
|
+
if external_ip:
|
|
41
|
+
# Get external IP address
|
|
42
|
+
response = requests.get(ip_resolver_url, timeout=30)
|
|
43
|
+
return 0, response.text
|
|
44
|
+
else:
|
|
45
|
+
# Get internal IP address
|
|
46
|
+
hostname = socket.gethostname()
|
|
47
|
+
return 0, socket.gethostbyname(hostname)
|
|
48
|
+
except (requests.exceptions.RequestException, socket.gaierror) as e:
|
|
49
|
+
xpk_print(f"Error getting IP address: {e}")
|
|
50
|
+
return 1, None
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def is_ip_in_any_network(ip_address, cidrs):
|
|
54
|
+
"""
|
|
55
|
+
Checks if an IP address is within any of the provided CIDR ranges.
|
|
56
|
+
|
|
57
|
+
Args:
|
|
58
|
+
ip_address: The IP address to check (as a string).
|
|
59
|
+
cidrs: A list of CIDR strings.
|
|
60
|
+
|
|
61
|
+
Returns:
|
|
62
|
+
True if the IP address is found in any of the CIDRs, False otherwise.
|
|
63
|
+
"""
|
|
64
|
+
|
|
65
|
+
try:
|
|
66
|
+
if not are_cidrs_valid(cidrs):
|
|
67
|
+
return False
|
|
68
|
+
|
|
69
|
+
if cidrs is None:
|
|
70
|
+
return False
|
|
71
|
+
|
|
72
|
+
current_ip = ipaddress.ip_address(ip_address)
|
|
73
|
+
for cidr in cidrs:
|
|
74
|
+
network = ipaddress.ip_network(cidr)
|
|
75
|
+
if current_ip in network:
|
|
76
|
+
return True
|
|
77
|
+
except ValueError as e:
|
|
78
|
+
xpk_print(f"Error: {e}")
|
|
79
|
+
return False
|
|
80
|
+
return False
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def is_current_machine_in_any_network(cidrs, external_ip=True):
|
|
84
|
+
"""
|
|
85
|
+
Checks if the current machine's IP address is within any of the provided CIDR ranges.
|
|
86
|
+
|
|
87
|
+
Args:
|
|
88
|
+
cidrs: A list of CIDR strings.
|
|
89
|
+
external_ip: If True (default), checks the external IP. If False, checks the internal IP.
|
|
90
|
+
|
|
91
|
+
Returns:
|
|
92
|
+
True if the IP address is found in any of the CIDRs, False otherwise.
|
|
93
|
+
"""
|
|
94
|
+
if not are_cidrs_valid(cidrs):
|
|
95
|
+
return 1, False
|
|
96
|
+
|
|
97
|
+
if cidrs is None:
|
|
98
|
+
return 0, False
|
|
99
|
+
|
|
100
|
+
return_code, ip_address = get_current_machine_ip(external_ip)
|
|
101
|
+
if return_code > 0:
|
|
102
|
+
return return_code, False
|
|
103
|
+
else:
|
|
104
|
+
return return_code, is_ip_in_any_network(ip_address, cidrs)
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def add_current_machine_to_networks(cidrs, external_ip=True):
|
|
108
|
+
"""
|
|
109
|
+
Adds the current machine's IP address to the list of CIDRs if it's not already present.
|
|
110
|
+
|
|
111
|
+
Args:
|
|
112
|
+
cidrs: A list of CIDR strings.
|
|
113
|
+
external_ip: If True, uses the external IP. If False (default), uses the internal IP.
|
|
114
|
+
|
|
115
|
+
Returns:
|
|
116
|
+
The updated list of CIDRs with the current machine's IP added (if necessary).
|
|
117
|
+
"""
|
|
118
|
+
if not are_cidrs_valid(cidrs):
|
|
119
|
+
return 1, None
|
|
120
|
+
|
|
121
|
+
return_code, ip_address = get_current_machine_ip(external_ip)
|
|
122
|
+
if return_code > 0:
|
|
123
|
+
return return_code, None
|
|
124
|
+
|
|
125
|
+
if not is_ip_in_any_network(ip_address, cidrs):
|
|
126
|
+
cidrs.append(f"{ip_address}/32")
|
|
127
|
+
|
|
128
|
+
return 0, cidrs
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
def is_cidr_valid(cidr):
|
|
132
|
+
"""
|
|
133
|
+
Validates a CIDR string.
|
|
134
|
+
|
|
135
|
+
Args:
|
|
136
|
+
cidr: The CIDR string to validate.
|
|
137
|
+
|
|
138
|
+
Returns:
|
|
139
|
+
True if the CIDR string is valid, False otherwise.
|
|
140
|
+
"""
|
|
141
|
+
try:
|
|
142
|
+
ipaddress.ip_network(cidr)
|
|
143
|
+
return True
|
|
144
|
+
except ValueError:
|
|
145
|
+
return False
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
def are_cidrs_valid(cidrs):
|
|
149
|
+
"""
|
|
150
|
+
Validates a list of CIDR strings.
|
|
151
|
+
|
|
152
|
+
Args:
|
|
153
|
+
cidrs: A list of CIDR strings to validate.
|
|
154
|
+
|
|
155
|
+
Returns:
|
|
156
|
+
True if all CIDR strings in the list are valid, False otherwise.
|
|
157
|
+
"""
|
|
158
|
+
if cidrs is None:
|
|
159
|
+
return True
|
|
160
|
+
|
|
161
|
+
all_cidrs_are_valid = True
|
|
162
|
+
|
|
163
|
+
for cidr in cidrs:
|
|
164
|
+
if not is_cidr_valid(cidr):
|
|
165
|
+
all_cidrs_are_valid = False
|
|
166
|
+
xpk_print(f"Error: the string '{cidr}' is not a valid CIDR.")
|
|
167
|
+
|
|
168
|
+
return all_cidrs_are_valid
|
xpk/utils/objects.py
ADDED
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright 2024 Google LLC
|
|
3
|
+
|
|
4
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
you may not use this file except in compliance with the License.
|
|
6
|
+
You may obtain a copy of the License at
|
|
7
|
+
|
|
8
|
+
https://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
|
|
10
|
+
Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
See the License for the specific language governing permissions and
|
|
14
|
+
limitations under the License.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
import hashlib
|
|
18
|
+
from .console import xpk_print
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def chunks(lst: list, n: int):
|
|
22
|
+
"""Return a list of n-sized chunks from lst.
|
|
23
|
+
|
|
24
|
+
Args:
|
|
25
|
+
lst: input list to get chunks from.
|
|
26
|
+
n: size of each chunk.
|
|
27
|
+
|
|
28
|
+
Returns:
|
|
29
|
+
List of n-sized chunks for lst.
|
|
30
|
+
"""
|
|
31
|
+
return [lst[i : i + n] for i in range(0, len(lst), n)]
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def get_value_from_map(key: str, map_to_search: dict) -> tuple[int, str | None]:
|
|
35
|
+
"""Helper function to get value from a map if the key exists.
|
|
36
|
+
|
|
37
|
+
Args:
|
|
38
|
+
key: The key to look for in the map
|
|
39
|
+
map_to_search: The map to look in for the value
|
|
40
|
+
|
|
41
|
+
Returns:
|
|
42
|
+
Tuple of int, str where
|
|
43
|
+
int is the return code
|
|
44
|
+
str is the value if found
|
|
45
|
+
"""
|
|
46
|
+
value = map_to_search.get(key)
|
|
47
|
+
if value:
|
|
48
|
+
return 0, value
|
|
49
|
+
else:
|
|
50
|
+
xpk_print(
|
|
51
|
+
f'Unable to find key: {key} in map: {map_to_search}.'
|
|
52
|
+
f'The map has the following keys: {map_to_search.keys()}'
|
|
53
|
+
)
|
|
54
|
+
return 1, value
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def is_text_true(text: str) -> bool:
|
|
58
|
+
return text.strip().lower() == 'true'
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def hash_string(input_string: str, length: int | None = None):
|
|
62
|
+
"""
|
|
63
|
+
Generates a hash of a string using characters 0-9 and a-z.
|
|
64
|
+
|
|
65
|
+
Args:
|
|
66
|
+
input_string: The string to hash.
|
|
67
|
+
length: The desired length of the hash (optional).
|
|
68
|
+
If not provided, or less than 0, the full SHA256 hash is returned.
|
|
69
|
+
|
|
70
|
+
Returns:
|
|
71
|
+
A hash string of the specified length or the full SHA256 hash,
|
|
72
|
+
using characters 0-9 and a-z.
|
|
73
|
+
"""
|
|
74
|
+
hash_value = int(hashlib.sha256(input_string.encode()).hexdigest(), 16)
|
|
75
|
+
|
|
76
|
+
if length is None or length < 0:
|
|
77
|
+
length = 64 # Use the full length of the SHA256 hash
|
|
78
|
+
|
|
79
|
+
charset = '0123456789abcdefghijklmnopqrstuvwxyz'
|
|
80
|
+
result = ''
|
|
81
|
+
while hash_value > 0 and len(result) < length:
|
|
82
|
+
hash_value, index = divmod(hash_value, 36) # Get quotient and remainder
|
|
83
|
+
result = charset[index] + result
|
|
84
|
+
|
|
85
|
+
return result.zfill(length).lower() # Pad with zeros if necessary
|
xpk/utils/yaml.py
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright 2024 Google LLC
|
|
3
|
+
|
|
4
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
you may not use this file except in compliance with the License.
|
|
6
|
+
You may obtain a copy of the License at
|
|
7
|
+
|
|
8
|
+
https://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
|
|
10
|
+
Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
See the License for the specific language governing permissions and
|
|
14
|
+
limitations under the License.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
import yaml
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class literal_string(str):
|
|
21
|
+
pass
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def literal_string_representer(
|
|
25
|
+
dumper: yaml.Dumper, data
|
|
26
|
+
) -> yaml.nodes.ScalarNode:
|
|
27
|
+
return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
yaml.add_representer(literal_string, literal_string_representer)
|