zenml-nightly 0.74.0.dev20250224__py3-none-any.whl → 0.74.0.dev20250225__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.74.0.dev20250224
+0.74.0.dev20250225

zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py

@@ -71,6 +71,7 @@ class SagemakerOrchestratorSettings(BaseSettings):
             For processor_args.instance_type, check
             https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-available-instance-types.html
             for a list of available instance types.
+        environment: Environment variables to pass to the container.
         estimator_args: Arguments that are directly passed to the SageMaker
             Estimator for a specific step, allowing for overriding the default
             settings provided when configuring the component. See
@@ -116,6 +117,7 @@ class SagemakerOrchestratorSettings(BaseSettings):
 
     processor_args: Dict[str, Any] = {}
     estimator_args: Dict[str, Any] = {}
+    environment: Dict[str, str] = {}
 
     input_data_s3_mode: str = "File"
     input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(

zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py

@@ -55,6 +55,7 @@ class SagemakerStepOperatorSettings(BaseSettings):
             For estimator_args.instance_type, check
             https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-available-instance-types.html
             for a list of available instance types.
+        environment: Environment variables to pass to the container.
 
     """
 
@@ -64,6 +65,7 @@ class SagemakerStepOperatorSettings(BaseSettings):
         default=None, union_mode="left_to_right"
     )
     estimator_args: Dict[str, Any] = {}
+    environment: Dict[str, str] = {}
 
     _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
         "instance_type"

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -323,6 +323,19 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                 ExecutionVariables.PIPELINE_EXECUTION_ARN
             )
 
+            if step_settings.environment:
+                step_environment = step_settings.environment.copy()
+                # Sagemaker does not allow environment variables longer than 256
+                # characters to be passed to Processor steps. If an environment variable
+                # is longer than 256 characters, we split it into multiple environment
+                # variables (chunks) and re-construct it on the other side using the
+                # custom entrypoint configuration.
+                split_environment_variables(
+                    size_limit=SAGEMAKER_PROCESSOR_STEP_ENV_VAR_SIZE_LIMIT,
+                    env=step_environment,
+                )
+                environment.update(step_environment)
+
             use_training_step = (
                 step_settings.use_training_step
                 if step_settings.use_training_step is not None
@@ -457,6 +470,11 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                         )
                     )
 
+            # Convert environment to a dict of strings
+            environment = {
+                key: str(value) for key, value in environment.items()
+            }
+
             if use_training_step:
                 # Create Estimator and TrainingStep
                 estimator = sagemaker.estimator.Estimator(

zenml/integrations/aws/step_operators/sagemaker_step_operator.py

@@ -181,6 +181,11 @@ class SagemakerStepOperator(BaseStepOperator):
                 self.name,
             )
 
+        settings = cast(SagemakerStepOperatorSettings, self.get_settings(info))
+
+        if settings.environment:
+            environment.update(settings.environment)
+
         # Sagemaker does not allow environment variables longer than 512
         # characters to be passed to Estimator steps. If an environment variable
         # is longer than 512 characters, we split it into multiple environment
@@ -194,8 +199,6 @@ class SagemakerStepOperator(BaseStepOperator):
         image_name = info.get_image(key=SAGEMAKER_DOCKER_IMAGE_KEY)
         environment[_ENTRYPOINT_ENV_VARIABLE] = " ".join(entrypoint_command)
 
-        settings = cast(SagemakerStepOperatorSettings, self.get_settings(info))
-
         # Get and default fill SageMaker estimator arguments for full ZenML support
         estimator_args = settings.estimator_args
 
@@ -221,6 +224,9 @@ class SagemakerStepOperator(BaseStepOperator):
             "instance_type", settings.instance_type or "ml.m5.large"
         )
 
+        # Convert environment to a dict of strings
+        environment = {key: str(value) for key, value in environment.items()}
+
         estimator_args["environment"] = environment
         estimator_args["instance_count"] = 1
         estimator_args["sagemaker_session"] = session

zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py

@@ -86,7 +86,7 @@ class VertexOrchestratorConfig(
             then a subdirectory of the artifact store will be used.
         encryption_spec_key_name: The Cloud KMS resource identifier of the
             customer managed encryption key used to protect the job. Has the form:
-            `projects/<PRJCT>/locations/<REGION>/keyRings/<KR>/cryptoKeys/<KEY>`
+            `projects/<PROJECT>/locations/<REGION>/keyRings/<KR>/cryptoKeys/<KEY>`
             . The key needs to be in the same region as where the compute
             resource is created.
         workload_service_account: the service account for workload run-as
@@ -124,13 +124,14 @@ class VertexOrchestratorConfig(
     pipeline_root: Optional[str] = None
     encryption_spec_key_name: Optional[str] = None
     workload_service_account: Optional[str] = None
-    function_service_account: Optional[str] = None
-    scheduler_service_account: Optional[str] = None
     network: Optional[str] = None
 
+    # Deprecated
     cpu_limit: Optional[str] = None
     memory_limit: Optional[str] = None
     gpu_limit: Optional[int] = None
+    function_service_account: Optional[str] = None
+    scheduler_service_account: Optional[str] = None
 
     _resource_deprecation = deprecation_utils.deprecate_pydantic_attributes(
         "cpu_limit",

zenml/zen_stores/sql_zen_store.py

@@ -986,6 +986,7 @@ class SqlZenStore(BaseZenStore):
             RuntimeError: if the schema does not have a `to_model` method.
         """
         query = filter_model.apply_filter(query=query, table=table)
+        query = filter_model.apply_sorting(query=query, table=table)
         query = query.distinct()
 
         # Get the total amount of items in the database for a given query
@@ -1005,9 +1006,6 @@ class SqlZenStore(BaseZenStore):
             else:
                 total = 0
 
-        # Sorting
-        query = filter_model.apply_sorting(query=query, table=table)
-
         # Get the total amount of pages in the database for a given query
         if total == 0:
             total_pages = 1

{zenml_nightly-0.74.0.dev20250224.dist-info → zenml_nightly-0.74.0.dev20250225.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.74.0.dev20250224
+Version: 0.74.0.dev20250225
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops

{zenml_nightly-0.74.0.dev20250224.dist-info → zenml_nightly-0.74.0.dev20250225.dist-info}/RECORD

@@ -1,5 +1,5 @@
 zenml/README.md,sha256=827dekbOWAs1BpW7VF1a4d7EbwPbjwccX-2zdXBENZo,1777
-zenml/VERSION,sha256=IyEgTi3qkZbE5T-L9C0C2yJSvaE5wb9untVYsQrfdo0,19
+zenml/VERSION,sha256=zSzvVswwIzD3tSsA78NP_8AZNJtWkvKKmLed4ZMMYj0,19
 zenml/__init__.py,sha256=SkMObQA41ajqdZqGErN00S1Vf3KAxpLvbZ-OBy5uYoo,2130
 zenml/actions/__init__.py,sha256=mrt6wPo73iKRxK754_NqsGyJ3buW7RnVeIGXr1xEw8Y,681
 zenml/actions/base_action.py,sha256=UcaHev6BTuLDwuswnyaPjdA8AgUqB5xPZ-lRtuvf2FU,25553
@@ -141,17 +141,17 @@ zenml/integrations/aws/container_registries/aws_container_registry.py,sha256=SCx
 zenml/integrations/aws/flavors/__init__.py,sha256=XYL9fpwKzeFfHCjakU0iJ3SyHVRJk63QT7luOy9Giek,1480
 zenml/integrations/aws/flavors/aws_container_registry_flavor.py,sha256=GIDLOySz1zF08YSkmKIj89TxBqSLWueXLAHyxYwm-NI,4169
 zenml/integrations/aws/flavors/aws_image_builder_flavor.py,sha256=XobJOw5ymbY22i7YHWGYOKDQoJQAqTeMIfvkADt-DDc,5343
-zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py,sha256=Shcvegidrxq0ogue4ya4Hm7fVjp6_hBd7O0LGdnnBio,13498
-zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py,sha256=OnNokixzGvOTBoZJQ5TDG3k4FFsP1pJmxbiij2VLW4s,5978
+zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py,sha256=aPBuzQjkroSZOdGWynXYu_A1c0YjK5dB96zb0NvHkBs,13604
+zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py,sha256=e3locb2OnF7bqV3iafIUB_tHhDE8-i7eyB4H6Hyqj1Y,6084
 zenml/integrations/aws/image_builders/__init__.py,sha256=91hgH1OphG2i-vk-G8N4yKBFIzK89Wu7BK4-T5yOA7E,786
 zenml/integrations/aws/image_builders/aws_image_builder.py,sha256=UcPYYYjJjfsicY3hV4OZeJt552AVmwPZPlv-AsG1g1I,11489
 zenml/integrations/aws/orchestrators/__init__.py,sha256=Wh0Fhtt_uo6YrkvXY9kL0M478FL7XpapjoFreUZbgUg,794
-zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py,sha256=wEMUL1kFYlhUVELbVgJTvHPh0WdgtZvKn6943PqnZn8,36348
+zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py,sha256=ZkZ9uH7FeN0J_0NyxAlTk1G_OAzTNwGiFj4sMcdbc4w,37254
 zenml/integrations/aws/orchestrators/sagemaker_orchestrator_entrypoint_config.py,sha256=WXCWdVojIZxx5_3-g1T95S2vsJ-RLNGcp-V409wgme0,1555
 zenml/integrations/aws/service_connectors/__init__.py,sha256=w2Md40yG89PwmU9eBceh6dGy3XYZ3MKusNAZ51sGOgE,783
 zenml/integrations/aws/service_connectors/aws_service_connector.py,sha256=7H69IoOYmyn5QcXEfL1-OmC0UaQ54TfNNhv2t8A6Daw,92107
 zenml/integrations/aws/step_operators/__init__.py,sha256=HK5oIqLixYKUKoiN4LnqTyGjAZJUdUqWYGqJhFSWyo0,817
-zenml/integrations/aws/step_operators/sagemaker_step_operator.py,sha256=-y1zqfIDJm2wwy_KETJ1nHdKuSbIHPFK3bT-0i9tpoA,10113
+zenml/integrations/aws/step_operators/sagemaker_step_operator.py,sha256=8g9p9eHvZkg2LY2lnWPVkriB7Jr_LbQcw2HQFjVgZ6E,10330
 zenml/integrations/aws/step_operators/sagemaker_step_operator_entrypoint_config.py,sha256=2cXroe6-bXyHVkO5yPnZagNpqx6MrMDcvuvNr8J2j-A,1581
 zenml/integrations/azure/__init__.py,sha256=99cefcfBlEXl5lYCVUzWN0uvYn_TFGpBsEATvn1d3c4,2938
 zenml/integrations/azure/artifact_stores/__init__.py,sha256=dlIwbpgjE0Hy4rhMbelNJHVKm4t8tj_hRu9mQ_cEIAg,820
@@ -263,7 +263,7 @@ zenml/integrations/gcp/flavors/__init__.py,sha256=GcB8EvYjXM_VSku16jnDSNyJYMgKc2
 zenml/integrations/gcp/flavors/gcp_artifact_store_flavor.py,sha256=Ts2jvR3IgGH8YyaFBT6OQtx2kCKD9dgZgceKiRAv_QI,3560
 zenml/integrations/gcp/flavors/gcp_image_builder_flavor.py,sha256=K6sE9D-okbdlctNwNeDYEfhWmMXXW-S92x342dnhjqY,4451
 zenml/integrations/gcp/flavors/vertex_experiment_tracker_flavor.py,sha256=icexIPoqyJ_tsO6M5-Vncd1TAUaKTGbdUG0cDOYC6Kc,6834
-zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py,sha256=m01jsJbyX9i29l4VmCNKRo6skwyLi_RQgTJM5oKWMMQ,9907
+zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py,sha256=zJBSs7sY0dT_oJFt1PCFVD3usav6LANfXEyavYg8usE,9926
 zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py,sha256=VHg2coeKoBJjSYkiVSTtbfKRg9qW9o8_QZ_xXskfzI4,6653
 zenml/integrations/gcp/google_credentials_mixin.py,sha256=bPy3JYCCcyuTmPiVFqbY81YJ5g1yRdzHLlBlokvbeqg,4026
 zenml/integrations/gcp/image_builders/__init__.py,sha256=2IvTL6U2YpUoxGQXeXew-6WFoL5hHIxkqr4DaA5Ez9w,786
@@ -1277,11 +1277,11 @@ zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py,sha256=NfW1EHIA99lseb
 zenml/zen_stores/secrets_stores/secrets_store_interface.py,sha256=Q2Jbnt2Pp7NGlR-u1YBfRZV2g8su2Fd0ArBMdksAE-Q,2819
 zenml/zen_stores/secrets_stores/service_connector_secrets_store.py,sha256=kPYX-Z_OOhZCI1CP77ncfV7IsV4e8brklnTXmKxZYNc,7078
 zenml/zen_stores/secrets_stores/sql_secrets_store.py,sha256=nEO0bAPlULBLxLVk-UTRIZiUeVpATggo8qCsKmgEU1E,8788
-zenml/zen_stores/sql_zen_store.py,sha256=tKOghX2Wa0f0xSYu_ReqHgCwslvHrBhE3RePJG2tcA0,415921
+zenml/zen_stores/sql_zen_store.py,sha256=L3PtBT-VULJVh4udD4eY4TbU_o1iEC23VyUh26uBo5Y,415902
 zenml/zen_stores/template_utils.py,sha256=EKYBgmDLTS_PSMWaIO5yvHPLiQvMqHcsAe6NUCrv-i4,9068
 zenml/zen_stores/zen_store_interface.py,sha256=vf2gKBWfUUPtcGZC35oQB6pPNVzWVyQC8nWxVLjfrxM,92692
-zenml_nightly-0.74.0.dev20250224.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
-zenml_nightly-0.74.0.dev20250224.dist-info/METADATA,sha256=hlyV5SEOUeUqvgz0VaEP8O4XXHE6-HnAPrRMo-A_oGY,21643
-zenml_nightly-0.74.0.dev20250224.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-zenml_nightly-0.74.0.dev20250224.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
-zenml_nightly-0.74.0.dev20250224.dist-info/RECORD,,
+zenml_nightly-0.74.0.dev20250225.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
+zenml_nightly-0.74.0.dev20250225.dist-info/METADATA,sha256=m6rLsnw-E4Kl_0cIfYAMQiKk3_pd5TRUIWeqqYL06AA,21643
+zenml_nightly-0.74.0.dev20250225.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+zenml_nightly-0.74.0.dev20250225.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
+zenml_nightly-0.74.0.dev20250225.dist-info/RECORD,,