From 26d2608df76c5a8dfec87bf3a78f85960c19d8e3 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Mon, 12 Oct 2026 12:08:16 -0500 Subject: [PATCH] fix(mlops): call LocalPipelineSession's pipeline methods in local mode In v3 the local pipeline methods (create_pipeline, update_pipeline, describe_pipeline, delete_pipeline, start_pipeline_execution) moved from LocalSagemakerClient to sagemaker.mlops.local.LocalPipelineSession, but Pipeline's local-mode branches still called them on sagemaker_session.sagemaker_client. Pipeline.create(), upsert(), update() and start() with a LocalPipelineSession therefore always failed with AttributeError: 'LocalSagemakerClient' object has no attribute 'create_pipeline' (describe() and delete() hit the real client). Call the session's methods in local mode. The new test runs the Pipeline API against a real LocalPipelineSession; the existing local-mode unit tests asserted the client calls and now assert the session calls. --- .../src/sagemaker/mlops/workflow/pipeline.py | 13 +++++--- .../unit/local/test_local_pipeline_session.py | 33 +++++++++++++++++++ .../tests/unit/workflow/test_pipeline.py | 10 +++--- .../unit/workflow/test_pipeline_class.py | 10 +++--- 4 files changed, 52 insertions(+), 14 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py index 05622e9ca0..273326ffdf 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py @@ -212,8 +212,8 @@ def create( if self.sagemaker_session.local_mode: if parallelism_config: logger.warning("Pipeline parallelism config is not supported in the local mode.") - # TODO: replace with sagemaker-core methods - return self.sagemaker_session.sagemaker_client.create_pipeline(self, description) + # Local pipelines are managed by LocalPipelineSession itself, not by its client. + return self.sagemaker_session.create_pipeline(self, description) tags = format_tags(tags) tags = _append_project_tags(tags) tags = self.sagemaker_session._append_sagemaker_config_tags(tags, PIPELINE_TAGS_PATH) @@ -299,6 +299,9 @@ def describe(self, pipeline_version_id: int = None) -> Dict[str, Any]: `_ """ + if self.sagemaker_session.local_mode: + return self.sagemaker_session.describe_pipeline(PipelineName=self.name) + kwargs = dict(PipelineName=self.name) if pipeline_version_id: kwargs["PipelineVersionId"] = pipeline_version_id @@ -337,7 +340,7 @@ def update( if self.sagemaker_session.local_mode: if parallelism_config: logger.warning("Pipeline parallelism config is not supported in the local mode.") - return self.sagemaker_session.sagemaker_client.update_pipeline(self, description) + return self.sagemaker_session.update_pipeline(self, description) kwargs = self._create_args(role_arn, description, parallelism_config) return self.sagemaker_session.sagemaker_client.update_pipeline(**kwargs) @@ -412,6 +415,8 @@ def delete(self) -> Dict[str, Any]: "If triggers have been setup for this target, they will become orphaned." "You will need to clean them up manually via the CLI or EventBridge console." ) + if self.sagemaker_session.local_mode: + return self.sagemaker_session.delete_pipeline(PipelineName=self.name) return self.sagemaker_session.sagemaker_client.delete_pipeline(PipelineName=self.name) @_telemetry_emitter( @@ -477,7 +482,7 @@ def start( ) if self.sagemaker_session.local_mode: update_args(kwargs, PipelineParameters=parameters) - return self.sagemaker_session.sagemaker_client.start_pipeline_execution(**kwargs) + return self.sagemaker_session.start_pipeline_execution(**kwargs) update_args(kwargs, PipelineParameters=format_start_parameters(parameters)) # retry on AccessDeniedException to cover case of tag propagation delay diff --git a/sagemaker-mlops/tests/unit/local/test_local_pipeline_session.py b/sagemaker-mlops/tests/unit/local/test_local_pipeline_session.py index 44a646f1bd..cdc6ca043c 100644 --- a/sagemaker-mlops/tests/unit/local/test_local_pipeline_session.py +++ b/sagemaker-mlops/tests/unit/local/test_local_pipeline_session.py @@ -274,3 +274,36 @@ def test_start_pipeline_execution_not_found(local_session): error = exc_info.value assert error.response["Error"]["Code"] == "ResourceNotFound" assert "nonexistent-pipeline" in error.response["Error"]["Message"] + + +@patch( + "sagemaker.mlops.workflow.pipeline.resolve_and_validate_role", + return_value="arn:aws:iam::123456789012:role/role", +) +def test_pipeline_api_in_local_mode(_mock_role): + """Pipeline's create/upsert/describe/start/delete work with a real LocalPipelineSession.""" + import boto3 + from sagemaker.mlops.workflow.fail_step import FailStep + from sagemaker.mlops.workflow.pipeline import Pipeline + + boto_session = boto3.Session( + region_name="us-west-2", aws_access_key_id="testing", aws_secret_access_key="testing" + ) + session = LocalPipelineSession(boto_session=boto_session, default_bucket="bucket") + pipeline = Pipeline( + name="local-pipeline", + steps=[FailStep("fail", error_message="boom")], + sagemaker_session=session, + ) + + assert pipeline.create(description="d") == {"PipelineArn": "local-pipeline"} + assert pipeline.upsert(description="d2") == {"PipelineArn": "local-pipeline"} + assert pipeline.update(description="d3") == {"PipelineArn": "local-pipeline"} + assert pipeline.describe()["PipelineArn"] == "local-pipeline" + + execution = pipeline.start() + assert execution.describe()["PipelineExecutionStatus"] == "Failed" + + pipeline.delete() + with pytest.raises(ClientError): + pipeline.describe() diff --git a/sagemaker-mlops/tests/unit/workflow/test_pipeline.py b/sagemaker-mlops/tests/unit/workflow/test_pipeline.py index 759b906c02..28ba37ee6c 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_pipeline.py +++ b/sagemaker-mlops/tests/unit/workflow/test_pipeline.py @@ -63,7 +63,7 @@ def test_pipeline_create_without_role_auto_resolves(mock_session, mock_step): role_type="pipeline", sagemaker_session=mock_session, ) - mock_session.sagemaker_client.create_pipeline.assert_called_once() + mock_session.create_pipeline.assert_called_once() def test_format_start_parameters(): @@ -104,7 +104,7 @@ def test_pipeline_create_local_mode(mock_session, mock_step): description="test", parallelism_config={"MaxParallelExecutionSteps": 2}, ) - mock_session.sagemaker_client.create_pipeline.assert_called_once() + mock_session.create_pipeline.assert_called_once() def test_pipeline_create_large_definition(mock_session, mock_step): @@ -147,7 +147,7 @@ def test_pipeline_update_without_role_auto_resolves(mock_session, mock_step): role_type="pipeline", sagemaker_session=mock_session, ) - mock_session.sagemaker_client.update_pipeline.assert_called_once() + mock_session.update_pipeline.assert_called_once() def test_pipeline_update_local_mode(mock_session, mock_step): @@ -165,7 +165,7 @@ def test_pipeline_update_local_mode(mock_session, mock_step): description="test", parallelism_config={"MaxParallelExecutionSteps": 2}, ) - mock_session.sagemaker_client.update_pipeline.assert_called_once() + mock_session.update_pipeline.assert_called_once() def test_pipeline_upsert_without_role_auto_resolves(mock_session, mock_step): @@ -245,7 +245,7 @@ def test_pipeline_start_local_mode(mock_session, mock_step): pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session) pipeline.start(parameters={"param": "value"}) - mock_session.sagemaker_client.start_pipeline_execution.assert_called_once() + mock_session.start_pipeline_execution.assert_called_once() def test_create_args_converts_parallelism_config_to_dict(mock_session, mock_step): diff --git a/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py b/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py index cab734f1d5..f2e3173a10 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py +++ b/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py @@ -272,7 +272,7 @@ def test_create_with_parallelism_config(self, mock_session): def test_create_local_mode(self, mock_session): """Test create pipeline in local mode.""" mock_session.local_mode = True - mock_session.sagemaker_client.create_pipeline = Mock( + mock_session.create_pipeline = Mock( return_value={"PipelineArn": "test-pipeline"} ) @@ -288,7 +288,7 @@ def test_create_local_mode(self, mock_session): result = pipeline.create(role_arn="arn:aws:iam::123:role/SageMakerRole") assert result["PipelineArn"] == "test-pipeline" - mock_session.sagemaker_client.create_pipeline.assert_called_once() + mock_session.create_pipeline.assert_called_once() class TestPipelineUpdate: @@ -334,7 +334,7 @@ def test_update_without_role_raises_error(self, mock_session): def test_update_local_mode(self, mock_session): """Test update pipeline in local mode.""" mock_session.local_mode = True - mock_session.sagemaker_client.update_pipeline = Mock( + mock_session.update_pipeline = Mock( return_value={"PipelineArn": "test-pipeline"} ) @@ -665,14 +665,14 @@ def test_start_with_execution_display_name(self, mock_session): def test_start_local_mode(self, mock_session): """Test start pipeline in local mode.""" mock_session.local_mode = True - mock_session.sagemaker_client.start_pipeline_execution = Mock(return_value=Mock()) + mock_session.start_pipeline_execution = Mock(return_value=Mock()) pipeline = Pipeline(name="test-pipeline", sagemaker_session=mock_session) result = pipeline.start(parameters={"param1": "value1"}) assert result is not None - mock_session.sagemaker_client.start_pipeline_execution.assert_called_once() + mock_session.start_pipeline_execution.assert_called_once() class TestPipelineDefinition: