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..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 @@ -33,84 +33,149 @@ 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 +### 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. {% 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 %} -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. +#### 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 │ +│ │ Impersonation │ Service Account Key JSON │ +│ │ │ optionally base64 encoded.) │ +│ │ │ target_principal (GCP Service │ +│ │ │ Account Email to impersonate) │ +│ │ │ │ +└─────────┴────────────────────────────────┴────────────────────────────────┘ +``` +{% endcode %} + +### Defining cloud components 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://************** │ -└────────┴─────────────────────────────────────────────────────────────┘ + Available GCP storages +┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ Choice ┃ Storage ┃ +┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ [0] │ gs://*************************** │ +├───────────────┼───────────────────────────────────────────────────────────┤ +│ [1] │ gs://*************************** │ +└───────────────┴───────────────────────────────────────────────────────────┘ ``` {% endcode %} +### Final steps + Based on your selection, ZenML will create the stack component and ultimately register the stack for you. diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index 791f5687d87..133142beabe 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, @@ -378,6 +391,7 @@ def register_stack( if provider: labels["zenml:provider"] = provider service_connector_resource_model = None + can_generate_long_tokens = False # create components needed_components = ( (StackComponentType.ARTIFACT_STORE, artifact_store), @@ -420,26 +434,38 @@ 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, + existing_service_connector_info = ( + client.get_service_connector( + service_connector + ) ) - ) - if service_connector_resource_model is None: - cli_utils.error( - f"Failed to validate service connector {service_connector}..." + can_generate_long_tokens = not existing_service_connector_info.configuration.get( + "generate_temporary_tokens", True ) + 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, + ) + ) + can_generate_long_tokens = True + 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, @@ -456,6 +482,7 @@ def register_stack( cloud_provider=provider, service_connector_resource_models=service_connector_resource_model.resources, service_connector_index=0, + can_generate_long_tokens=can_generate_long_tokens, ) component_name = stack_name created_objects.add(component_type.value) @@ -471,6 +498,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 @@ -2284,7 +2323,7 @@ def _get_service_connector_info( """ from rich.prompt import Prompt - if cloud_provider not in {"aws"}: + if cloud_provider not in {"aws", "gcp"}: raise ValueError(f"Unknown cloud provider {cloud_provider}") client = Client() @@ -2313,7 +2352,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.") @@ -2343,7 +2382,6 @@ 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, @@ -2358,6 +2396,7 @@ def _get_stack_component_info( service_connector_resource_models: List[ ServiceConnectorTypedResourcesModel ], + can_generate_long_tokens: bool, service_connector_index: Optional[int] = None, ) -> ComponentInfo: """Get a stack component info with given type and service connector. @@ -2366,6 +2405,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. + can_generate_long_tokens: Whether connector can generate long-living tokens. service_connector_index: The index of the service connector to use. Returns: @@ -2383,6 +2423,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 @@ -2406,7 +2449,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", @@ -2422,12 +2477,29 @@ def _get_stack_component_info( config = {"path": selected_storage} service_connector_resource_id = selected_storage elif component_type == "orchestrator": + + def query_gcp_region(compute_type: str) -> str: + region = Prompt.ask( + 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 + if cloud_provider == "aws": available_orchestrators = [] for each in service_connector_resource_models: types = [] if each.resource_type == "aws-generic": - types = ["Sagemaker", "Skypilot (EC2)"] + types = ["Sagemaker"] + if can_generate_long_tokens: + types.append("Skypilot (EC2)") if each.resource_type == "kubernetes-cluster": types = ["Kubernetes"] @@ -2448,7 +2520,32 @@ 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"] + if can_generate_long_tokens: + types.append("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 @@ -2465,25 +2562,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("Skypilot cluster") + elif selected_orchestrator[0] == "Vertex AI": + flavor = "vertex" + config["location"] = query_gcp_region("Vertex AI job") 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 [] @@ -2492,14 +2595,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}" ) - elif cloud_provider == "azure": - flavor = "azure" - elif cloud_provider == "gcp": + return available_registries + + if cloud_provider == "aws": + flavor = "aws" + available_registries = _get_registries( + "ECR", f"{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", diff --git a/src/zenml/constants.py b/src/zenml/constants.py index 81be1e5bc53..b2adfdf3f3c 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 6d584a4bb18..2dd200f887e 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -92,6 +92,7 @@ SERVICE_CONNECTOR_RESOURCES, SERVICE_CONNECTOR_TYPES, SERVICE_CONNECTOR_VERIFY, + SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT, SERVICE_CONNECTORS, SERVICES, STACK, @@ -2477,6 +2478,10 @@ def verify_service_connector_config( f"{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_VERIFY}", body=service_connector, params={"list_resources": list_resources}, + timeout=max( + self.config.http_timeout, + SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT, + ), ) resources = ServiceConnectorResourcesModel.model_validate( @@ -2514,6 +2519,10 @@ def verify_service_connector( response_body = self.put( f"{SERVICE_CONNECTORS}/{str(service_connector_id)}{SERVICE_CONNECTOR_VERIFY}", params=params, + timeout=max( + self.config.http_timeout, + SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT, + ), ) resources = ServiceConnectorResourcesModel.model_validate( @@ -4073,6 +4082,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. @@ -4081,6 +4091,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: @@ -4103,7 +4114,7 @@ def _request( url, params=params, verify=self.config.verify_ssl, - timeout=self.config.http_timeout, + timeout=timeout or self.config.http_timeout, **kwargs, ) ) @@ -4132,13 +4143,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: @@ -4146,17 +4162,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: @@ -4167,6 +4192,7 @@ def delete( "DELETE", self.url + API + VERSION_1 + path, params=params, + timeout=timeout, **kwargs, ) @@ -4175,6 +4201,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. @@ -4183,6 +4210,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: @@ -4194,6 +4222,7 @@ def post( self.url + API + VERSION_1 + path, data=body.model_dump_json(), params=params, + timeout=timeout, **kwargs, ) @@ -4202,6 +4231,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. @@ -4210,6 +4240,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: @@ -4222,6 +4253,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 eae4dfe1f8d..987e92499d7 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -7007,7 +7007,6 @@ def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse: # Stack Components components_mapping: Dict[StackComponentType, List[UUID]] = {} - for ( component_type, component_info,