More changes

Author: schustmi

src/zenml/client.py

@@ -1211,12 +1211,7 @@ def create_stack(
             )
             stack_components[c_type] = [component.id]
 
-        secret_ids = []
-        for secret in secrets or []:
-            if isinstance(secret, UUID):
-                secret_ids.append(secret)
-            else:
-                secret_ids.append(self.get_secret(secret).id)
+        secret_ids = [self.get_secret(secret).id for secret in secrets or []]
 
         stack = StackRequest(
             name=name,
@@ -1323,6 +1318,8 @@ def update_stack(
         component_updates: Optional[
             Dict[StackComponentType, List[Union[UUID, str]]]
         ] = None,
+        add_secrets: Optional[List[Union[UUID, str]]] = None,
+        remove_secrets: Optional[List[Union[UUID, str]]] = None,
     ) -> StackResponse:
         """Updates a stack and its components.
 
@@ -1334,6 +1331,8 @@ def update_stack(
             description: the new description of the stack.
             component_updates: dictionary which maps stack component types to
                 lists of new stack component names or ids.
+            add_secrets: The secrets to add to the stack.
+            remove_secrets: The secrets to remove from the stack.
 
         Returns:
             The model of the updated stack.
@@ -1391,6 +1390,16 @@ def update_stack(
             }
             update_model.labels = existing_labels
 
+        if add_secrets:
+            update_model.add_secrets = [
+                self.get_secret(secret).id for secret in add_secrets
+            ]
+
+        if remove_secrets:
+            update_model.remove_secrets = [
+                self.get_secret(secret).id for secret in remove_secrets
+            ]
+
         updated_stack = self.zen_store.update_stack(
             stack_id=stack.id,
             stack_update=update_model,
@@ -2000,6 +2009,7 @@ def create_stack_component(
         component_type: StackComponentType,
         configuration: Dict[str, str],
         labels: Optional[Dict[str, Any]] = None,
+        secrets: Optional[List[Union[UUID, str]]] = None,
     ) -> "ComponentResponse":
         """Registers a stack component.
 
@@ -2009,7 +2019,7 @@ def create_stack_component(
             component_type: The type of the stack component.
             configuration: The configuration of the stack component.
             labels: The labels of the stack component.
-
+            secrets: The secrets of the stack component.
         Returns:
             The model of the registered component.
         """
@@ -2030,12 +2040,15 @@ def create_stack_component(
         assert validated_config is not None
         warn_if_config_server_mismatch(validated_config)
 
+        secret_ids = [self.get_secret(secret).id for secret in secrets or []]
+
         create_component_model = ComponentRequest(
             name=name,
             type=component_type,
             flavor=flavor,
             configuration=configuration,
             labels=labels,
+            secrets=secret_ids,
         )
 
         # Register the new model
@@ -2053,6 +2066,8 @@ def update_stack_component(
         disconnect: Optional[bool] = None,
         connector_id: Optional[UUID] = None,
         connector_resource_id: Optional[str] = None,
+        add_secrets: Optional[List[Union[UUID, str]]] = None,
+        remove_secrets: Optional[List[Union[UUID, str]]] = None,
     ) -> ComponentResponse:
         """Updates a stack component.
 
@@ -2068,6 +2083,8 @@ def update_stack_component(
             connector_id: The new connector id of the stack component.
             connector_resource_id: The new connector resource id of the
                 stack component.
+            add_secrets: The secrets to add to the stack component.
+            remove_secrets: The secrets to remove from the stack component.
 
         Returns:
             The updated stack component.
@@ -2150,6 +2167,16 @@ def update_stack_component(
                     existing_component.connector_resource_id
                 )
 
+        if add_secrets:
+            update_model.add_secrets = [
+                self.get_secret(secret).id for secret in add_secrets
+            ]
+
+        if remove_secrets:
+            update_model.remove_secrets = [
+                self.get_secret(secret).id for secret in remove_secrets
+            ]
+
         # Send the updated component to the ZenStore
         return self.zen_store.update_stack_component(
             component_id=component.id,

src/zenml/config/compiler.py

@@ -37,16 +37,11 @@
     StepConfigurationUpdate,
     StepSpec,
 )
-from zenml.constants import (
-    ENV_ZENML_ACTIVE_PROJECT_ID,
-    ENV_ZENML_ACTIVE_STACK_ID,
-    ENV_ZENML_STORE_PREFIX,
-)
 from zenml.environment import get_run_environment_dict
 from zenml.exceptions import StackValidationError
 from zenml.models import PipelineDeploymentBase
 from zenml.pipelines.run_utils import get_default_run_name
-from zenml.utils import pydantic_utils, settings_utils
+from zenml.utils import pydantic_utils, secret_utils, settings_utils
 
 if TYPE_CHECKING:
     from zenml.config.source import Source
@@ -115,7 +110,9 @@ def compile(
         pipeline_environment = finalize_environment_variables(
             pipeline.configuration.environment
         )
-        pipeline_secrets = pipeline.configuration.secrets.copy()
+        pipeline_secrets = secret_utils.resolve_and_verify_secrets(
+            pipeline.configuration.secrets
+        )
         pipeline_settings = self._filter_and_validate_settings(
             settings=pipeline.configuration.settings,
             configuration_level=ConfigurationLevel.PIPELINE,
@@ -509,10 +506,10 @@ def _compile_step_invocation(
             step.configuration.settings, stack=stack
         )
         step_spec = self._get_step_spec(invocation=invocation)
-        step_environment = finalize_environment_variables(
-            step.configuration.environment
+        step_environment = step.configuration.environment
+        step_secrets = secret_utils.resolve_and_verify_secrets(
+            step.configuration.secrets
         )
-        step_secrets = step.configuration.secrets
         step_settings = self._filter_and_validate_settings(
             settings=step.configuration.settings,
             configuration_level=ConfigurationLevel.STEP,
@@ -722,17 +719,4 @@ def finalize_environment_variables(
             else:
                 environment[key_without_prefix] = value
 
-    finalized_env = {}
-
-    for key, value in environment.items():
-        if key.upper().startswith(ENV_ZENML_STORE_PREFIX) or key.upper() in [
-            ENV_ZENML_ACTIVE_PROJECT_ID,
-            ENV_ZENML_ACTIVE_STACK_ID,
-        ]:
-            logger.warning(
-                "Not allowed to set `%s` config environment variable.", key
-            )
-            continue
-        finalized_env[key] = str(value)
-
-    return finalized_env
+    return environment

src/zenml/config/pipeline_configurations.py

@@ -15,6 +15,7 @@
 
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from uuid import UUID
 
 from pydantic import SerializeAsAny, field_validator
 
@@ -42,7 +43,7 @@ class PipelineConfigurationUpdate(StrictBaseModel):
     enable_artifact_visualization: Optional[bool] = None
     enable_step_logs: Optional[bool] = None
     environment: Dict[str, Any] = {}
-    secrets: List[str] = []
+    secrets: List[Union[str, UUID]] = []
     enable_pipeline_logs: Optional[bool] = None
     settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
     tags: Optional[List[Union[str, "Tag"]]] = None

src/zenml/config/pipeline_run_configuration.py

@@ -47,7 +47,7 @@ class PipelineRunConfiguration(
     )
     steps: Dict[str, StepConfigurationUpdate] = {}
     environment: Dict[str, Any] = {}
-    secrets: List[str] = []
+    secrets: List[Union[str, UUID]] = []
     settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
     tags: Optional[List[Union[str, Tag]]] = None
     extra: Dict[str, Any] = {}

src/zenml/orchestrators/step_runner.py

@@ -189,9 +189,14 @@ def run(
             )
 
             step_failed = False
-            environment = env_utils.gather_step_environment(
+            environment = env_utils.get_step_environment(
                 step_config=step_run.config, stack=self._stack
             )
+            secret_environment = env_utils.get_step_secret_environment(
+                step_config=step_run.config, stack=self._stack
+            )
+            environment.update(secret_environment)
+
             with env_utils.temporary_environment(environment):
                 try:
                     return_values = step_instance.call_entrypoint(

src/zenml/steps/base_step.py

@@ -60,7 +60,6 @@
     materializer_utils,
     notebook_utils,
     pydantic_utils,
-    secret_utils,
     settings_utils,
     source_code_utils,
     source_utils,
@@ -716,9 +715,6 @@ def _convert_to_tuple(value: Any) -> Tuple[Source, ...]:
         if merge and secrets and self._configuration.secrets:
             secrets = self._configuration.secrets + secrets
 
-        if secrets:
-            secrets = secret_utils.convert_to_secret_ids(secrets)
-
         values = dict_utils.remove_none_values(
             {
                 "enable_cache": enable_cache,

src/zenml/utils/env_utils.py

@@ -203,10 +203,10 @@ def temporary_environment(environment: Dict[str, str]) -> Iterator[None]:
                     os.environ[key] = previous_value
 
 
-def gather_step_environment(
+def get_step_environment(
     step_config: "StepConfiguration", stack: "Stack"
 ) -> Dict[str, str]:
-    """Gather the environment variables for a step.
+    """Get the environment variables for a step.
 
     Args:
         step_config: The step configuration.
@@ -216,28 +216,46 @@ def gather_step_environment(
         A dictionary of environment variables.
     """
     environment = {}
-    secrets = []
     for component in stack.components.values():
         environment.update(component.environment)
-        secrets.extend(component.secrets)
 
     environment.update(stack.environment)
+    environment.update(step_config.environment)
+
+    return environment
+
+
+def get_step_secret_environment(
+    step_config: "StepConfiguration", stack: "Stack"
+) -> Dict[str, str]:
+    """Get the environment variables for a step.
+
+    Args:
+        step_config: The step configuration.
+        stack: The stack on which the step will run.
+
+    Returns:
+        A dictionary of environment variables.
+    """
+    secrets = step_config.secrets
     secrets.extend(stack.secrets)
 
-    environment.update(step_config.environment)
-    secrets.extend(step_config.secrets)
+    for component in stack.components.values():
+        secrets.extend(component.secrets)
 
-    # Remove duplicates while preserving order, only the last occurrence of
-    # each secret will be used to handle overrides
-    secrets = list(reversed(dict.fromkeys(reversed(secrets))))
+    # Removes duplicates while preserving order, only the first occurrence of
+    # each secret will be kept. We then reverse the list to make sure the
+    # overrides are applied in the correct order.
+    secrets = list(reversed(dict.fromkeys(secrets)))
 
+    environment = {}
     for secret_name_or_id in secrets:
         try:
             secret = Client().get_secret(secret_name_or_id)
         except Exception as e:
             logger.warning(
                 "Failed to get secret `%s` with error: %s. Skipping setting "
-                "environment variable for this secret.",
+                "environment variables for this secret.",
                 secret_name_or_id,
                 e,
             )
@@ -247,7 +265,7 @@ def gather_step_environment(
             logger.warning(
                 "Did not find any secret values for secret `%s`. This might be "
                 "because you do not have permissions to read the secret "
-                "values. Skipping setting environment variable for this "
+                "values. Skipping setting environment variables for this "
                 "secret.",
                 secret_name_or_id,
             )

src/zenml/utils/secret_utils.py

@@ -14,18 +14,21 @@
 """Utility functions for secrets and secret references."""
 
 import re
-from typing import TYPE_CHECKING, Any, List, NamedTuple, Union
+from typing import TYPE_CHECKING, Any, List, NamedTuple, Optional, Union
 
 from pydantic import Field, PlainSerializer, SecretStr
 from typing_extensions import Annotated
 
 from zenml.logger import get_logger
+from zenml.utils import uuid_utils
 
 if TYPE_CHECKING:
     from uuid import UUID
 
     from pydantic.fields import FieldInfo
 
+    from zenml.zen_stores.zen_store_interface import ZenStoreInterface
+
 _secret_reference_expression = re.compile(r"\{\{\s*\S+?\.\S+\s*\}\}")
 
 PYDANTIC_SENSITIVE_FIELD_MARKER = "sensitive"
@@ -186,22 +189,47 @@ def is_clear_text_field(field: "FieldInfo") -> bool:
     return False
 
 
-def convert_to_secret_ids(secrets: List[Union[str, "UUID"]]) -> List["UUID"]:
+def resolve_and_verify_secrets(
+    secrets: List[Union[str, "UUID"]],
+    zen_store: Optional["ZenStoreInterface"] = None,
+) -> List["UUID"]:
     """Convert a list of secret names or IDs to a list of secret IDs.
 
     Args:
         secrets: A list of secret names or IDs.
+        zen_store: The ZenML store to use to resolve the secrets.
 
     Returns:
         A list of secret IDs.
     """
-    from zenml.client import Client
+    if zen_store:
+        from zenml.models import SecretFilter
 
-    secret_ids = []
-    for secret in secrets:
-        if isinstance(secret, str):
-            secret_ids.append(Client().get_secret(secret, hydrate=False).id)
-        else:
-            secret_ids.append(secret)
+        resolved_secrets = []
 
-    return secret_ids
+        for secret_name_or_id in secrets:
+            if uuid_utils.is_valid_uuid(secret_name_or_id):
+                secret_id = (
+                    secret_name_or_id
+                    if isinstance(secret_name_or_id, UUID)
+                    else UUID(secret_name_or_id)
+                )
+                secret = zen_store.get_secret(secret_id, hydrate=False)
+                resolved_secrets.append(secret.id)
+            else:
+                filter_model = SecretFilter(name=secret_name_or_id)
+                secrets = zen_store.list_secrets(filter_model=filter_model)
+                if not secrets:
+                    raise KeyError(
+                        f"Secret with name {secret_name_or_id} not found."
+                    )
+
+                resolved_secrets.append(secrets[0].id)
+
+        return resolved_secrets
+    else:
+        from zenml.client import Client
+
+        return [
+            Client().get_secret(secret, hydrate=False).id for secret in secrets
+        ]