Repository: aws/sagemaker-python-sdk

Issue:
Wrong notebook being run in pipeline when creating NotebookJobStep with shared environment_variables dict
**Describe the bug**
When I create a Pipeline with two NotebookJobStep steps and both steps were created using the same dict as environment_variables parameter the first step is run with the second input notebook isntead of its own input notebook.

**To reproduce**
```
env_vars = {
    'test': 'test',
}
steps = [
    NotebookJobStep(
        image_uri="885854791233.dkr.ecr.us-east-1.amazonaws.com/sagemaker-distribution-prod:1-cpu",
        kernel_name="python3",
        input_notebook="job1.ipynb",
        initialization_script="setup.sh",
        environment_variables=env_vars,
    ),
    NotebookJobStep(
        image_uri="885854791233.dkr.ecr.us-east-1.amazonaws.com/sagemaker-distribution-prod:1-cpu",
        kernel_name="python3",
        input_notebook="job2.ipynb",
        initialization_script="setup.sh",
        environment_variables=env_vars,
    ),
]
pipeline = Pipeline(
    name="pipeline",
    steps=steps,
)
pipeline.upsert(role_arn=role)
execution = pipeline.start()
```

The problem seems the env vars for each step:

print(json.loads(pipeline.definition())["Steps"][0]["Arguments"]["Environment"])
{
 'test': 'test',
 'AWS_DEFAULT_REGION': 'us-east-1',
 'SM_JOB_DEF_VERSION': '1.0',
 'SM_ENV_NAME': 'sagemaker-default-env',
 'SM_SKIP_EFS_SIMULATION': 'true',
 'SM_EXECUTION_INPUT_PATH': '/opt/ml/input/data/sagemaker_headless_execution_pipelinestep',
 'SM_KERNEL_NAME': 'python3',
 'SM_INPUT_NOTEBOOK_NAME': 'job2.ipynb',  **<<==== wrong input**
 'SM_OUTPUT_NOTEBOOK_NAME': 'job2-ipynb-2024-08-29-15-04-49-575.ipynb',
 'SM_INIT_SCRIPT': 'setup.sh'
}

print(json.loads(pipeline.definition())["Steps"][1]["Arguments"]["Environment"])
{
 'test': 'test',
 'AWS_DEFAULT_REGION': 'us-east-1',
 'SM_JOB_DEF_VERSION': '1.0',
 'SM_ENV_NAME': 'sagemaker-default-env',
 'SM_SKIP_EFS_SIMULATION': 'true',
 'SM_EXECUTION_INPUT_PATH': '/opt/ml/input/data/sagemaker_headless_execution_pipelinestep',
 'SM_KERNEL_NAME': 'python3',
 'SM_INPUT_NOTEBOOK_NAME': 'job2.ipynb',
 'SM_OUTPUT_NOTEBOOK_NAME': 'job2-ipynb-2024-08-29-15-04-49-575.ipynb',
 'SM_INIT_SCRIPT': 'setup.sh'
}

**Expected behavior**
Run job1.ipynb and job2.ipynb in each step.

**Screenshots or logs**

Screenshot of notebook jobs in Studio UI:

![image](https://github.com/user-attachments/assets/15865368-d16c-4cb0-8286-f3e079a52dfb)


**System information**
A description of your system. Please provide:
- **SageMaker Python SDK version**: 2.226.1
- **Framework name (eg. PyTorch) or algorithm (eg. KMeans)**:
- **Framework version**:
- **Python version**: 3.8.18
- **CPU or GPU**: CPU
- **Custom Docker image (Y/N)**: N

Discussion:
@qidewenwhen: Hi @fakio, thanks for the good finding! And thanks for the investigation.

This is a bug for sure. 

The issue is here https://github.com/aws/sagemaker-python-sdk/blob/1e679b415301ade807653308701b7fed4eefd9ec/src/sagemaker/workflow/notebook_job_step.py#L377-L391

The `environment_variables` from user is directly updated with the `system_envs`. Thus if the same `environment_variables` is used across multiple notebook job steps, they could override each other.

The fix is to simply copy the `environment_variables` from users and add `system_envs` to the copy.

However, as a workaround, you can create separate `environment_variables` for different steps until we release a fix.

I'll creating a backlog for the fix of this bug in our service queue, and will update this issue once the fix PR is published.

Solution patch:
diff --git a/src/sagemaker/workflow/notebook_job_step.py b/src/sagemaker/workflow/notebook_job_step.py
index 8a1dd6bc..ca0ecac1 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 @@ class NotebookJobStep(ConfigurableRetryStep):
         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",
diff --git a/src/sagemaker/workflow/steps.py b/src/sagemaker/workflow/steps.py
index f49e457b..dbc37371 100644
--- a/src/sagemaker/workflow/steps.py
+++ b/src/sagemaker/workflow/steps.py
@@ -18,7 +18,6 @@ import warnings
 
 from enum import Enum
 from typing import Dict, List, Set, Union, Optional, Any, TYPE_CHECKING
-from urllib.parse import urlparse
 
 import attr
 
@@ -465,6 +464,7 @@ class TrainingStep(ConfigurableRetryStep):
         self.step_args = step_args
         self.estimator = estimator
         self.inputs = inputs
+        self.job_name = None
 
         self._properties = Properties(
             step_name=name, step=self, shape_name="DescribeTrainingJobResponse"
@@ -493,19 +493,6 @@ class TrainingStep(ConfigurableRetryStep):
                 DeprecationWarning,
             )
 
-        self.job_name = None
-        if estimator and (estimator.source_dir or estimator.entry_point):
-            # By default, `Estimator` will upload the local code to an S3 path
-            # containing a timestamp. This causes cache misses whenever a
-            # pipeline is updated, even if the underlying script hasn't changed.
-            # To avoid this, hash the contents of the training script and include it
-            # in the `job_name` passed to the `Estimator`, which will be used
-            # instead of the timestamped path.
-            if not is_pipeline_variable(estimator.source_dir) and not is_pipeline_variable(
-                estimator.entry_point
-            ):
-                self.job_name = self._generate_code_upload_path()
-
     @property
     def arguments(self) -> RequestType:
         """The arguments dictionary that is used to call `create_training_job`.
@@ -554,26 +541,6 @@ class TrainingStep(ConfigurableRetryStep):
 
         return request_dict
 
-    def _generate_code_upload_path(self) -> str or None:
-        """Generate an upload path for local training scripts based on their content."""
-        from sagemaker.workflow.utilities import hash_files_or_dirs
-
-        if self.estimator.source_dir:
-            source_dir_url = urlparse(self.estimator.source_dir)
-            if source_dir_url.scheme == "" or source_dir_url.scheme == "file":
-                code_hash = hash_files_or_dirs(
-                    [self.estimator.source_dir] + self.estimator.dependencies
-                )
-                return f"{self.name}-{code_hash}"[:1024]
-        elif self.estimator.entry_point:
-            entry_point_url = urlparse(self.estimator.entry_point)
-            if entry_point_url.scheme == "" or entry_point_url.scheme == "file":
-                code_hash = hash_files_or_dirs(
-                    [self.estimator.entry_point] + self.estimator.dependencies
-                )
-                return f"{self.name}-{code_hash}"[:1024]
-        return None
-
 
 class CreateModelStep(ConfigurableRetryStep):
     """`CreateModelStep` for SageMaker Pipelines Workflows."""
@@ -895,16 +862,6 @@ class ProcessingStep(ConfigurableRetryStep):
                         "code argument has to be a valid S3 URI or local file path "
                         + "rather than a pipeline variable"
                     )
-                code_url = urlparse(code)
-                if code_url.scheme == "" or code_url.scheme == "file":
-                    # By default, `Processor` will upload the local code to an S3 path
-                    # containing a timestamp. This causes cache misses whenever a
-                    # pipeline is updated, even if the underlying script hasn't changed.
-                    # To avoid this, hash the contents of the script and include it
-                    # in the `job_name` passed to the `Processor`, which will be used
-                    # instead of the timestamped path.
-                    self.job_name = self._generate_code_upload_path()
-
             warnings.warn(
                 (
                     'We are deprecating the instantiation of ProcessingStep using "processor".'


Test patch:
diff --git a/tests/unit/sagemaker/workflow/test_notebook_job_step.py b/tests/unit/sagemaker/workflow/test_notebook_job_step.py
index 9cc34ee2..6a5bb20d 100644
--- a/tests/unit/sagemaker/workflow/test_notebook_job_step.py
+++ b/tests/unit/sagemaker/workflow/test_notebook_job_step.py
@@ -12,11 +12,13 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import os
 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 +575,62 @@ class TestNotebookJobStep(unittest.TestCase):
             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",
+        )
diff --git a/tests/unit/sagemaker/workflow/test_steps.py b/tests/unit/sagemaker/workflow/test_steps.py
index 248fee65..84906ce6 100644
--- a/tests/unit/sagemaker/workflow/test_steps.py
+++ b/tests/unit/sagemaker/workflow/test_steps.py
@@ -671,7 +671,7 @@ def test_processing_step_normalizes_args_with_local_code(mock_normalize_args, sc
     mock_normalize_args.return_value = [step.inputs, step.outputs]
     step.to_request()
     mock_normalize_args.assert_called_with(
-        job_name="MyProcessingStep-a22fc59b38f13da26f6a40b18687ba598cf669f74104b793cefd9c63eddf4ac7",
+        job_name=None,
         arguments=step.job_arguments,
         inputs=step.inputs,
         outputs=step.outputs,
diff --git a/tests/unit/sagemaker/workflow/test_utils.py b/tests/unit/sagemaker/workflow/test_utils.py
index e16293a1..b18ed71f 100644
--- a/tests/unit/sagemaker/workflow/test_utils.py
+++ b/tests/unit/sagemaker/workflow/test_utils.py
@@ -80,11 +80,11 @@ def test_repack_model_step(estimator):
     assert hyperparameters["inference_script"] == '"dummy_script.py"'
     assert hyperparameters["model_archive"] == '"s3://my-bucket/model.tar.gz"'
     assert hyperparameters["sagemaker_program"] == f'"{REPACK_SCRIPT_LAUNCHER}"'
-    assert (
-        hyperparameters["sagemaker_submit_directory"]
-        == '"s3://my-bucket/MyRepackModelStep-717d7bdd388168c27e9ad2938ff0314e35be50b3157cf2498688c7525ea27e1e\
-/source/sourcedir.tar.gz"'
-    )
+
+    # ex: "gits3://my-bucket/sagemaker-scikit-learn-2025-04-07-20-39-38-854/source/sourcedir.tar.gz"
+    sagemaker_submit_directory = hyperparameters["sagemaker_submit_directory"]
+    assert sagemaker_submit_directory.startswith('"s3://my-bucket/sagemaker-scikit-learn-')
+    assert sagemaker_submit_directory.endswith('/source/sourcedir.tar.gz"')
 
     del request_dict["Arguments"]["HyperParameters"]
     del request_dict["Arguments"]["AlgorithmSpecification"]["TrainingImage"]