From 4af4307cc5c3998135cfdabaae5fc9fde93a8fa4 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 26 Jun 2024 09:53:07 +0200 Subject: [PATCH 01/71] initial codes --- src/zenml/cli/__init__.py | 28 +++++ src/zenml/cli/stack.py | 256 +++++++++++++++++++++++++++++++++++++- 2 files changed, 278 insertions(+), 6 deletions(-) diff --git a/src/zenml/cli/__init__.py b/src/zenml/cli/__init__.py index 1fa35f6168e..5b21e1bbc5b 100644 --- a/src/zenml/cli/__init__.py +++ b/src/zenml/cli/__init__.py @@ -1411,6 +1411,34 @@ Each corresponding argument should be the name, id or even the first few letters of the id that uniquely identify the artifact store or orchestrator. +To create a new stack using the new service connector with a set of minimal components, +use the following command: + +```bash +zenml stack register STACK_NAME \ + -cp CLOUD_PROVIDER +``` + +To create a new stack using the existing service connector with a set of minimal components, +use the following command: + +```bash +zenml stack register STACK_NAME \ + -sc SERVICE_CONNECTOR_NAME +``` + +To create a new stack using the existing service connector with existing components ( +important, that the components are already registered in the service connector), use the +following command: + +```bash +zenml stack register STACK_NAME \ + -sc SERVICE_CONNECTOR_NAME \ + -a ARTIFACT_STORE_NAME \ + -o ORCHESTRATOR_NAME \ + ... +``` + If you want to immediately set this newly created stack as your active stack, simply pass along the `--set` flag. diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 7229d216500..26638955853 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -53,6 +53,8 @@ from zenml.io.fileio import rmtree from zenml.logger import get_logger from zenml.models import StackFilter +from zenml.models.v2.core.component import ComponentResponse +from zenml.models.v2.core.service_connector import ServiceConnectorResponse from zenml.utils.dashboard_utils import get_stack_url from zenml.utils.io_utils import create_dir_recursive_if_not_exists from zenml.utils.mlstacks_utils import ( @@ -93,7 +95,7 @@ def stack() -> None: "artifact_store", help="Name of the artifact store for this stack.", type=str, - required=True, + required=False, ) @click.option( "-o", @@ -101,7 +103,7 @@ def stack() -> None: "orchestrator", help="Name of the orchestrator for this stack.", type=str, - required=True, + required=False, ) @click.option( "-c", @@ -190,10 +192,24 @@ def stack() -> None: help="Immediately set this stack as active.", type=click.BOOL, ) +@click.option( + "-cp", + "--cloud", + help="Name of the cloud provider for this stack.", + type=click.Choice(["aws", "azure", "gcp"]), + required=False, +) +@click.option( + "-sc", + "--connector", + help="Name of the service connector for this stack.", + type=str, + required=False, +) def register_stack( stack_name: str, - artifact_store: str, - orchestrator: str, + artifact_store: Optional[str] = None, + orchestrator: Optional[str] = None, container_registry: Optional[str] = None, model_registry: Optional[str] = None, step_operator: Optional[str] = None, @@ -205,6 +221,8 @@ def register_stack( data_validator: Optional[str] = None, image_builder: Optional[str] = None, set_stack: bool = False, + cloud: Optional[str] = None, + connector: Optional[str] = None, ) -> None: """Register a stack. @@ -223,10 +241,139 @@ def register_stack( data_validator: Name of the data validator for this stack. image_builder: Name of the new image builder for this stack. set_stack: Immediately set this stack as active. + cloud: Name of the cloud provider for this stack. + connector: Name of the service connector for this stack. """ - with console.status(f"Registering stack '{stack_name}'...\n"): - client = Client() + if (cloud is None and connector is None) and ( + artifact_store is None or orchestrator is None + ): + cli_utils.error( + "Only stack using service connector can be registered " + "without specifying an artifact store and an orchestrator. " + "Please specify the artifact store and the orchestrator or " + "the service connector or cloud type settings." + ) + client = Client() + + # cloud flow + service_connector = None + if cloud is not None and connector is None: + # if more than 100 service connectors of given type exist this might be an issue + existing_connectors = client.list_service_connectors( + connector_type=cloud, size=100 + ) + selected_connector = 0 + if existing_connectors.total: + selected_connector = int( + click.prompt( + f"We found following {cloud.upper()} service connectors. " + "Do you want to create a new one or use one of the existing ones?\n" + "[0] - Create a new service connector\n" + + "\n".join( + [ + f"[{i+1}] - {sc.name}" + for i, sc in enumerate(existing_connectors.items) + ] + ) + + "\n", + type=click.Choice( + [ + str(i) + for i in range( + 0, len(existing_connectors.items) + 1 + ) + ] + ), + default="0", + show_choices=False, + ) + ) + + if selected_connector != 0: + service_connector = existing_connectors.items[ + selected_connector - 1 + ] + else: + service_connector = _create_service_connector(cloud_provider=cloud) + elif connector is not None: + service_connector = client.get_service_connector(connector) + if service_connector.type != cloud: + cli_utils.warning( + f"The service connector `{connector}` is not of type `{cloud}`." + ) + if service_connector and _verify_service_connector(service_connector): + # create components + needed_components = { + ("artifact_store", artifact_store), + ("container_registry", container_registry), + ("orchestrator", orchestrator), # for azure only k8s orchestrator + } + created_components: List[ComponentResponse] = [] + for component_type, preset_name in needed_components: + if preset_name is not None: + component_response = client.get_stack_component( + component_type, preset_name + ) + if component_response.connector.id != service_connector.id: + cli_utils.error( + f"The {component_type.replace('_', ' ')} and service connector " + "do not match. Please check your inputs and try again." + ) + else: + # find existing components under same connector + existing_components = client.list_stack_components( + type=component_type, connector_id=service_connector.id + ) + # if some existing components are found - prompt user what to do + selected_component = 0 + if existing_components.total > 0: + # explore rich for interactive components + selected_component = int( + click.prompt( + f"We found following {component_type.replace('_', ' ')} " + "connected using the current service connector. Do you " + "want to create a new one or use existing one?\n" + f"[0] - Create a new {component_type.replace('_', ' ')}\n" + + "\n".join( + [ + f"[{i+1}] - {as_.name}" + for i, as_ in enumerate( + existing_components.items + ) + ] + ) + + "\n", + type=click.Choice( + [ + str(i) + for i in range( + 0, len(existing_components.items) + 1 + ) + ] + ), + default="0", + show_choices=False, + ) + ) + if selected_component != 0: + component_response = existing_components.items[ + selected_component - 1 + ] + else: + component_response = _create_stack_component( + component_type, service_connector + ) + + if component_type == "orchestrator": + orchestrator = component_response.name + elif component_type == "artifact_store": + artifact_store = component_response.name + elif component_type == "container_registry": + container_registry = component_response.name + + # normal flow once all components are defined + with console.status(f"Registering stack '{stack_name}'...\n"): components: Dict[StackComponentType, Union[str, UUID]] = dict() components[StackComponentType.ARTIFACT_STORE] = artifact_store @@ -279,6 +426,8 @@ def register_stack( print_model_url(get_stack_url(created_stack)) + # TODO: print how to delete stack and how to run a pipeline on it + @stack.command( "update", @@ -1727,3 +1876,98 @@ def connect_stack( interactive=interactive, no_verify=no_verify, ) + + +def _create_service_connector(cloud_provider: str) -> ServiceConnectorResponse: + """Create a service connector with given cloud provider. + + Args: + cloud_provider: The cloud provider to use. + + Returns: + The model of the created service connector. + """ + # TODO: here we question user and fill the needed details: + # - which auth you going to use? + # - depends on auth type, ask for specific credentials + if cloud_provider == "aws": + return ServiceConnectorResponse() + elif cloud_provider == "azure": + return ServiceConnectorResponse() + elif cloud_provider == "gcp": + return ServiceConnectorResponse() + else: + raise ValueError(f"Unknown cloud provider {cloud_provider}") + + +def _create_stack_component( + component_type: str, service_connector: ServiceConnectorResponse +) -> ComponentResponse: + """Create a stack component with given type and service connector. + + Args: + component_type: The type of component to create. + service_connector: The service connector to use. + + Returns: + The model of the created component. + """ + from zenml.cli.stack_components import ( + connect_stack_component_with_service_connector, + ) + + client = Client() + + # TODO: here we question user and fill the needed details + # - generic question will be the name + if component_type == "artifact_store": + # - here we list the available resources (buckets) to offer them to the user + # - user can select one of them + # - user can add suffix to the path, e.g. s3://bucket/suffix/path + component = ComponentResponse() + elif component_type == "container_registry": + # - here we list available registries and ask to pick one + component = ComponentResponse() + elif component_type == "orchestrator": + # - here we list available regions/locations and ask to pick one + # - ask for extra mandatory params, like Role ARN for AWS + component = ComponentResponse() + elif component_type == "image_builder": + # - only relevant for GCP and is fully optional + # - if on GCP - offer what we found, if none, go for local silently + component = ComponentResponse() + else: + raise ValueError(f"Unknown component type {component_type}") + + connect_stack_component_with_service_connector( + component_type=component.type, + name_id_or_prefix=component.id, + connector=service_connector, + interactive=False, + no_verify=False, + ) + + +def _verify_service_connector( + service_connector: ServiceConnectorResponse, +) -> bool: + """Verifies if a service connector has access to one or more resources. + + Args: + service_connector: The service connector to verify. + + Returns: + True if the service connector has proper permissions, False + otherwise. + + Raises: + ValueError: If the service connector has unexpected type. + """ + if service_connector.type == "aws": + return True + elif service_connector.type == "azure": + return True + elif service_connector.type == "gcp": + return True + else: + raise ValueError(f"Unknown cloud provider {service_connector.type}") From 7461deca6fc96eabc34a7a37ccce244e0022ba1a Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 26 Jun 2024 15:06:17 +0200 Subject: [PATCH 02/71] prettify a bit --- src/zenml/cli/stack.py | 208 +++++++++++++++++++++++++---------------- 1 file changed, 125 insertions(+), 83 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 26638955853..7e865b18d1b 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -244,6 +244,84 @@ def register_stack( cloud: Name of the cloud provider for this stack. connector: Name of the service connector for this stack. """ + from rich import print + from rich.console import Console + from rich.prompt import Prompt + from rich.table import Table + + def show_status( + cloud: str = None, + connector: str = None, + artifact_store: str = None, + orchestrator: str = None, + container_registry: str = None, + ) -> None: + status = [] + for each in [ + cloud, + connector, + artifact_store, + orchestrator, + container_registry, + ]: + if not each: + each = ":x:" + status.append(each) + + status_table = Table( + title="New cloud stack registration progress", + show_header=True, + expand=True, + ) + for c in [ + "Cloud", + "Service Connector", + "Artifact Store", + "Orchestrator", + "Container Registry", + ]: + status_table.add_column(c, justify="center", width=1) + + status_table.add_row(*status) + Console().clear() + print(status_table) + + def multi_choice_prompt( + object_type: str, choices_nameable: List[Any], prompt_text: str + ) -> int: + table = Table( + title=f"Available {object_type}", + show_header=False, + border_style=None, + expand=True, + ) + table.add_column("", justify="left", width=1) + for _ in range(min(len(choices_nameable) // 10, 3)): + table.add_column("", justify="left", width=1) + + ins = [f"[bold][0] - Create a new {object_type}[/bold]"] + for i, one_choice in enumerate(choices_nameable): + if len(ins) == len(table.columns): + table.add_row(*ins) + ins = [] + if len(ins) < len(table.columns): + ins.append(f"[{i+1}] - {one_choice.name}") + if ins: + while len(ins) < len(table.columns): + ins.append("") + table.add_row(*ins) + + print(table) + + return int( + Prompt.ask( + prompt_text, + choices=[str(i) for i in range(0, len(choices_nameable) + 1)], + default="0", + show_choices=False, + ) + ) + if (cloud is None and connector is None) and ( artifact_store is None or orchestrator is None ): @@ -259,37 +337,24 @@ def register_stack( # cloud flow service_connector = None if cloud is not None and connector is None: - # if more than 100 service connectors of given type exist this might be an issue + show_status( + cloud=cloud, + connector=connector, + artifact_store=artifact_store, + orchestrator=orchestrator, + container_registry=container_registry, + ) existing_connectors = client.list_service_connectors( connector_type=cloud, size=100 ) selected_connector = 0 if existing_connectors.total: - selected_connector = int( - click.prompt( - f"We found following {cloud.upper()} service connectors. " - "Do you want to create a new one or use one of the existing ones?\n" - "[0] - Create a new service connector\n" - + "\n".join( - [ - f"[{i+1}] - {sc.name}" - for i, sc in enumerate(existing_connectors.items) - ] - ) - + "\n", - type=click.Choice( - [ - str(i) - for i in range( - 0, len(existing_connectors.items) + 1 - ) - ] - ), - default="0", - show_choices=False, - ) + selected_connector = multi_choice_prompt( + object_type=f"{cloud.upper()} service connectors", + choices_nameable=existing_connectors.items, + prompt_text=f"We found these {cloud.upper()} service connectors. " + "Do you want to create a new one or use one of the existing ones?", ) - if selected_connector != 0: service_connector = existing_connectors.items[ selected_connector - 1 @@ -297,19 +362,33 @@ def register_stack( else: service_connector = _create_service_connector(cloud_provider=cloud) elif connector is not None: + show_status( + cloud=cloud, + connector=connector, + artifact_store=artifact_store, + orchestrator=orchestrator, + container_registry=container_registry, + ) service_connector = client.get_service_connector(connector) if service_connector.type != cloud: cli_utils.warning( f"The service connector `{connector}` is not of type `{cloud}`." ) - if service_connector and _verify_service_connector(service_connector): + show_status( + cloud=cloud, + connector=service_connector.name, + artifact_store=artifact_store, + orchestrator=orchestrator, + container_registry=container_registry, + ) + + if service_connector: # create components - needed_components = { + needed_components = ( ("artifact_store", artifact_store), - ("container_registry", container_registry), ("orchestrator", orchestrator), # for azure only k8s orchestrator - } - created_components: List[ComponentResponse] = [] + ("container_registry", container_registry), + ) for component_type, preset_name in needed_components: if preset_name is not None: component_response = client.get_stack_component( @@ -323,38 +402,19 @@ def register_stack( else: # find existing components under same connector existing_components = client.list_stack_components( - type=component_type, connector_id=service_connector.id + type=component_type, + connector_id=service_connector.id, + size=100, ) # if some existing components are found - prompt user what to do selected_component = 0 if existing_components.total > 0: - # explore rich for interactive components - selected_component = int( - click.prompt( - f"We found following {component_type.replace('_', ' ')} " - "connected using the current service connector. Do you " - "want to create a new one or use existing one?\n" - f"[0] - Create a new {component_type.replace('_', ' ')}\n" - + "\n".join( - [ - f"[{i+1}] - {as_.name}" - for i, as_ in enumerate( - existing_components.items - ) - ] - ) - + "\n", - type=click.Choice( - [ - str(i) - for i in range( - 0, len(existing_components.items) + 1 - ) - ] - ), - default="0", - show_choices=False, - ) + selected_component = multi_choice_prompt( + object_type=component_type.replace("_", " "), + choices_nameable=existing_components.items, + prompt_text=f"We found these {component_type.replace('_', ' ')} " + "connected using the current service connector. Do you " + "want to create a new one or use existing one?", ) if selected_component != 0: component_response = existing_components.items[ @@ -371,6 +431,13 @@ def register_stack( artifact_store = component_response.name elif component_type == "container_registry": container_registry = component_response.name + show_status( + cloud=cloud, + connector=service_connector.name, + artifact_store=artifact_store, + orchestrator=orchestrator, + container_registry=container_registry, + ) # normal flow once all components are defined with console.status(f"Registering stack '{stack_name}'...\n"): @@ -1946,28 +2013,3 @@ def _create_stack_component( interactive=False, no_verify=False, ) - - -def _verify_service_connector( - service_connector: ServiceConnectorResponse, -) -> bool: - """Verifies if a service connector has access to one or more resources. - - Args: - service_connector: The service connector to verify. - - Returns: - True if the service connector has proper permissions, False - otherwise. - - Raises: - ValueError: If the service connector has unexpected type. - """ - if service_connector.type == "aws": - return True - elif service_connector.type == "azure": - return True - elif service_connector.type == "gcp": - return True - else: - raise ValueError(f"Unknown cloud provider {service_connector.type}") From d03804a36f3a2dd557b48352fded75fc97906483 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 26 Jun 2024 17:59:52 +0200 Subject: [PATCH 03/71] some aws pieces --- src/zenml/cli/stack.py | 241 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 221 insertions(+), 20 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 7e865b18d1b..100e8ea687e 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -55,6 +55,9 @@ from zenml.models import StackFilter from zenml.models.v2.core.component import ComponentResponse from zenml.models.v2.core.service_connector import ServiceConnectorResponse +from zenml.models.v2.misc.service_connector_type import ( + ServiceConnectorResourcesModel, +) from zenml.utils.dashboard_utils import get_stack_url from zenml.utils.io_utils import create_dir_recursive_if_not_exists from zenml.utils.mlstacks_utils import ( @@ -383,6 +386,7 @@ def multi_choice_prompt( ) if service_connector: + service_connector_resource_model = None # create components needed_components = ( ("artifact_store", artifact_store), @@ -421,8 +425,16 @@ def multi_choice_prompt( selected_component - 1 ] else: + if service_connector_resource_model is None: + service_connector_resource_model = ( + client.verify_service_connector( + service_connector.id + ) + ) component_response = _create_stack_component( - component_type, service_connector + component_type, + service_connector, + service_connector_resource_model, ) if component_type == "orchestrator": @@ -1954,21 +1966,89 @@ def _create_service_connector(cloud_provider: str) -> ServiceConnectorResponse: Returns: The model of the created service connector. """ - # TODO: here we question user and fill the needed details: - # - which auth you going to use? - # - depends on auth type, ask for specific credentials - if cloud_provider == "aws": - return ServiceConnectorResponse() - elif cloud_provider == "azure": - return ServiceConnectorResponse() - elif cloud_provider == "gcp": - return ServiceConnectorResponse() - else: + from rich import print + from rich.console import Console + from rich.markdown import Markdown + from rich.prompt import Confirm, Prompt + from rich.table import Table + + if cloud_provider not in {"aws", "azure", "gcp"}: raise ValueError(f"Unknown cloud provider {cloud_provider}") + client = Client() + auth_methods = client.get_service_connector_type( + cloud_provider + ).auth_method_dict + auth_methods_table = Table( + title=f"Available authentication methods for {cloud_provider}", + expand=True, + show_lines=True, + ) + auth_methods_table.add_column("Choice", justify="left", width=1) + auth_methods_table.add_column("Name", justify="left", width=10) + auth_methods_table.add_column("Required", justify="left", width=10) + + fixed_auth_methods = list(enumerate(auth_methods.items())) + for i, (_, value) in fixed_auth_methods: + schema = value.config_schema + required = "" + for each_req in schema["required"]: + field = schema["properties"][each_req] + required += f"[bold]{each_req}[/bold] [italic]({field.get('title','no description')})[/italic]\n" + auth_methods_table.add_row(str(i), value.name, required) + auth_selected = False + selected_auth_model = None + while not auth_selected: + Console().clear() + print(auth_methods_table) + selected_auth_idx = int( + Prompt.ask( + "Please choose one of the authentication option above to see detailed description:", + choices=[str(i) for i in range(len(auth_methods))], + ) + ) + selected_auth_model = fixed_auth_methods[selected_auth_idx][1][1] + print( + Markdown( + f"## {selected_auth_model.name}\n" + + selected_auth_model.description + ) + ) + + auth_selected = Confirm.ask( + "Do you want to continue or go back to authentication methods selection?", + default=True, + ) + if not selected_auth_model: + raise ValueError("No authentication method selected") + required_fields = selected_auth_model.config_schema["required"] + data_entered = False + while not data_entered: + answers = {} + for req_field in required_fields: + answers[req_field] = Prompt.ask( + f"Please enter value for `{req_field}`:" + ) + data_entered = Confirm.ask( + "Please confirm the values you entered:\n" + str(answers), + default=True, + ) + + connector_name = Prompt.ask( + "Please enter a name for the service connector:" + ) + return client.create_service_connector( + name=connector_name, + connector_type=cloud_provider, + auth_method=fixed_auth_methods[selected_auth_idx][1][0], + configuration=answers, + )[0] + def _create_stack_component( - component_type: str, service_connector: ServiceConnectorResponse + component_type: str, + service_connector: ServiceConnectorResponse, + service_connector_resource_model: ServiceConnectorResourcesModel, ) -> ComponentResponse: """Create a stack component with given type and service connector. @@ -1979,6 +2059,10 @@ def _create_stack_component( Returns: The model of the created component. """ + from rich import print + from rich.prompt import Confirm, Prompt + from rich.table import Table + from zenml.cli.stack_components import ( connect_stack_component_with_service_connector, ) @@ -1988,17 +2072,132 @@ def _create_stack_component( # TODO: here we question user and fill the needed details # - generic question will be the name if component_type == "artifact_store": - # - here we list the available resources (buckets) to offer them to the user - # - user can select one of them - # - user can add suffix to the path, e.g. s3://bucket/suffix/path - component = ComponentResponse() + config_confirmed = False + while not config_confirmed: + if service_connector.type == "aws": + for each in service_connector_resource_model.resources: + if each.resource_type == "s3-bucket": + available_buckets = each.resource_ids + available_buckets_table = Table( + title="Available S3 buckets:", expand=True + ) + available_buckets_table.add_column( + "Choice", justify="left", width=1 + ) + available_buckets_table.add_column( + "Bucket", justify="left", width=10 + ) + for i, bucket in enumerate(available_buckets): + available_buckets_table.add_row(str(i), bucket) + print(available_buckets_table) + selected_bucket = available_buckets[ + int( + Prompt.ask( + "Please choose one of the buckets for new artifact store:", + choices=[ + str(i) for i in range(len(available_buckets)) + ], + ) + ) + ] + extra_path = Prompt.ask( + f"Please enter any further path inside the bucket, if needed ({selected_bucket}/...):", + default="", + ) + flavor = "s3" + config = { + "path": f"{selected_bucket.strip('/')}/{extra_path.strip('/')}" + } + + as_name = Prompt.ask( + "Please enter a name for the artifact store:" + ) + + print({"name": as_name, "config": config}) + config_confirmed = Confirm.ask( + "Please confirm the values you entered:", default=True + ) + component = client.create_stack_component( + name=as_name, + flavor=flavor, + component_type=StackComponentType.ARTIFACT_STORE, + configuration=config, + ) elif component_type == "container_registry": # - here we list available registries and ask to pick one component = ComponentResponse() elif component_type == "orchestrator": - # - here we list available regions/locations and ask to pick one - # - ask for extra mandatory params, like Role ARN for AWS - component = ComponentResponse() + config_confirmed = False + while not config_confirmed: + if service_connector.type == "aws": + available_orchestrators = {} + for each in service_connector_resource_model.resources: + if each.resource_type == "aws-generic": + available_orchestrators["Sagemaker"] = ( + each.resource_ids or [] + ) + + if each.resource_type == "kubernetes-cluster": + available_orchestrators["K8S"] = ( + each.resource_ids or [] + ) + available_orchestrators_table = Table( + title="Available orchestrators on AWS:", expand=True + ) + available_orchestrators_table.add_column( + "Choice", justify="left", width=1 + ) + available_orchestrators_table.add_column( + "Orchestrator details", justify="left", width=10 + ) + choice_number = 0 + choices_mapper = {} + for type_ in available_orchestrators: + for i, orchestrator in enumerate( + available_orchestrators[type_] + ): + available_orchestrators_table.add_row( + str(choice_number), f"{type_} - {orchestrator}" + ) + choices_mapper[choice_number] = (type_, i) + choice_number += 1 + print(available_orchestrators_table) + orchestrator_choice = int( + Prompt.ask( + "Please choose one of the options for the new orchestrator:", + choices=[str(i) for i in range(choice_number)], + ) + ) + selected_orchestrator = available_orchestrators[ + choices_mapper[orchestrator_choice][0] + ][choices_mapper[orchestrator_choice][1]] + + if choices_mapper[orchestrator_choice][0] == "Sagemaker": + flavor = "sagemaker" + execution_role = Prompt.ask( + "Please enter an execution role ARN:" + ) + config = {"execution_role": execution_role} + elif choices_mapper[orchestrator_choice][0] == "K8S": + flavor = "kubernetes" + config = {} + else: + raise ValueError( + f"Unknown orchestrator type {choices_mapper[orchestrator_choice][0]}" + ) + orchestrator_name = Prompt.ask( + "Please enter a name for the orchestrator:" + ) + print({"name": orchestrator_name, "config": config}) + config_confirmed = Confirm.ask( + "Please confirm the values you entered:", default=True + ) + component = client.create_stack_component( + name=orchestrator_name, + flavor=flavor, + component_type=StackComponentType.ORCHESTRATOR, + configuration=config, + ) elif component_type == "image_builder": # - only relevant for GCP and is fully optional # - if on GCP - offer what we found, if none, go for local silently @@ -2009,7 +2208,9 @@ def _create_stack_component( connect_stack_component_with_service_connector( component_type=component.type, name_id_or_prefix=component.id, - connector=service_connector, + connector=service_connector.id, interactive=False, no_verify=False, ) + + return component From 7399ba803ed8bd55c564b515a7f57d64a07b97c7 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 27 Jun 2024 10:39:27 +0200 Subject: [PATCH 04/71] add container registry --- src/zenml/cli/stack.py | 238 ++++++++++++++++++++++++++--------------- 1 file changed, 151 insertions(+), 87 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 100e8ea687e..d4400029739 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -2067,51 +2067,56 @@ def _create_stack_component( connect_stack_component_with_service_connector, ) + if service_connector.type not in {"aws", "azure", "gcp"}: + raise ValueError(f"Unknown cloud provider {service_connector.type}") + client = Client() - # TODO: here we question user and fill the needed details - # - generic question will be the name if component_type == "artifact_store": config_confirmed = False while not config_confirmed: if service_connector.type == "aws": for each in service_connector_resource_model.resources: if each.resource_type == "s3-bucket": - available_buckets = each.resource_ids - available_buckets_table = Table( - title="Available S3 buckets:", expand=True - ) - available_buckets_table.add_column( - "Choice", justify="left", width=1 - ) - available_buckets_table.add_column( - "Bucket", justify="left", width=10 - ) - for i, bucket in enumerate(available_buckets): - available_buckets_table.add_row(str(i), bucket) - print(available_buckets_table) - selected_bucket = available_buckets[ - int( - Prompt.ask( - "Please choose one of the buckets for new artifact store:", - choices=[ - str(i) for i in range(len(available_buckets)) - ], - ) + available_storages = each.resource_ids + flavor = "s3" + elif service_connector.type == "azure": + flavor = "azure" + elif service_connector.type == "gcp": + flavor = "gcs" + + available_storages_table = Table( + title=f"Available {service_connector.type.upper()} storages:", + expand=True, + ) + available_storages_table.add_column( + "Choice", justify="left", width=1 + ) + available_storages_table.add_column( + "Storage", justify="left", width=10 + ) + for i, storage in enumerate(available_storages): + available_storages_table.add_row(str(i), storage) + print(available_storages_table) + selected_storage = available_storages[ + int( + Prompt.ask( + "Please choose one of the storages for the new artifact store:", + choices=[ + str(i) for i in range(len(available_storages)) + ], ) - ] - extra_path = Prompt.ask( - f"Please enter any further path inside the bucket, if needed ({selected_bucket}/...):", - default="", ) - flavor = "s3" - config = { - "path": f"{selected_bucket.strip('/')}/{extra_path.strip('/')}" - } + ] + extra_path = Prompt.ask( + f"Please enter any further path inside the storage, if needed ({selected_storage}/...):", + default="", + ) + config = { + "path": f"{selected_storage.strip('/')}/{extra_path.strip('/')}" + } - as_name = Prompt.ask( - "Please enter a name for the artifact store:" - ) + as_name = Prompt.ask("Please enter a name for the artifact store:") print({"name": as_name, "config": config}) config_confirmed = Confirm.ask( @@ -2123,9 +2128,6 @@ def _create_stack_component( component_type=StackComponentType.ARTIFACT_STORE, configuration=config, ) - elif component_type == "container_registry": - # - here we list available registries and ask to pick one - component = ComponentResponse() elif component_type == "orchestrator": config_confirmed = False while not config_confirmed: @@ -2136,55 +2138,67 @@ def _create_stack_component( available_orchestrators["Sagemaker"] = ( each.resource_ids or [] ) + available_orchestrators["VM AWS"] = ( + each.resource_ids or [] + ) if each.resource_type == "kubernetes-cluster": available_orchestrators["K8S"] = ( each.resource_ids or [] ) - available_orchestrators_table = Table( - title="Available orchestrators on AWS:", expand=True - ) - available_orchestrators_table.add_column( - "Choice", justify="left", width=1 + elif service_connector.type == "gcp": + pass + elif service_connector.type == "azure": + pass + + available_orchestrators_table = Table( + title=f"Available orchestrators on {service_connector.type.upper()}:", + expand=True, + ) + available_orchestrators_table.add_column( + "Choice", justify="left", width=1 + ) + available_orchestrators_table.add_column( + "Orchestrator details", justify="left", width=10 + ) + choice_number = 0 + choices_mapper = {} + for type_ in available_orchestrators: + for i, orchestrator in enumerate( + available_orchestrators[type_] + ): + available_orchestrators_table.add_row( + str(choice_number), f"{type_} - {orchestrator}" + ) + choices_mapper[choice_number] = (type_, i) + choice_number += 1 + print(available_orchestrators_table) + orchestrator_choice = int( + Prompt.ask( + "Please choose one of the options for the new orchestrator:", + choices=[str(i) for i in range(choice_number)], ) - available_orchestrators_table.add_column( - "Orchestrator details", justify="left", width=10 + ) + selected_orchestrator = available_orchestrators[ + choices_mapper[orchestrator_choice][0] + ][choices_mapper[orchestrator_choice][1]] + + if choices_mapper[orchestrator_choice][0] == "Sagemaker": + flavor = "sagemaker" + execution_role = Prompt.ask( + "Please enter an execution role ARN:" ) - choice_number = 0 - choices_mapper = {} - for type_ in available_orchestrators: - for i, orchestrator in enumerate( - available_orchestrators[type_] - ): - available_orchestrators_table.add_row( - str(choice_number), f"{type_} - {orchestrator}" - ) - choices_mapper[choice_number] = (type_, i) - choice_number += 1 - print(available_orchestrators_table) - orchestrator_choice = int( - Prompt.ask( - "Please choose one of the options for the new orchestrator:", - choices=[str(i) for i in range(choice_number)], - ) + config = {"execution_role": execution_role} + elif choices_mapper[orchestrator_choice][0] == "VM AWS": + flavor = "vm_aws" + config = {} + elif choices_mapper[orchestrator_choice][0] == "K8S": + flavor = "kubernetes" + config = {} + else: + raise ValueError( + f"Unknown orchestrator type {choices_mapper[orchestrator_choice][0]}" ) - selected_orchestrator = available_orchestrators[ - choices_mapper[orchestrator_choice][0] - ][choices_mapper[orchestrator_choice][1]] - - if choices_mapper[orchestrator_choice][0] == "Sagemaker": - flavor = "sagemaker" - execution_role = Prompt.ask( - "Please enter an execution role ARN:" - ) - config = {"execution_role": execution_role} - elif choices_mapper[orchestrator_choice][0] == "K8S": - flavor = "kubernetes" - config = {} - else: - raise ValueError( - f"Unknown orchestrator type {choices_mapper[orchestrator_choice][0]}" - ) orchestrator_name = Prompt.ask( "Please enter a name for the orchestrator:" ) @@ -2192,16 +2206,66 @@ def _create_stack_component( config_confirmed = Confirm.ask( "Please confirm the values you entered:", default=True ) - component = client.create_stack_component( - name=orchestrator_name, - flavor=flavor, - component_type=StackComponentType.ORCHESTRATOR, - configuration=config, + component = client.create_stack_component( + name=orchestrator_name, + flavor=flavor, + component_type=StackComponentType.ORCHESTRATOR, + configuration=config, + ) + elif component_type == "container_registry": + config_confirmed = False + while not config_confirmed: + if service_connector.type == "aws": + for each in service_connector_resource_model.resources: + if each.resource_type == "docker-registry": + available_registries = each.resource_ids + flavor = "aws" + elif service_connector.type == "azure": + flavor = "azure" + available_registries = [] + elif service_connector.type == "gcp": + flavor = "gcp" + available_registries = [] + + available_registries_table = Table( + title=f"Available {service_connector.type.upper()} registries:", + expand=True, + ) + available_registries_table.add_column( + "Choice", justify="left", width=1 + ) + available_registries_table.add_column( + "Container Registry", justify="left", width=10 ) - elif component_type == "image_builder": - # - only relevant for GCP and is fully optional - # - if on GCP - offer what we found, if none, go for local silently - component = ComponentResponse() + for i, registry in enumerate(available_registries): + available_registries_table.add_row(str(i), registry) + print(available_registries_table) + selected_storage = available_registries[ + int( + Prompt.ask( + "Please choose one of the registries for the new container registry:", + choices=[ + str(i) for i in range(len(available_registries)) + ], + ) + ) + ] + config = {"uri": selected_storage} + + cr_name = Prompt.ask( + "Please enter a name for the container registry:" + ) + + print({"name": cr_name, "config": config}) + config_confirmed = Confirm.ask( + "Please confirm the values you entered:", default=True + ) + component = client.create_stack_component( + name=cr_name, + flavor=flavor, + component_type=StackComponentType.CONTAINER_REGISTRY, + configuration=config, + ) else: raise ValueError(f"Unknown component type {component_type}") From 0d5b64643695b1a01a987bb470a5ba4309b8fcba Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 27 Jun 2024 16:05:00 +0200 Subject: [PATCH 05/71] refactor --- src/zenml/cli/stack.py | 286 +++++++++++------------------------------ src/zenml/cli/utils.py | 109 +++++++++++++++- 2 files changed, 177 insertions(+), 218 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index d4400029739..da47db4d2fc 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -247,84 +247,6 @@ def register_stack( cloud: Name of the cloud provider for this stack. connector: Name of the service connector for this stack. """ - from rich import print - from rich.console import Console - from rich.prompt import Prompt - from rich.table import Table - - def show_status( - cloud: str = None, - connector: str = None, - artifact_store: str = None, - orchestrator: str = None, - container_registry: str = None, - ) -> None: - status = [] - for each in [ - cloud, - connector, - artifact_store, - orchestrator, - container_registry, - ]: - if not each: - each = ":x:" - status.append(each) - - status_table = Table( - title="New cloud stack registration progress", - show_header=True, - expand=True, - ) - for c in [ - "Cloud", - "Service Connector", - "Artifact Store", - "Orchestrator", - "Container Registry", - ]: - status_table.add_column(c, justify="center", width=1) - - status_table.add_row(*status) - Console().clear() - print(status_table) - - def multi_choice_prompt( - object_type: str, choices_nameable: List[Any], prompt_text: str - ) -> int: - table = Table( - title=f"Available {object_type}", - show_header=False, - border_style=None, - expand=True, - ) - table.add_column("", justify="left", width=1) - for _ in range(min(len(choices_nameable) // 10, 3)): - table.add_column("", justify="left", width=1) - - ins = [f"[bold][0] - Create a new {object_type}[/bold]"] - for i, one_choice in enumerate(choices_nameable): - if len(ins) == len(table.columns): - table.add_row(*ins) - ins = [] - if len(ins) < len(table.columns): - ins.append(f"[{i+1}] - {one_choice.name}") - if ins: - while len(ins) < len(table.columns): - ins.append("") - table.add_row(*ins) - - print(table) - - return int( - Prompt.ask( - prompt_text, - choices=[str(i) for i in range(0, len(choices_nameable) + 1)], - default="0", - show_choices=False, - ) - ) - if (cloud is None and connector is None) and ( artifact_store is None or orchestrator is None ): @@ -340,7 +262,7 @@ def multi_choice_prompt( # cloud flow service_connector = None if cloud is not None and connector is None: - show_status( + cli_utils.show_status_from_kwargs( cloud=cloud, connector=connector, artifact_store=artifact_store, @@ -350,22 +272,25 @@ def multi_choice_prompt( existing_connectors = client.list_service_connectors( connector_type=cloud, size=100 ) - selected_connector = 0 + connector_selected: Optional[int] = None if existing_connectors.total: - selected_connector = multi_choice_prompt( + connector_selected = cli_utils.multi_choice_prompt( object_type=f"{cloud.upper()} service connectors", - choices_nameable=existing_connectors.items, + choices=[ + [connector.name] for connector in existing_connectors.items + ], + headers=["Name"], prompt_text=f"We found these {cloud.upper()} service connectors. " "Do you want to create a new one or use one of the existing ones?", + default_choice="0", + allow_zero_be_a_new_object=True, ) - if selected_connector != 0: - service_connector = existing_connectors.items[ - selected_connector - 1 - ] - else: + if connector_selected is None: service_connector = _create_service_connector(cloud_provider=cloud) + else: + service_connector = existing_connectors.items[connector_selected] elif connector is not None: - show_status( + cli_utils.show_status_from_kwargs( cloud=cloud, connector=connector, artifact_store=artifact_store, @@ -377,7 +302,7 @@ def multi_choice_prompt( cli_utils.warning( f"The service connector `{connector}` is not of type `{cloud}`." ) - show_status( + cli_utils.show_status_from_kwargs( cloud=cloud, connector=service_connector.name, artifact_store=artifact_store, @@ -411,20 +336,22 @@ def multi_choice_prompt( size=100, ) # if some existing components are found - prompt user what to do - selected_component = 0 + component_selected: Optional[int] = None if existing_components.total > 0: - selected_component = multi_choice_prompt( + component_selected = cli_utils.multi_choice_prompt( object_type=component_type.replace("_", " "), - choices_nameable=existing_components.items, + choices=[ + [component.name] + for component in existing_components.items + ], + headers=["Name"], prompt_text=f"We found these {component_type.replace('_', ' ')} " "connected using the current service connector. Do you " "want to create a new one or use existing one?", + default_choice="0", + allow_zero_be_a_new_object=True, ) - if selected_component != 0: - component_response = existing_components.items[ - selected_component - 1 - ] - else: + if component_selected is None: if service_connector_resource_model is None: service_connector_resource_model = ( client.verify_service_connector( @@ -436,6 +363,10 @@ def multi_choice_prompt( service_connector, service_connector_resource_model, ) + else: + component_response = existing_components.items[ + component_selected + ] if component_type == "orchestrator": orchestrator = component_response.name @@ -443,7 +374,7 @@ def multi_choice_prompt( artifact_store = component_response.name elif component_type == "container_registry": container_registry = component_response.name - show_status( + cli_utils.show_status_from_kwargs( cloud=cloud, connector=service_connector.name, artifact_store=artifact_store, @@ -1967,10 +1898,8 @@ def _create_service_connector(cloud_provider: str) -> ServiceConnectorResponse: The model of the created service connector. """ from rich import print - from rich.console import Console from rich.markdown import Markdown from rich.prompt import Confirm, Prompt - from rich.table import Table if cloud_provider not in {"aws", "azure", "gcp"}: raise ValueError(f"Unknown cloud provider {cloud_provider}") @@ -1979,35 +1908,26 @@ def _create_service_connector(cloud_provider: str) -> ServiceConnectorResponse: auth_methods = client.get_service_connector_type( cloud_provider ).auth_method_dict - auth_methods_table = Table( - title=f"Available authentication methods for {cloud_provider}", - expand=True, - show_lines=True, - ) - auth_methods_table.add_column("Choice", justify="left", width=1) - auth_methods_table.add_column("Name", justify="left", width=10) - auth_methods_table.add_column("Required", justify="left", width=10) - - fixed_auth_methods = list(enumerate(auth_methods.items())) - for i, (_, value) in fixed_auth_methods: + fixed_auth_methods = list(auth_methods.items()) + choices = [] + headers = ["Name", "Required"] + for _, value in fixed_auth_methods: schema = value.config_schema required = "" for each_req in schema["required"]: field = schema["properties"][each_req] required += f"[bold]{each_req}[/bold] [italic]({field.get('title','no description')})[/italic]\n" - auth_methods_table.add_row(str(i), value.name, required) + choices.append([value.name, required]) + auth_selected = False - selected_auth_model = None while not auth_selected: - Console().clear() - print(auth_methods_table) - selected_auth_idx = int( - Prompt.ask( - "Please choose one of the authentication option above to see detailed description:", - choices=[str(i) for i in range(len(auth_methods))], - ) + selected_auth_idx = cli_utils.multi_choice_prompt( + object_type=f"authentication methods for {cloud_provider}", + choices=choices, + headers=headers, + prompt_text="Please choose one of the authentication option above to see detailed description:", ) - selected_auth_model = fixed_auth_methods[selected_auth_idx][1][1] + selected_auth_model = fixed_auth_methods[selected_auth_idx][1] print( Markdown( f"## {selected_auth_model.name}\n" @@ -2040,7 +1960,7 @@ def _create_service_connector(cloud_provider: str) -> ServiceConnectorResponse: return client.create_service_connector( name=connector_name, connector_type=cloud_provider, - auth_method=fixed_auth_methods[selected_auth_idx][1][0], + auth_method=fixed_auth_methods[selected_auth_idx][0], configuration=answers, )[0] @@ -2061,7 +1981,6 @@ def _create_stack_component( """ from rich import print from rich.prompt import Confirm, Prompt - from rich.table import Table from zenml.cli.stack_components import ( connect_stack_component_with_service_connector, @@ -2085,29 +2004,14 @@ def _create_stack_component( elif service_connector.type == "gcp": flavor = "gcs" - available_storages_table = Table( - title=f"Available {service_connector.type.upper()} storages:", - expand=True, + selected_storage_idx = cli_utils.multi_choice_prompt( + object_type=f"{service_connector.type.upper()} storages", + choices=[[st] for st in available_storages], + headers=["Storage"], + prompt_text="Please choose one of the storages for the new artifact store:", ) - available_storages_table.add_column( - "Choice", justify="left", width=1 - ) - available_storages_table.add_column( - "Storage", justify="left", width=10 - ) - for i, storage in enumerate(available_storages): - available_storages_table.add_row(str(i), storage) - print(available_storages_table) - selected_storage = available_storages[ - int( - Prompt.ask( - "Please choose one of the storages for the new artifact store:", - choices=[ - str(i) for i in range(len(available_storages)) - ], - ) - ) - ] + selected_storage = available_storages[selected_storage_idx] + extra_path = Prompt.ask( f"Please enter any further path inside the storage, if needed ({selected_storage}/...):", default="", @@ -2132,72 +2036,48 @@ def _create_stack_component( config_confirmed = False while not config_confirmed: if service_connector.type == "aws": - available_orchestrators = {} + available_orchestrators = [] for each in service_connector_resource_model.resources: + types = [] if each.resource_type == "aws-generic": - available_orchestrators["Sagemaker"] = ( - each.resource_ids or [] - ) - available_orchestrators["VM AWS"] = ( - each.resource_ids or [] - ) - + types = ["Sagemaker", "VM AWS"] if each.resource_type == "kubernetes-cluster": - available_orchestrators["K8S"] = ( - each.resource_ids or [] - ) + types = ["K8S"] + + for orchestrator in each.resource_ids: + for t in types: + available_orchestrators.append([t, orchestrator]) elif service_connector.type == "gcp": pass elif service_connector.type == "azure": pass - available_orchestrators_table = Table( - title=f"Available orchestrators on {service_connector.type.upper()}:", - expand=True, - ) - available_orchestrators_table.add_column( - "Choice", justify="left", width=1 - ) - available_orchestrators_table.add_column( - "Orchestrator details", justify="left", width=10 - ) - choice_number = 0 - choices_mapper = {} - for type_ in available_orchestrators: - for i, orchestrator in enumerate( - available_orchestrators[type_] - ): - available_orchestrators_table.add_row( - str(choice_number), f"{type_} - {orchestrator}" - ) - choices_mapper[choice_number] = (type_, i) - choice_number += 1 - print(available_orchestrators_table) - orchestrator_choice = int( - Prompt.ask( - "Please choose one of the options for the new orchestrator:", - choices=[str(i) for i in range(choice_number)], - ) + selected_orchestrator_idx = cli_utils.multi_choice_prompt( + object_type=f"orchestrators on {service_connector.type.upper()}", + choices=available_orchestrators, + headers=["Orchestrator Type", "Orchestrator details"], + prompt_text="Please choose one of the orchestrators for the new orchestrator:", ) + selected_orchestrator = available_orchestrators[ - choices_mapper[orchestrator_choice][0] - ][choices_mapper[orchestrator_choice][1]] + selected_orchestrator_idx + ] - if choices_mapper[orchestrator_choice][0] == "Sagemaker": + if selected_orchestrator[0] == "Sagemaker": flavor = "sagemaker" execution_role = Prompt.ask( "Please enter an execution role ARN:" ) config = {"execution_role": execution_role} - elif choices_mapper[orchestrator_choice][0] == "VM AWS": + elif selected_orchestrator[0] == "VM AWS": flavor = "vm_aws" config = {} - elif choices_mapper[orchestrator_choice][0] == "K8S": + elif selected_orchestrator[0] == "K8S": flavor = "kubernetes" config = {} else: raise ValueError( - f"Unknown orchestrator type {choices_mapper[orchestrator_choice][0]}" + f"Unknown orchestrator type {selected_orchestrator[0]}" ) orchestrator_name = Prompt.ask( "Please enter a name for the orchestrator:" @@ -2227,29 +2107,13 @@ def _create_stack_component( flavor = "gcp" available_registries = [] - available_registries_table = Table( - title=f"Available {service_connector.type.upper()} registries:", - expand=True, - ) - available_registries_table.add_column( - "Choice", justify="left", width=1 + selected_registry_idx = cli_utils.multi_choice_prompt( + object_type=f"{service_connector.type.upper()} registries", + choices=[[st] for st in available_registries], + headers=["Container Registry"], + prompt_text="Please choose one of the registries for the new container registry:", ) - available_registries_table.add_column( - "Container Registry", justify="left", width=10 - ) - for i, registry in enumerate(available_registries): - available_registries_table.add_row(str(i), registry) - print(available_registries_table) - selected_storage = available_registries[ - int( - Prompt.ask( - "Please choose one of the registries for the new container registry:", - choices=[ - str(i) for i in range(len(available_registries)) - ], - ) - ) - ] + selected_storage = available_registries[selected_registry_idx] config = {"uri": selected_storage} cr_name = Prompt.ask( diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index e38f8b12f68..7b76983205b 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -46,8 +46,9 @@ from rich.emoji import Emoji, NoEmoji from rich.markdown import Markdown from rich.markup import escape -from rich.prompt import Confirm +from rich.prompt import Confirm, Prompt from rich.style import Style +from rich.table import Table from zenml.client import Client from zenml.console import console, zenml_style_defaults @@ -1276,9 +1277,11 @@ def pretty_print_model_version_table( "NAME": model_version.registered_model.name, "MODEL_VERSION": model_version.version, "VERSION_DESCRIPTION": model_version.description, - "METADATA": model_version.metadata.model_dump() - if model_version.metadata - else {}, + "METADATA": ( + model_version.metadata.model_dump() + if model_version.metadata + else {} + ), } for model_version in model_versions ] @@ -1318,9 +1321,11 @@ def pretty_print_model_version_details( if model_version.last_updated_at else "N/A" ), - "METADATA": model_version.metadata.model_dump() - if model_version.metadata - else {}, + "METADATA": ( + model_version.metadata.model_dump() + if model_version.metadata + else {} + ), "MODEL_SOURCE_URI": model_version.model_source_uri, "STAGE": model_version.stage.value, } @@ -2748,3 +2753,93 @@ def is_jupyter_installed() -> bool: return True except pkg_resources.DistributionNotFound: return False + + +def show_status_from_kwargs(default_value: str = ":x:", **kwargs: Any) -> None: + """Show status from kwargs. + + Args: + default_value: The default value to show status from. + **kwargs: The kwargs to show status from. If value is passed, + but `None` a default value is used. + """ + from rich import print + + status = [] + names = [] + for name, each in kwargs.items(): + if not each: + each = ":x:" + status.append(each) + names.append(name.replace("_", " ").capitalize()) + + status_table = Table( + title="New cloud stack registration progress", + show_header=True, + expand=True, + ) + for c in names: + status_table.add_column(c, justify="center", width=1) + + status_table.add_row(*status) + print(status_table) + + +def multi_choice_prompt( + object_type: str, + choices: List[List[Any]], + headers: List[str], + prompt_text: str, + allow_zero_be_a_new_object: bool = False, + default_choice: Optional[str] = None, +) -> Optional[int]: + """Prompts the user to select a choice from a list of choices. + + Args: + object_type: The type of the object + choices: The list of choices + prompt_text: The prompt text + selector_from_choices: The list of selectors to use + allow_zero_be_a_new_object: Whether to allow zero as a new object + default_choice: The default choice + + Returns: + The selected choice index or None for new object + """ + from rich import print + + table = Table( + title=f"Available {object_type}", + show_header=True, + border_style=None, + expand=True, + show_lines=True, + ) + table.add_column("Choice", justify="left", width=1) + for h in headers: + table.add_column( + h.replace("_", " ").capitalize(), justify="left", width=10 + ) + + i_shift = 0 + if allow_zero_be_a_new_object: + i_shift = 1 + table.add_row( + "[0]", + *([f"Create a new {object_type}"] * len(headers)), + ) + for i, one_choice in enumerate(choices): + table.add_row(f"[{i+i_shift}]", *[str(x) for x in one_choice]) + print(table) + + selected = Prompt.ask( + prompt_text, + choices=[str(i) for i in range(0, len(choices) + 1)], + default=default_choice, + show_choices=False, + ) + + if selected == "0" and allow_zero_be_a_new_object: + return None + else: + return int(selected) - i_shift From 2cf6ad52a0bb0813e665c9990a07ddaff74ff6f7 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 09:15:18 +0200 Subject: [PATCH 06/71] new models --- src/zenml/models/__init__.py | 8 +++ src/zenml/models/v2/misc/full_stack.py | 71 ++++++++++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 src/zenml/models/v2/misc/full_stack.py diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index 0724fdc89b0..53ded3453b4 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -327,6 +327,10 @@ ResourceTypeModel, ) from zenml.models.v2.misc.server_models import ServerDatabaseType, ServerModel +from zenml.models.v2.misc.full_stack import ( + FullStackRequest, + FullStackScopedRequest, +) from zenml.models.v2.core.trigger import ( TriggerRequest, TriggerFilter, @@ -405,6 +409,8 @@ EventSourceResponseResources.model_rebuild() FlavorResponseBody.model_rebuild() FlavorResponseMetadata.model_rebuild() +FullStackRequest.model_rebuild() +FullStackScopedRequest.model_rebuild() LazyArtifactVersionResponse.model_rebuild() LazyRunMetadataResponse.model_rebuild() ModelResponseBody.model_rebuild() @@ -706,6 +712,8 @@ "ServiceConnectorTypedResourcesModel", "ServiceConnectorRequirements", "ResourceTypeModel", + "FullStackRequest", + "FullStackScopedRequest", "UserAuthModel", "ExternalUserModel", "BuildItem", diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py new file mode 100644 index 00000000000..d9493c826af --- /dev/null +++ b/src/zenml/models/v2/misc/full_stack.py @@ -0,0 +1,71 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Models representing full stack requests.""" + +from typing import Any, Dict, Optional, Union +from uuid import UUID + +from pydantic import BaseModel, Field + +from zenml.constants import STR_FIELD_MAX_LENGTH +from zenml.enums import StackComponentType +from zenml.models.v2.base.scoped import WorkspaceScopedRequest + + +class ServiceConnectorInfo(BaseModel): + """Information about the service connector when creating a full stack.""" + + name: str + type: str + auth_type: str + configuration: Dict[str, Any] = {} + + +class ComponentInfo(BaseModel): + """Information about each stack components when creating a full stack.""" + + name: str + type: StackComponentType + flavor: str + service_connector: Optional[str] = None + configuration: Dict[str, Any] = {} + + +class FullStackRequest(BaseModel): + """Request model for a full-stack.""" + + name: str = Field( + title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH + ) + description: str = Field( + default="", + title="The description of the stack", + max_length=STR_FIELD_MAX_LENGTH, + ) + service_connector: Union[UUID, ServiceConnectorInfo] = Field( + title="The service connector for the full stack registration.", + description="The UUID of an already existing service connector or " + "request information to create a service connector from " + "scratch.", + ) + components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = Field( + title="The mapping for the components of the full stack registration.", + description="The mapping from component types to either UUIDs of " + "existing components or request information for brand new " + "components.", + ) + + +class FullStackScopedRequest(FullStackRequest, WorkspaceScopedRequest): + """Request model for a full-stack scoped by a user and workspace.""" From e60489e74c90126e846866ea95a3cdc0dd8ce599 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 09:27:30 +0200 Subject: [PATCH 07/71] new models --- src/zenml/models/__init__.py | 7 +------ src/zenml/models/v2/misc/full_stack.py | 8 +++----- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index 53ded3453b4..24efc00a984 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -327,10 +327,7 @@ ResourceTypeModel, ) from zenml.models.v2.misc.server_models import ServerDatabaseType, ServerModel -from zenml.models.v2.misc.full_stack import ( - FullStackRequest, - FullStackScopedRequest, -) +from zenml.models.v2.misc.full_stack import FullStackRequest from zenml.models.v2.core.trigger import ( TriggerRequest, TriggerFilter, @@ -410,7 +407,6 @@ FlavorResponseBody.model_rebuild() FlavorResponseMetadata.model_rebuild() FullStackRequest.model_rebuild() -FullStackScopedRequest.model_rebuild() LazyArtifactVersionResponse.model_rebuild() LazyRunMetadataResponse.model_rebuild() ModelResponseBody.model_rebuild() @@ -713,7 +709,6 @@ "ServiceConnectorRequirements", "ResourceTypeModel", "FullStackRequest", - "FullStackScopedRequest", "UserAuthModel", "ExternalUserModel", "BuildItem", diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index d9493c826af..ff295ccad16 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -20,7 +20,6 @@ from zenml.constants import STR_FIELD_MAX_LENGTH from zenml.enums import StackComponentType -from zenml.models.v2.base.scoped import WorkspaceScopedRequest class ServiceConnectorInfo(BaseModel): @@ -45,6 +44,9 @@ class ComponentInfo(BaseModel): class FullStackRequest(BaseModel): """Request model for a full-stack.""" + user_id: Optional[UUID] = None + workspace_id: Optional[UUID] = None + name: str = Field( title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH ) @@ -65,7 +67,3 @@ class FullStackRequest(BaseModel): "existing components or request information for brand new " "components.", ) - - -class FullStackScopedRequest(FullStackRequest, WorkspaceScopedRequest): - """Request model for a full-stack scoped by a user and workspace.""" From a293dc6a12efff0420aad62f391dcc858e232195 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 09:27:47 +0200 Subject: [PATCH 08/71] new constant --- src/zenml/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zenml/constants.py b/src/zenml/constants.py index 2f3ccc43015..00591662d58 100644 --- a/src/zenml/constants.py +++ b/src/zenml/constants.py @@ -346,6 +346,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int: EVENT_FLAVORS = "/event-flavors" EVENT_SOURCES = "/event-sources" FLAVORS = "/flavors" +FULL_STACK = "/full-stack" GET_OR_CREATE = "/get-or-create" GRAPH = "/graph" HEALTH = "/health" @@ -372,7 +373,6 @@ def handle_int_env_var(var: str, default: int = 0) -> int: SERVICE_CONNECTOR_RESOURCES = "/resources" SERVICE_CONNECTOR_TYPES = "/service_connector_types" SERVICE_CONNECTOR_VERIFY = "/verify" -SERVICE_CONNECTOR_RESOURCES = "/resources" MODELS = "/models" MODEL_VERSIONS = "/model_versions" MODEL_VERSION_ARTIFACTS = "/model_version_artifacts" From be18c347cf2d861f6c236ab29aca74ad8de21430 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 09:31:40 +0200 Subject: [PATCH 09/71] converting it into a request --- src/zenml/models/v2/misc/full_stack.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index ff295ccad16..3ad66597c7b 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -20,6 +20,7 @@ from zenml.constants import STR_FIELD_MAX_LENGTH from zenml.enums import StackComponentType +from zenml.models.v2.base.base import BaseRequest class ServiceConnectorInfo(BaseModel): @@ -41,7 +42,7 @@ class ComponentInfo(BaseModel): configuration: Dict[str, Any] = {} -class FullStackRequest(BaseModel): +class FullStackRequest(BaseRequest): """Request model for a full-stack.""" user_id: Optional[UUID] = None From 37e664f36a4bb4cd8051d1d4800adc2751f975d5 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 09:34:21 +0200 Subject: [PATCH 10/71] defining interface --- src/zenml/zen_stores/zen_store_interface.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index 61cac8050a6..06ab716325f 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -54,6 +54,7 @@ FlavorRequest, FlavorResponse, FlavorUpdate, + FullStackRequest, LogsResponse, ModelFilter, ModelRequest, @@ -2166,6 +2167,24 @@ def create_stack(self, stack: StackRequest) -> StackResponse: by this user in this workspace. """ + @abstractmethod + def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: + """Create a full stack. + + Args: + full_stack: The full stack configuration. + + Returns: + The created stack. + + Raises: + EntityExistsError: If a service connector with the same name + already exists. + StackComponentExistsError: If a stack component with the same name + already exists. + StackExistsError: If a stack with the same name already exists. + """ + @abstractmethod def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse: """Get a stack by its unique ID. From 092ea2469b81cedb1ec486b5ae09fa566a9b77dd Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 09:34:49 +0200 Subject: [PATCH 11/71] rest zen store --- src/zenml/zen_stores/rest_zen_store.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 406101fd269..b5146809d75 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -65,6 +65,7 @@ ENV_ZENML_DISABLE_CLIENT_SERVER_MISMATCH_WARNING, EVENT_SOURCES, FLAVORS, + FULL_STACK, GET_OR_CREATE, INFO, LOGIN, @@ -147,6 +148,7 @@ FlavorRequest, FlavorResponse, FlavorUpdate, + FullStackRequest, LogsResponse, ModelFilter, ModelRequest, @@ -2744,6 +2746,21 @@ def create_stack(self, stack: StackRequest) -> StackResponse: response_model=StackResponse, ) + def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: + """Register a full-stack. + + Args: + full_stack: The full stack configuration. + + Returns: + The registered stack. + """ + return self._create_resource( + resource=full_stack, + route=FULL_STACK, + response_model=StackResponse, + ) + def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse: """Get a stack by its unique ID. From 524e4aa05006234cc4f1482cdd3204f43c46598e Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 09:38:04 +0200 Subject: [PATCH 12/71] checking permissions and scoping the request in the endpoint --- .../routers/workspaces_endpoints.py | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/src/zenml/zen_server/routers/workspaces_endpoints.py b/src/zenml/zen_server/routers/workspaces_endpoints.py index 06820655ef3..f9cebe09c0e 100644 --- a/src/zenml/zen_server/routers/workspaces_endpoints.py +++ b/src/zenml/zen_server/routers/workspaces_endpoints.py @@ -22,6 +22,7 @@ API, ARTIFACTS, CODE_REPOSITORIES, + FULL_STACK, GET_OR_CREATE, MODEL_VERSIONS, MODELS, @@ -51,6 +52,7 @@ ComponentFilter, ComponentRequest, ComponentResponse, + FullStackRequest, ModelRequest, ModelResponse, ModelVersionArtifactRequest, @@ -349,6 +351,72 @@ def create_stack( ) +@router.post( + WORKSPACES + "/{workspace_name_or_id}" + FULL_STACK, + response_model=StackResponse, + responses={401: error_response, 409: error_response, 422: error_response}, +) +@handle_exceptions +def create_full_stack( + workspace_name_or_id: Union[str, UUID], + full_stack: FullStackRequest, + auth_context: AuthContext = Security(authorize), +) -> StackResponse: + """Creates a stack for a particular workspace. + + Args: + workspace_name_or_id: Name or ID of the workspace. + full_stack: Stack to register. + + Returns: + The created stack. + + Raises: + IllegalOperationError: If the workspace specified in the stack + does not match the current workspace. + """ + workspace = zen_store().get_workspace(workspace_name_or_id) + + if isinstance(full_stack.service_connector, UUID): + service_connector = zen_store().get_service_connector( + full_stack.service_connector + ) + verify_permission_for_model( + model=service_connector, action=Action.READ + ) + else: + verify_permission( + resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE + ) + + for component_type, component_id_or_request in full_stack.components: + if isinstance(component_id_or_request, UUID): + component = zen_store().get_stack_component( + full_stack.component_id_or_request + ) + verify_permission_for_model(model=component, action=Action.READ) + else: + verify_permission( + resource_type=ResourceType.STACK_COMPONENT, + action=Action.CREATE, + ) + + if "service_connector" in component_id_or_request: + verify_permission( + resource_type=ResourceType.STACK_COMPONENT, + action=Action.UPDATE, + ) + + verify_permission( + resource_type=ResourceType.STACK_COMPONENT, action=Action.CREATE + ) + + full_stack.user_id = auth_context.user.id + full_stack.workspace_id = workspace.id + + return zen_store().create_full_stack(full_stack) + + @router.get( WORKSPACES + "/{workspace_name_or_id}" + STACK_COMPONENTS, response_model=Page[ComponentResponse], From fa437faf6f3ac7f52feb4cf3939c661b3d5545b4 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 09:40:59 +0200 Subject: [PATCH 13/71] first version of the sql zen store --- src/zenml/zen_stores/sql_zen_store.py | 140 ++++++++++++++++++++++++++ 1 file changed, 140 insertions(+) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 867d80fc1f2..593ea15e8f9 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -172,6 +172,7 @@ FlavorRequest, FlavorResponse, FlavorUpdate, + FullStackRequest, LogsResponse, ModelFilter, ModelRequest, @@ -6919,6 +6920,145 @@ def create_stack(self, stack: StackRequest) -> StackResponse: return new_stack_schema.to_model(include_metadata=True) + @track_decorator(AnalyticsEvent.REGISTERED_STACK) + def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: + """Register a full stack. + + Args: + full_stack: The full stack configuration. + + Returns: + The registered stack. + """ + validate_name(full_stack) + + if isinstance(full_stack.service_connector, UUID): + service_connector = self.get_service_connector( + full_stack.service_connector + ) + else: + service_connector_request = ServiceConnectorRequest( + name=full_stack.service_connector.name, + flavor=full_stack.service_connector.flavor, + auth_type=full_stack.service_connector.auth_type, + configuration=full_stack.service_connector.configuration, + ) + service_connector = self.create_service_connector( + service_connector=service_connector_request + ) + + for component_type, component_info in full_stack.components: + if isinstance(component_info, UUID): + component = self.get_stack_component( + component_id=component_info + ) + else: + component_request = ComponentRequest( + name=component_info.name, + type=component_info.type, + flavor=component_info.flavor, + configuration=component_info.configuration, + ) + component = self.create_stack_component( + component=component_request + ) + if component_info.service_connector is not None: + flavor_list = self.list_flavors( + flavor_filter_model=FlavorFilter( + name=component_info.flavor, + type=component_info.type, + ) + ) + assert len(flavor_list) == 1 + + flavor_model = flavor_list[0] + + requirements = flavor_model.connector_requirements + + if not requirements: + raise RuntimeError( + f"The '{component_info.name}' implementation " + "does not support using a service connector to " + "connect to resources." + ) + + resource_id = None + resource_type = requirements.resource_type + if requirements.resource_id_attr is not None: + resource_id = component_info.configuration.get( + requirements.resource_id_attr + ) + + satisfied, msg = requirements.is_satisfied_by( + connector=service_connector, + component=component, + ) + + if not satisfied: + raise RuntimeError( + "Please pick a connector that is compatible with " + "the component flavor and try again, or use the " + "interactive mode to select a compatible connector." + ) + + if not resource_id: + if service_connector.resource_id: + resource_id = service_connector.resource_id + elif service_connector.supports_instances: + raise RuntimeError( + f"Multiple {resource_type} resources are " + "available for the selected connector. Please " + "use a `resource_id` to configure a " + f"{resource_type} resource." + ) + + try: + connector_resources = self.verify_service_connector( + service_connector_id=service_connector.id, + resource_type=requirements.resource_type, + resource_id=resource_id, + ) + except ( + KeyError, + ValueError, + IllegalOperationError, + NotImplementedError, + AuthorizationException, + ) as e: + raise RuntimeError( + f"Access to the resource could not be verified: {e}" + ) + + resources = connector_resources.resources[0] + if resources.resource_ids: + if len(resources.resource_ids) > 1: + raise RuntimeError( + f"Multiple {resource_type} resources are " + f"available for the selected connector. Please " + "use the a specific resource-id to configure a " + f"{resource_type} resource." + ) + else: + resource_id = resources.resource_ids[0] + + component_update = ComponentUpdate( + connector=service_connector.id, + connector_resource_id=resource_id, + ) + self.update_stack_component( + component_update=component_update + ) + + full_stack.components[component_type] = component.id + + stack_request = StackRequest( + name=full_stack.name, + description=full_stack.description, + component=full_stack.components, + ) + + return self.create_stack(stack_request) + def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse: """Get a stack by its unique ID. From a8bb05ceff0580237ae7a072fb91a977a4e7d0f7 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 28 Jun 2024 10:21:55 +0200 Subject: [PATCH 14/71] rbac fixes --- .../zen_server/routers/workspaces_endpoints.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/src/zenml/zen_server/routers/workspaces_endpoints.py b/src/zenml/zen_server/routers/workspaces_endpoints.py index f9cebe09c0e..f0f72a99ec8 100644 --- a/src/zenml/zen_server/routers/workspaces_endpoints.py +++ b/src/zenml/zen_server/routers/workspaces_endpoints.py @@ -379,7 +379,7 @@ def create_full_stack( if isinstance(full_stack.service_connector, UUID): service_connector = zen_store().get_service_connector( - full_stack.service_connector + full_stack.service_connector, hydrate=False ) verify_permission_for_model( model=service_connector, action=Action.READ @@ -389,10 +389,10 @@ def create_full_stack( resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE ) - for component_type, component_id_or_request in full_stack.components: + for component_id_or_request in full_stack.components.values(): if isinstance(component_id_or_request, UUID): component = zen_store().get_stack_component( - full_stack.component_id_or_request + component_id_or_request, hydrate=False ) verify_permission_for_model(model=component, action=Action.READ) else: @@ -401,15 +401,7 @@ def create_full_stack( action=Action.CREATE, ) - if "service_connector" in component_id_or_request: - verify_permission( - resource_type=ResourceType.STACK_COMPONENT, - action=Action.UPDATE, - ) - - verify_permission( - resource_type=ResourceType.STACK_COMPONENT, action=Action.CREATE - ) + verify_permission(resource_type=ResourceType.STACK, action=Action.CREATE) full_stack.user_id = auth_context.user.id full_stack.workspace_id = workspace.id From 62635a8434592dd2642f6bc7fac6dab2861e2451 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 28 Jun 2024 10:52:48 +0200 Subject: [PATCH 15/71] misc fixes --- src/zenml/cli/utils.py | 8 +++----- src/zenml/models/v2/misc/full_stack.py | 8 ++++---- src/zenml/zen_stores/sql_zen_store.py | 21 +++++++++++++++------ 3 files changed, 22 insertions(+), 15 deletions(-) diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index ad885f97350..93683a51e19 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -44,6 +44,7 @@ import yaml from pydantic import BaseModel, SecretStr from rich import box, table +from rich.console import Console from rich.emoji import Emoji, NoEmoji from rich.markdown import Markdown from rich.markup import escape @@ -2764,8 +2765,6 @@ def show_status_from_kwargs(default_value: str = ":x:", **kwargs: Any) -> None: **kwargs: The kwargs to show status from. If value is passed, but `None` a default value is used. """ - from rich import print - status = [] names = [] for name, each in kwargs.items(): @@ -2783,7 +2782,7 @@ def show_status_from_kwargs(default_value: str = ":x:", **kwargs: Any) -> None: status_table.add_column(c, justify="center", width=1) status_table.add_row(*status) - print(status_table) + Console().print(status_table) def multi_choice_prompt( @@ -2807,7 +2806,6 @@ def multi_choice_prompt( Returns: The selected choice index or None for new object """ - from rich import print table = Table( title=f"Available {object_type}", @@ -2831,7 +2829,7 @@ def multi_choice_prompt( ) for i, one_choice in enumerate(choices): table.add_row(f"[{i+i_shift}]", *[str(x) for x in one_choice]) - print(table) + Console().print(table) selected = Prompt.ask( prompt_text, diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 3ad66597c7b..0e7ac25e766 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -27,7 +27,7 @@ class ServiceConnectorInfo(BaseModel): """Information about the service connector when creating a full stack.""" name: str - type: str + connector_type: str auth_type: str configuration: Dict[str, Any] = {} @@ -36,7 +36,6 @@ class ComponentInfo(BaseModel): """Information about each stack components when creating a full stack.""" name: str - type: StackComponentType flavor: str service_connector: Optional[str] = None configuration: Dict[str, Any] = {} @@ -51,12 +50,13 @@ class FullStackRequest(BaseRequest): name: str = Field( title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH ) - description: str = Field( + description: Optional[str] = Field( default="", title="The description of the stack", max_length=STR_FIELD_MAX_LENGTH, ) - service_connector: Union[UUID, ServiceConnectorInfo] = Field( + service_connector: Optional[Union[UUID, ServiceConnectorInfo]] = Field( + default=None, title="The service connector for the full stack registration.", description="The UUID of an already existing service connector or " "request information to create a service connector from " diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 593ea15e8f9..6520fa59903 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6939,15 +6939,18 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: else: service_connector_request = ServiceConnectorRequest( name=full_stack.service_connector.name, - flavor=full_stack.service_connector.flavor, + connector_type=full_stack.service_connector.connector_type, auth_type=full_stack.service_connector.auth_type, configuration=full_stack.service_connector.configuration, + user=full_stack.user_id, + workspace=full_stack.workspace_id, ) service_connector = self.create_service_connector( service_connector=service_connector_request ) - for component_type, component_info in full_stack.components: + components_mapping: Dict[StackComponentType, List[UUID]] = {} + for component_type, component_info in full_stack.components.items(): if isinstance(component_info, UUID): component = self.get_stack_component( component_id=component_info @@ -6955,9 +6958,11 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: else: component_request = ComponentRequest( name=component_info.name, - type=component_info.type, + type=component_type, flavor=component_info.flavor, configuration=component_info.configuration, + user=full_stack.user_id, + workspace=full_stack.workspace_id, ) component = self.create_stack_component( component=component_request @@ -6966,7 +6971,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: flavor_list = self.list_flavors( flavor_filter_model=FlavorFilter( name=component_info.flavor, - type=component_info.type, + type=component_type, ) ) assert len(flavor_list) == 1 @@ -7049,12 +7054,16 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: component_update=component_update ) - full_stack.components[component_type] = component.id + components_mapping[component_type] = [ + component.id, + ] stack_request = StackRequest( + user=full_stack.user_id, + workspace=full_stack.workspace_id, name=full_stack.name, description=full_stack.description, - component=full_stack.components, + components=components_mapping, ) return self.create_stack(stack_request) From 8c0b71d8a1a02987bc38cdb813a781e11b4c5a22 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 28 Jun 2024 10:57:26 +0200 Subject: [PATCH 16/71] remove names duplicates --- src/zenml/models/v2/misc/full_stack.py | 2 -- src/zenml/zen_stores/sql_zen_store.py | 4 ++-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 0e7ac25e766..1ee253f7dac 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -26,7 +26,6 @@ class ServiceConnectorInfo(BaseModel): """Information about the service connector when creating a full stack.""" - name: str connector_type: str auth_type: str configuration: Dict[str, Any] = {} @@ -35,7 +34,6 @@ class ServiceConnectorInfo(BaseModel): class ComponentInfo(BaseModel): """Information about each stack components when creating a full stack.""" - name: str flavor: str service_connector: Optional[str] = None configuration: Dict[str, Any] = {} diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 6520fa59903..990fdff7ba0 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6938,7 +6938,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: ) else: service_connector_request = ServiceConnectorRequest( - name=full_stack.service_connector.name, + name=full_stack.name, # try and fail, then randomize connector_type=full_stack.service_connector.connector_type, auth_type=full_stack.service_connector.auth_type, configuration=full_stack.service_connector.configuration, @@ -6957,7 +6957,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: ) else: component_request = ComponentRequest( - name=component_info.name, + name=full_stack.name, # try and fail, then randomize type=component_type, flavor=component_info.flavor, configuration=component_info.configuration, From 3402bf86c7144598893d6c81756a9d44a440143b Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 28 Jun 2024 11:20:39 +0200 Subject: [PATCH 17/71] multiple service connectors support --- src/zenml/models/v2/misc/full_stack.py | 23 +++++++---- .../routers/workspaces_endpoints.py | 37 ++++++++++------- src/zenml/zen_stores/sql_zen_store.py | 41 +++++++++++-------- 3 files changed, 61 insertions(+), 40 deletions(-) diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 1ee253f7dac..d309d5a5dd6 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Models representing full stack requests.""" -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, List, Optional, Union from uuid import UUID from pydantic import BaseModel, Field @@ -35,7 +35,12 @@ class ComponentInfo(BaseModel): """Information about each stack components when creating a full stack.""" flavor: str - service_connector: Optional[str] = None + service_connector_index: Optional[int] = Field( + default=None, + title="The id of the service connector from the list `service_connectors`.", + description="The id of the service connector from the list " + "`service_connectors` from `FullStackRequest`.", + ) configuration: Dict[str, Any] = {} @@ -53,12 +58,14 @@ class FullStackRequest(BaseRequest): title="The description of the stack", max_length=STR_FIELD_MAX_LENGTH, ) - service_connector: Optional[Union[UUID, ServiceConnectorInfo]] = Field( - default=None, - title="The service connector for the full stack registration.", - description="The UUID of an already existing service connector or " - "request information to create a service connector from " - "scratch.", + service_connectors: Optional[List[Union[UUID, ServiceConnectorInfo]]] = ( + Field( + default=[], + title="The service connectors dictionary for the full stack registration.", + description="The UUID of an already existing service connector or " + "request information to create a service connector from " + "scratch.", + ) ) components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = Field( title="The mapping for the components of the full stack registration.", diff --git a/src/zenml/zen_server/routers/workspaces_endpoints.py b/src/zenml/zen_server/routers/workspaces_endpoints.py index f0f72a99ec8..219e0a9bc0a 100644 --- a/src/zenml/zen_server/routers/workspaces_endpoints.py +++ b/src/zenml/zen_server/routers/workspaces_endpoints.py @@ -377,29 +377,36 @@ def create_full_stack( """ workspace = zen_store().get_workspace(workspace_name_or_id) - if isinstance(full_stack.service_connector, UUID): - service_connector = zen_store().get_service_connector( - full_stack.service_connector, hydrate=False - ) - verify_permission_for_model( - model=service_connector, action=Action.READ - ) - else: + is_connector_create_needed = False + for connector_id_or_info in full_stack.service_connectors: + if isinstance(connector_id_or_info, UUID): + service_connector = zen_store().get_service_connector( + connector_id_or_info, hydrate=False + ) + verify_permission_for_model( + model=service_connector, action=Action.READ + ) + else: + is_connector_create_needed = True + if is_connector_create_needed: verify_permission( resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE ) - for component_id_or_request in full_stack.components.values(): - if isinstance(component_id_or_request, UUID): + is_component_create_needed = False + for component_id_or_info in full_stack.components.values(): + if isinstance(component_id_or_info, UUID): component = zen_store().get_stack_component( - component_id_or_request, hydrate=False + component_id_or_info, hydrate=False ) verify_permission_for_model(model=component, action=Action.READ) else: - verify_permission( - resource_type=ResourceType.STACK_COMPONENT, - action=Action.CREATE, - ) + is_component_create_needed = True + if is_component_create_needed: + verify_permission( + resource_type=ResourceType.STACK_COMPONENT, + action=Action.CREATE, + ) verify_permission(resource_type=ResourceType.STACK, action=Action.CREATE) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 990fdff7ba0..ca13b3ded14 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6932,22 +6932,26 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: """ validate_name(full_stack) - if isinstance(full_stack.service_connector, UUID): - service_connector = self.get_service_connector( - full_stack.service_connector - ) - else: - service_connector_request = ServiceConnectorRequest( - name=full_stack.name, # try and fail, then randomize - connector_type=full_stack.service_connector.connector_type, - auth_type=full_stack.service_connector.auth_type, - configuration=full_stack.service_connector.configuration, - user=full_stack.user_id, - workspace=full_stack.workspace_id, - ) - service_connector = self.create_service_connector( - service_connector=service_connector_request - ) + service_connectors: List[ServiceConnectorResponse] = [] + for connector_id_or_info in full_stack.service_connectors: + if isinstance(connector_id_or_info, UUID): + service_connectors.append( + self.get_service_connector(connector_id_or_info) + ) + else: + service_connector_request = ServiceConnectorRequest( + name=full_stack.name, # try and fail, then randomize + connector_type=connector_id_or_info.connector_type, + auth_type=connector_id_or_info.auth_type, + configuration=connector_id_or_info.configuration, + user=full_stack.user_id, + workspace=full_stack.workspace_id, + ) + service_connectors.append( + self.create_service_connector( + service_connector=service_connector_request + ) + ) components_mapping: Dict[StackComponentType, List[UUID]] = {} for component_type, component_info in full_stack.components.items(): @@ -6967,7 +6971,10 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: component = self.create_stack_component( component=component_request ) - if component_info.service_connector is not None: + if component_info.service_connector_index is not None: + service_connector = service_connectors[ + component_info.service_connector_index + ] flavor_list = self.list_flavors( flavor_filter_model=FlavorFilter( name=component_info.flavor, From 8cd5c0297005e0cbef838e919c37485bcc48d5c9 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 28 Jun 2024 12:03:25 +0200 Subject: [PATCH 18/71] maybe working --- src/zenml/cli/stack.py | 329 +++++++++++++++++++++-------------------- 1 file changed, 167 insertions(+), 162 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index da47db4d2fc..aa292055bac 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -53,10 +53,13 @@ from zenml.io.fileio import rmtree from zenml.logger import get_logger from zenml.models import StackFilter -from zenml.models.v2.core.component import ComponentResponse -from zenml.models.v2.core.service_connector import ServiceConnectorResponse +from zenml.models.v2.misc.full_stack import ( + ComponentInfo, + FullStackRequest, + ServiceConnectorInfo, +) from zenml.models.v2.misc.service_connector_type import ( - ServiceConnectorResourcesModel, + ServiceConnectorTypedResourcesModel, ) from zenml.utils.dashboard_utils import get_stack_url from zenml.utils.io_utils import create_dir_recursive_if_not_exists @@ -259,6 +262,18 @@ def register_stack( client = Client() + try: + client.get_stack(name_id_or_prefix=stack_name) + cli_utils.error( + f"A stack with name `{stack_name}` already exists, " + "please use a different name." + ) + except KeyError: + pass + except Exception as e: + raise e + + components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = {} # cloud flow service_connector = None if cloud is not None and connector is None: @@ -286,9 +301,17 @@ def register_stack( allow_zero_be_a_new_object=True, ) if connector_selected is None: - service_connector = _create_service_connector(cloud_provider=cloud) + service_connector = _get_service_connector_info( + cloud_provider=cloud + ) else: - service_connector = existing_connectors.items[connector_selected] + selected_connector = existing_connectors.items[connector_selected] + service_connector = selected_connector.id + connector = selected_connector.name + if isinstance(selected_connector.connector_type, str): + cloud = selected_connector.connector_type + else: + cloud = selected_connector.connector_type.connector_type elif connector is not None: cli_utils.show_status_from_kwargs( cloud=cloud, @@ -304,7 +327,7 @@ def register_stack( ) cli_utils.show_status_from_kwargs( cloud=cloud, - connector=service_connector.name, + connector=connector, artifact_store=artifact_store, orchestrator=orchestrator, container_registry=container_registry, @@ -314,69 +337,84 @@ def register_stack( service_connector_resource_model = None # create components needed_components = ( - ("artifact_store", artifact_store), - ("orchestrator", orchestrator), # for azure only k8s orchestrator - ("container_registry", container_registry), + (StackComponentType.ARTIFACT_STORE, artifact_store), + (StackComponentType.ORCHESTRATOR, orchestrator), + (StackComponentType.CONTAINER_REGISTRY, container_registry), ) for component_type, preset_name in needed_components: if preset_name is not None: - component_response = client.get_stack_component( + component_info = client.get_stack_component( component_type, preset_name ) - if component_response.connector.id != service_connector.id: - cli_utils.error( - f"The {component_type.replace('_', ' ')} and service connector " - "do not match. Please check your inputs and try again." - ) + component_info = component_info.id else: - # find existing components under same connector - existing_components = client.list_stack_components( - type=component_type, - connector_id=service_connector.id, - size=100, - ) - # if some existing components are found - prompt user what to do - component_selected: Optional[int] = None - if existing_components.total > 0: - component_selected = cli_utils.multi_choice_prompt( - object_type=component_type.replace("_", " "), - choices=[ - [component.name] - for component in existing_components.items - ], - headers=["Name"], - prompt_text=f"We found these {component_type.replace('_', ' ')} " - "connected using the current service connector. Do you " - "want to create a new one or use existing one?", - default_choice="0", - allow_zero_be_a_new_object=True, + if isinstance(service_connector, UUID): + # find existing components under same connector + existing_components = client.list_stack_components( + type=component_type.value, + connector_id=service_connector, + size=100, ) + # if some existing components are found - prompt user what to do + component_selected: Optional[int] = None + if existing_components.total > 0: + component_selected = cli_utils.multi_choice_prompt( + object_type=component_type.value.replace("_", " "), + choices=[ + [component.name] + for component in existing_components.items + ], + headers=["Name"], + prompt_text=f"We found these {component_type.value.replace('_', ' ')} " + "connected using the current service connector. Do you " + "want to create a new one or use existing one?", + default_choice="0", + allow_zero_be_a_new_object=True, + ) + else: + component_selected = None + if component_selected is None: if service_connector_resource_model is None: - service_connector_resource_model = ( - client.verify_service_connector( - service_connector.id + if isinstance(service_connector, UUID): + service_connector_resource_model = ( + client.verify_service_connector( + service_connector + ) ) - ) - component_response = _create_stack_component( - component_type, - service_connector, - service_connector_resource_model, + else: + service_connector_resource_model = client.create_service_connector( + name=service_connector.name, + type=service_connector.connector_type, + auth_method=service_connector.auth_method, + configuration=service_connector.configuration, + register=False, + ) + + component_info = _get_stack_component_info( + component_type=component_type.value, + cloud_provider=cloud, + service_connector_resource_models=service_connector_resource_model.resources, + service_connector_index=0, ) + component_name = stack_name else: - component_response = existing_components.items[ + selected_component = existing_components.items[ component_selected ] - - if component_type == "orchestrator": - orchestrator = component_response.name - elif component_type == "artifact_store": - artifact_store = component_response.name - elif component_type == "container_registry": - container_registry = component_response.name + component_info = selected_component.id + component_name = selected_component.name + + components[component_type] = component_info + if component_type == StackComponentType.ARTIFACT_STORE: + artifact_store = component_name + if component_type == StackComponentType.ORCHESTRATOR: + orchestrator = component_name + if component_type == StackComponentType.CONTAINER_REGISTRY: + container_registry = component_name cli_utils.show_status_from_kwargs( cloud=cloud, - connector=service_connector.name, + connector=connector, artifact_store=artifact_store, orchestrator=orchestrator, container_registry=container_registry, @@ -384,40 +422,36 @@ def register_stack( # normal flow once all components are defined with console.status(f"Registering stack '{stack_name}'...\n"): - components: Dict[StackComponentType, Union[str, UUID]] = dict() - - components[StackComponentType.ARTIFACT_STORE] = artifact_store - components[StackComponentType.ORCHESTRATOR] = orchestrator - - if alerter: - components[StackComponentType.ALERTER] = alerter - if annotator: - components[StackComponentType.ANNOTATOR] = annotator - if data_validator: - components[StackComponentType.DATA_VALIDATOR] = data_validator - if feature_store: - components[StackComponentType.FEATURE_STORE] = feature_store - if image_builder: - components[StackComponentType.IMAGE_BUILDER] = image_builder - if model_deployer: - components[StackComponentType.MODEL_DEPLOYER] = model_deployer - if model_registry: - components[StackComponentType.MODEL_REGISTRY] = model_registry - if step_operator: - components[StackComponentType.STEP_OPERATOR] = step_operator - if experiment_tracker: - components[StackComponentType.EXPERIMENT_TRACKER] = ( - experiment_tracker - ) - if container_registry: - components[StackComponentType.CONTAINER_REGISTRY] = ( - container_registry - ) + for component_type, component_name in [ + (StackComponentType.ARTIFACT_STORE, artifact_store), + (StackComponentType.ORCHESTRATOR, orchestrator), + (StackComponentType.ALERTER, alerter), + (StackComponentType.ANNOTATOR, annotator), + (StackComponentType.DATA_VALIDATOR, data_validator), + (StackComponentType.FEATURE_STORE, feature_store), + (StackComponentType.IMAGE_BUILDER, image_builder), + (StackComponentType.MODEL_DEPLOYER, model_deployer), + (StackComponentType.MODEL_REGISTRY, model_registry), + (StackComponentType.STEP_OPERATOR, step_operator), + (StackComponentType.EXPERIMENT_TRACKER, experiment_tracker), + (StackComponentType.CONTAINER_REGISTRY, container_registry), + ]: + if component_name and component_type not in components: + components[component_type] = client.get_stack_component( + component_type, component_name + ).id try: - created_stack = client.create_stack( - name=stack_name, - components=components, + created_stack = client.zen_store.create_full_stack( + full_stack=FullStackRequest( + user_id=client.active_user.id, + workspace_id=client.active_workspace.id, + name=stack_name, + components=components, + service_connectors=[ + service_connector, + ], + ) ) except (KeyError, IllegalOperationError) as err: cli_utils.error(str(err)) @@ -1888,14 +1922,14 @@ def connect_stack( ) -def _create_service_connector(cloud_provider: str) -> ServiceConnectorResponse: - """Create a service connector with given cloud provider. +def _get_service_connector_info(cloud_provider: str) -> ServiceConnectorInfo: + """Get a service connector info with given cloud provider. Args: cloud_provider: The cloud provider to use. Returns: - The model of the created service connector. + The info model of the created service connector. """ from rich import print from rich.markdown import Markdown @@ -1954,58 +1988,59 @@ def _create_service_connector(cloud_provider: str) -> ServiceConnectorResponse: default=True, ) - connector_name = Prompt.ask( - "Please enter a name for the service connector:" - ) - return client.create_service_connector( - name=connector_name, + return ServiceConnectorInfo( connector_type=cloud_provider, auth_method=fixed_auth_methods[selected_auth_idx][0], configuration=answers, - )[0] + ) -def _create_stack_component( +def _get_stack_component_info( component_type: str, - service_connector: ServiceConnectorResponse, - service_connector_resource_model: ServiceConnectorResourcesModel, -) -> ComponentResponse: - """Create a stack component with given type and service connector. + cloud_provider: str, + service_connector_resource_models: List[ + ServiceConnectorTypedResourcesModel + ], + service_connector_index: Optional[int] = None, +) -> ComponentInfo: + """Get a stack component info with given type and service connector. Args: component_type: The type of component to create. - service_connector: The service connector to use. + cloud_provider: The cloud provider to use. + service_connector_resource_models: The list of the available service connector resource models. + service_connector_index: The index of the service connector to use. Returns: - The model of the created component. + The info model of the stack component. + + Raises: + ValueError: If the cloud provider is not supported. + ValueError: If the component type is not supported. """ from rich import print from rich.prompt import Confirm, Prompt - from zenml.cli.stack_components import ( - connect_stack_component_with_service_connector, - ) - - if service_connector.type not in {"aws", "azure", "gcp"}: - raise ValueError(f"Unknown cloud provider {service_connector.type}") - - client = Client() + if cloud_provider not in {"aws", "azure", "gcp"}: + raise ValueError(f"Unknown cloud provider {cloud_provider}") + flavor = "undefined" + config = {} if component_type == "artifact_store": config_confirmed = False while not config_confirmed: - if service_connector.type == "aws": - for each in service_connector_resource_model.resources: + if cloud_provider == "aws": + for each in service_connector_resource_models: if each.resource_type == "s3-bucket": available_storages = each.resource_ids flavor = "s3" - elif service_connector.type == "azure": + elif cloud_provider == "azure": flavor = "azure" - elif service_connector.type == "gcp": + elif cloud_provider == "gcp": flavor = "gcs" selected_storage_idx = cli_utils.multi_choice_prompt( - object_type=f"{service_connector.type.upper()} storages", + object_type=f"{cloud_provider.upper()} storages", choices=[[st] for st in available_storages], headers=["Storage"], prompt_text="Please choose one of the storages for the new artifact store:", @@ -2020,24 +2055,16 @@ def _create_stack_component( "path": f"{selected_storage.strip('/')}/{extra_path.strip('/')}" } - as_name = Prompt.ask("Please enter a name for the artifact store:") - - print({"name": as_name, "config": config}) + print(config) config_confirmed = Confirm.ask( "Please confirm the values you entered:", default=True ) - component = client.create_stack_component( - name=as_name, - flavor=flavor, - component_type=StackComponentType.ARTIFACT_STORE, - configuration=config, - ) elif component_type == "orchestrator": config_confirmed = False while not config_confirmed: - if service_connector.type == "aws": + if cloud_provider == "aws": available_orchestrators = [] - for each in service_connector_resource_model.resources: + for each in service_connector_resource_models: types = [] if each.resource_type == "aws-generic": types = ["Sagemaker", "VM AWS"] @@ -2047,13 +2074,13 @@ def _create_stack_component( for orchestrator in each.resource_ids: for t in types: available_orchestrators.append([t, orchestrator]) - elif service_connector.type == "gcp": + elif cloud_provider == "gcp": pass - elif service_connector.type == "azure": + elif cloud_provider == "azure": pass selected_orchestrator_idx = cli_utils.multi_choice_prompt( - object_type=f"orchestrators on {service_connector.type.upper()}", + object_type=f"orchestrators on {cloud_provider.upper()}", choices=available_orchestrators, headers=["Orchestrator Type", "Orchestrator details"], prompt_text="Please choose one of the orchestrators for the new orchestrator:", @@ -2079,36 +2106,27 @@ def _create_stack_component( raise ValueError( f"Unknown orchestrator type {selected_orchestrator[0]}" ) - orchestrator_name = Prompt.ask( - "Please enter a name for the orchestrator:" - ) - print({"name": orchestrator_name, "config": config}) + print(config) config_confirmed = Confirm.ask( "Please confirm the values you entered:", default=True ) - component = client.create_stack_component( - name=orchestrator_name, - flavor=flavor, - component_type=StackComponentType.ORCHESTRATOR, - configuration=config, - ) elif component_type == "container_registry": config_confirmed = False while not config_confirmed: - if service_connector.type == "aws": - for each in service_connector_resource_model.resources: + if cloud_provider == "aws": + for each in service_connector_resource_models: if each.resource_type == "docker-registry": available_registries = each.resource_ids flavor = "aws" - elif service_connector.type == "azure": + elif cloud_provider == "azure": flavor = "azure" available_registries = [] - elif service_connector.type == "gcp": + elif cloud_provider == "gcp": flavor = "gcp" available_registries = [] selected_registry_idx = cli_utils.multi_choice_prompt( - object_type=f"{service_connector.type.upper()} registries", + object_type=f"{cloud_provider.upper()} registries", choices=[[st] for st in available_registries], headers=["Container Registry"], prompt_text="Please choose one of the registries for the new container registry:", @@ -2116,29 +2134,16 @@ def _create_stack_component( selected_storage = available_registries[selected_registry_idx] config = {"uri": selected_storage} - cr_name = Prompt.ask( - "Please enter a name for the container registry:" - ) - - print({"name": cr_name, "config": config}) + print(config) config_confirmed = Confirm.ask( "Please confirm the values you entered:", default=True ) - component = client.create_stack_component( - name=cr_name, - flavor=flavor, - component_type=StackComponentType.CONTAINER_REGISTRY, - configuration=config, - ) else: raise ValueError(f"Unknown component type {component_type}") + service_connector_index - connect_stack_component_with_service_connector( - component_type=component.type, - name_id_or_prefix=component.id, - connector=service_connector.id, - interactive=False, - no_verify=False, + return ComponentInfo( + flavor=flavor, + configuration=config, + service_connector_index=service_connector_index, ) - - return component From a0530b82102e88f69251ae240aa84f46261b89b1 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 28 Jun 2024 13:01:12 +0200 Subject: [PATCH 19/71] likely to work --- src/zenml/cli/stack.py | 18 ++++++++++-------- src/zenml/zen_stores/sql_zen_store.py | 2 +- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index aa292055bac..7cf4ef101ab 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -320,7 +320,7 @@ def register_stack( orchestrator=orchestrator, container_registry=container_registry, ) - service_connector = client.get_service_connector(connector) + service_connector = client.get_service_connector(connector).id if service_connector.type != cloud: cli_utils.warning( f"The service connector `{connector}` is not of type `{cloud}`." @@ -383,12 +383,14 @@ def register_stack( ) ) else: - service_connector_resource_model = client.create_service_connector( - name=service_connector.name, - type=service_connector.connector_type, - auth_method=service_connector.auth_method, - configuration=service_connector.configuration, - register=False, + _, service_connector_resource_model = ( + client.create_service_connector( + name=stack_name, + connector_type=service_connector.connector_type, + auth_method=service_connector.auth_type, + configuration=service_connector.configuration, + register=False, + ) ) component_info = _get_stack_component_info( @@ -1990,7 +1992,7 @@ def _get_service_connector_info(cloud_provider: str) -> ServiceConnectorInfo: return ServiceConnectorInfo( connector_type=cloud_provider, - auth_method=fixed_auth_methods[selected_auth_idx][0], + auth_type=fixed_auth_methods[selected_auth_idx][0], configuration=answers, ) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index ca13b3ded14..73440e49885 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6942,7 +6942,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: service_connector_request = ServiceConnectorRequest( name=full_stack.name, # try and fail, then randomize connector_type=connector_id_or_info.connector_type, - auth_type=connector_id_or_info.auth_type, + auth_method=connector_id_or_info.auth_type, configuration=connector_id_or_info.configuration, user=full_stack.user_id, workspace=full_stack.workspace_id, From 4c2e37fbcc3eca559e1b4805758935446b169e66 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 14:32:21 +0200 Subject: [PATCH 20/71] some fixes and naming checks --- src/zenml/cli/utils.py | 3 +- src/zenml/models/v2/misc/full_stack.py | 6 +- .../routers/workspaces_endpoints.py | 1 + src/zenml/zen_stores/sql_zen_store.py | 71 ++++++++++++------- 4 files changed, 52 insertions(+), 29 deletions(-) diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index 93683a51e19..c17049ef731 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -2799,14 +2799,13 @@ def multi_choice_prompt( object_type: The type of the object choices: The list of choices prompt_text: The prompt text - selector_from_choices: The list of selectors to use + headers: The list of headers. allow_zero_be_a_new_object: Whether to allow zero as a new object default_choice: The default choice Returns: The selected choice index or None for new object """ - table = Table( title=f"Available {object_type}", show_header=True, diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index d309d5a5dd6..90b16e93e91 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -37,7 +37,8 @@ class ComponentInfo(BaseModel): flavor: str service_connector_index: Optional[int] = Field( default=None, - title="The id of the service connector from the list `service_connectors`.", + title="The id of the service connector from the list " + "`service_connectors`.", description="The id of the service connector from the list " "`service_connectors` from `FullStackRequest`.", ) @@ -61,7 +62,8 @@ class FullStackRequest(BaseRequest): service_connectors: Optional[List[Union[UUID, ServiceConnectorInfo]]] = ( Field( default=[], - title="The service connectors dictionary for the full stack registration.", + title="The service connectors dictionary for the full stack " + "registration.", description="The UUID of an already existing service connector or " "request information to create a service connector from " "scratch.", diff --git a/src/zenml/zen_server/routers/workspaces_endpoints.py b/src/zenml/zen_server/routers/workspaces_endpoints.py index 219e0a9bc0a..be858eff467 100644 --- a/src/zenml/zen_server/routers/workspaces_endpoints.py +++ b/src/zenml/zen_server/routers/workspaces_endpoints.py @@ -367,6 +367,7 @@ def create_full_stack( Args: workspace_name_or_id: Name or ID of the workspace. full_stack: Stack to register. + auth_context: Authentication context. Returns: The created stack. diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 73440e49885..6364ad6c9ec 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6939,19 +6939,29 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: self.get_service_connector(connector_id_or_info) ) else: - service_connector_request = ServiceConnectorRequest( - name=full_stack.name, # try and fail, then randomize - connector_type=connector_id_or_info.connector_type, - auth_method=connector_id_or_info.auth_type, - configuration=connector_id_or_info.configuration, - user=full_stack.user_id, - workspace=full_stack.workspace_id, - ) - service_connectors.append( - self.create_service_connector( - service_connector=service_connector_request - ) - ) + connector_name = full_stack.name + while True: + try: + service_connector_request = ServiceConnectorRequest( + name=connector_name, + connector_type=connector_id_or_info.connector_type, + auth_method=connector_id_or_info.auth_type, + configuration=connector_id_or_info.configuration, + user=full_stack.user_id, + workspace=full_stack.workspace_id, + ) + service_connectors.append( + self.create_service_connector( + service_connector=service_connector_request + ) + ) + break + except EntityExistsError: + connector_name = ( + f"{full_stack.name}-{random_str(4)}".lower() + ) + + continue components_mapping: Dict[StackComponentType, List[UUID]] = {} for component_type, component_info in full_stack.components.items(): @@ -6960,17 +6970,27 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: component_id=component_info ) else: - component_request = ComponentRequest( - name=full_stack.name, # try and fail, then randomize - type=component_type, - flavor=component_info.flavor, - configuration=component_info.configuration, - user=full_stack.user_id, - workspace=full_stack.workspace_id, - ) - component = self.create_stack_component( - component=component_request - ) + component_name = full_stack.name + while True: + try: + component_request = ComponentRequest( + name=component_name, + type=component_type, + flavor=component_info.flavor, + configuration=component_info.configuration, + user=full_stack.user_id, + workspace=full_stack.workspace_id, + ) + component = self.create_stack_component( + component=component_request + ) + break + except EntityExistsError: + component_name = ( + f"{full_stack.name}-{random_str(4)}".lower() + ) + continue + if component_info.service_connector_index is not None: service_connector = service_connectors[ component_info.service_connector_index @@ -7058,7 +7078,8 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: connector_resource_id=resource_id, ) self.update_stack_component( - component_update=component_update + component_id=component.id, + component_update=component_update, ) components_mapping[component_type] = [ From 8b847baea48ed8e55e60afb567bb05145bf29538 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 14:40:38 +0200 Subject: [PATCH 21/71] linting --- src/zenml/models/v2/misc/full_stack.py | 16 +++++++--------- src/zenml/zen_stores/sql_zen_store.py | 2 +- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 90b16e93e91..5d21ec6c5af 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -59,15 +59,13 @@ class FullStackRequest(BaseRequest): title="The description of the stack", max_length=STR_FIELD_MAX_LENGTH, ) - service_connectors: Optional[List[Union[UUID, ServiceConnectorInfo]]] = ( - Field( - default=[], - title="The service connectors dictionary for the full stack " - "registration.", - description="The UUID of an already existing service connector or " - "request information to create a service connector from " - "scratch.", - ) + service_connectors: List[Union[UUID, ServiceConnectorInfo]] = Field( + default=[], + title="The service connectors dictionary for the full stack " + "registration.", + description="The UUID of an already existing service connector or " + "request information to create a service connector from " + "scratch.", ) components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = Field( title="The mapping for the components of the full stack registration.", diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 6364ad6c9ec..7c9ecaa1e47 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -7009,7 +7009,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: if not requirements: raise RuntimeError( - f"The '{component_info.name}' implementation " + f"The '{flavor_model.name}' implementation " "does not support using a service connector to " "connect to resources." ) From 0801c5925e8d25eb7018e0297fe71f43b9576d5f Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 16:17:23 +0200 Subject: [PATCH 22/71] analytics, name checks for stacks and cleanup --- src/zenml/zen_stores/sql_zen_store.py | 342 +++++++++++++++----------- 1 file changed, 195 insertions(+), 147 deletions(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 7c9ecaa1e47..60d7653f618 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6920,7 +6920,6 @@ def create_stack(self, stack: StackRequest) -> StackResponse: return new_stack_schema.to_model(include_metadata=True) - @track_decorator(AnalyticsEvent.REGISTERED_STACK) def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: """Register a full stack. @@ -6930,171 +6929,220 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: Returns: The registered stack. """ - validate_name(full_stack) + # We can not use the decorator here, as we need custom metadata + with track_handler( + event=AnalyticsEvent.REGISTERED_STACK, + metadata={"generated_by_wizard": True}, + ) as handler: + # For clean-up purposes, each created entity is tracked here + try: + # Validate the name of the new stack + validate_name(full_stack) - service_connectors: List[ServiceConnectorResponse] = [] - for connector_id_or_info in full_stack.service_connectors: - if isinstance(connector_id_or_info, UUID): - service_connectors.append( - self.get_service_connector(connector_id_or_info) - ) - else: - connector_name = full_stack.name - while True: - try: - service_connector_request = ServiceConnectorRequest( - name=connector_name, - connector_type=connector_id_or_info.connector_type, - auth_method=connector_id_or_info.auth_type, - configuration=connector_id_or_info.configuration, - user=full_stack.user_id, - workspace=full_stack.workspace_id, - ) + # Service Connectors + service_connectors: List[ServiceConnectorResponse] = [] + service_connectors_created_ids: List[UUID] = [] + + for connector_id_or_info in full_stack.service_connectors: + # Fetch an existing service connector + if isinstance(connector_id_or_info, UUID): service_connectors.append( - self.create_service_connector( - service_connector=service_connector_request - ) + self.get_service_connector(connector_id_or_info) ) - break - except EntityExistsError: - connector_name = ( - f"{full_stack.name}-{random_str(4)}".lower() + # Create a new service connector + else: + connector_name = full_stack.name + while True: + try: + service_connector_request = ServiceConnectorRequest( + name=connector_name, + connector_type=connector_id_or_info.connector_type, + auth_method=connector_id_or_info.auth_type, + configuration=connector_id_or_info.configuration, + user=full_stack.user_id, + workspace=full_stack.workspace_id, + ) + service_connector_response = self.create_service_connector( + service_connector=service_connector_request + ) + service_connectors.append( + service_connector_response + ) + service_connectors_created_ids.append( + service_connector_response.id + ) + break + except EntityExistsError: + connector_name = f"{full_stack.name}-{random_str(4)}".lower() + continue + + # Stack Components + components_mapping: Dict[StackComponentType, List[UUID]] = {} + components_created_ids: List[UUID] = [] + + for ( + component_type, + component_info, + ) in full_stack.components.items(): + # Fetch an existing component + if isinstance(component_info, UUID): + component = self.get_stack_component( + component_id=component_info ) + # Create a new component + else: + component_name = full_stack.name + while True: + try: + component_request = ComponentRequest( + name=component_name, + type=component_type, + flavor=component_info.flavor, + configuration=component_info.configuration, + user=full_stack.user_id, + workspace=full_stack.workspace_id, + ) + component = self.create_stack_component( + component=component_request + ) + components_created_ids.append(component.id) + break + except EntityExistsError: + component_name = f"{full_stack.name}-{random_str(4)}".lower() + continue + + if component_info.service_connector_index is not None: + service_connector = service_connectors[ + component_info.service_connector_index + ] + flavor_list = self.list_flavors( + flavor_filter_model=FlavorFilter( + name=component_info.flavor, + type=component_type, + ) + ) + assert len(flavor_list) == 1 - continue + flavor_model = flavor_list[0] - components_mapping: Dict[StackComponentType, List[UUID]] = {} - for component_type, component_info in full_stack.components.items(): - if isinstance(component_info, UUID): - component = self.get_stack_component( - component_id=component_info - ) - else: - component_name = full_stack.name + requirements = flavor_model.connector_requirements + + if not requirements: + raise RuntimeError( + f"The '{flavor_model.name}' implementation " + "does not support using a service connector to " + "connect to resources." + ) + + resource_id = None + resource_type = requirements.resource_type + if requirements.resource_id_attr is not None: + resource_id = component_info.configuration.get( + requirements.resource_id_attr + ) + + satisfied, msg = requirements.is_satisfied_by( + connector=service_connector, + component=component, + ) + + if not satisfied: + raise RuntimeError( + "Please pick a connector that is " + "compatible with the component flavor and " + "try again.." + ) + + if not resource_id: + if service_connector.resource_id: + resource_id = service_connector.resource_id + elif service_connector.supports_instances: + raise RuntimeError( + f"Multiple {resource_type} resources " + "are available for the selected " + "connector. Please use a `resource_id` " + "to configure a " + f"{resource_type} resource." + ) + + try: + connector_resources = self.verify_service_connector( + service_connector_id=service_connector.id, + resource_type=requirements.resource_type, + resource_id=resource_id, + ) + except ( + KeyError, + ValueError, + IllegalOperationError, + NotImplementedError, + AuthorizationException, + ) as e: + raise RuntimeError( + f"Access to the resource could not be " + f"verified: {e}" + ) + + resources = connector_resources.resources[0] + if resources.resource_ids: + if len(resources.resource_ids) > 1: + raise RuntimeError( + f"Multiple {resource_type} resources are " + "available for the selected connector. " + "Please use the a specific resource-id to " + f"configure a {resource_type} resource." + ) + else: + resource_id = resources.resource_ids[0] + + component_update = ComponentUpdate( + connector=service_connector.id, + connector_resource_id=resource_id, + ) + self.update_stack_component( + component_id=component.id, + component_update=component_update, + ) + + components_mapping[component_type] = [ + component.id, + ] + + # Stck + stack_name = full_stack.name while True: try: - component_request = ComponentRequest( - name=component_name, - type=component_type, - flavor=component_info.flavor, - configuration=component_info.configuration, + stack_request = StackRequest( user=full_stack.user_id, workspace=full_stack.workspace_id, + name=stack_name, + description=full_stack.description, + components=components_mapping, ) - component = self.create_stack_component( - component=component_request - ) + stack_response = self.create_stack(stack_request) break except EntityExistsError: - component_name = ( + stack_name = ( f"{full_stack.name}-{random_str(4)}".lower() ) continue - if component_info.service_connector_index is not None: - service_connector = service_connectors[ - component_info.service_connector_index - ] - flavor_list = self.list_flavors( - flavor_filter_model=FlavorFilter( - name=component_info.flavor, - type=component_type, - ) - ) - assert len(flavor_list) == 1 - - flavor_model = flavor_list[0] - - requirements = flavor_model.connector_requirements - - if not requirements: - raise RuntimeError( - f"The '{flavor_model.name}' implementation " - "does not support using a service connector to " - "connect to resources." - ) - - resource_id = None - resource_type = requirements.resource_type - if requirements.resource_id_attr is not None: - resource_id = component_info.configuration.get( - requirements.resource_id_attr - ) - - satisfied, msg = requirements.is_satisfied_by( - connector=service_connector, - component=component, - ) - - if not satisfied: - raise RuntimeError( - "Please pick a connector that is compatible with " - "the component flavor and try again, or use the " - "interactive mode to select a compatible connector." - ) - - if not resource_id: - if service_connector.resource_id: - resource_id = service_connector.resource_id - elif service_connector.supports_instances: - raise RuntimeError( - f"Multiple {resource_type} resources are " - "available for the selected connector. Please " - "use a `resource_id` to configure a " - f"{resource_type} resource." - ) - - try: - connector_resources = self.verify_service_connector( - service_connector_id=service_connector.id, - resource_type=requirements.resource_type, - resource_id=resource_id, - ) - except ( - KeyError, - ValueError, - IllegalOperationError, - NotImplementedError, - AuthorizationException, - ) as e: - raise RuntimeError( - f"Access to the resource could not be verified: {e}" - ) - - resources = connector_resources.resources[0] - if resources.resource_ids: - if len(resources.resource_ids) > 1: - raise RuntimeError( - f"Multiple {resource_type} resources are " - f"available for the selected connector. Please " - "use the a specific resource-id to configure a " - f"{resource_type} resource." - ) - else: - resource_id = resources.resource_ids[0] - - component_update = ComponentUpdate( - connector=service_connector.id, - connector_resource_id=resource_id, - ) - self.update_stack_component( - component_id=component.id, - component_update=component_update, + handler.metadata.update( + stack_response.get_analytics_metadata() + ) + except Exception as e: + for service_connector_id in service_connectors_created_ids: + self.delete_service_connector( + service_connector_id=service_connector_id ) + for component_id in components_created_ids: + self.delete_stack_component(component_id=component_id) - components_mapping[component_type] = [ - component.id, - ] - - stack_request = StackRequest( - user=full_stack.user_id, - workspace=full_stack.workspace_id, - name=full_stack.name, - description=full_stack.description, - components=components_mapping, - ) - - return self.create_stack(stack_request) + raise RuntimeError( + f"Full Stack creation has failed {e}. Cleaning up the " + f"created entities." + ) + return stack_response def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse: """Get a stack by its unique ID. From 63d9a35b618e0a82fd8024da579cf5472d3a3ff4 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Fri, 28 Jun 2024 16:22:10 +0200 Subject: [PATCH 23/71] typo --- src/zenml/zen_stores/sql_zen_store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 60d7653f618..80bf7518137 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -7108,7 +7108,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: component.id, ] - # Stck + # Stack stack_name = full_stack.name while True: try: From 5e4235e7c47dfed79658e96261cec9d09da5c923 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 28 Jun 2024 16:50:23 +0200 Subject: [PATCH 24/71] hamza's feedback --- src/zenml/cli/stack.py | 422 +++++++++++++++++++++-------------------- src/zenml/cli/utils.py | 28 --- 2 files changed, 218 insertions(+), 232 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 7cf4ef101ab..daacad60e81 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -20,6 +20,8 @@ from uuid import UUID import click +from rich.console import Console +from rich.syntax import Syntax import zenml from zenml.analytics.enums import AnalyticsEvent @@ -199,8 +201,8 @@ def stack() -> None: type=click.BOOL, ) @click.option( - "-cp", - "--cloud", + "-p", + "--provider", help="Name of the cloud provider for this stack.", type=click.Choice(["aws", "azure", "gcp"]), required=False, @@ -227,7 +229,7 @@ def register_stack( data_validator: Optional[str] = None, image_builder: Optional[str] = None, set_stack: bool = False, - cloud: Optional[str] = None, + provider: Optional[str] = None, connector: Optional[str] = None, ) -> None: """Register a stack. @@ -247,10 +249,10 @@ def register_stack( data_validator: Name of the data validator for this stack. image_builder: Name of the new image builder for this stack. set_stack: Immediately set this stack as active. - cloud: Name of the cloud provider for this stack. + provider: Name of the cloud provider for this stack. connector: Name of the service connector for this stack. """ - if (cloud is None and connector is None) and ( + if (provider is None and connector is None) and ( artifact_store is None or orchestrator is None ): cli_utils.error( @@ -275,63 +277,58 @@ def register_stack( components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = {} # cloud flow + created_objects = set() service_connector = None - if cloud is not None and connector is None: - cli_utils.show_status_from_kwargs( - cloud=cloud, - connector=connector, - artifact_store=artifact_store, - orchestrator=orchestrator, - container_registry=container_registry, - ) - existing_connectors = client.list_service_connectors( - connector_type=cloud, size=100 + if provider is not None and connector is None: + use_implicit = cli_utils.Confirm.ask( + f"[bold]{provider.upper()} cloud service connector[/bold] " + "can use the Implicit Authentication by accessing connection " + "configuration of the environment or use one of " + "the authentications methods supported.\n" + "The implicit authentication is great to quickstart, but " + "actual credentials may vary system to system impacting reproducibility.\n" + "Would you like to use the Implicit Authentication method?", + default=False, + show_choices=True, + show_default=True, ) - connector_selected: Optional[int] = None - if existing_connectors.total: - connector_selected = cli_utils.multi_choice_prompt( - object_type=f"{cloud.upper()} service connectors", - choices=[ - [connector.name] for connector in existing_connectors.items - ], - headers=["Name"], - prompt_text=f"We found these {cloud.upper()} service connectors. " - "Do you want to create a new one or use one of the existing ones?", - default_choice="0", - allow_zero_be_a_new_object=True, + if not use_implicit: + existing_connectors = client.list_service_connectors( + connector_type=provider, size=100 ) - if connector_selected is None: + connector_selected: Optional[int] = None + if existing_connectors.total: + connector_selected = cli_utils.multi_choice_prompt( + object_type=f"{provider.upper()} service connectors", + choices=[ + [connector.name] + for connector in existing_connectors.items + ], + headers=["Name"], + prompt_text=f"We found these {provider.upper()} service connectors. " + "Do you want to create a new one or use one of the existing ones?", + default_choice="0", + allow_zero_be_a_new_object=True, + ) + if use_implicit or connector_selected is None: service_connector = _get_service_connector_info( - cloud_provider=cloud + cloud_provider=provider, use_implicit=use_implicit ) + created_objects.add("service_connector") else: selected_connector = existing_connectors.items[connector_selected] service_connector = selected_connector.id connector = selected_connector.name if isinstance(selected_connector.connector_type, str): - cloud = selected_connector.connector_type + provider = selected_connector.connector_type else: - cloud = selected_connector.connector_type.connector_type + provider = selected_connector.connector_type.connector_type elif connector is not None: - cli_utils.show_status_from_kwargs( - cloud=cloud, - connector=connector, - artifact_store=artifact_store, - orchestrator=orchestrator, - container_registry=container_registry, - ) service_connector = client.get_service_connector(connector).id - if service_connector.type != cloud: + if service_connector.type != provider: cli_utils.warning( - f"The service connector `{connector}` is not of type `{cloud}`." + f"The service connector `{connector}` is not of type `{provider}`." ) - cli_utils.show_status_from_kwargs( - cloud=cloud, - connector=connector, - artifact_store=artifact_store, - orchestrator=orchestrator, - container_registry=container_registry, - ) if service_connector: service_connector_resource_model = None @@ -395,11 +392,12 @@ def register_stack( component_info = _get_stack_component_info( component_type=component_type.value, - cloud_provider=cloud, + cloud_provider=provider, service_connector_resource_models=service_connector_resource_model.resources, service_connector_index=0, ) component_name = stack_name + created_objects.add(component_type.value) else: selected_component = existing_components.items[ component_selected @@ -414,13 +412,6 @@ def register_stack( orchestrator = component_name if component_type == StackComponentType.CONTAINER_REGISTRY: container_registry = component_name - cli_utils.show_status_from_kwargs( - cloud=cloud, - connector=connector, - artifact_store=artifact_store, - orchestrator=orchestrator, - container_registry=container_registry, - ) # normal flow once all components are defined with console.status(f"Registering stack '{stack_name}'...\n"): @@ -461,6 +452,10 @@ def register_stack( cli_utils.declare( f"Stack '{created_stack.name}' successfully registered!" ) + cli_utils.print_stack_configuration( + stack=created_stack, + active=created_stack.id == client.active_stack_model.id, + ) if set_stack: client.activate_stack(created_stack.id) @@ -470,9 +465,28 @@ def register_stack( f"Active {scope} stack set to:'{created_stack.name}'" ) - print_model_url(get_stack_url(created_stack)) + delete_commands = [] + if "service_connector" in created_objects: + created_objects.remove("service_connector") + connectors = set() + for each in created_objects: + connectors.add(created_stack.components[each][0].connector.name) + for connector in connectors: + delete_commands.append( + "zenml service-connector delete " + connector + ) + for each in created_objects: + delete_commands.append( + f"zenml {each.replace('_', '-')} delete {created_stack.components[each][0].name}" + ) + delete_commands.append("zenml stack delete -y " + created_stack.name) + + Console().print( + "To delete the objects created by this command run, please run in a sequence:\n" + ) + Console().print(Syntax("\n".join(delete_commands[::-1]), "bash")) - # TODO: print how to delete stack and how to run a pipeline on it + print_model_url(get_stack_url(created_stack)) @stack.command( @@ -1924,18 +1938,19 @@ def connect_stack( ) -def _get_service_connector_info(cloud_provider: str) -> ServiceConnectorInfo: +def _get_service_connector_info( + cloud_provider: str, use_implicit: bool +) -> ServiceConnectorInfo: """Get a service connector info with given cloud provider. Args: cloud_provider: The cloud provider to use. + use_implicit: Whether to use implicit credentials. Returns: The info model of the created service connector. """ - from rich import print - from rich.markdown import Markdown - from rich.prompt import Confirm, Prompt + from rich.prompt import Prompt if cloud_provider not in {"aws", "azure", "gcp"}: raise ValueError(f"Unknown cloud provider {cloud_provider}") @@ -1944,55 +1959,50 @@ def _get_service_connector_info(cloud_provider: str) -> ServiceConnectorInfo: auth_methods = client.get_service_connector_type( cloud_provider ).auth_method_dict - fixed_auth_methods = list(auth_methods.items()) - choices = [] - headers = ["Name", "Required"] - for _, value in fixed_auth_methods: - schema = value.config_schema - required = "" - for each_req in schema["required"]: - field = schema["properties"][each_req] - required += f"[bold]{each_req}[/bold] [italic]({field.get('title','no description')})[/italic]\n" - choices.append([value.name, required]) - - auth_selected = False - while not auth_selected: + if not use_implicit: + fixed_auth_methods = list( + [ + (key, value) + for key, value in auth_methods.items() + if key != "implicit" + ] + ) + choices = [] + headers = ["Name", "Required"] + for _, value in fixed_auth_methods: + schema = value.config_schema + required = "" + for each_req in schema["required"]: + field = schema["properties"][each_req] + required += f"[bold]{each_req}[/bold] [italic]({field.get('title','no description')})[/italic]\n" + choices.append([value.name, required]) + selected_auth_idx = cli_utils.multi_choice_prompt( object_type=f"authentication methods for {cloud_provider}", choices=choices, headers=headers, - prompt_text="Please choose one of the authentication option above to see detailed description:", - ) - selected_auth_model = fixed_auth_methods[selected_auth_idx][1] - print( - Markdown( - f"## {selected_auth_model.name}\n" - + selected_auth_model.description - ) + prompt_text="Please choose one of the authentication option above.", ) + auth_type = fixed_auth_methods[selected_auth_idx][0] + else: + auth_type = "implicit" + + selected_auth_model = auth_methods[auth_type] - auth_selected = Confirm.ask( - "Do you want to continue or go back to authentication methods selection?", - default=True, - ) - if not selected_auth_model: - raise ValueError("No authentication method selected") required_fields = selected_auth_model.config_schema["required"] - data_entered = False - while not data_entered: - answers = {} - for req_field in required_fields: - answers[req_field] = Prompt.ask( - f"Please enter value for `{req_field}`:" - ) - data_entered = Confirm.ask( - "Please confirm the values you entered:\n" + str(answers), - default=True, + properties = selected_auth_model.config_schema["properties"] + + answers = {} + for req_field in required_fields: + answers[req_field] = Prompt.ask( + f"Please enter value for `{req_field}`:", + password="format" in properties[req_field] + and properties[req_field]["format"] == "password", ) return ServiceConnectorInfo( connector_type=cloud_provider, - auth_type=fixed_auth_methods[selected_auth_idx][0], + auth_type=auth_type, configuration=answers, ) @@ -2020,129 +2030,133 @@ def _get_stack_component_info( ValueError: If the cloud provider is not supported. ValueError: If the component type is not supported. """ - from rich import print - from rich.prompt import Confirm, Prompt + from rich.prompt import Prompt if cloud_provider not in {"aws", "azure", "gcp"}: raise ValueError(f"Unknown cloud provider {cloud_provider}") + AWS_DOCS = ( + "https://docs.zenml.io/how-to/auth-management/aws-service-connector" + ) + flavor = "undefined" config = {} if component_type == "artifact_store": - config_confirmed = False - while not config_confirmed: - if cloud_provider == "aws": - for each in service_connector_resource_models: - if each.resource_type == "s3-bucket": - available_storages = each.resource_ids - flavor = "s3" - elif cloud_provider == "azure": - flavor = "azure" - elif cloud_provider == "gcp": - flavor = "gcs" - - selected_storage_idx = cli_utils.multi_choice_prompt( - object_type=f"{cloud_provider.upper()} storages", - choices=[[st] for st in available_storages], - headers=["Storage"], - prompt_text="Please choose one of the storages for the new artifact store:", - ) - selected_storage = available_storages[selected_storage_idx] - - extra_path = Prompt.ask( - f"Please enter any further path inside the storage, if needed ({selected_storage}/...):", - default="", - ) - config = { - "path": f"{selected_storage.strip('/')}/{extra_path.strip('/')}" - } + if cloud_provider == "aws": + for each in service_connector_resource_models: + if each.resource_type == "s3-bucket": + available_storages = each.resource_ids + flavor = "s3" + if not available_storages: + cli_utils.error( + "We were unable to find any S3 buckets available " + "to configured service connector. Please, verify " + "that needed permission are granted for the " + "service connector.\nDocumentation for the S3 " + "Buckets configuration can be found at " + f"{AWS_DOCS}#s3-bucket" + ) + elif cloud_provider == "azure": + flavor = "azure" + elif cloud_provider == "gcp": + flavor = "gcs" + + selected_storage_idx = cli_utils.multi_choice_prompt( + object_type=f"{cloud_provider.upper()} storages", + choices=[[st] for st in available_storages], + headers=["Storage"], + prompt_text="Please choose one of the storages for the new artifact store:", + ) + selected_storage = available_storages[selected_storage_idx] - print(config) - config_confirmed = Confirm.ask( - "Please confirm the values you entered:", default=True - ) + config = {"path": selected_storage} elif component_type == "orchestrator": - config_confirmed = False - while not config_confirmed: - if cloud_provider == "aws": - available_orchestrators = [] - for each in service_connector_resource_models: - types = [] - if each.resource_type == "aws-generic": - types = ["Sagemaker", "VM AWS"] - if each.resource_type == "kubernetes-cluster": - types = ["K8S"] - + if cloud_provider == "aws": + available_orchestrators = [] + for each in service_connector_resource_models: + types = [] + if each.resource_type == "aws-generic": + types = ["Sagemaker", "VM AWS"] + if each.resource_type == "kubernetes-cluster": + types = ["K8S"] + + if each.resource_ids: for orchestrator in each.resource_ids: for t in types: available_orchestrators.append([t, orchestrator]) - elif cloud_provider == "gcp": - pass - elif cloud_provider == "azure": - pass - - selected_orchestrator_idx = cli_utils.multi_choice_prompt( - object_type=f"orchestrators on {cloud_provider.upper()}", - choices=available_orchestrators, - headers=["Orchestrator Type", "Orchestrator details"], - prompt_text="Please choose one of the orchestrators for the new orchestrator:", - ) + if not available_orchestrators: + cli_utils.error( + "We were unable to find any orchestrator engines " + "available to the service connector. Please, verify " + "that needed permission are granted for the " + "service connector.\nDocumentation for the Generic " + "AWS resource configuration can be found at " + f"{AWS_DOCS}#generic-aws-resource\n" + "Documentation for the Kubernetes resource " + "configuration can be found at " + f"{AWS_DOCS}#eks-kubernetes-cluster" + ) + elif cloud_provider == "gcp": + pass + elif cloud_provider == "azure": + pass + + selected_orchestrator_idx = cli_utils.multi_choice_prompt( + object_type=f"orchestrators on {cloud_provider.upper()}", + choices=available_orchestrators, + headers=["Orchestrator Type", "Orchestrator details"], + prompt_text="Please choose one of the orchestrators for the new orchestrator:", + ) - selected_orchestrator = available_orchestrators[ - selected_orchestrator_idx - ] + selected_orchestrator = available_orchestrators[ + selected_orchestrator_idx + ] - if selected_orchestrator[0] == "Sagemaker": - flavor = "sagemaker" - execution_role = Prompt.ask( - "Please enter an execution role ARN:" - ) - config = {"execution_role": execution_role} - elif selected_orchestrator[0] == "VM AWS": - flavor = "vm_aws" - config = {} - elif selected_orchestrator[0] == "K8S": - flavor = "kubernetes" - config = {} - else: - raise ValueError( - f"Unknown orchestrator type {selected_orchestrator[0]}" - ) - print(config) - config_confirmed = Confirm.ask( - "Please confirm the values you entered:", default=True + if selected_orchestrator[0] == "Sagemaker": + flavor = "sagemaker" + execution_role = Prompt.ask("Please enter an execution role ARN:") + config = {"execution_role": execution_role} + elif selected_orchestrator[0] == "VM AWS": + flavor = "vm_aws" + config = {} + elif selected_orchestrator[0] == "K8S": + flavor = "kubernetes" + config = {} + else: + raise ValueError( + f"Unknown orchestrator type {selected_orchestrator[0]}" ) elif component_type == "container_registry": - config_confirmed = False - while not config_confirmed: - if cloud_provider == "aws": - for each in service_connector_resource_models: - if each.resource_type == "docker-registry": - available_registries = each.resource_ids - flavor = "aws" - elif cloud_provider == "azure": - flavor = "azure" - available_registries = [] - elif cloud_provider == "gcp": - flavor = "gcp" - available_registries = [] - - selected_registry_idx = cli_utils.multi_choice_prompt( - object_type=f"{cloud_provider.upper()} registries", - choices=[[st] for st in available_registries], - headers=["Container Registry"], - prompt_text="Please choose one of the registries for the new container registry:", - ) - selected_storage = available_registries[selected_registry_idx] - config = {"uri": selected_storage} - - print(config) - config_confirmed = Confirm.ask( - "Please confirm the values you entered:", default=True - ) + available_registries = [] + if cloud_provider == "aws": + flavor = "aws" + for each in service_connector_resource_models: + if each.resource_type == "docker-registry": + available_registries = each.resource_ids + if not available_registries: + cli_utils.error( + "We were unable to find any container registries " + "available to the service connector. Please, verify " + "that needed permission are granted for the " + "service connector.\nDocumentation for the ECR " + "container registry resource configuration can " + f"be found at {AWS_DOCS}#ecr-container-registry" + ) + elif cloud_provider == "azure": + flavor = "azure" + elif cloud_provider == "gcp": + flavor = "gcp" + + selected_registry_idx = cli_utils.multi_choice_prompt( + object_type=f"{cloud_provider.upper()} registries", + choices=[[st] for st in available_registries], + headers=["Container Registry"], + prompt_text="Please choose one of the registries for the new container registry:", + ) + selected_storage = available_registries[selected_registry_idx] + config = {"uri": selected_storage} else: raise ValueError(f"Unknown component type {component_type}") - service_connector_index return ComponentInfo( flavor=flavor, diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index c17049ef731..47abc1b1228 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -2757,34 +2757,6 @@ def is_jupyter_installed() -> bool: return False -def show_status_from_kwargs(default_value: str = ":x:", **kwargs: Any) -> None: - """Show status from kwargs. - - Args: - default_value: The default value to show status from. - **kwargs: The kwargs to show status from. If value is passed, - but `None` a default value is used. - """ - status = [] - names = [] - for name, each in kwargs.items(): - if not each: - each = ":x:" - status.append(each) - names.append(name.replace("_", " ").capitalize()) - - status_table = Table( - title="New cloud stack registration progress", - show_header=True, - expand=True, - ) - for c in names: - status_table.add_column(c, justify="center", width=1) - - status_table.add_row(*status) - Console().print(status_table) - - def multi_choice_prompt( object_type: str, choices: List[List[Any]], From e73590d06c1d5e0419396356a1549811cadf1e79 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 28 Jun 2024 16:51:24 +0200 Subject: [PATCH 25/71] minor docs update --- src/zenml/cli/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zenml/cli/__init__.py b/src/zenml/cli/__init__.py index 5b21e1bbc5b..7c670be0420 100644 --- a/src/zenml/cli/__init__.py +++ b/src/zenml/cli/__init__.py @@ -1416,7 +1416,7 @@ ```bash zenml stack register STACK_NAME \ - -cp CLOUD_PROVIDER + -p CLOUD_PROVIDER ``` To create a new stack using the existing service connector with a set of minimal components, From b13c09ea4631cc88afd2b0a2398704818d5870be Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Mon, 1 Jul 2024 14:22:29 +0200 Subject: [PATCH 26/71] various fixes from PR comments --- src/zenml/cli/stack.py | 8 ++++---- src/zenml/models/v2/misc/full_stack.py | 4 ++-- src/zenml/zen_stores/rest_zen_store.py | 4 ++-- src/zenml/zen_stores/sql_zen_store.py | 7 +++++-- 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index daacad60e81..ab51e2e15fe 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -383,8 +383,8 @@ def register_stack( _, service_connector_resource_model = ( client.create_service_connector( name=stack_name, - connector_type=service_connector.connector_type, - auth_method=service_connector.auth_type, + connector_type=service_connector.type, + auth_method=service_connector.auth_method, configuration=service_connector.configuration, register=False, ) @@ -2001,8 +2001,8 @@ def _get_service_connector_info( ) return ServiceConnectorInfo( - connector_type=cloud_provider, - auth_type=auth_type, + type=cloud_provider, + auth_method=auth_type, configuration=answers, ) diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 5d21ec6c5af..4d0d8baad02 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -26,8 +26,8 @@ class ServiceConnectorInfo(BaseModel): """Information about the service connector when creating a full stack.""" - connector_type: str - auth_type: str + type: str + auth_method: str configuration: Dict[str, Any] = {} diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index b5146809d75..f3a739e2537 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -2677,9 +2677,9 @@ def list_service_connector_types( } for connector in local_connector_types: - if connector.connector_type in connector_types_map: + if connector.type in connector_types_map: connector.remote = True - connector_types_map[connector.connector_type] = connector + connector_types_map[connector.type] = connector return list(connector_types_map.values()) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 80bf7518137..239161840c8 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6928,6 +6928,9 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: Returns: The registered stack. + + Raises: + RuntimeError: If the full stack creation fails """ # We can not use the decorator here, as we need custom metadata with track_handler( @@ -6956,8 +6959,8 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: try: service_connector_request = ServiceConnectorRequest( name=connector_name, - connector_type=connector_id_or_info.connector_type, - auth_method=connector_id_or_info.auth_type, + connector_type=connector_id_or_info.type, + auth_method=connector_id_or_info.auth_method, configuration=connector_id_or_info.configuration, user=full_stack.user_id, workspace=full_stack.workspace_id, From c97afe930278f2772c6da5a2ea31249eb12342f0 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Mon, 1 Jul 2024 14:32:34 +0200 Subject: [PATCH 27/71] several more fixes --- src/zenml/zen_stores/sql_zen_store.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 239161840c8..3186f1aea85 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6930,7 +6930,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: The registered stack. Raises: - RuntimeError: If the full stack creation fails + RuntimeError: If the full stack creation fails. """ # We can not use the decorator here, as we need custom metadata with track_handler( @@ -6938,13 +6938,15 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: metadata={"generated_by_wizard": True}, ) as handler: # For clean-up purposes, each created entity is tracked here + service_connectors_created_ids: List[UUID] = [] + components_created_ids: List[UUID] = [] + try: # Validate the name of the new stack validate_name(full_stack) # Service Connectors service_connectors: List[ServiceConnectorResponse] = [] - service_connectors_created_ids: List[UUID] = [] for connector_id_or_info in full_stack.service_connectors: # Fetch an existing service connector @@ -6981,7 +6983,6 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: # Stack Components components_mapping: Dict[StackComponentType, List[UUID]] = {} - components_created_ids: List[UUID] = [] for ( component_type, @@ -7144,7 +7145,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: raise RuntimeError( f"Full Stack creation has failed {e}. Cleaning up the " f"created entities." - ) + ) from e return stack_response def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse: From 81632cd8df1d63f6e188359ad07d1d67e35634e9 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Mon, 1 Jul 2024 14:39:04 +0200 Subject: [PATCH 28/71] reversing the order of the cleanup --- src/zenml/zen_stores/sql_zen_store.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 3186f1aea85..a02482874ac 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -7135,13 +7135,12 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: stack_response.get_analytics_metadata() ) except Exception as e: + for component_id in components_created_ids: + self.delete_stack_component(component_id=component_id) for service_connector_id in service_connectors_created_ids: self.delete_service_connector( service_connector_id=service_connector_id ) - for component_id in components_created_ids: - self.delete_stack_component(component_id=component_id) - raise RuntimeError( f"Full Stack creation has failed {e}. Cleaning up the " f"created entities." From 6c7e9d30c6cb23b0844e04b4cf6fdb8f853e4e33 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Mon, 1 Jul 2024 15:49:28 +0200 Subject: [PATCH 29/71] darglint --- src/zenml/cli/stack.py | 5 +++-- src/zenml/zen_server/routers/workspaces_endpoints.py | 4 ---- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index ab51e2e15fe..8c04fde73cc 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -272,8 +272,6 @@ def register_stack( ) except KeyError: pass - except Exception as e: - raise e components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = {} # cloud flow @@ -1949,6 +1947,9 @@ def _get_service_connector_info( Returns: The info model of the created service connector. + + Raises: + ValueError: If the cloud provider is not supported. """ from rich.prompt import Prompt diff --git a/src/zenml/zen_server/routers/workspaces_endpoints.py b/src/zenml/zen_server/routers/workspaces_endpoints.py index be858eff467..a5dee6f55d5 100644 --- a/src/zenml/zen_server/routers/workspaces_endpoints.py +++ b/src/zenml/zen_server/routers/workspaces_endpoints.py @@ -371,10 +371,6 @@ def create_full_stack( Returns: The created stack. - - Raises: - IllegalOperationError: If the workspace specified in the stack - does not match the current workspace. """ workspace = zen_store().get_workspace(workspace_name_or_id) From 37eec5126b934b43c9fcd1cac1050bbfe68e660a Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Mon, 1 Jul 2024 15:52:11 +0200 Subject: [PATCH 30/71] ruff --- src/zenml/utils/function_utils.py | 2 +- src/zenml/zen_stores/migrations/utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/zenml/utils/function_utils.py b/src/zenml/utils/function_utils.py index 4a9a7ba06e1..ebf309340a9 100644 --- a/src/zenml/utils/function_utils.py +++ b/src/zenml/utils/function_utils.py @@ -146,7 +146,7 @@ def _cli_wrapped_function(func: F) -> F: if _is_valid_optional_arg(arg_type): arg_type = arg_type.__args__[0] arg_name = _cli_arg_name(arg_name) - if arg_type == bool: + if arg_type is bool: options.append( click.option( f"--{arg_name}", diff --git a/src/zenml/zen_stores/migrations/utils.py b/src/zenml/zen_stores/migrations/utils.py index 02d3be03d97..b79a0b72d83 100644 --- a/src/zenml/zen_stores/migrations/utils.py +++ b/src/zenml/zen_stores/migrations/utils.py @@ -364,7 +364,7 @@ def restore_database_from_storage( # Convert column values to the correct type for column in table.columns: # Blob columns are stored as binary strings - if column.type.python_type == bytes and isinstance( + if column.type.python_type is bytes and isinstance( row[column.name], str ): # Convert the string to bytes From 19ddc4206e7e85735a899898b2a85a025914a054 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Mon, 1 Jul 2024 16:13:50 +0200 Subject: [PATCH 31/71] mypy --- src/zenml/cli/stack.py | 78 ++++++++++++++++++++++++++++++------------ 1 file changed, 56 insertions(+), 22 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 8c04fde73cc..0d7222a4ac8 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -16,11 +16,12 @@ import getpass import os from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union from uuid import UUID import click from rich.console import Console +from rich.prompt import Confirm from rich.syntax import Syntax import zenml @@ -252,6 +253,7 @@ def register_stack( provider: Name of the cloud provider for this stack. connector: Name of the service connector for this stack. """ + if (provider is None and connector is None) and ( artifact_store is None or orchestrator is None ): @@ -275,10 +277,10 @@ def register_stack( components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = {} # cloud flow - created_objects = set() - service_connector = None + created_objects: Set[str] = set() + service_connector: Optional[Union[UUID, ServiceConnectorInfo]] = None if provider is not None and connector is None: - use_implicit = cli_utils.Confirm.ask( + use_implicit = Confirm.ask( f"[bold]{provider.upper()} cloud service connector[/bold] " "can use the Implicit Authentication by accessing connection " "configuration of the environment or use one of " @@ -322,11 +324,15 @@ def register_stack( else: provider = selected_connector.connector_type.connector_type elif connector is not None: - service_connector = client.get_service_connector(connector).id - if service_connector.type != provider: - cli_utils.warning( - f"The service connector `{connector}` is not of type `{provider}`." - ) + service_connector_response = client.get_service_connector(connector) + service_connector = service_connector_response.id + if provider: + if service_connector_response.type != provider: + cli_utils.warning( + f"The service connector `{connector}` is not of type `{provider}`." + ) + else: + provider = service_connector_response.type if service_connector: service_connector_resource_model = None @@ -337,11 +343,12 @@ def register_stack( (StackComponentType.CONTAINER_REGISTRY, container_registry), ) for component_type, preset_name in needed_components: + component_info: Optional[Union[UUID, ComponentInfo]] = None if preset_name is not None: - component_info = client.get_stack_component( + component_response = client.get_stack_component( component_type, preset_name ) - component_info = component_info.id + component_info = component_response.id else: if isinstance(service_connector, UUID): # find existing components under same connector @@ -387,6 +394,20 @@ def register_stack( register=False, ) ) + if service_connector_resource_model is None: + cli_utils.error( + f"Failed to validate service connector {service_connector}..." + ) + if provider is None: + if isinstance( + service_connector_resource_model.connector_type, + str, + ): + provider = ( + service_connector_resource_model.connector_type + ) + else: + provider = service_connector_resource_model.connector_type.connector_type component_info = _get_stack_component_info( component_type=component_type.value, @@ -413,7 +434,7 @@ def register_stack( # normal flow once all components are defined with console.status(f"Registering stack '{stack_name}'...\n"): - for component_type, component_name in [ + for component_type_, component_name_ in [ (StackComponentType.ARTIFACT_STORE, artifact_store), (StackComponentType.ORCHESTRATOR, orchestrator), (StackComponentType.ALERTER, alerter), @@ -427,9 +448,9 @@ def register_stack( (StackComponentType.EXPERIMENT_TRACKER, experiment_tracker), (StackComponentType.CONTAINER_REGISTRY, container_registry), ]: - if component_name and component_type not in components: - components[component_type] = client.get_stack_component( - component_type, component_name + if component_name_ and component_type_ not in components: + components[component_type_] = client.get_stack_component( + component_type_, component_name_ ).id try: @@ -468,15 +489,18 @@ def register_stack( created_objects.remove("service_connector") connectors = set() for each in created_objects: - connectors.add(created_stack.components[each][0].connector.name) + if comps_ := created_stack.components[StackComponentType(each)]: + if conn_ := comps_[0].connector: + connectors.add(conn_.name) for connector in connectors: delete_commands.append( "zenml service-connector delete " + connector ) for each in created_objects: - delete_commands.append( - f"zenml {each.replace('_', '-')} delete {created_stack.components[each][0].name}" - ) + if comps_ := created_stack.components[StackComponentType(each)]: + delete_commands.append( + f"zenml {each.replace('_', '-')} delete {comps_[0].name}" + ) delete_commands.append("zenml stack delete -y " + created_stack.name) Console().print( @@ -1984,6 +2008,8 @@ def _get_service_connector_info( headers=headers, prompt_text="Please choose one of the authentication option above.", ) + if selected_auth_idx is None: + cli_utils.error("No authentication method selected.") auth_type = fixed_auth_methods[selected_auth_idx][0] else: auth_type = "implicit" @@ -2043,10 +2069,11 @@ def _get_stack_component_info( flavor = "undefined" config = {} if component_type == "artifact_store": + available_storages: List[str] = [] if cloud_provider == "aws": for each in service_connector_resource_models: if each.resource_type == "s3-bucket": - available_storages = each.resource_ids + available_storages = each.resource_ids or [] flavor = "s3" if not available_storages: cli_utils.error( @@ -2068,6 +2095,9 @@ def _get_stack_component_info( headers=["Storage"], prompt_text="Please choose one of the storages for the new artifact store:", ) + if selected_storage_idx is None: + cli_utils.error("No storage selected.") + selected_storage = available_storages[selected_storage_idx] config = {"path": selected_storage} @@ -2108,6 +2138,8 @@ def _get_stack_component_info( headers=["Orchestrator Type", "Orchestrator details"], prompt_text="Please choose one of the orchestrators for the new orchestrator:", ) + if selected_orchestrator_idx is None: + cli_utils.error("No orchestrator selected.") selected_orchestrator = available_orchestrators[ selected_orchestrator_idx @@ -2128,12 +2160,12 @@ def _get_stack_component_info( f"Unknown orchestrator type {selected_orchestrator[0]}" ) elif component_type == "container_registry": - available_registries = [] + available_registries: List[str] = [] if cloud_provider == "aws": flavor = "aws" for each in service_connector_resource_models: if each.resource_type == "docker-registry": - available_registries = each.resource_ids + available_registries = each.resource_ids or [] if not available_registries: cli_utils.error( "We were unable to find any container registries " @@ -2154,6 +2186,8 @@ def _get_stack_component_info( headers=["Container Registry"], prompt_text="Please choose one of the registries for the new container registry:", ) + if selected_registry_idx is None: + cli_utils.error("No container registry selected.") selected_storage = available_registries[selected_registry_idx] config = {"uri": selected_storage} else: From 00823f6c4a197ab12e6e3d28b2ce597c28fddcf3 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Mon, 1 Jul 2024 16:18:53 +0200 Subject: [PATCH 32/71] ruff --- src/zenml/cli/stack.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 0d7222a4ac8..56c0746d42b 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -253,7 +253,6 @@ def register_stack( provider: Name of the cloud provider for this stack. connector: Name of the service connector for this stack. """ - if (provider is None and connector is None) and ( artifact_store is None or orchestrator is None ): From 5437a19b13f49331f7a8e3c0b0407034919f1d98 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Mon, 1 Jul 2024 16:22:58 +0200 Subject: [PATCH 33/71] ruff --- tests/unit/materializers/test_built_in_materializer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/materializers/test_built_in_materializer.py b/tests/unit/materializers/test_built_in_materializer.py index 90d822a5820..31741bafbc1 100644 --- a/tests/unit/materializers/test_built_in_materializer.py +++ b/tests/unit/materializers/test_built_in_materializer.py @@ -34,7 +34,7 @@ def test_basic_type_materialization(): result = _test_materializer( step_output_type=type_, step_output=example, - expected_metadata_size=1 if type_ == str else 2, + expected_metadata_size=1 if type_ is str else 2, ) assert result == example From 79f835fbeb4b17d9aed8449312e2b6c09f6a40b7 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Mon, 1 Jul 2024 16:26:04 +0200 Subject: [PATCH 34/71] removing the extra verification --- src/zenml/zen_stores/sql_zen_store.py | 30 --------------------------- 1 file changed, 30 deletions(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index a02482874ac..555a66a6b6b 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -7069,36 +7069,6 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: f"{resource_type} resource." ) - try: - connector_resources = self.verify_service_connector( - service_connector_id=service_connector.id, - resource_type=requirements.resource_type, - resource_id=resource_id, - ) - except ( - KeyError, - ValueError, - IllegalOperationError, - NotImplementedError, - AuthorizationException, - ) as e: - raise RuntimeError( - f"Access to the resource could not be " - f"verified: {e}" - ) - - resources = connector_resources.resources[0] - if resources.resource_ids: - if len(resources.resource_ids) > 1: - raise RuntimeError( - f"Multiple {resource_type} resources are " - "available for the selected connector. " - "Please use the a specific resource-id to " - f"configure a {resource_type} resource." - ) - else: - resource_id = resources.resource_ids[0] - component_update = ComponentUpdate( connector=service_connector.id, connector_resource_id=resource_id, From 8574d86818ae80a12046059e6d72b947427062dd Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Mon, 1 Jul 2024 16:30:54 +0200 Subject: [PATCH 35/71] fetch stack by full name --- src/zenml/cli/stack.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 56c0746d42b..7c541d09331 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -266,7 +266,10 @@ def register_stack( client = Client() try: - client.get_stack(name_id_or_prefix=stack_name) + client.get_stack( + name_id_or_prefix=stack_name, + allow_name_prefix_match=False, + ) cli_utils.error( f"A stack with name `{stack_name}` already exists, " "please use a different name." From 35a36b4c9231fd589d648812954169a4b28d16aa Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Mon, 1 Jul 2024 16:44:48 +0200 Subject: [PATCH 36/71] mypy --- src/zenml/cli/utils.py | 5 +++++ src/zenml/zen_stores/rest_zen_store.py | 4 ++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index 47abc1b1228..2864f9a25f7 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -2777,6 +2777,9 @@ def multi_choice_prompt( Returns: The selected choice index or None for new object + + Raises: + RuntimeError: If no choice is made. """ table = Table( title=f"Available {object_type}", @@ -2808,6 +2811,8 @@ def multi_choice_prompt( default=default_choice, show_choices=False, ) + if selected is None: + raise RuntimeError(f"No {object_type} was selected") if selected == "0" and allow_zero_be_a_new_object: return None diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index f3a739e2537..b5146809d75 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -2677,9 +2677,9 @@ def list_service_connector_types( } for connector in local_connector_types: - if connector.type in connector_types_map: + if connector.connector_type in connector_types_map: connector.remote = True - connector_types_map[connector.type] = connector + connector_types_map[connector.connector_type] = connector return list(connector_types_map.values()) From dcd8e5d3b51d70123a79cac719ad8db6f1074903 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Tue, 2 Jul 2024 13:10:32 +0200 Subject: [PATCH 37/71] workspaced full stack request --- src/zenml/cli/stack.py | 4 ++-- src/zenml/cli/utils.py | 3 ++- src/zenml/models/v2/misc/full_stack.py | 4 ++-- src/zenml/zen_server/routers/workspaces_endpoints.py | 4 ++-- src/zenml/zen_stores/rest_zen_store.py | 2 +- src/zenml/zen_stores/sql_zen_store.py | 12 ++++++------ 6 files changed, 15 insertions(+), 14 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 7c541d09331..c1a239f3180 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -458,8 +458,8 @@ def register_stack( try: created_stack = client.zen_store.create_full_stack( full_stack=FullStackRequest( - user_id=client.active_user.id, - workspace_id=client.active_workspace.id, + user=client.active_user.id, + workspace=client.active_workspace.id, name=stack_name, components=components, service_connectors=[ diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index 2864f9a25f7..3edd52d05e3 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -2830,6 +2830,8 @@ def requires_mac_env_var_warning() -> bool: Returns: bool: True if a warning needs to be shown, False otherwise. """ + breakpoint() + if mac_version := platform.mac_ver()[0]: try: major, minor, _ = mac_version.split(".") @@ -2840,7 +2842,6 @@ def requires_mac_env_var_warning() -> bool: return True else: mac_version_tuple = (0, 0) - return ( not os.getenv("OBJC_DISABLE_INITIALIZE_FORK_SAFETY") and sys.platform == "darwin" diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 4d0d8baad02..b9881dec33b 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -48,8 +48,8 @@ class ComponentInfo(BaseModel): class FullStackRequest(BaseRequest): """Request model for a full-stack.""" - user_id: Optional[UUID] = None - workspace_id: Optional[UUID] = None + user: Optional[UUID] = None + workspace: Optional[UUID] = None name: str = Field( title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH diff --git a/src/zenml/zen_server/routers/workspaces_endpoints.py b/src/zenml/zen_server/routers/workspaces_endpoints.py index a5dee6f55d5..743b3556af5 100644 --- a/src/zenml/zen_server/routers/workspaces_endpoints.py +++ b/src/zenml/zen_server/routers/workspaces_endpoints.py @@ -407,8 +407,8 @@ def create_full_stack( verify_permission(resource_type=ResourceType.STACK, action=Action.CREATE) - full_stack.user_id = auth_context.user.id - full_stack.workspace_id = workspace.id + full_stack.user = auth_context.user.id + full_stack.workspace = workspace.id return zen_store().create_full_stack(full_stack) diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index b5146809d75..660ecb12463 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -2755,7 +2755,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: Returns: The registered stack. """ - return self._create_resource( + return self._create_workspace_scoped_resource( resource=full_stack, route=FULL_STACK, response_model=StackResponse, diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 555a66a6b6b..e295480ae27 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6964,8 +6964,8 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: connector_type=connector_id_or_info.type, auth_method=connector_id_or_info.auth_method, configuration=connector_id_or_info.configuration, - user=full_stack.user_id, - workspace=full_stack.workspace_id, + user=full_stack.user, + workspace=full_stack.workspace, ) service_connector_response = self.create_service_connector( service_connector=service_connector_request @@ -7003,8 +7003,8 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: type=component_type, flavor=component_info.flavor, configuration=component_info.configuration, - user=full_stack.user_id, - workspace=full_stack.workspace_id, + user=full_stack.user, + workspace=full_stack.workspace, ) component = self.create_stack_component( component=component_request @@ -7087,8 +7087,8 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: while True: try: stack_request = StackRequest( - user=full_stack.user_id, - workspace=full_stack.workspace_id, + user=full_stack.user, + workspace=full_stack.workspace, name=stack_name, description=full_stack.description, components=components_mapping, From 63ffd29caed9d8872896f885a1b63f7de8a8ff1b Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Tue, 2 Jul 2024 14:02:46 +0200 Subject: [PATCH 38/71] fix normal stack register --- src/zenml/cli/stack.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index c1a239f3180..1b3956c44a9 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -462,9 +462,9 @@ def register_stack( workspace=client.active_workspace.id, name=stack_name, components=components, - service_connectors=[ - service_connector, - ], + service_connectors=[service_connector] + if service_connector + else [], ) ) except (KeyError, IllegalOperationError) as err: From 9a9dd3fb6e15e5577bc4ad7f83e507016d3eb80f Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Tue, 2 Jul 2024 17:58:42 +0200 Subject: [PATCH 39/71] add `service_connector_resource_id` --- src/zenml/cli/stack.py | 11 ++++++++--- src/zenml/models/v2/misc/full_stack.py | 1 + 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 1b3956c44a9..3fa2941341d 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -2069,6 +2069,7 @@ def _get_stack_component_info( ) flavor = "undefined" + service_connector_resource_id = None config = {} if component_type == "artifact_store": available_storages: List[str] = [] @@ -2103,6 +2104,7 @@ def _get_stack_component_info( selected_storage = available_storages[selected_storage_idx] config = {"path": selected_storage} + service_connector_resource_id = selected_storage elif component_type == "orchestrator": if cloud_provider == "aws": available_orchestrators = [] @@ -2153,7 +2155,7 @@ def _get_stack_component_info( config = {"execution_role": execution_role} elif selected_orchestrator[0] == "VM AWS": flavor = "vm_aws" - config = {} + config = {"region": selected_orchestrator[1]} elif selected_orchestrator[0] == "K8S": flavor = "kubernetes" config = {} @@ -2161,6 +2163,7 @@ def _get_stack_component_info( raise ValueError( f"Unknown orchestrator type {selected_orchestrator[0]}" ) + service_connector_resource_id = selected_orchestrator[1] elif component_type == "container_registry": available_registries: List[str] = [] if cloud_provider == "aws": @@ -2190,8 +2193,9 @@ def _get_stack_component_info( ) if selected_registry_idx is None: cli_utils.error("No container registry selected.") - selected_storage = available_registries[selected_registry_idx] - config = {"uri": selected_storage} + selected_registry = available_registries[selected_registry_idx] + config = {"uri": selected_registry} + service_connector_resource_id = selected_registry else: raise ValueError(f"Unknown component type {component_type}") @@ -2199,4 +2203,5 @@ def _get_stack_component_info( flavor=flavor, configuration=config, service_connector_index=service_connector_index, + service_connector_resource_id=service_connector_resource_id, ) diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index b9881dec33b..9299dd37fd4 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -42,6 +42,7 @@ class ComponentInfo(BaseModel): description="The id of the service connector from the list " "`service_connectors` from `FullStackRequest`.", ) + service_connector_resource_id: Optional[str] = None configuration: Dict[str, Any] = {} From 10a781b5da2c4a28f9003c6d2174c946b2ab52ac Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Tue, 2 Jul 2024 18:02:10 +0200 Subject: [PATCH 40/71] tiny renames --- src/zenml/cli/stack.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 3fa2941341d..8c0f7f37666 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -2111,9 +2111,9 @@ def _get_stack_component_info( for each in service_connector_resource_models: types = [] if each.resource_type == "aws-generic": - types = ["Sagemaker", "VM AWS"] + types = ["Sagemaker", "Skypilot (EC2)"] if each.resource_type == "kubernetes-cluster": - types = ["K8S"] + types = ["Kubernetes"] if each.resource_ids: for orchestrator in each.resource_ids: @@ -2153,10 +2153,10 @@ def _get_stack_component_info( flavor = "sagemaker" execution_role = Prompt.ask("Please enter an execution role ARN:") config = {"execution_role": execution_role} - elif selected_orchestrator[0] == "VM AWS": + elif selected_orchestrator[0] == "Skypilot (EC2)": flavor = "vm_aws" config = {"region": selected_orchestrator[1]} - elif selected_orchestrator[0] == "K8S": + elif selected_orchestrator[0] == "Kubernetes": flavor = "kubernetes" config = {} else: From dc66278dd1225a43e055667e096e8c5a098564fd Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Tue, 2 Jul 2024 18:06:24 +0200 Subject: [PATCH 41/71] use passed `resource_id` --- src/zenml/zen_stores/sql_zen_store.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index e295480ae27..0b1e5e073c3 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -7038,12 +7038,17 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: "connect to resources." ) - resource_id = None - resource_type = requirements.resource_type - if requirements.resource_id_attr is not None: - resource_id = component_info.configuration.get( - requirements.resource_id_attr - ) + if component_info.service_connector_resource_id: + resource_id = component_info.service_connector_resource_id + else: + resource_id = None + resource_type = requirements.resource_type + if requirements.resource_id_attr is not None: + resource_id = ( + component_info.configuration.get( + requirements.resource_id_attr + ) + ) satisfied, msg = requirements.is_satisfied_by( connector=service_connector, From 19cd0e370328e116f1b16aaf9ce2bb43c9101e84 Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Wed, 3 Jul 2024 10:31:14 +0200 Subject: [PATCH 42/71] Fix connector resource type in `zenml describe` CLI command --- src/zenml/cli/stack_components.py | 7 +++++++ src/zenml/cli/utils.py | 16 ++++++++++++++-- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/src/zenml/cli/stack_components.py b/src/zenml/cli/stack_components.py index 58789c9c9c3..10c6dcbd377 100644 --- a/src/zenml/cli/stack_components.py +++ b/src/zenml/cli/stack_components.py @@ -136,9 +136,16 @@ def describe_stack_component_command(name_id_or_prefix: str) -> None: if active_components: active_component_id = active_components[0].id + if component_.connector: + # We also need the flavor to get the connector requirements + flavor = client.get_flavor_by_name_and_type( + name=component_.flavor, component_type=component_type + ) + cli_utils.print_stack_component_configuration( component=component_, active_status=component_.id == active_component_id, + connector_requirements=flavor.connector_requirements, ) print_model_url(get_component_url(component_)) diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index 3edd52d05e3..d8df9e584ca 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -72,6 +72,7 @@ BoolFilter, NumericFilter, Page, + ServiceConnectorRequirements, StrFilter, UUIDFilter, ) @@ -616,13 +617,18 @@ def print_flavor_list(flavors: Page["FlavorResponse"]) -> None: def print_stack_component_configuration( - component: "ComponentResponse", active_status: bool + component: "ComponentResponse", + active_status: bool, + connector_requirements: Optional[ServiceConnectorRequirements] = None, ) -> None: """Prints the configuration options of a stack component. Args: component: The stack component to print. active_status: Whether the stack component is active. + connector_requirements: Connector requirements for the component, taken + from the component flavor. Only needed if the component has a + connector. """ if component.user: user_name = component.user.name @@ -693,11 +699,17 @@ def print_stack_component_configuration( rich_table.add_column("PROPERTY") rich_table.add_column("VALUE", overflow="fold") + resource_type = ( + connector_requirements.resource_type + if connector_requirements + else component.connector.resource_types[0] + ) + connector_dict = { "ID": str(component.connector.id), "NAME": component.connector.name, "TYPE": component.connector.type, - "RESOURCE TYPE": component.connector.resource_types[0], + "RESOURCE TYPE": resource_type, "RESOURCE NAME": component.connector_resource_id or component.connector.resource_id or "N/A", From 00a1023bc57e2ea2337006eae1562ea5974f0500 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 3 Jul 2024 11:37:24 +0200 Subject: [PATCH 43/71] use auto config option instead of implicit --- src/zenml/cli/stack.py | 79 +++++++++++++++++++++++++++++------------- 1 file changed, 54 insertions(+), 25 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 8c0f7f37666..9872e9fe21a 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -56,6 +56,10 @@ from zenml.io.fileio import rmtree from zenml.logger import get_logger from zenml.models import StackFilter +from zenml.models.v2.core.service_connector import ( + ServiceConnectorRequest, + ServiceConnectorResponse, +) from zenml.models.v2.misc.full_stack import ( ComponentInfo, FullStackRequest, @@ -282,23 +286,34 @@ def register_stack( created_objects: Set[str] = set() service_connector: Optional[Union[UUID, ServiceConnectorInfo]] = None if provider is not None and connector is None: - use_implicit = Confirm.ask( - f"[bold]{provider.upper()} cloud service connector[/bold] " - "can use the Implicit Authentication by accessing connection " - "configuration of the environment or use one of " - "the authentications methods supported.\n" - "The implicit authentication is great to quickstart, but " - "actual credentials may vary system to system impacting reproducibility.\n" - "Would you like to use the Implicit Authentication method?", - default=False, - show_choices=True, - show_default=True, - ) - if not use_implicit: + service_connector_response = None + use_auto_configure = False + try: + service_connector_response, _ = client.create_service_connector( + name=stack_name, + connector_type=provider, + register=False, + auto_configure=True, + verify=False, + ) + except Exception: + pass + if service_connector_response: + use_auto_configure = Confirm.ask( + f"[bold]{provider.upper()} cloud service connector[/bold] " + "has detected connection credentials in your environment.\n" + "Would you like to use these credentials or create a new " + "configuration by providing connection details?", + default=True, + show_choices=True, + show_default=True, + ) + + connector_selected: Optional[int] = None + if not use_auto_configure: existing_connectors = client.list_service_connectors( connector_type=provider, size=100 ) - connector_selected: Optional[int] = None if existing_connectors.total: connector_selected = cli_utils.multi_choice_prompt( object_type=f"{provider.upper()} service connectors", @@ -312,9 +327,10 @@ def register_stack( default_choice="0", allow_zero_be_a_new_object=True, ) - if use_implicit or connector_selected is None: + if use_auto_configure or connector_selected is None: service_connector = _get_service_connector_info( - cloud_provider=provider, use_implicit=use_implicit + cloud_provider=provider, + connector_details=service_connector_response, ) created_objects.add("service_connector") else: @@ -1963,13 +1979,16 @@ def connect_stack( def _get_service_connector_info( - cloud_provider: str, use_implicit: bool + cloud_provider: str, + connector_details: Optional[ + Union[ServiceConnectorResponse, ServiceConnectorRequest] + ], ) -> ServiceConnectorInfo: """Get a service connector info with given cloud provider. Args: cloud_provider: The cloud provider to use. - use_implicit: Whether to use implicit credentials. + connector_details: Whether to use implicit credentials. Returns: The info model of the created service connector. @@ -1986,7 +2005,7 @@ def _get_service_connector_info( auth_methods = client.get_service_connector_type( cloud_provider ).auth_method_dict - if not use_implicit: + if not connector_details: fixed_auth_methods = list( [ (key, value) @@ -2014,7 +2033,7 @@ def _get_service_connector_info( cli_utils.error("No authentication method selected.") auth_type = fixed_auth_methods[selected_auth_idx][0] else: - auth_type = "implicit" + auth_type = connector_details.auth_method selected_auth_model = auth_methods[auth_type] @@ -2023,11 +2042,21 @@ def _get_service_connector_info( answers = {} for req_field in required_fields: - answers[req_field] = Prompt.ask( - f"Please enter value for `{req_field}`:", - password="format" in properties[req_field] - and properties[req_field]["format"] == "password", - ) + if connector_details: + if conf_value := connector_details.configuration.get( + req_field, None + ): + answers[req_field] = conf_value + elif secret_value := connector_details.secrets.get( + req_field, None + ): + answers[req_field] = secret_value.get_secret_value() + if req_field not in answers: + answers[req_field] = Prompt.ask( + f"Please enter value for `{req_field}`:", + password="format" in properties[req_field] + and properties[req_field]["format"] == "password", + ) return ServiceConnectorInfo( type=cloud_provider, From eeb8bea961946d1594fa1edcf20637a9c36423df Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 3 Jul 2024 13:46:06 +0200 Subject: [PATCH 44/71] add connector index validator --- src/zenml/models/v2/misc/full_stack.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 9299dd37fd4..3f95e1bba39 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -16,7 +16,7 @@ from typing import Any, Dict, List, Optional, Union from uuid import UUID -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator from zenml.constants import STR_FIELD_MAX_LENGTH from zenml.enums import StackComponentType @@ -74,3 +74,20 @@ class FullStackRequest(BaseRequest): "existing components or request information for brand new " "components.", ) + + @model_validator(mode="after") + def _validate_indexes_in_components(self) -> "FullStackRequest": + for component in self.components.values(): + if isinstance(component, ComponentInfo): + if component.service_connector_index is not None: + if ( + component.service_connector_index < 0 + or component.service_connector_index + >= len(self.service_connectors) + ): + raise ValueError( + f"Service connector index {component.service_connector_index} " + "is out of range. Please provide a valid index referring to " + "the position in the list of service connectors." + ) + return self From 8bf0a3426f0a7d8ac03b534336345be556fd4ec6 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 3 Jul 2024 14:42:01 +0200 Subject: [PATCH 45/71] suggestions and mypy --- src/zenml/models/v2/misc/full_stack.py | 4 ++-- src/zenml/zen_stores/sql_zen_store.py | 9 +++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 3f95e1bba39..6af321a508e 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -49,8 +49,8 @@ class ComponentInfo(BaseModel): class FullStackRequest(BaseRequest): """Request model for a full-stack.""" - user: Optional[UUID] = None - workspace: Optional[UUID] = None + user: UUID + workspace: UUID name: str = Field( title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 0b1e5e073c3..16601769ed8 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6930,7 +6930,8 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: The registered stack. Raises: - RuntimeError: If the full stack creation fails. + ValueError: If the full stack creation fails, due to the corrupted input. + RuntimeError: If the full stack creation fails, due to unforeseen errors. """ # We can not use the decorator here, as we need custom metadata with track_handler( @@ -7032,7 +7033,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: requirements = flavor_model.connector_requirements if not requirements: - raise RuntimeError( + raise ValueError( f"The '{flavor_model.name}' implementation " "does not support using a service connector to " "connect to resources." @@ -7056,7 +7057,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: ) if not satisfied: - raise RuntimeError( + raise ValueError( "Please pick a connector that is " "compatible with the component flavor and " "try again.." @@ -7066,7 +7067,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: if service_connector.resource_id: resource_id = service_connector.resource_id elif service_connector.supports_instances: - raise RuntimeError( + raise ValueError( f"Multiple {resource_type} resources " "are available for the selected " "connector. Please use a `resource_id` " From a1cf86208ac68353876bdd0854132c6eb292f14b Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 3 Jul 2024 15:30:52 +0200 Subject: [PATCH 46/71] ignore local, if rejected --- src/zenml/cli/stack.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 9872e9fe21a..f5d156da731 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -311,6 +311,7 @@ def register_stack( connector_selected: Optional[int] = None if not use_auto_configure: + service_connector_response = None existing_connectors = client.list_service_connectors( connector_type=provider, size=100 ) From b577cc81007d254010657d773015ddbfca92385d Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 3 Jul 2024 16:46:37 +0200 Subject: [PATCH 47/71] mypy --- src/zenml/models/v2/misc/full_stack.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 6af321a508e..5e5a96996be 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -20,7 +20,7 @@ from zenml.constants import STR_FIELD_MAX_LENGTH from zenml.enums import StackComponentType -from zenml.models.v2.base.base import BaseRequest +from zenml.models.v2.base.scoped import WorkspaceScopedRequest class ServiceConnectorInfo(BaseModel): @@ -46,12 +46,9 @@ class ComponentInfo(BaseModel): configuration: Dict[str, Any] = {} -class FullStackRequest(BaseRequest): +class FullStackRequest(WorkspaceScopedRequest): """Request model for a full-stack.""" - user: UUID - workspace: UUID - name: str = Field( title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH ) From 7f280280f9d59b6f9d0938afeb32d3cf057f713d Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Wed, 3 Jul 2024 14:48:40 +0000 Subject: [PATCH 48/71] Auto-update of LLM Finetuning template --- examples/llm_finetuning/steps/finetune.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/examples/llm_finetuning/steps/finetune.py b/examples/llm_finetuning/steps/finetune.py index 70a837d4269..eddf19a3292 100644 --- a/examples/llm_finetuning/steps/finetune.py +++ b/examples/llm_finetuning/steps/finetune.py @@ -102,8 +102,12 @@ def finetune( if should_print: logger.info("Loading datasets...") tokenizer = load_tokenizer(base_model_id, use_fast=use_fast) - tokenized_train_dataset = load_from_disk(dataset_dir / "train") - tokenized_val_dataset = load_from_disk(dataset_dir / "val") + tokenized_train_dataset = load_from_disk( + str((dataset_dir / "train").absolute()) + ) + tokenized_val_dataset = load_from_disk( + str((dataset_dir / "val").absolute()) + ) if should_print: logger.info("Loading base model...") @@ -162,7 +166,6 @@ def finetune( if should_print: logger.info("Saving model...") - ft_model_dir = Path(ft_model_dir) if not use_accelerate or accelerator.is_main_process: ft_model_dir.mkdir(parents=True, exist_ok=True) if not use_accelerate: From d1c039449d6a4e63338821429f356b4bb63c12c6 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 3 Jul 2024 16:58:33 +0200 Subject: [PATCH 49/71] revert fullstack model changes --- src/zenml/cli/stack.py | 4 +++- src/zenml/cli/utils.py | 2 -- src/zenml/models/v2/misc/full_stack.py | 7 +++++-- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index f5d156da731..523fd0c6115 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -404,6 +404,7 @@ def register_stack( ) ) else: + breakpoint() _, service_connector_resource_model = ( client.create_service_connector( name=stack_name, @@ -1999,7 +2000,7 @@ def _get_service_connector_info( """ from rich.prompt import Prompt - if cloud_provider not in {"aws", "azure", "gcp"}: + if cloud_provider not in {"aws"}: raise ValueError(f"Unknown cloud provider {cloud_provider}") client = Client() @@ -2058,6 +2059,7 @@ def _get_service_connector_info( password="format" in properties[req_field] and properties[req_field]["format"] == "password", ) + Console().print("All mandatory configuration parameters received!") return ServiceConnectorInfo( type=cloud_provider, diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index d8df9e584ca..ee202aff3a7 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -2842,8 +2842,6 @@ def requires_mac_env_var_warning() -> bool: Returns: bool: True if a warning needs to be shown, False otherwise. """ - breakpoint() - if mac_version := platform.mac_ver()[0]: try: major, minor, _ = mac_version.split(".") diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 5e5a96996be..3f95e1bba39 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -20,7 +20,7 @@ from zenml.constants import STR_FIELD_MAX_LENGTH from zenml.enums import StackComponentType -from zenml.models.v2.base.scoped import WorkspaceScopedRequest +from zenml.models.v2.base.base import BaseRequest class ServiceConnectorInfo(BaseModel): @@ -46,9 +46,12 @@ class ComponentInfo(BaseModel): configuration: Dict[str, Any] = {} -class FullStackRequest(WorkspaceScopedRequest): +class FullStackRequest(BaseRequest): """Request model for a full-stack.""" + user: Optional[UUID] = None + workspace: Optional[UUID] = None + name: str = Field( title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH ) From 15661e040e0d6208574c1ae9c9979b853b14dc79 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Thu, 4 Jul 2024 12:08:13 +0200 Subject: [PATCH 50/71] solving the linting issue --- src/zenml/zen_stores/rest_zen_store.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 660ecb12463..485c86e4c79 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -2755,10 +2755,12 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: Returns: The registered stack. """ - return self._create_workspace_scoped_resource( + assert full_stack.workspace is not None + + return self._create_resource( resource=full_stack, - route=FULL_STACK, response_model=StackResponse, + route=f"{WORKSPACES}/{str(full_stack.workspace)}{FULL_STACK}", ) def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse: From 28a6db015c36bf754ff83c6ccf33ed62d6e93867 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 4 Jul 2024 13:59:54 +0200 Subject: [PATCH 51/71] GCP stack wizard --- src/zenml/cli/stack.py | 144 ++++++++++++++++++++------ src/zenml/zen_stores/sql_zen_store.py | 8 +- 2 files changed, 120 insertions(+), 32 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index f5d156da731..8aec3b0c355 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -80,6 +80,7 @@ verify_spec_and_tf_files_exist, ) from zenml.utils.yaml_utils import read_yaml, write_yaml +from zenml.zen_stores.rest_zen_store import RestZenStore if TYPE_CHECKING: from zenml.models import StackResponse @@ -397,26 +398,39 @@ def register_stack( if component_selected is None: if service_connector_resource_model is None: - if isinstance(service_connector, UUID): - service_connector_resource_model = ( - client.verify_service_connector( - service_connector + with console.status( + "Exploring resources available to the service connector...\n" + ): + if isinstance(service_connector, UUID): + service_connector_resource_model = ( + client.verify_service_connector( + service_connector + ) ) - ) - else: - _, service_connector_resource_model = ( - client.create_service_connector( - name=stack_name, - connector_type=service_connector.type, - auth_method=service_connector.auth_method, - configuration=service_connector.configuration, - register=False, + else: + default_http_timeout = 30 + if isinstance(client.zen_store, RestZenStore): + default_http_timeout = ( + client.zen_store.config.http_timeout + ) + client.zen_store.config.http_timeout = 120 + _, service_connector_resource_model = ( + client.create_service_connector( + name=stack_name, + connector_type=service_connector.type, + auth_method=service_connector.auth_method, + configuration=service_connector.configuration, + register=False, + ) ) + if isinstance(client.zen_store, RestZenStore): + client.zen_store.config.http_timeout = ( + default_http_timeout + ) + if service_connector_resource_model is None: + cli_utils.error( + f"Failed to validate service connector {service_connector}..." ) - if service_connector_resource_model is None: - cli_utils.error( - f"Failed to validate service connector {service_connector}..." - ) if provider is None: if isinstance( service_connector_resource_model.connector_type, @@ -1999,7 +2013,7 @@ def _get_service_connector_info( """ from rich.prompt import Prompt - if cloud_provider not in {"aws", "azure", "gcp"}: + if cloud_provider not in {"aws", "gcp"}: raise ValueError(f"Unknown cloud provider {cloud_provider}") client = Client() @@ -2028,7 +2042,7 @@ def _get_service_connector_info( object_type=f"authentication methods for {cloud_provider}", choices=choices, headers=headers, - prompt_text="Please choose one of the authentication option above.", + prompt_text="Please choose one of the authentication option above", ) if selected_auth_idx is None: cli_utils.error("No authentication method selected.") @@ -2058,6 +2072,11 @@ def _get_service_connector_info( password="format" in properties[req_field] and properties[req_field]["format"] == "password", ) + if cloud_provider == "gcp" and auth_type in { + "service-account", + "external-account", + }: + answers["generate_temporary_tokens"] = False return ServiceConnectorInfo( type=cloud_provider, @@ -2097,6 +2116,9 @@ def _get_stack_component_info( AWS_DOCS = ( "https://docs.zenml.io/how-to/auth-management/aws-service-connector" ) + GCP_DOCS = ( + "https://docs.zenml.io/how-to/auth-management/gcp-service-connector" + ) flavor = "undefined" service_connector_resource_id = None @@ -2120,7 +2142,19 @@ def _get_stack_component_info( elif cloud_provider == "azure": flavor = "azure" elif cloud_provider == "gcp": - flavor = "gcs" + flavor = "gcp" + for each in service_connector_resource_models: + if each.resource_type == "gcs-bucket": + available_storages = each.resource_ids or [] + if not available_storages: + cli_utils.error( + "We were unable to find any GCS buckets available " + "to configured service connector. Please, verify " + "that needed permission are granted for the " + "service connector.\nDocumentation for the GCS " + "Buckets configuration can be found at " + f"{GCP_DOCS}#gcs-bucket" + ) selected_storage_idx = cli_utils.multi_choice_prompt( object_type=f"{cloud_provider.upper()} storages", @@ -2136,6 +2170,17 @@ def _get_stack_component_info( config = {"path": selected_storage} service_connector_resource_id = selected_storage elif component_type == "orchestrator": + + def query_gcp_region() -> str: + from google.cloud.aiplatform.constants import base as constants + + region = Prompt.ask( + "Select a location for your Vertex AI jobs:", + choices=sorted(list(constants.SUPPORTED_REGIONS)), + show_choices=True, + ) + return region + if cloud_provider == "aws": available_orchestrators = [] for each in service_connector_resource_models: @@ -2162,7 +2207,30 @@ def _get_stack_component_info( f"{AWS_DOCS}#eks-kubernetes-cluster" ) elif cloud_provider == "gcp": - pass + available_orchestrators = [] + for each in service_connector_resource_models: + types = [] + if each.resource_type == "gcp-generic": + types = ["Vertex AI", "Skypilot (Compute)"] + if each.resource_type == "kubernetes-cluster": + types = ["Kubernetes"] + + if each.resource_ids: + for orchestrator in each.resource_ids: + for t in types: + available_orchestrators.append([t, orchestrator]) + if not available_orchestrators: + cli_utils.error( + "We were unable to find any orchestrator engines " + "available to the service connector. Please, verify " + "that needed permission are granted for the " + "service connector.\nDocumentation for the Generic " + "GCP resource configuration can be found at " + f"{GCP_DOCS}#generic-gcp-resource\n" + "Documentation for the GKE Kubernetes resource " + "configuration can be found at " + f"{GCP_DOCS}#gke-kubernetes-cluster" + ) elif cloud_provider == "azure": pass @@ -2179,25 +2247,31 @@ def _get_stack_component_info( selected_orchestrator_idx ] + config = {} if selected_orchestrator[0] == "Sagemaker": flavor = "sagemaker" - execution_role = Prompt.ask("Please enter an execution role ARN:") - config = {"execution_role": execution_role} + execution_role = Prompt.ask("Enter an execution role ARN:") + config["execution_role"] = execution_role elif selected_orchestrator[0] == "Skypilot (EC2)": flavor = "vm_aws" - config = {"region": selected_orchestrator[1]} + config["region"] = selected_orchestrator[1] + elif selected_orchestrator[0] == "Skypilot (Compute)": + flavor = "vm_gcp" + config["region"] = query_gcp_region() + elif selected_orchestrator[0] == "Vertex AI": + flavor = "vertex" + config["location"] = query_gcp_region() elif selected_orchestrator[0] == "Kubernetes": flavor = "kubernetes" - config = {} else: raise ValueError( f"Unknown orchestrator type {selected_orchestrator[0]}" ) service_connector_resource_id = selected_orchestrator[1] elif component_type == "container_registry": - available_registries: List[str] = [] - if cloud_provider == "aws": - flavor = "aws" + + def _get_registries(registry_name: str, docs_link: str) -> List[str]: + available_registries: List[str] = [] for each in service_connector_resource_models: if each.resource_type == "docker-registry": available_registries = each.resource_ids or [] @@ -2206,14 +2280,24 @@ def _get_stack_component_info( "We were unable to find any container registries " "available to the service connector. Please, verify " "that needed permission are granted for the " - "service connector.\nDocumentation for the ECR " + f"service connector.\nDocumentation for the {registry_name} " "container registry resource configuration can " - f"be found at {AWS_DOCS}#ecr-container-registry" + f"be found at {docs_link}" ) + return available_registries + + if cloud_provider == "aws": + flavor = "aws" + available_registries = _get_registries( + "ECR", f"{AWS_DOCS}#ecr-container-registry" + ) elif cloud_provider == "azure": flavor = "azure" elif cloud_provider == "gcp": flavor = "gcp" + available_registries = _get_registries( + "GCR", f"{GCP_DOCS}#gcr-container-registry" + ) selected_registry_idx = cli_utils.multi_choice_prompt( object_type=f"{cloud_provider.upper()} registries", diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 16601769ed8..25e4d9d9d21 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -7026,7 +7026,10 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: type=component_type, ) ) - assert len(flavor_list) == 1 + assert len(flavor_list) == 1, ( + f"{len(flavor_list)} flavors found for " + f"{component_info.flavor}/{component_type}" + ) flavor_model = flavor_list[0] @@ -7111,6 +7114,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: stack_response.get_analytics_metadata() ) except Exception as e: + breakpoint() for component_id in components_created_ids: self.delete_stack_component(component_id=component_id) for service_connector_id in service_connectors_created_ids: @@ -7118,7 +7122,7 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: service_connector_id=service_connector_id ) raise RuntimeError( - f"Full Stack creation has failed {e}. Cleaning up the " + f"Full Stack creation has failed:\n{e}\n Cleaning up the " f"created entities." ) from e return stack_response From 5e574795bf14dcce4f8d2df7e5bbb0ee642c0318 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Thu, 4 Jul 2024 15:23:13 +0200 Subject: [PATCH 52/71] changing the models to include the new labels as well --- src/zenml/models/v2/core/component.py | 18 +++++++++++ src/zenml/models/v2/core/service_connector.py | 17 ++++++++++ src/zenml/models/v2/core/stack.py | 31 +++++++++++++++++++ src/zenml/models/v2/misc/full_stack.py | 4 +++ 4 files changed, 70 insertions(+) diff --git a/src/zenml/models/v2/core/component.py b/src/zenml/models/v2/core/component.py index 086bb4c8358..46edf5a7372 100644 --- a/src/zenml/models/v2/core/component.py +++ b/src/zenml/models/v2/core/component.py @@ -238,6 +238,24 @@ class ComponentResponse( max_length=STR_FIELD_MAX_LENGTH, ) + def get_analytics_metadata(self) -> Dict[str, Any]: + """Add the component labels to analytics metadata. + + Returns: + Dict of analytics metadata. + """ + metadata = super().get_analytics_metadata() + + if self.labels is not None: + metadata.update( + { + label: value + for label, value in self.labels.items() + if label.startswith("zenml:") + } + ) + return metadata + def get_hydrated_version(self) -> "ComponentResponse": """Get the hydrated version of this component. diff --git a/src/zenml/models/v2/core/service_connector.py b/src/zenml/models/v2/core/service_connector.py index f788242e023..825ef1bc7e5 100644 --- a/src/zenml/models/v2/core/service_connector.py +++ b/src/zenml/models/v2/core/service_connector.py @@ -493,6 +493,23 @@ class ServiceConnectorResponse( max_length=STR_FIELD_MAX_LENGTH, ) + def get_analytics_metadata(self) -> Dict[str, Any]: + """Add the service connector labels to analytics metadata. + + Returns: + Dict of analytics metadata. + """ + metadata = super().get_analytics_metadata() + + metadata.update( + { + label: value + for label, value in self.labels.items() + if label.startswith("zenml:") + } + ) + return metadata + def get_hydrated_version(self) -> "ServiceConnectorResponse": """Get the hydrated version of this service connector. diff --git a/src/zenml/models/v2/core/stack.py b/src/zenml/models/v2/core/stack.py index 390b88ef149..f7fe723fdde 100644 --- a/src/zenml/models/v2/core/stack.py +++ b/src/zenml/models/v2/core/stack.py @@ -36,6 +36,7 @@ if TYPE_CHECKING: from sqlalchemy.sql.elements import ColumnElement + # ------------------ Request Model ------------------ @@ -59,6 +60,10 @@ class StackRequest(WorkspaceScopedRequest): title="A mapping of stack component types to the actual" "instances of components of this type.", ) + labels: Optional[Dict[str, Any]] = Field( + default=None, + title="The stack labels.", + ) @property def is_valid(self) -> bool: @@ -109,6 +114,10 @@ class StackUpdate(BaseUpdate): "instances of components of this type.", default=None, ) + labels: Optional[Dict[str, Any]] = Field( + default=None, + title="The stack labels.", + ) # ------------------ Response Model ------------------ @@ -134,6 +143,10 @@ class StackResponseMetadata(WorkspaceScopedResponseMetadata): default=None, title="The path to the stack spec used for mlstacks deployments.", ) + labels: Optional[Dict[str, Any]] = Field( + default=None, + title="The stack labels.", + ) class StackResponseResources(WorkspaceScopedResponseResources): @@ -214,6 +227,15 @@ def get_analytics_metadata(self) -> Dict[str, Any]: """ metadata = super().get_analytics_metadata() metadata.update({ct: c[0].flavor for ct, c in self.components.items()}) + + if self.labels is not None: + metadata.update( + { + label: value + for label, value in self.labels.items() + if label.startswith("zenml:") + } + ) return metadata @property @@ -243,6 +265,15 @@ def components(self) -> Dict[StackComponentType, List[ComponentResponse]]: """ return self.get_metadata().components + @property + def labels(self) -> Optional[Dict[str, Any]]: + """The `labels` property. + + Returns: + the value of the property. + """ + return self.get_metadata().labels + # ------------------ Filter Model ------------------ diff --git a/src/zenml/models/v2/misc/full_stack.py b/src/zenml/models/v2/misc/full_stack.py index 3f95e1bba39..ccbc3a8fcad 100644 --- a/src/zenml/models/v2/misc/full_stack.py +++ b/src/zenml/models/v2/misc/full_stack.py @@ -60,6 +60,10 @@ class FullStackRequest(BaseRequest): title="The description of the stack", max_length=STR_FIELD_MAX_LENGTH, ) + labels: Optional[Dict[str, Any]] = Field( + default=None, + title="The stack labels.", + ) service_connectors: List[Union[UUID, ServiceConnectorInfo]] = Field( default=[], title="The service connectors dictionary for the full stack " From ad3a438799b2a47b7a32f9895f70c8cf28207c39 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Thu, 4 Jul 2024 15:26:08 +0200 Subject: [PATCH 53/71] adding the labels to stacks in the client --- src/zenml/client.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/zenml/client.py b/src/zenml/client.py index 6972a90d860..9e299916a94 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -1153,6 +1153,7 @@ def create_stack( name: str, components: Mapping[StackComponentType, Union[str, UUID]], stack_spec_file: Optional[str] = None, + labels: Optional[Dict[str, Any]] = None, ) -> StackResponse: """Registers a stack and its components. @@ -1160,6 +1161,7 @@ def create_stack( name: The name of the stack to register. components: dictionary which maps component types to component names stack_spec_file: path to the stack spec file + labels: The labels of the stack. Returns: The model of the registered stack. @@ -1184,6 +1186,7 @@ def create_stack( stack_spec_path=stack_spec_file, workspace=self.active_workspace.id, user=self.active_user.id, + labels=labels, ) self._validate_stack_configuration(stack=stack) @@ -1279,6 +1282,7 @@ def update_stack( name_id_or_prefix: Optional[Union[UUID, str]] = None, name: Optional[str] = None, stack_spec_file: Optional[str] = None, + labels: Optional[Dict[str, Any]] = None, description: Optional[str] = None, component_updates: Optional[ Dict[StackComponentType, List[Union[UUID, str]]] @@ -1289,7 +1293,8 @@ def update_stack( Args: name_id_or_prefix: The name, id or prefix of the stack to update. name: the new name of the stack. - stack_spec_file: path to the stack spec file + stack_spec_file: path to the stack spec file. + labels: The new labels of the stack component. description: the new description of the stack. component_updates: dictionary which maps stack component types to lists of new stack component names or ids. @@ -1343,6 +1348,15 @@ def update_stack( for c_type, c_list in components_dict.items() } + if labels is not None: + existing_labels = stack.labels or {} + existing_labels.update(labels) + + existing_labels = { + k: v for k, v in existing_labels.items() if v is not None + } + update_model.labels = existing_labels + updated_stack = self.zen_store.update_stack( stack_id=stack.id, stack_update=update_model, From 6a1d35d3cbca598cafbb6be94b7a9b093a517374 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Thu, 4 Jul 2024 15:27:17 +0200 Subject: [PATCH 54/71] changing the schema --- src/zenml/zen_stores/schemas/stack_schemas.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/zenml/zen_stores/schemas/stack_schemas.py b/src/zenml/zen_stores/schemas/stack_schemas.py index f538fe4aac0..4307d2cb6ff 100644 --- a/src/zenml/zen_stores/schemas/stack_schemas.py +++ b/src/zenml/zen_stores/schemas/stack_schemas.py @@ -13,6 +13,8 @@ # permissions and limitations under the License. """SQL Model Implementations for Stacks.""" +import base64 +import json from datetime import datetime from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -75,6 +77,7 @@ class StackSchema(NamedSchema, table=True): __tablename__ = "stack" stack_spec_path: Optional[str] + labels: Optional[bytes] workspace_id: UUID = build_foreign_key_field( source=__tablename__, @@ -124,6 +127,10 @@ def update( ).items(): if field == "components": self.components = components + elif field == "labels": + self.labels = base64.b64encode( + json.dumps(stack_update.labels).encode("utf-8") + ) else: setattr(self, field, value) @@ -158,6 +165,9 @@ def to_model( workspace=self.workspace.to_model(), components={c.type: [c.to_model()] for c in self.components}, stack_spec_path=self.stack_spec_path, + labels=json.loads(base64.b64decode(self.labels).decode()) + if self.labels + else None, ) return StackResponse( From 0c32fd365f306b2b4ed336ae0a9fd06b7bd1b562 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Thu, 4 Jul 2024 15:30:10 +0200 Subject: [PATCH 55/71] adding the necessary migration script --- .../6d2ebf1a2017_adding_labels_to_stacks.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 src/zenml/zen_stores/migrations/versions/6d2ebf1a2017_adding_labels_to_stacks.py diff --git a/src/zenml/zen_stores/migrations/versions/6d2ebf1a2017_adding_labels_to_stacks.py b/src/zenml/zen_stores/migrations/versions/6d2ebf1a2017_adding_labels_to_stacks.py new file mode 100644 index 00000000000..14cc41fc4de --- /dev/null +++ b/src/zenml/zen_stores/migrations/versions/6d2ebf1a2017_adding_labels_to_stacks.py @@ -0,0 +1,27 @@ +"""adding labels to stacks [6d2ebf1a2017]. + +Revision ID: 6d2ebf1a2017 +Revises: 0.60.0 +Create Date: 2024-07-04 15:29:45.744811 + +""" + +# revision identifiers, used by Alembic. +revision = "6d2ebf1a2017" +down_revision = "0.60.0" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + """Upgrade database schema and/or data, creating a new revision.""" + # ### commands auto generated by Alembic - please adjust! ### + pass + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade database schema and/or data back to the previous revision.""" + # ### commands auto generated by Alembic - please adjust! ### + pass + # ### end Alembic commands ### From c7bfc24b2c1a1cfadf2d3436b5e9913021248577 Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Thu, 4 Jul 2024 15:30:23 +0200 Subject: [PATCH 56/71] adjusting the sql zen store after the labels change --- src/zenml/zen_stores/sql_zen_store.py | 346 +++++++++++++------------- 1 file changed, 177 insertions(+), 169 deletions(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 16601769ed8..4c41499e82a 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -6912,6 +6912,9 @@ def create_stack(self, stack: StackRequest) -> StackResponse: name=stack.name, description=stack.description, components=defined_components, + labels=base64.b64encode( + json.dumps(stack.labels).encode("utf-8") + ), ) session.add(new_stack_schema) @@ -6930,199 +6933,204 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: The registered stack. Raises: - ValueError: If the full stack creation fails, due to the corrupted input. - RuntimeError: If the full stack creation fails, due to unforeseen errors. + ValueError: If the full stack creation fails, due to the corrupted + input. + RuntimeError: If the full stack creation fails, due to unforeseen + errors. """ - # We can not use the decorator here, as we need custom metadata - with track_handler( - event=AnalyticsEvent.REGISTERED_STACK, - metadata={"generated_by_wizard": True}, - ) as handler: - # For clean-up purposes, each created entity is tracked here - service_connectors_created_ids: List[UUID] = [] - components_created_ids: List[UUID] = [] + # For clean-up purposes, each created entity is tracked here + service_connectors_created_ids: List[UUID] = [] + components_created_ids: List[UUID] = [] - try: - # Validate the name of the new stack - validate_name(full_stack) + try: + # Validate the name of the new stack + validate_name(full_stack) - # Service Connectors - service_connectors: List[ServiceConnectorResponse] = [] + if full_stack.labels is None: + full_stack.labels = {} - for connector_id_or_info in full_stack.service_connectors: - # Fetch an existing service connector - if isinstance(connector_id_or_info, UUID): - service_connectors.append( - self.get_service_connector(connector_id_or_info) - ) - # Create a new service connector - else: - connector_name = full_stack.name - while True: - try: - service_connector_request = ServiceConnectorRequest( - name=connector_name, - connector_type=connector_id_or_info.type, - auth_method=connector_id_or_info.auth_method, - configuration=connector_id_or_info.configuration, - user=full_stack.user, - workspace=full_stack.workspace, - ) - service_connector_response = self.create_service_connector( + full_stack.labels.update({"zenml:full_stack": True}) + + # Service Connectors + service_connectors: List[ServiceConnectorResponse] = [] + + for connector_id_or_info in full_stack.service_connectors: + # Fetch an existing service connector + if isinstance(connector_id_or_info, UUID): + service_connectors.append( + self.get_service_connector(connector_id_or_info) + ) + # Create a new service connector + else: + connector_name = full_stack.name + while True: + try: + service_connector_request = ServiceConnectorRequest( + name=connector_name, + connector_type=connector_id_or_info.type, + auth_method=connector_id_or_info.auth_method, + configuration=connector_id_or_info.configuration, + user=full_stack.user, + workspace=full_stack.workspace, + ) + service_connector_response = ( + self.create_service_connector( service_connector=service_connector_request ) - service_connectors.append( - service_connector_response - ) - service_connectors_created_ids.append( - service_connector_response.id - ) - break - except EntityExistsError: - connector_name = f"{full_stack.name}-{random_str(4)}".lower() - continue - - # Stack Components - components_mapping: Dict[StackComponentType, List[UUID]] = {} - - for ( - component_type, - component_info, - ) in full_stack.components.items(): - # Fetch an existing component - if isinstance(component_info, UUID): - component = self.get_stack_component( - component_id=component_info - ) - # Create a new component - else: - component_name = full_stack.name - while True: - try: - component_request = ComponentRequest( - name=component_name, - type=component_type, - flavor=component_info.flavor, - configuration=component_info.configuration, - user=full_stack.user, - workspace=full_stack.workspace, - ) - component = self.create_stack_component( - component=component_request - ) - components_created_ids.append(component.id) - break - except EntityExistsError: - component_name = f"{full_stack.name}-{random_str(4)}".lower() - continue - - if component_info.service_connector_index is not None: - service_connector = service_connectors[ - component_info.service_connector_index - ] - flavor_list = self.list_flavors( - flavor_filter_model=FlavorFilter( - name=component_info.flavor, - type=component_type, - ) ) - assert len(flavor_list) == 1 + service_connectors.append( + service_connector_response + ) + service_connectors_created_ids.append( + service_connector_response.id + ) + break + except EntityExistsError: + connector_name = ( + f"{full_stack.name}-{random_str(4)}".lower() + ) + continue + + # Stack Components + components_mapping: Dict[StackComponentType, List[UUID]] = {} - flavor_model = flavor_list[0] + for ( + component_type, + component_info, + ) in full_stack.components.items(): + # Fetch an existing component + if isinstance(component_info, UUID): + component = self.get_stack_component( + component_id=component_info + ) + # Create a new component + else: + component_name = full_stack.name + while True: + try: + component_request = ComponentRequest( + name=component_name, + type=component_type, + flavor=component_info.flavor, + configuration=component_info.configuration, + user=full_stack.user, + workspace=full_stack.workspace, + ) + component = self.create_stack_component( + component=component_request + ) + components_created_ids.append(component.id) + break + except EntityExistsError: + component_name = ( + f"{full_stack.name}-{random_str(4)}".lower() + ) + continue + + if component_info.service_connector_index is not None: + service_connector = service_connectors[ + component_info.service_connector_index + ] + flavor_list = self.list_flavors( + flavor_filter_model=FlavorFilter( + name=component_info.flavor, + type=component_type, + ) + ) + assert len(flavor_list) == 1 - requirements = flavor_model.connector_requirements + flavor_model = flavor_list[0] - if not requirements: - raise ValueError( - f"The '{flavor_model.name}' implementation " - "does not support using a service connector to " - "connect to resources." - ) + requirements = flavor_model.connector_requirements - if component_info.service_connector_resource_id: - resource_id = component_info.service_connector_resource_id - else: - resource_id = None - resource_type = requirements.resource_type - if requirements.resource_id_attr is not None: - resource_id = ( - component_info.configuration.get( - requirements.resource_id_attr - ) - ) - - satisfied, msg = requirements.is_satisfied_by( - connector=service_connector, - component=component, + if not requirements: + raise ValueError( + f"The '{flavor_model.name}' implementation " + "does not support using a service connector to " + "connect to resources." ) - if not satisfied: - raise ValueError( - "Please pick a connector that is " - "compatible with the component flavor and " - "try again.." + if component_info.service_connector_resource_id: + resource_id = ( + component_info.service_connector_resource_id + ) + else: + resource_id = None + resource_type = requirements.resource_type + if requirements.resource_id_attr is not None: + resource_id = component_info.configuration.get( + requirements.resource_id_attr ) - if not resource_id: - if service_connector.resource_id: - resource_id = service_connector.resource_id - elif service_connector.supports_instances: - raise ValueError( - f"Multiple {resource_type} resources " - "are available for the selected " - "connector. Please use a `resource_id` " - "to configure a " - f"{resource_type} resource." - ) - - component_update = ComponentUpdate( - connector=service_connector.id, - connector_resource_id=resource_id, - ) - self.update_stack_component( - component_id=component.id, - component_update=component_update, + satisfied, msg = requirements.is_satisfied_by( + connector=service_connector, + component=component, + ) + + if not satisfied: + raise ValueError( + "Please pick a connector that is " + "compatible with the component flavor and " + "try again.." ) - components_mapping[component_type] = [ - component.id, - ] + if not resource_id: + if service_connector.resource_id: + resource_id = service_connector.resource_id + elif service_connector.supports_instances: + raise ValueError( + f"Multiple {resource_type} resources " + "are available for the selected " + "connector. Please use a `resource_id` " + "to configure a " + f"{resource_type} resource." + ) - # Stack - stack_name = full_stack.name - while True: - try: - stack_request = StackRequest( - user=full_stack.user, - workspace=full_stack.workspace, - name=stack_name, - description=full_stack.description, - components=components_mapping, + component_update = ComponentUpdate( + connector=service_connector.id, + connector_resource_id=resource_id, ) - stack_response = self.create_stack(stack_request) - break - except EntityExistsError: - stack_name = ( - f"{full_stack.name}-{random_str(4)}".lower() + self.update_stack_component( + component_id=component.id, + component_update=component_update, ) - continue - handler.metadata.update( - stack_response.get_analytics_metadata() - ) - except Exception as e: - for component_id in components_created_ids: - self.delete_stack_component(component_id=component_id) - for service_connector_id in service_connectors_created_ids: - self.delete_service_connector( - service_connector_id=service_connector_id + components_mapping[component_type] = [ + component.id, + ] + + # Stack + stack_name = full_stack.name + while True: + try: + stack_request = StackRequest( + user=full_stack.user, + workspace=full_stack.workspace, + name=stack_name, + description=full_stack.description, + components=components_mapping, + labels=full_stack.labels, ) - raise RuntimeError( - f"Full Stack creation has failed {e}. Cleaning up the " - f"created entities." - ) from e + stack_response = self.create_stack(stack_request) + + break + except EntityExistsError: + stack_name = f"{full_stack.name}-{random_str(4)}".lower() + return stack_response + except Exception as e: + for component_id in components_created_ids: + self.delete_stack_component(component_id=component_id) + for service_connector_id in service_connectors_created_ids: + self.delete_service_connector( + service_connector_id=service_connector_id + ) + raise RuntimeError( + f"Full Stack creation has failed {e}. Cleaning up the " + f"created entities." + ) from e + def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse: """Get a stack by its unique ID. From 893a1fc16f15a1689db9126b2d40cd137715314d Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Thu, 4 Jul 2024 16:10:38 +0200 Subject: [PATCH 57/71] the correct migration script --- ...> 0d707865f404_adding_labels_to_stacks.py} | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) rename src/zenml/zen_stores/migrations/versions/{6d2ebf1a2017_adding_labels_to_stacks.py => 0d707865f404_adding_labels_to_stacks.py} (53%) diff --git a/src/zenml/zen_stores/migrations/versions/6d2ebf1a2017_adding_labels_to_stacks.py b/src/zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py similarity index 53% rename from src/zenml/zen_stores/migrations/versions/6d2ebf1a2017_adding_labels_to_stacks.py rename to src/zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py index 14cc41fc4de..527940ca2cd 100644 --- a/src/zenml/zen_stores/migrations/versions/6d2ebf1a2017_adding_labels_to_stacks.py +++ b/src/zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py @@ -1,13 +1,16 @@ -"""adding labels to stacks [6d2ebf1a2017]. +"""adding labels to stacks [0d707865f404]. -Revision ID: 6d2ebf1a2017 +Revision ID: 0d707865f404 Revises: 0.60.0 -Create Date: 2024-07-04 15:29:45.744811 +Create Date: 2024-07-04 16:10:07.709184 """ +import sqlalchemy as sa +from alembic import op + # revision identifiers, used by Alembic. -revision = "6d2ebf1a2017" +revision = "0d707865f404" down_revision = "0.60.0" branch_labels = None depends_on = None @@ -16,12 +19,18 @@ def upgrade() -> None: """Upgrade database schema and/or data, creating a new revision.""" # ### commands auto generated by Alembic - please adjust! ### - pass + with op.batch_alter_table("stack", schema=None) as batch_op: + batch_op.add_column( + sa.Column("labels", sa.LargeBinary(), nullable=True) + ) + # ### end Alembic commands ### def downgrade() -> None: """Downgrade database schema and/or data back to the previous revision.""" # ### commands auto generated by Alembic - please adjust! ### - pass + with op.batch_alter_table("stack", schema=None) as batch_op: + batch_op.drop_column("labels") + # ### end Alembic commands ### From 7ccdb351989e990635c7a43327a83e62c7bcc55c Mon Sep 17 00:00:00 2001 From: Baris Can Durak Date: Thu, 4 Jul 2024 16:12:48 +0200 Subject: [PATCH 58/71] removed autogenerated comments --- .../versions/0d707865f404_adding_labels_to_stacks.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py b/src/zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py index 527940ca2cd..777d4c4c2ff 100644 --- a/src/zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py +++ b/src/zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py @@ -18,19 +18,13 @@ def upgrade() -> None: """Upgrade database schema and/or data, creating a new revision.""" - # ### commands auto generated by Alembic - please adjust! ### with op.batch_alter_table("stack", schema=None) as batch_op: batch_op.add_column( sa.Column("labels", sa.LargeBinary(), nullable=True) ) - # ### end Alembic commands ### - def downgrade() -> None: """Downgrade database schema and/or data back to the previous revision.""" - # ### commands auto generated by Alembic - please adjust! ### with op.batch_alter_table("stack", schema=None) as batch_op: batch_op.drop_column("labels") - - # ### end Alembic commands ### From 8742967d4c6fe65527225be9ce18743d2b507223 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 5 Jul 2024 15:40:54 +0200 Subject: [PATCH 59/71] fix after merge mess --- src/zenml/cli/stack.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 1fe35e6b979..29f133f91e4 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -2291,13 +2291,13 @@ def _get_registries(registry_name: str, docs_link: str) -> List[str]: available_registries = _get_registries( "ECR", f"{AWS_DOCS}#ecr-container-registry" ) - "service connector.\nDocumentation for the ECR " - "container registry resource configuration can " - f"be found at {AWS_DOCS}#ecr-container-registry" - ) + if cloud_provider == "gcp": + flavor = "gcp" available_registries = _get_registries( "GCR", f"{GCP_DOCS}#gcr-container-registry" ) + if cloud_provider == "azure": + flavor = "azure" selected_registry_idx = cli_utils.multi_choice_prompt( object_type=f"{cloud_provider.upper()} registries", From 2cc1441d40265f93dc37e8206df27f39ad5a449a Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Mon, 8 Jul 2024 14:27:59 +0200 Subject: [PATCH 60/71] lint --- .../kubeflow/flavors/kubeflow_orchestrator_flavor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py b/src/zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py index 6a11cc254fa..7f3b0b912b6 100644 --- a/src/zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py +++ b/src/zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py @@ -149,7 +149,7 @@ class KubeflowOrchestratorConfig( kubeflow_hostname: Optional[str] = None kubeflow_namespace: str = "kubeflow" - kubernetes_context: Optional[str] = None # TODO: Potential setting + kubernetes_context: Optional[str] = None # TODO: Potential setting @model_validator(mode="before") @classmethod From 5f487935f88fad206dc596f2100b73c0dee7307d Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 10 Jul 2024 15:06:35 +0200 Subject: [PATCH 61/71] move timeout closer to rest --- src/zenml/cli/stack.py | 14 ++------ src/zenml/client.py | 7 ++++ src/zenml/zen_stores/rest_zen_store.py | 37 ++++++++++++++++++--- src/zenml/zen_stores/sql_zen_store.py | 5 +++ src/zenml/zen_stores/zen_store_interface.py | 5 +++ 5 files changed, 52 insertions(+), 16 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 544d766b1f4..418e39d08cf 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -98,7 +98,6 @@ verify_spec_and_tf_files_exist, ) from zenml.utils.yaml_utils import read_yaml, write_yaml -from zenml.zen_stores.rest_zen_store import RestZenStore if TYPE_CHECKING: from zenml.models import StackResponse @@ -426,16 +425,10 @@ def register_stack( if isinstance(service_connector, UUID): service_connector_resource_model = ( client.verify_service_connector( - service_connector + service_connector, timeout=120 ) ) else: - default_http_timeout = 30 - if isinstance(client.zen_store, RestZenStore): - default_http_timeout = ( - client.zen_store.config.http_timeout - ) - client.zen_store.config.http_timeout = 120 _, service_connector_resource_model = ( client.create_service_connector( name=stack_name, @@ -443,12 +436,9 @@ def register_stack( auth_method=service_connector.auth_method, configuration=service_connector.configuration, register=False, + timeout=120, ) ) - if isinstance(client.zen_store, RestZenStore): - client.zen_store.config.http_timeout = ( - default_http_timeout - ) if service_connector_resource_model is None: cli_utils.error( f"Failed to validate service connector {service_connector}..." diff --git a/src/zenml/client.py b/src/zenml/client.py index 869010024ab..0568ff0c4d3 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -46,6 +46,7 @@ from zenml.config.pipeline_run_configuration import PipelineRunConfiguration from zenml.config.source import Source from zenml.constants import ( + DEFAULT_HTTP_TIMEOUT, ENV_ZENML_ACTIVE_STACK_ID, ENV_ZENML_ACTIVE_WORKSPACE_ID, ENV_ZENML_ENABLE_REPO_INIT_WARNINGS, @@ -4974,6 +4975,7 @@ def create_service_connector( verify: bool = True, list_resources: bool = True, register: bool = True, + timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> Tuple[ Optional[ Union[ @@ -5006,6 +5008,7 @@ def create_service_connector( list_resources: Whether to also list the resources that the service connector can give access to (if verify is True). register: Whether to register the service connector or not. + timeout: The timeout in seconds for the HTTP request. Returns: The model of the registered service connector and the resources @@ -5138,6 +5141,7 @@ def create_service_connector( self.zen_store.verify_service_connector_config( connector_request, list_resources=list_resources, + timeout=timeout, ) ) else: @@ -5577,6 +5581,7 @@ def verify_service_connector( resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True, + timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> "ServiceConnectorResourcesModel": """Verifies if a service connector has access to one or more resources. @@ -5591,6 +5596,7 @@ def verify_service_connector( configuration will be used. list_resources: Whether to list the resources that the service connector has access to. + timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector has access to, @@ -5625,6 +5631,7 @@ def verify_service_connector( resource_type=resource_type, resource_id=resource_id, list_resources=list_resources, + timeout=timeout, ) else: connector_instance = ( diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 835534e5379..e0bd8d27490 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -2459,6 +2459,7 @@ def verify_service_connector_config( self, service_connector: ServiceConnectorRequest, list_resources: bool = True, + timeout: Optional[int] = None, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector configuration has access to resources. @@ -2467,6 +2468,7 @@ def verify_service_connector_config( list_resources: If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned. + timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector configuration has @@ -2476,6 +2478,7 @@ def verify_service_connector_config( f"{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_VERIFY}", body=service_connector, params={"list_resources": list_resources}, + timeout=timeout, ) resources = ServiceConnectorResourcesModel.model_validate( @@ -2490,6 +2493,7 @@ def verify_service_connector( resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True, + timeout: Optional[int] = None, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector instance has access to one or more resources. @@ -2500,6 +2504,7 @@ def verify_service_connector( list_resources: If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned. + timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector has access to, @@ -2513,6 +2518,7 @@ def verify_service_connector( response_body = self.put( f"{SERVICE_CONNECTORS}/{str(service_connector_id)}{SERVICE_CONNECTOR_VERIFY}", params=params, + timeout=timeout, ) resources = ServiceConnectorResourcesModel.model_validate( @@ -4081,6 +4087,7 @@ def _request( method: str, url: str, params: Optional[Dict[str, Any]] = None, + timeout: Optional[int] = None, **kwargs: Any, ) -> Json: """Make a request to the REST API. @@ -4089,6 +4096,7 @@ def _request( method: The HTTP method to use. url: The URL to request. params: The query parameters to pass to the endpoint. + timeout: The request timeout in seconds. kwargs: Additional keyword arguments to pass to the request. Returns: @@ -4111,7 +4119,7 @@ def _request( url, params=params, verify=self.config.verify_ssl, - timeout=self.config.http_timeout, + timeout=timeout or self.config.http_timeout, **kwargs, ) ) @@ -4140,13 +4148,18 @@ def _request( raise def get( - self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any + self, + path: str, + params: Optional[Dict[str, Any]] = None, + timeout: Optional[int] = None, + **kwargs: Any, ) -> Json: """Make a GET request to the given endpoint path. Args: path: The path to the endpoint. params: The query parameters to pass to the endpoint. + timeout: The request timeout in seconds. kwargs: Additional keyword arguments to pass to the request. Returns: @@ -4154,17 +4167,26 @@ def get( """ logger.debug(f"Sending GET request to {path}...") return self._request( - "GET", self.url + API + VERSION_1 + path, params=params, **kwargs + "GET", + self.url + API + VERSION_1 + path, + params=params, + timeout=timeout, + **kwargs, ) def delete( - self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any + self, + path: str, + params: Optional[Dict[str, Any]] = None, + timeout: Optional[int] = None, + **kwargs: Any, ) -> Json: """Make a DELETE request to the given endpoint path. Args: path: The path to the endpoint. params: The query parameters to pass to the endpoint. + timeout: The request timeout in seconds. kwargs: Additional keyword arguments to pass to the request. Returns: @@ -4175,6 +4197,7 @@ def delete( "DELETE", self.url + API + VERSION_1 + path, params=params, + timeout=timeout, **kwargs, ) @@ -4183,6 +4206,7 @@ def post( path: str, body: BaseModel, params: Optional[Dict[str, Any]] = None, + timeout: Optional[int] = None, **kwargs: Any, ) -> Json: """Make a POST request to the given endpoint path. @@ -4191,6 +4215,7 @@ def post( path: The path to the endpoint. body: The body to send. params: The query parameters to pass to the endpoint. + timeout: The request timeout in seconds. kwargs: Additional keyword arguments to pass to the request. Returns: @@ -4202,6 +4227,7 @@ def post( self.url + API + VERSION_1 + path, data=body.model_dump_json(), params=params, + timeout=timeout, **kwargs, ) @@ -4210,6 +4236,7 @@ def put( path: str, body: Optional[BaseModel] = None, params: Optional[Dict[str, Any]] = None, + timeout: Optional[int] = None, **kwargs: Any, ) -> Json: """Make a PUT request to the given endpoint path. @@ -4218,6 +4245,7 @@ def put( path: The path to the endpoint. body: The body to send. params: The query parameters to pass to the endpoint. + timeout: The request timeout in seconds. kwargs: Additional keyword arguments to pass to the request. Returns: @@ -4230,6 +4258,7 @@ def put( self.url + API + VERSION_1 + path, data=data, params=params, + timeout=timeout, **kwargs, ) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index f17950f1b8d..b3c1aed987d 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -86,6 +86,7 @@ from zenml.config.server_config import ServerConfiguration from zenml.config.store_config import StoreConfiguration from zenml.constants import ( + DEFAULT_HTTP_TIMEOUT, DEFAULT_PASSWORD, DEFAULT_STACK_AND_COMPONENT_NAME, DEFAULT_USERNAME, @@ -6644,6 +6645,7 @@ def verify_service_connector_config( self, service_connector: ServiceConnectorRequest, list_resources: bool = True, + timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector configuration has access to resources. @@ -6651,6 +6653,7 @@ def verify_service_connector_config( service_connector: The service connector configuration to verify. list_resources: If True, the list of all resources accessible through the service connector is returned. + timeout: not used. Returns: The list of resources that the service connector configuration has @@ -6667,6 +6670,7 @@ def verify_service_connector( resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True, + timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector instance has access to one or more resources. @@ -6677,6 +6681,7 @@ def verify_service_connector( list_resources: If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned. + timeout: not used Returns: The list of resources that the service connector has access to, diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index 46865035572..78f68fd937b 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -19,6 +19,7 @@ from uuid import UUID from zenml.config.pipeline_run_configuration import PipelineRunConfiguration +from zenml.constants import DEFAULT_HTTP_TIMEOUT from zenml.enums import StackDeploymentProvider from zenml.models import ( ActionFilter, @@ -2029,6 +2030,7 @@ def verify_service_connector_config( self, service_connector: ServiceConnectorRequest, list_resources: bool = True, + timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector configuration has access to resources. @@ -2036,6 +2038,7 @@ def verify_service_connector_config( service_connector: The service connector configuration to verify. list_resources: If True, the list of all resources accessible through the service connector is returned. + timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector configuration has @@ -2053,6 +2056,7 @@ def verify_service_connector( resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True, + timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector instance has access to one or more resources. @@ -2063,6 +2067,7 @@ def verify_service_connector( list_resources: If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned. + timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector has access to, From 25e32eecd40187fb530c11acef4463e7c3744ef6 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Wed, 10 Jul 2024 15:32:32 +0200 Subject: [PATCH 62/71] generate_temporary_tokens only in skypilot cases --- src/zenml/cli/stack.py | 40 +++++++++++++++++++++++++++++++++------- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 418e39d08cf..c1a1f949a45 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -377,6 +377,7 @@ def register_stack( if provider: labels["zenml:provider"] = provider service_connector_resource_model = None + generate_temporary_tokens = True # create components needed_components = ( (StackComponentType.ARTIFACT_STORE, artifact_store), @@ -428,6 +429,14 @@ def register_stack( service_connector, timeout=120 ) ) + existing_service_connector_info = ( + client.get_service_connector( + service_connector + ) + ) + generate_temporary_tokens = existing_service_connector_info.configuration.get( + "generate_temporary_tokens", True + ) else: _, service_connector_resource_model = ( client.create_service_connector( @@ -439,6 +448,7 @@ def register_stack( timeout=120, ) ) + generate_temporary_tokens = False if service_connector_resource_model is None: cli_utils.error( f"Failed to validate service connector {service_connector}..." @@ -459,6 +469,7 @@ def register_stack( cloud_provider=provider, service_connector_resource_models=service_connector_resource_model.resources, service_connector_index=0, + generate_temporary_tokens=generate_temporary_tokens, ) component_name = stack_name created_objects.add(component_type.value) @@ -474,6 +485,18 @@ def register_stack( artifact_store = component_name if component_type == StackComponentType.ORCHESTRATOR: orchestrator = component_name + if not isinstance( + component_info, UUID + ) and component_info.flavor.startswith("vm"): + if isinstance( + service_connector, ServiceConnectorInfo + ) and service_connector.auth_method in { + "service-account", + "external-account", + }: + service_connector.configuration[ + "generate_temporary_tokens" + ] = False if component_type == StackComponentType.CONTAINER_REGISTRY: container_registry = component_name @@ -2311,11 +2334,6 @@ def _get_service_connector_info( password="format" in properties[req_field] and properties[req_field]["format"] == "password", ) - if cloud_provider == "gcp" and auth_type in { - "service-account", - "external-account", - }: - answers["generate_temporary_tokens"] = False return ServiceConnectorInfo( type=cloud_provider, @@ -2330,6 +2348,7 @@ def _get_stack_component_info( service_connector_resource_models: List[ ServiceConnectorTypedResourcesModel ], + generate_temporary_tokens: bool, service_connector_index: Optional[int] = None, ) -> ComponentInfo: """Get a stack component info with given type and service connector. @@ -2338,6 +2357,7 @@ def _get_stack_component_info( component_type: The type of component to create. cloud_provider: The cloud provider to use. service_connector_resource_models: The list of the available service connector resource models. + generate_temporary_tokens: Whether to generate temporary tokens in connector. service_connector_index: The index of the service connector to use. Returns: @@ -2425,7 +2445,10 @@ def query_gcp_region() -> str: for each in service_connector_resource_models: types = [] if each.resource_type == "aws-generic": - types = ["Sagemaker", "Skypilot (EC2)"] + if generate_temporary_tokens: + types = ["Sagemaker"] + else: + types = ["Sagemaker", "Skypilot (EC2)"] if each.resource_type == "kubernetes-cluster": types = ["Kubernetes"] @@ -2450,7 +2473,10 @@ def query_gcp_region() -> str: for each in service_connector_resource_models: types = [] if each.resource_type == "gcp-generic": - types = ["Vertex AI", "Skypilot (Compute)"] + if generate_temporary_tokens: + types = ["Vertex AI"] + else: + types = ["Vertex AI", "Skypilot (Compute)"] if each.resource_type == "kubernetes-cluster": types = ["Kubernetes"] From 8461da936c790183f2e704e866fb4f31eb62977d Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 11 Jul 2024 09:30:36 +0200 Subject: [PATCH 63/71] rework timeout as suggested --- src/zenml/cli/stack.py | 2 +- src/zenml/client.py | 4 ---- src/zenml/constants.py | 1 + src/zenml/zen_stores/rest_zen_store.py | 15 +++++++++------ src/zenml/zen_stores/sql_zen_store.py | 5 ----- src/zenml/zen_stores/zen_store_interface.py | 5 ----- 6 files changed, 11 insertions(+), 21 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index c1a1f949a45..98f6a74180d 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -426,7 +426,7 @@ def register_stack( if isinstance(service_connector, UUID): service_connector_resource_model = ( client.verify_service_connector( - service_connector, timeout=120 + service_connector ) ) existing_service_connector_info = ( diff --git a/src/zenml/client.py b/src/zenml/client.py index 0568ff0c4d3..5bbe1a879a1 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -5141,7 +5141,6 @@ def create_service_connector( self.zen_store.verify_service_connector_config( connector_request, list_resources=list_resources, - timeout=timeout, ) ) else: @@ -5581,7 +5580,6 @@ def verify_service_connector( resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True, - timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> "ServiceConnectorResourcesModel": """Verifies if a service connector has access to one or more resources. @@ -5596,7 +5594,6 @@ def verify_service_connector( configuration will be used. list_resources: Whether to list the resources that the service connector has access to. - timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector has access to, @@ -5631,7 +5628,6 @@ def verify_service_connector( resource_type=resource_type, resource_id=resource_id, list_resources=list_resources, - timeout=timeout, ) else: connector_instance = ( diff --git a/src/zenml/constants.py b/src/zenml/constants.py index 5e61166c5ef..285b207a23f 100644 --- a/src/zenml/constants.py +++ b/src/zenml/constants.py @@ -263,6 +263,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int: DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT = 60 * 5 # 5 minutes DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING = 5 # seconds DEFAULT_HTTP_TIMEOUT = 30 +SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT = 120 # seconds ZENML_API_KEY_PREFIX = "ZENKEY_" DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW = 60 * 48 # 48 hours DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE = 5 diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index e0bd8d27490..3049092139d 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -91,6 +91,7 @@ SERVICE_CONNECTOR_RESOURCES, SERVICE_CONNECTOR_TYPES, SERVICE_CONNECTOR_VERIFY, + SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT, SERVICE_CONNECTORS, SERVICES, STACK, @@ -2459,7 +2460,6 @@ def verify_service_connector_config( self, service_connector: ServiceConnectorRequest, list_resources: bool = True, - timeout: Optional[int] = None, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector configuration has access to resources. @@ -2468,7 +2468,6 @@ def verify_service_connector_config( list_resources: If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned. - timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector configuration has @@ -2478,7 +2477,10 @@ def verify_service_connector_config( f"{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_VERIFY}", body=service_connector, params={"list_resources": list_resources}, - timeout=timeout, + timeout=max( + self.config.http_timeout, + SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT, + ), ) resources = ServiceConnectorResourcesModel.model_validate( @@ -2493,7 +2495,6 @@ def verify_service_connector( resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True, - timeout: Optional[int] = None, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector instance has access to one or more resources. @@ -2504,7 +2505,6 @@ def verify_service_connector( list_resources: If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned. - timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector has access to, @@ -2518,7 +2518,10 @@ def verify_service_connector( response_body = self.put( f"{SERVICE_CONNECTORS}/{str(service_connector_id)}{SERVICE_CONNECTOR_VERIFY}", params=params, - timeout=timeout, + timeout=max( + self.config.http_timeout, + SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT, + ), ) resources = ServiceConnectorResourcesModel.model_validate( diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index b3c1aed987d..f17950f1b8d 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -86,7 +86,6 @@ from zenml.config.server_config import ServerConfiguration from zenml.config.store_config import StoreConfiguration from zenml.constants import ( - DEFAULT_HTTP_TIMEOUT, DEFAULT_PASSWORD, DEFAULT_STACK_AND_COMPONENT_NAME, DEFAULT_USERNAME, @@ -6645,7 +6644,6 @@ def verify_service_connector_config( self, service_connector: ServiceConnectorRequest, list_resources: bool = True, - timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector configuration has access to resources. @@ -6653,7 +6651,6 @@ def verify_service_connector_config( service_connector: The service connector configuration to verify. list_resources: If True, the list of all resources accessible through the service connector is returned. - timeout: not used. Returns: The list of resources that the service connector configuration has @@ -6670,7 +6667,6 @@ def verify_service_connector( resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True, - timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector instance has access to one or more resources. @@ -6681,7 +6677,6 @@ def verify_service_connector( list_resources: If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned. - timeout: not used Returns: The list of resources that the service connector has access to, diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index 78f68fd937b..46865035572 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -19,7 +19,6 @@ from uuid import UUID from zenml.config.pipeline_run_configuration import PipelineRunConfiguration -from zenml.constants import DEFAULT_HTTP_TIMEOUT from zenml.enums import StackDeploymentProvider from zenml.models import ( ActionFilter, @@ -2030,7 +2029,6 @@ def verify_service_connector_config( self, service_connector: ServiceConnectorRequest, list_resources: bool = True, - timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector configuration has access to resources. @@ -2038,7 +2036,6 @@ def verify_service_connector_config( service_connector: The service connector configuration to verify. list_resources: If True, the list of all resources accessible through the service connector is returned. - timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector configuration has @@ -2056,7 +2053,6 @@ def verify_service_connector( resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True, - timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> ServiceConnectorResourcesModel: """Verifies if a service connector instance has access to one or more resources. @@ -2067,7 +2063,6 @@ def verify_service_connector( list_resources: If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned. - timeout: The timeout in seconds for the HTTP request. Returns: The list of resources that the service connector has access to, From 740ca42203150416871dfae504c8db7bf5f237ec Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 11 Jul 2024 09:36:10 +0200 Subject: [PATCH 64/71] rename --- src/zenml/cli/stack.py | 26 ++++++++++++-------------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 98f6a74180d..ea820ee584b 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -377,7 +377,7 @@ def register_stack( if provider: labels["zenml:provider"] = provider service_connector_resource_model = None - generate_temporary_tokens = True + can_generate_long_tokens = False # create components needed_components = ( (StackComponentType.ARTIFACT_STORE, artifact_store), @@ -434,7 +434,7 @@ def register_stack( service_connector ) ) - generate_temporary_tokens = existing_service_connector_info.configuration.get( + can_generate_long_tokens = not existing_service_connector_info.configuration.get( "generate_temporary_tokens", True ) else: @@ -448,7 +448,7 @@ def register_stack( timeout=120, ) ) - generate_temporary_tokens = False + can_generate_long_tokens = True if service_connector_resource_model is None: cli_utils.error( f"Failed to validate service connector {service_connector}..." @@ -469,7 +469,7 @@ def register_stack( cloud_provider=provider, service_connector_resource_models=service_connector_resource_model.resources, service_connector_index=0, - generate_temporary_tokens=generate_temporary_tokens, + can_generate_long_tokens=can_generate_long_tokens, ) component_name = stack_name created_objects.add(component_type.value) @@ -2348,7 +2348,7 @@ def _get_stack_component_info( service_connector_resource_models: List[ ServiceConnectorTypedResourcesModel ], - generate_temporary_tokens: bool, + can_generate_long_tokens: bool, service_connector_index: Optional[int] = None, ) -> ComponentInfo: """Get a stack component info with given type and service connector. @@ -2357,7 +2357,7 @@ def _get_stack_component_info( component_type: The type of component to create. cloud_provider: The cloud provider to use. service_connector_resource_models: The list of the available service connector resource models. - generate_temporary_tokens: Whether to generate temporary tokens in connector. + can_generate_long_tokens: Whether connector can generate long-living tokens. service_connector_index: The index of the service connector to use. Returns: @@ -2445,10 +2445,9 @@ def query_gcp_region() -> str: for each in service_connector_resource_models: types = [] if each.resource_type == "aws-generic": - if generate_temporary_tokens: - types = ["Sagemaker"] - else: - types = ["Sagemaker", "Skypilot (EC2)"] + types = ["Sagemaker"] + if can_generate_long_tokens: + types.append("Skypilot (EC2)") if each.resource_type == "kubernetes-cluster": types = ["Kubernetes"] @@ -2473,10 +2472,9 @@ def query_gcp_region() -> str: for each in service_connector_resource_models: types = [] if each.resource_type == "gcp-generic": - if generate_temporary_tokens: - types = ["Vertex AI"] - else: - types = ["Vertex AI", "Skypilot (Compute)"] + types = ["Vertex AI"] + if can_generate_long_tokens: + types.append("Skypilot (Compute)") if each.resource_type == "kubernetes-cluster": types = ["Kubernetes"] From 4c77f8ab3527b42d158e51aefe42620cc1948b15 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 11 Jul 2024 11:27:27 +0200 Subject: [PATCH 65/71] get regions from zen_store --- src/zenml/cli/stack.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index c5e332e986d..7bbc1363161 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -2466,12 +2466,16 @@ def _get_stack_component_info( service_connector_resource_id = selected_storage elif component_type == "orchestrator": - def query_gcp_region() -> str: - from google.cloud.aiplatform.constants import base as constants - + def query_gcp_region(compute_type: str) -> str: region = Prompt.ask( - "Select a location for your Vertex AI jobs:", - choices=sorted(list(constants.SUPPORTED_REGIONS)), + f"Select the location for your {compute_type}:", + choices=sorted( + Client() + .zen_store.get_stack_deployment_info( + StackDeploymentProvider.GCP + ) + .locations.values() + ), show_choices=True, ) return region @@ -2556,10 +2560,10 @@ def query_gcp_region() -> str: config["region"] = selected_orchestrator[1] elif selected_orchestrator[0] == "Skypilot (Compute)": flavor = "vm_gcp" - config["region"] = query_gcp_region() + config["region"] = query_gcp_region("Skypilot cluster") elif selected_orchestrator[0] == "Vertex AI": flavor = "vertex" - config["location"] = query_gcp_region() + config["location"] = query_gcp_region("Vertex AI job") elif selected_orchestrator[0] == "Kubernetes": flavor = "kubernetes" else: From b35f2b8afe8604a7718782088e49cf4b6a796808 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 11 Jul 2024 11:45:01 +0200 Subject: [PATCH 66/71] disable wizard with SC for local ZenMLs --- src/zenml/cli/stack.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 7bbc1363161..93bd4093213 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -288,6 +288,19 @@ def register_stack( client = Client() + if provider is not None or connector is not None: + if client.zen_store.is_local_store(): + cli_utils.error( + "You are registering a stack using a service connector, but " + "this feature cannot be used with a local ZenML deployment. " + "ZenML needs to be accessible from the cloud provider to allow the " + "stack and its components to be registered automatically. " + "Please deploy ZenML in a remote environment as described in the " + "documentation: https://docs.zenml.io/getting-started/deploying-zenml " + "or use a managed ZenML Pro server instance for quick access to " + "this feature and more: https://www.zenml.io/pro" + ) + try: client.get_stack( name_id_or_prefix=stack_name, From dae3ad7e7a48dfa777f6ab9036139189db8e358e Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 11 Jul 2024 11:49:33 +0200 Subject: [PATCH 67/71] left-over timeout --- src/zenml/cli/stack.py | 1 - src/zenml/client.py | 3 --- 2 files changed, 4 deletions(-) diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 93bd4093213..133142beabe 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -459,7 +459,6 @@ def register_stack( auth_method=service_connector.auth_method, configuration=service_connector.configuration, register=False, - timeout=120, ) ) can_generate_long_tokens = True diff --git a/src/zenml/client.py b/src/zenml/client.py index 5bbe1a879a1..869010024ab 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -46,7 +46,6 @@ from zenml.config.pipeline_run_configuration import PipelineRunConfiguration from zenml.config.source import Source from zenml.constants import ( - DEFAULT_HTTP_TIMEOUT, ENV_ZENML_ACTIVE_STACK_ID, ENV_ZENML_ACTIVE_WORKSPACE_ID, ENV_ZENML_ENABLE_REPO_INIT_WARNINGS, @@ -4975,7 +4974,6 @@ def create_service_connector( verify: bool = True, list_resources: bool = True, register: bool = True, - timeout: int = DEFAULT_HTTP_TIMEOUT, ) -> Tuple[ Optional[ Union[ @@ -5008,7 +5006,6 @@ def create_service_connector( list_resources: Whether to also list the resources that the service connector can give access to (if verify is True). register: Whether to register the service connector or not. - timeout: The timeout in seconds for the HTTP request. Returns: The model of the registered service connector and the resources From dbb8e513e9d3a337afbc48586aa0bf5e9077b2c6 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 11 Jul 2024 16:32:12 +0200 Subject: [PATCH 68/71] add docs --- .../register-a-cloud-stack.md | 73 +++++++++++++++++-- 1 file changed, 67 insertions(+), 6 deletions(-) diff --git a/docs/book/how-to/stack-deployment/register-a-cloud-stack.md b/docs/book/how-to/stack-deployment/register-a-cloud-stack.md index a002d6c864d..bdb737c6da9 100644 --- a/docs/book/how-to/stack-deployment/register-a-cloud-stack.md +++ b/docs/book/how-to/stack-deployment/register-a-cloud-stack.md @@ -33,21 +33,21 @@ In order to register a remote stack over the CLI with the stack wizard, you can use the following command: ```shell -zenml stack register -p aws +zenml stack register -p {aws|gcp} ``` To register the cloud stack, the first thing that the wizard needs is a service -connector. You can either use an existing connector by providing its ID -`-c ` or the wizard will create one for you. +connector. You can either use an existing connector by providing its ID or name +`-sc ` or the wizard will create one for you. Similar to the service connector, you can also use existing stack components. -However, this is only possible if these component are already configured with +However, this is only possible if these components are already configured with the same service connector that you provided through the parameter described above. {% hint style="warning" %} -Currently, the stack wizard only works on AWS. We are working on bringing -support to GCP and Azure as well. Stay in touch for further updates. +Currently, the stack wizard only works on AWS and GCP. We are working on bringing +support to Azure as well. Stay in touch for further updates. {% endhint %} ### AWS @@ -111,6 +111,67 @@ you as follows: ``` {% endcode %} +### GCP + +If you select `gcp` as your cloud provider, and you haven't selected a connector +yet, you will be prompted to select an authentication method for your stack. + +{% code title="Example Command Output" %} +``` + Available authentication methods for gcp +┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ Choice ┃ Name ┃ Required ┃ +┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ [0] │ GCP User Account │ user_account_json (GCP User Account Credentials │ +│ │ │ JSON optionally base64 encoded.) │ +│ │ │ project_id (GCP Project ID where the target │ +│ │ │ resource is located.) │ +│ │ │ │ +├────────┼──────────────────────┼───────────────────────────────────────────────────┤ +│ [1] │ GCP Service Account │ service_account_json (GCP Service Account Key │ +│ │ │ JSON optionally base64 encoded.) │ +│ │ │ │ +├────────┼──────────────────────┼───────────────────────────────────────────────────┤ +│ [2] │ GCP External Account │ external_account_json (GCP External Account │ +│ │ │ JSON optionally base64 encoded.) │ +│ │ │ project_id (GCP Project ID where the target │ +│ │ │ resource is located.) │ +│ │ │ │ +├────────┼──────────────────────┼───────────────────────────────────────────────────┤ +│ [3] │ GCP Oauth 2.0 Token │ token (GCP OAuth 2.0 Token) │ +│ │ │ project_id (GCP Project ID where the target │ +│ │ │ resource is located.) │ +│ │ │ │ +├────────┼──────────────────────┼───────────────────────────────────────────────────┤ +│ [4] │ GCP Service Account │ service_account_json (GCP Service Account Key │ +│ │ │ JSON optionally base64 encoded.) │ +│ │ Impersonation │ target_principal (GCP Service Account Email to │ +│ │ │ impersonate) │ +│ │ │ │ +└────────┴──────────────────────┴───────────────────────────────────────────────────┘ +``` +{% endcode %} + +Based on your selection, you will have to provide the required parameters listed +above. This will allow ZenML to create a Service Connector and +authenticate you to use your cloud resources. + +Next, for each missing component, the available resources will be listed to +you as follows: + +{% code title="Example Command Output for Artifact Stores" %} +``` + Available GCP storages +┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ Choice ┃ Storage ┃ +┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━┩ +│ [0] │ gs://************** │ +├─────────┼───────────────────────┤ +│ [1] │ gs://************** │ +└─────────┴───────────────────────┘ +``` +{% endcode %} + Based on your selection, ZenML will create the stack component and ultimately register the stack for you. From 1cde618bff1545943e01597bf4242890b211e7ba Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 11 Jul 2024 16:44:59 +0200 Subject: [PATCH 69/71] formatting --- .../register-a-cloud-stack.md | 75 ++++++++++--------- 1 file changed, 41 insertions(+), 34 deletions(-) diff --git a/docs/book/how-to/stack-deployment/register-a-cloud-stack.md b/docs/book/how-to/stack-deployment/register-a-cloud-stack.md index bdb737c6da9..3b487dd1387 100644 --- a/docs/book/how-to/stack-deployment/register-a-cloud-stack.md +++ b/docs/book/how-to/stack-deployment/register-a-cloud-stack.md @@ -57,7 +57,7 @@ yet, you will be prompted to select an authentication method for your stack. {% code title="Example Command Output" %} ``` - Available authentication methods for aws + Available authentication methods for aws ┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Choice ┃ Name ┃ Required ┃ ┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ @@ -100,7 +100,7 @@ you as follows: {% code title="Example Command Output for Artifact Stores" %} ``` - Available AWS storages + Available AWS storages ┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Choice ┃ Storage ┃ ┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ @@ -118,37 +118,44 @@ yet, you will be prompted to select an authentication method for your stack. {% code title="Example Command Output" %} ``` - Available authentication methods for gcp -┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ -┃ Choice ┃ Name ┃ Required ┃ -┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ -│ [0] │ GCP User Account │ user_account_json (GCP User Account Credentials │ -│ │ │ JSON optionally base64 encoded.) │ -│ │ │ project_id (GCP Project ID where the target │ -│ │ │ resource is located.) │ -│ │ │ │ -├────────┼──────────────────────┼───────────────────────────────────────────────────┤ -│ [1] │ GCP Service Account │ service_account_json (GCP Service Account Key │ -│ │ │ JSON optionally base64 encoded.) │ -│ │ │ │ -├────────┼──────────────────────┼───────────────────────────────────────────────────┤ -│ [2] │ GCP External Account │ external_account_json (GCP External Account │ -│ │ │ JSON optionally base64 encoded.) │ -│ │ │ project_id (GCP Project ID where the target │ -│ │ │ resource is located.) │ -│ │ │ │ -├────────┼──────────────────────┼───────────────────────────────────────────────────┤ -│ [3] │ GCP Oauth 2.0 Token │ token (GCP OAuth 2.0 Token) │ -│ │ │ project_id (GCP Project ID where the target │ -│ │ │ resource is located.) │ -│ │ │ │ -├────────┼──────────────────────┼───────────────────────────────────────────────────┤ -│ [4] │ GCP Service Account │ service_account_json (GCP Service Account Key │ -│ │ │ JSON optionally base64 encoded.) │ -│ │ Impersonation │ target_principal (GCP Service Account Email to │ -│ │ │ impersonate) │ -│ │ │ │ -└────────┴──────────────────────┴───────────────────────────────────────────────────┘ + Available authentication methods for gcp +┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ Choice ┃ Name ┃ Required ┃ +┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ [0] │ GCP User Account │ *user_account_json* (GCP User │ +│ │ │ Account Credentials JSON │ +│ │ │ optionally base64 encoded.) │ +│ │ │ *project_id* (GCP Project ID │ +│ │ │ where the target resource is │ +│ │ │ located.) │ +│ │ │ │ +├─────────┼──────────────────────┼────────────────────────────────┤ +│ [1] │ GCP Service Account │ *service_account_json* (GCP │ +│ │ │ Service Account Key JSON │ +│ │ │ optionally base64 encoded.) │ +│ │ │ │ +├─────────┼──────────────────────┼────────────────────────────────┤ +│ [2] │ GCP External Account │ *external_account_json* (GCP │ +│ │ │ External Account JSON │ +│ │ │ optionally base64 encoded.) │ +│ │ │ *project_id* (GCP Project ID │ +│ │ │ where the target resource is │ +│ │ │ located.) │ +│ │ │ │ +├─────────┼──────────────────────┼────────────────────────────────┤ +│ [3] │ GCP Oauth 2.0 Token │ *token* (GCP OAuth 2.0 Token) │ +│ │ │ *project_id* (GCP Project ID │ +│ │ │ where the target resource is │ +│ │ │ located.) │ +│ │ │ │ +├─────────┼──────────────────────┼────────────────────────────────┤ +│ [4] │ GCP Service Account │ *service_account_json* (GCP │ +│ │ Impersonation │ Service Account Key JSON │ +│ │ │ optionally base64 encoded.) │ +│ │ │ *target_principal* (GCP Service │ +│ │ │ Account Email to impersonate) │ +│ │ │ │ +└─────────┴──────────────────────┴────────────────────────────────┘ ``` {% endcode %} @@ -161,7 +168,7 @@ you as follows: {% code title="Example Command Output for Artifact Stores" %} ``` - Available GCP storages + Available GCP storages ┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Choice ┃ Storage ┃ ┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━┩ From c45b73aa4da5f49a1f7b6c123c2272060f3dd7ce Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 11 Jul 2024 16:49:41 +0200 Subject: [PATCH 70/71] formatting --- .../register-a-cloud-stack.md | 181 ++++++++++-------- 1 file changed, 96 insertions(+), 85 deletions(-) diff --git a/docs/book/how-to/stack-deployment/register-a-cloud-stack.md b/docs/book/how-to/stack-deployment/register-a-cloud-stack.md index 3b487dd1387..5b065221904 100644 --- a/docs/book/how-to/stack-deployment/register-a-cloud-stack.md +++ b/docs/book/how-to/stack-deployment/register-a-cloud-stack.md @@ -57,37 +57,48 @@ yet, you will be prompted to select an authentication method for your stack. {% code title="Example Command Output" %} ``` - Available authentication methods for aws -┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ -┃ Choice ┃ Name ┃ Required ┃ -┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ -│ [0] │ AWS Secret Key │ aws_access_key_id (AWS Access Key ID) │ -│ │ │ aws_secret_access_key (AWS Secret Access Key) │ -│ │ │ region (AWS Region) │ -│ │ │ │ -├────────┼──────────────────────┼────────────────────────────────────────────────┤ -│ [1] │ AWS STS Token │ aws_access_key_id (AWS Access Key ID) │ -│ │ │ aws_secret_access_key (AWS Secret Access Key) │ -│ │ │ aws_session_token (AWS Session Token) │ -│ │ │ region (AWS Region) │ -│ │ │ │ -├────────┼──────────────────────┼────────────────────────────────────────────────┤ -│ [2] │ AWS IAM Role │ aws_access_key_id (AWS Access Key ID) │ -│ │ │ aws_secret_access_key (AWS Secret Access Key) │ -│ │ │ region (AWS Region) │ -│ │ │ role_arn (AWS IAM Role ARN) │ -│ │ │ │ -├────────┼──────────────────────┼────────────────────────────────────────────────┤ -│ [3] │ AWS Session Token │ aws_access_key_id (AWS Access Key ID) │ -│ │ │ aws_secret_access_key (AWS Secret Access Key) │ -│ │ │ region (AWS Region) │ -│ │ │ │ -├────────┼──────────────────────┼────────────────────────────────────────────────┤ -│ [4] │ AWS Federation Token │ aws_access_key_id (AWS Access Key ID) │ -│ │ │ aws_secret_access_key (AWS Secret Access Key) │ -│ │ │ region (AWS Region) │ -│ │ │ │ -└────────┴──────────────────────┴────────────────────────────────────────────────┘ + Available authentication methods for aws +┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ Choice ┃ Name ┃ Required ┃ +┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ [0] │ AWS Secret Key │ aws_access_key_id (AWS Access │ +│ │ │ Key ID) │ +│ │ │ aws_secret_access_key (AWS │ +│ │ │ Secret Access Key) │ +│ │ │ region (AWS Region) │ +│ │ │ │ +├─────────┼────────────────────────────────┼────────────────────────────────┤ +│ [1] │ AWS STS Token │ aws_access_key_id (AWS Access │ +│ │ │ Key ID) │ +│ │ │ aws_secret_access_key (AWS │ +│ │ │ Secret Access Key) │ +│ │ │ aws_session_token (AWS │ +│ │ │ Session Token) │ +│ │ │ region (AWS Region) │ +│ │ │ │ +├─────────┼────────────────────────────────┼────────────────────────────────┤ +│ [2] │ AWS IAM Role │ aws_access_key_id (AWS Access │ +│ │ │ Key ID) │ +│ │ │ aws_secret_access_key (AWS │ +│ │ │ Secret Access Key) │ +│ │ │ region (AWS Region) │ +│ │ │ role_arn (AWS IAM Role ARN) │ +│ │ │ │ +├─────────┼────────────────────────────────┼────────────────────────────────┤ +│ [3] │ AWS Session Token │ aws_access_key_id (AWS Access │ +│ │ │ Key ID) │ +│ │ │ aws_secret_access_key (AWS │ +│ │ │ Secret Access Key) │ +│ │ │ region (AWS Region) │ +│ │ │ │ +├─────────┼────────────────────────────────┼────────────────────────────────┤ +│ [4] │ AWS Federation Token │ aws_access_key_id (AWS Access │ +│ │ │ Key ID) │ +│ │ │ aws_secret_access_key (AWS │ +│ │ │ Secret Access Key) │ +│ │ │ region (AWS Region) │ +│ │ │ │ +└─────────┴────────────────────────────────┴────────────────────────────────┘ ``` {% endcode %} @@ -100,14 +111,14 @@ you as follows: {% code title="Example Command Output for Artifact Stores" %} ``` - Available AWS storages -┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ -┃ Choice ┃ Storage ┃ -┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ -│ [0] │ s3://************** │ -├────────┼─────────────────────────────────────────────────────────────┤ -│ [1] │ s3://************** │ -└────────┴─────────────────────────────────────────────────────────────┘ + Available AWS storages +┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ Choice ┃ Storage ┃ +┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ [0] │ s3://*************************** │ +├───────────────┼───────────────────────────────────────────────────────────┤ +│ [1] │ s3://*************************** │ +└───────────────┴───────────────────────────────────────────────────────────┘ ``` {% endcode %} @@ -118,44 +129,44 @@ yet, you will be prompted to select an authentication method for your stack. {% code title="Example Command Output" %} ``` - Available authentication methods for gcp -┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ -┃ Choice ┃ Name ┃ Required ┃ -┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ -│ [0] │ GCP User Account │ *user_account_json* (GCP User │ -│ │ │ Account Credentials JSON │ -│ │ │ optionally base64 encoded.) │ -│ │ │ *project_id* (GCP Project ID │ -│ │ │ where the target resource is │ -│ │ │ located.) │ -│ │ │ │ -├─────────┼──────────────────────┼────────────────────────────────┤ -│ [1] │ GCP Service Account │ *service_account_json* (GCP │ -│ │ │ Service Account Key JSON │ -│ │ │ optionally base64 encoded.) │ -│ │ │ │ -├─────────┼──────────────────────┼────────────────────────────────┤ -│ [2] │ GCP External Account │ *external_account_json* (GCP │ -│ │ │ External Account JSON │ -│ │ │ optionally base64 encoded.) │ -│ │ │ *project_id* (GCP Project ID │ -│ │ │ where the target resource is │ -│ │ │ located.) │ -│ │ │ │ -├─────────┼──────────────────────┼────────────────────────────────┤ -│ [3] │ GCP Oauth 2.0 Token │ *token* (GCP OAuth 2.0 Token) │ -│ │ │ *project_id* (GCP Project ID │ -│ │ │ where the target resource is │ -│ │ │ located.) │ -│ │ │ │ -├─────────┼──────────────────────┼────────────────────────────────┤ -│ [4] │ GCP Service Account │ *service_account_json* (GCP │ -│ │ Impersonation │ Service Account Key JSON │ -│ │ │ optionally base64 encoded.) │ -│ │ │ *target_principal* (GCP Service │ -│ │ │ Account Email to impersonate) │ -│ │ │ │ -└─────────┴──────────────────────┴────────────────────────────────┘ + Available authentication methods for gcp +┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ Choice ┃ Name ┃ Required ┃ +┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ [0] │ GCP User Account │ user_account_json (GCP User │ +│ │ │ Account Credentials JSON │ +│ │ │ optionally base64 encoded.) │ +│ │ │ project_id (GCP Project ID │ +│ │ │ where the target resource is │ +│ │ │ located.) │ +│ │ │ │ +├─────────┼────────────────────────────────┼────────────────────────────────┤ +│ [1] │ GCP Service Account │ service_account_json (GCP │ +│ │ │ Service Account Key JSON │ +│ │ │ optionally base64 encoded.) │ +│ │ │ │ +├─────────┼────────────────────────────────┼────────────────────────────────┤ +│ [2] │ GCP External Account │ external_account_json (GCP │ +│ │ │ External Account JSON │ +│ │ │ optionally base64 encoded.) │ +│ │ │ project_id (GCP Project ID │ +│ │ │ where the target resource is │ +│ │ │ located.) │ +│ │ │ │ +├─────────┼────────────────────────────────┼────────────────────────────────┤ +│ [3] │ GCP Oauth 2.0 Token │ token (GCP OAuth 2.0 Token) │ +│ │ │ project_id (GCP Project ID │ +│ │ │ where the target resource is │ +│ │ │ located.) │ +│ │ │ │ +├─────────┼────────────────────────────────┼────────────────────────────────┤ +│ [4] │ GCP Service Account │ service_account_json (GCP │ +│ │ Impersonation │ Service Account Key JSON │ +│ │ │ optionally base64 encoded.) │ +│ │ │ target_principal (GCP Service │ +│ │ │ Account Email to impersonate) │ +│ │ │ │ +└─────────┴────────────────────────────────┴────────────────────────────────┘ ``` {% endcode %} @@ -168,14 +179,14 @@ you as follows: {% code title="Example Command Output for Artifact Stores" %} ``` - Available GCP storages -┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┓ -┃ Choice ┃ Storage ┃ -┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━┩ -│ [0] │ gs://************** │ -├─────────┼───────────────────────┤ -│ [1] │ gs://************** │ -└─────────┴───────────────────────┘ + Available GCP storages +┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ Choice ┃ Storage ┃ +┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ [0] │ gs://*************************** │ +├───────────────┼───────────────────────────────────────────────────────────┤ +│ [1] │ gs://*************************** │ +└───────────────┴───────────────────────────────────────────────────────────┘ ``` {% endcode %} From 21c69fd6baec84850139ed0045e60bb0a880142b Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Thu, 11 Jul 2024 16:55:46 +0200 Subject: [PATCH 71/71] restructure a bit --- .../register-a-cloud-stack.md | 36 ++++++------------- 1 file changed, 11 insertions(+), 25 deletions(-) diff --git a/docs/book/how-to/stack-deployment/register-a-cloud-stack.md b/docs/book/how-to/stack-deployment/register-a-cloud-stack.md index 5b065221904..389b3d940ab 100644 --- a/docs/book/how-to/stack-deployment/register-a-cloud-stack.md +++ b/docs/book/how-to/stack-deployment/register-a-cloud-stack.md @@ -50,7 +50,13 @@ Currently, the stack wizard only works on AWS and GCP. We are working on bringin support to Azure as well. Stay in touch for further updates. {% endhint %} -### AWS +### Define Service Connector + +Below you will find cloud-specific selection options. Based on your selection, you will have to provide the required parameters listed +below. This will allow ZenML to create a Service Connector and +authenticate you to use your cloud resources. + +#### AWS If you select `aws` as your cloud provider, and you haven't selected a connector yet, you will be prompted to select an authentication method for your stack. @@ -102,27 +108,7 @@ yet, you will be prompted to select an authentication method for your stack. ``` {% endcode %} -Based on your selection, you will have to provide the required parameters listed -above. This will allow ZenML to create a Service Connector and -authenticate you to use your cloud resources. - -Next, for each missing component, the available resources will be listed to -you as follows: - -{% code title="Example Command Output for Artifact Stores" %} -``` - Available AWS storages -┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ -┃ Choice ┃ Storage ┃ -┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ -│ [0] │ s3://*************************** │ -├───────────────┼───────────────────────────────────────────────────────────┤ -│ [1] │ s3://*************************** │ -└───────────────┴───────────────────────────────────────────────────────────┘ -``` -{% endcode %} - -### GCP +#### GCP If you select `gcp` as your cloud provider, and you haven't selected a connector yet, you will be prompted to select an authentication method for your stack. @@ -170,9 +156,7 @@ yet, you will be prompted to select an authentication method for your stack. ``` {% endcode %} -Based on your selection, you will have to provide the required parameters listed -above. This will allow ZenML to create a Service Connector and -authenticate you to use your cloud resources. +### Defining cloud components Next, for each missing component, the available resources will be listed to you as follows: @@ -190,6 +174,8 @@ you as follows: ``` {% endcode %} +### Final steps + Based on your selection, ZenML will create the stack component and ultimately register the stack for you.