Skip to content

Commit 49f1ece

Browse files
authored
Split core.py into smaller files (#363)
1 parent ce7f3e2 commit 49f1ece

39 files changed

+3752
-3497
lines changed

.gitignore

+1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# editor and IDE paraphernalia
22
.idea/
3+
.vscode/
34

45
*__pycache__*
56
tmp/

pytype-conf.cfg

-4
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,8 @@ exclude =
1111
src/xpk/commands
1212
src/xpk/core/tests
1313
src/xpk/core/__init__.py
14-
src/xpk/core/app_profile.py
15-
src/xpk/core/blueprint.py
1614
src/xpk/core/cluster_private.py
1715
src/xpk/core/commands.py
18-
src/xpk/core/core.py
19-
src/xpk/core/job_template.py
2016
src/xpk/core/kjob.py
2117
src/xpk/core/kueue.py
2218
src/xpk/core/nap.py

src/xpk/commands/batch.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,12 @@
1616

1717
from argparse import Namespace
1818

19+
from ..core.commands import run_command_for_value
20+
from ..core.gcloud_context import add_zone_and_project
21+
from ..core.kjob import AppProfileDefaults
1922
from ..core.kueue import LOCAL_QUEUE_NAME
2023
from ..utils.console import xpk_exit, xpk_print
2124
from .common import set_cluster_command
22-
from ..core.core import add_zone_and_project
23-
from ..core.kjob import AppProfileDefaults
24-
from ..core.commands import run_command_for_value
2525
from .kind import set_local_cluster_command
2626

2727

src/xpk/commands/cluster.py

+22-24
Original file line numberDiff line numberDiff line change
@@ -14,58 +14,56 @@
1414
limitations under the License.
1515
"""
1616

17-
from ..core.commands import run_command_for_value, run_command_with_updates
18-
from ..core.core import (
19-
VERTEX_TENSORBOARD_FEATURE_FLAG,
20-
add_zone_and_project,
21-
create_cluster_configmaps,
22-
create_cluster_network_config,
23-
create_vertex_tensorboard,
24-
delete_cluster_subnets,
17+
from tabulate import tabulate
18+
19+
from ..core.capacity import H100_DEVICE_TYPE
20+
from ..core.cluster import (
2521
get_all_clusters_programmatic,
2622
get_cluster_credentials,
27-
get_gke_control_plane_version,
28-
get_gke_node_pool_version,
29-
get_gke_server_config,
30-
h100_device_type,
3123
install_nccl_on_cluster,
32-
run_gke_node_pool_create_command,
3324
set_jobset_on_cluster,
34-
set_up_cluster_network_for_gpu,
3525
setup_k8s_env,
3626
update_cluster_with_gcsfuse_driver_if_necessary,
3727
update_cluster_with_workload_identity_if_necessary,
38-
zone_to_region,
39-
get_user_input,
4028
)
4129
from ..core.cluster_private import authorize_private_cluster_access_if_necessary
42-
from ..core.kjob import (
43-
verify_kjob_installed,
44-
prepare_kjob,
45-
apply_kjob_crds,
30+
from ..core.commands import run_command_for_value, run_command_with_updates
31+
from ..core.config import VERTEX_TENSORBOARD_FEATURE_FLAG
32+
from ..core.gcloud_context import (
33+
add_zone_and_project,
34+
get_gke_control_plane_version,
35+
get_gke_server_config,
36+
zone_to_region,
4637
)
38+
from ..core.kjob import apply_kjob_crds, prepare_kjob, verify_kjob_installed
4739
from ..core.kueue import (
4840
cluster_preheat_yml,
4941
install_kueue_crs,
5042
install_kueue_on_cluster,
5143
wait_for_kueue_available,
5244
)
5345
from ..core.nap import enable_autoprovisioning_on_cluster
46+
from ..core.network import (
47+
create_cluster_network_config,
48+
delete_cluster_subnets,
49+
set_up_cluster_network_for_gpu,
50+
)
51+
from ..core.nodepool import get_gke_node_pool_version, run_gke_node_pool_create_command
5452
from ..core.ray import install_ray_cluster
53+
from ..core.resources import create_cluster_configmaps
5554
from ..core.storage import install_storage_crd
5655
from ..core.system_characteristics import (
5756
AcceleratorType,
5857
AcceleratorTypeToAcceleratorCharacteristics,
5958
SystemCharacteristics,
6059
get_system_characteristics,
6160
)
61+
from ..core.vertex import create_vertex_tensorboard
6262
from ..core.workload import get_workload_list
63+
from ..utils.console import get_user_input, xpk_exit, xpk_print
6364
from ..utils.file import write_tmp_file
64-
from ..utils.console import xpk_exit, xpk_print
6565
from . import cluster_gcluster
6666

67-
from tabulate import tabulate
68-
6967

7068
def cluster_create(args) -> None:
7169
"""Function around cluster creation.
@@ -148,7 +146,7 @@ def cluster_create(args) -> None:
148146
if set_up_cluster_network_code != 0:
149147
xpk_exit(set_up_cluster_network_code)
150148

151-
if system.device_type == h100_device_type:
149+
if system.device_type == H100_DEVICE_TYPE:
152150
xpk_print('Creating Network Config for cluster')
153151
create_cluster_network_config_code = create_cluster_network_config(args)
154152
if create_cluster_network_config_code != 0:

src/xpk/commands/cluster_gcluster.py

+12-4
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,24 @@
1414
limitations under the License.
1515
"""
1616

17-
from ..core.blueprint.blueprint_generator import BlueprintGenerator, BlueprintGeneratorOutput, supported_device_types, a3mega_device_type, a3ultra_device_type
17+
import os
18+
19+
from ..core.blueprint.blueprint_generator import (
20+
BlueprintGenerator,
21+
BlueprintGeneratorOutput,
22+
a3mega_device_type,
23+
a3ultra_device_type,
24+
supported_device_types,
25+
)
26+
from ..core.capacity import get_capacity_type
1827
from ..core.docker_manager import DockerManager
28+
from ..core.gcloud_context import zone_to_region
1929
from ..core.gcluster_manager import GclusterManager
20-
from ..core.core import zone_to_region, get_capacity_type
2130
from ..utils.console import xpk_exit, xpk_print
22-
from ..utils.network import all_IPs_cidr
2331
from ..utils.file import ensure_directory_exists
32+
from ..utils.network import all_IPs_cidr
2433
from ..utils.objects import hash_string
2534
from .common import set_cluster_command
26-
import os
2735

2836
blueprints_path = os.path.abspath('xpkclusters/blueprints')
2937
gcluster_working_dir = os.path.abspath('xpkclusters/gcluster-out')

src/xpk/commands/common.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
"""
1616

1717
from ..core.commands import run_command_with_updates_retry
18-
from ..core.core import zone_to_region
18+
from ..core.gcloud_context import zone_to_region
1919
from ..utils.console import xpk_print
2020

2121

src/xpk/commands/info.py

+8-10
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,17 @@
1414
limitations under the License.
1515
"""
1616

17-
from ..utils.console import xpk_exit, xpk_print
18-
from ..core.kueue import verify_kueuectl
19-
from .common import set_cluster_command
20-
from ..core.commands import (
21-
run_command_for_value,
22-
)
23-
from ..core.core import (
24-
add_zone_and_project,
25-
)
2617
import json
27-
from tabulate import tabulate
2818
from argparse import Namespace
2919

20+
from tabulate import tabulate
21+
22+
from ..core.commands import run_command_for_value
23+
from ..core.gcloud_context import add_zone_and_project
24+
from ..core.kueue import verify_kueuectl
25+
from ..utils.console import xpk_exit, xpk_print
26+
from .common import set_cluster_command
27+
3028
table_fmt = 'plain'
3129

3230

src/xpk/commands/inspector.py

+4-8
Original file line numberDiff line numberDiff line change
@@ -14,17 +14,13 @@
1414
limitations under the License.
1515
"""
1616

17+
from ..core.cluster import get_cluster_credentials
1718
from ..core.commands import run_command_for_value
18-
from ..core.core import (
19-
CLUSTER_METADATA_CONFIGMAP,
20-
CLUSTER_RESOURCES_CONFIGMAP,
21-
add_zone_and_project,
22-
get_cluster_credentials,
23-
zone_to_region,
24-
)
19+
from ..core.gcloud_context import add_zone_and_project, zone_to_region
2520
from ..core.kueue import CLUSTER_QUEUE_NAME, LOCAL_QUEUE_NAME
26-
from ..utils.file import append_tmp_file, write_tmp_file
21+
from ..core.resources import CLUSTER_METADATA_CONFIGMAP, CLUSTER_RESOURCES_CONFIGMAP
2722
from ..utils.console import xpk_exit, xpk_print
23+
from ..utils.file import append_tmp_file, write_tmp_file
2824
from .workload import get_workload_list
2925

3026

src/xpk/commands/job.py

+9-7
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,18 @@
1414
limitations under the License.
1515
"""
1616

17-
from .common import set_cluster_command
18-
from .kind import set_local_cluster_command
19-
from ..core.commands import run_command_for_value, run_command_with_updates
20-
from ..utils.console import xpk_exit, xpk_print
21-
from ..core.kjob import AppProfileDefaults
22-
from ..core.core import add_zone_and_project
23-
from ruamel.yaml import YAML
2417
import re
2518
import sys
2619

20+
from ruamel.yaml import YAML
21+
22+
from ..core.commands import run_command_for_value, run_command_with_updates
23+
from ..core.gcloud_context import add_zone_and_project
24+
from ..core.kjob import AppProfileDefaults
25+
from ..utils.console import xpk_exit, xpk_print
26+
from .common import set_cluster_command
27+
from .kind import set_local_cluster_command
28+
2729

2830
def job_info(args):
2931
"""Run commands obtaining information about a job given by name.

src/xpk/commands/kind.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
run_command_for_value,
1919
run_command_with_updates,
2020
)
21-
from ..core.core import (
21+
from ..core.cluster import (
2222
set_jobset_on_cluster,
2323
)
2424
from ..core.kjob import (

src/xpk/commands/run.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,12 @@
1616

1717
from argparse import Namespace
1818

19+
from ..core.commands import run_command_with_full_controls
20+
from ..core.gcloud_context import add_zone_and_project
21+
from ..core.kjob import AppProfileDefaults
1922
from ..core.kueue import LOCAL_QUEUE_NAME
2023
from ..utils.console import xpk_exit, xpk_print
2124
from .common import set_cluster_command
22-
from ..core.core import add_zone_and_project
23-
from ..core.kjob import AppProfileDefaults
24-
from ..core.commands import run_command_with_full_controls
2525
from .kind import set_local_cluster_command
2626

2727

src/xpk/commands/storage.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from kubernetes import client as k8s_client
2020
from kubernetes.client.rest import ApiException
2121

22-
from ..core.core import (
22+
from ..core.cluster import (
2323
setup_k8s_env,
2424
update_cluster_with_gcsfuse_driver_if_necessary,
2525
update_cluster_with_workload_identity_if_necessary,

src/xpk/commands/version.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
limitations under the License.
1515
"""
1616

17-
from ..core.core import __version__
17+
from ..core.config import __version__
1818
from ..utils.console import xpk_print
1919

2020

0 commit comments

Comments
 (0)