From 4c9b9452b75cfd5fddb8ed4f745fe9c223c19675 Mon Sep 17 00:00:00 2001 From: Nathan Park Date: Mon, 31 Mar 2025 16:24:59 -0700 Subject: [PATCH 1/3] Fix issue #4856 by copying environment variables --- src/sagemaker/workflow/notebook_job_step.py | 50 +++++++-------------- 1 file changed, 17 insertions(+), 33 deletions(-) diff --git a/src/sagemaker/workflow/notebook_job_step.py b/src/sagemaker/workflow/notebook_job_step.py index 8a1dd6bc53..ca0ecac15b 100644 --- a/src/sagemaker/workflow/notebook_job_step.py +++ b/src/sagemaker/workflow/notebook_job_step.py @@ -13,49 +13,33 @@ """The notebook job step definitions for workflow.""" from __future__ import absolute_import +import os import re import shutil -import os +from typing import Dict, List, Optional, Union -from typing import ( - List, - Optional, - Union, - Dict, +from sagemaker import vpc_utils +from sagemaker.config.config_schema import ( + NOTEBOOK_JOB_ROLE_ARN, + NOTEBOOK_JOB_S3_KMS_KEY_ID, + NOTEBOOK_JOB_S3_ROOT_URI, + NOTEBOOK_JOB_VOLUME_KMS_KEY_ID, + NOTEBOOK_JOB_VPC_CONFIG_SECURITY_GROUP_IDS, + NOTEBOOK_JOB_VPC_CONFIG_SUBNETS, ) - +from sagemaker.s3 import S3Uploader +from sagemaker.s3_utils import s3_path_join +from sagemaker.session import get_execution_role +from sagemaker.utils import Tags, _tmpdir, format_tags, name_from_base, resolve_value_from_config +from sagemaker.workflow.entities import PipelineVariable, RequestType from sagemaker.workflow.execution_variables import ExecutionVariables from sagemaker.workflow.functions import Join from sagemaker.workflow.properties import Properties from sagemaker.workflow.retry import RetryPolicy -from sagemaker.workflow.steps import ( - Step, - ConfigurableRetryStep, - StepTypeEnum, -) from sagemaker.workflow.step_collections import StepCollection from sagemaker.workflow.step_outputs import StepOutput - -from sagemaker.workflow.entities import ( - RequestType, - PipelineVariable, -) +from sagemaker.workflow.steps import ConfigurableRetryStep, Step, StepTypeEnum from sagemaker.workflow.utilities import _collect_parameters, load_step_compilation_context -from sagemaker.session import get_execution_role - -from sagemaker.s3_utils import s3_path_join -from sagemaker.s3 import S3Uploader -from sagemaker.utils import _tmpdir, name_from_base, resolve_value_from_config, format_tags, Tags -from sagemaker import vpc_utils - -from sagemaker.config.config_schema import ( - NOTEBOOK_JOB_ROLE_ARN, - NOTEBOOK_JOB_S3_ROOT_URI, - NOTEBOOK_JOB_S3_KMS_KEY_ID, - NOTEBOOK_JOB_VOLUME_KMS_KEY_ID, - NOTEBOOK_JOB_VPC_CONFIG_SUBNETS, - NOTEBOOK_JOB_VPC_CONFIG_SECURITY_GROUP_IDS, -) # disable E1101 as collect_parameters decorator sets the attributes @@ -374,7 +358,7 @@ def _prepare_env_variables(self): execution mechanism. """ - job_envs = self.environment_variables if self.environment_variables else {} + job_envs = dict(self.environment_variables or {}) system_envs = { "AWS_DEFAULT_REGION": self._region_from_session, "SM_JOB_DEF_VERSION": "1.0", From f5eeeac75dd9bc8124d1849a8890bc2b4e746a19 Mon Sep 17 00:00:00 2001 From: Nathan Park Date: Mon, 7 Apr 2025 15:39:29 -0700 Subject: [PATCH 2/3] Add unit test --- .../workflow/test_notebook_job_step.py | 62 ++++++++++++++++++- 1 file changed, 61 insertions(+), 1 deletion(-) diff --git a/tests/unit/sagemaker/workflow/test_notebook_job_step.py b/tests/unit/sagemaker/workflow/test_notebook_job_step.py index 9cc34ee243..387609a222 100644 --- a/tests/unit/sagemaker/workflow/test_notebook_job_step.py +++ b/tests/unit/sagemaker/workflow/test_notebook_job_step.py @@ -13,10 +13,11 @@ from __future__ import absolute_import import unittest + from mock import Mock, patch -from sagemaker.workflow.notebook_job_step import NotebookJobStep from sagemaker.workflow.functions import Join +from sagemaker.workflow.notebook_job_step import NotebookJobStep REGION = "us-west-2" PIPELINE_NAME = "test-pipeline-name" @@ -573,3 +574,62 @@ def _create_step_with_required_fields(self): image_uri=IMAGE_URI, kernel_name=KERNEL_NAME, ) + + def test_environment_variables_not_shared(self): + """Test that environment variables are not shared between NotebookJob steps""" + # Setup shared environment variables + shared_env_vars = {"test": "test"} + + # Create two steps with the same environment variables dictionary + step1 = NotebookJobStep( + name="step1", + input_notebook=INPUT_NOTEBOOK, + image_uri=IMAGE_URI, + kernel_name=KERNEL_NAME, + environment_variables=shared_env_vars, + ) + + step2 = NotebookJobStep( + name="step2", + input_notebook=INPUT_NOTEBOOK, + image_uri=IMAGE_URI, + kernel_name=KERNEL_NAME, + environment_variables=shared_env_vars, + ) + + # Get the arguments for both steps + step1_args = step1.arguments + step2_args = step2.arguments + + # Verify that the environment variables are different objects + self.assertIsNot( + step1_args["Environment"], + step2_args["Environment"], + "Environment dictionaries should be different objects", + ) + + # Verify that modifying one step's environment doesn't affect the other + step1_env = step1_args["Environment"] + step2_env = step2_args["Environment"] + + # Both should have the original test value + self.assertEqual(step1_env["test"], "test") + self.assertEqual(step2_env["test"], "test") + + # Modify step1's environment + step1_env["test"] = "modified" + + # Verify step2's environment remains unchanged + self.assertEqual(step2_env["test"], "test") + + # Verify notebook names are correct for each step + self.assertEqual( + step1_env["SM_INPUT_NOTEBOOK_NAME"], + os.path.basename(INPUT_NOTEBOOK), + "Step 1 should have its own notebook name", + ) + self.assertEqual( + step2_env["SM_INPUT_NOTEBOOK_NAME"], + os.path.basename(INPUT_NOTEBOOK), + "Step 2 should have its own notebook name", + ) From eb946776dc147b020a7b44b739018b8e31931ab6 Mon Sep 17 00:00:00 2001 From: Nathan Park Date: Mon, 7 Apr 2025 18:00:45 -0700 Subject: [PATCH 3/3] Fix an import --- tests/unit/sagemaker/workflow/test_notebook_job_step.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/sagemaker/workflow/test_notebook_job_step.py b/tests/unit/sagemaker/workflow/test_notebook_job_step.py index 387609a222..6a5bb20daa 100644 --- a/tests/unit/sagemaker/workflow/test_notebook_job_step.py +++ b/tests/unit/sagemaker/workflow/test_notebook_job_step.py @@ -12,6 +12,7 @@ # language governing permissions and limitations under the License. from __future__ import absolute_import +import os import unittest from mock import Mock, patch