diff --git a/src/codeflare_sdk/cluster/cluster.py b/src/codeflare_sdk/cluster/cluster.py index b98255e68..256c14d80 100644 --- a/src/codeflare_sdk/cluster/cluster.py +++ b/src/codeflare_sdk/cluster/cluster.py @@ -50,7 +50,7 @@ class Cluster: torchx_scheduler = "ray" - def __init__(self, config: ClusterConfiguration): + def __init__(self, config: ClusterConfiguration, generate_app_wrapper: bool = True): """ Create the resource cluster object by passing in a ClusterConfiguration (defined in the config sub-module). An AppWrapper will then be generated @@ -58,13 +58,17 @@ def __init__(self, config: ClusterConfiguration): request. """ self.config = config - self.app_wrapper_yaml = self.create_app_wrapper() - self.app_wrapper_name = self.app_wrapper_yaml.split(".")[0] + self.app_wrapper_yaml = None + self.app_wrapper_name = None + + if generate_app_wrapper: + self.app_wrapper_yaml = self.create_app_wrapper() + self.app_wrapper_name = self.app_wrapper_yaml.split(".")[0] def create_app_wrapper(self): """ - Called upon cluster object creation, creates an AppWrapper yaml based on - the specifications of the ClusterConfiguration. + Creates an AppWrapper yaml based on the specified cluster config + based on the specifications of the ClusterConfiguration. """ if self.config.namespace is None: @@ -115,6 +119,9 @@ def up(self): Applies the AppWrapper yaml, pushing the resource request onto the MCAD queue. """ + if self.app_wrapper_yaml is None: + self.app_wrapper_yaml = self.create_app_wrapper() + self.app_wrapper_name = self.app_wrapper_yaml.split(".")[0] namespace = self.config.namespace try: config_check() @@ -137,6 +144,9 @@ def down(self): associated with the cluster. """ namespace = self.config.namespace + if not self.config.name and not self.app_wrapper_name: + print("Error taking down cluster: missing name or AppWrapper") + return try: config_check() api_instance = client.CustomObjectsApi(api_config_handler()) @@ -145,7 +155,7 @@ def down(self): version="v1beta1", namespace=namespace, plural="appwrappers", - name=self.app_wrapper_name, + name=self.app_wrapper_name or self.config.name, ) except Exception as e: # pragma: no cover return _kube_api_error_handling(e) @@ -346,7 +356,7 @@ def from_k8_cluster_object(rc): ]["image"], local_interactive=local_interactive, ) - return Cluster(cluster_config) + return Cluster(cluster_config, False) def from_definition_yaml(yaml_path): try: diff --git a/tests/unit_test.py b/tests/unit_test.py index a9258017f..0c1e1f875 100644 --- a/tests/unit_test.py +++ b/tests/unit_test.py @@ -2518,7 +2518,6 @@ def test_cleanup(): os.remove("unit-test-default-cluster.yaml") os.remove("test.yaml") os.remove("raytest2.yaml") - os.remove("quicktest.yaml") os.remove("tls-cluster-namespace/ca.crt") os.remove("tls-cluster-namespace/tls.crt") os.remove("tls-cluster-namespace/tls.key")