@step decorator import doesn't work

[Rizhiy]

Describe the bug When running:

from sagemaker.workflow.function_step import step

I get:

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[9], line 1
----> 1 from sagemaker.workflow.function_step import step
      2 from sagemaker.workflow.pipeline import Pipeline
      3 from sagemaker.workflow.pipeline_context import LocalPipelineSession

File /opt/conda/lib/python3.10/site-packages/sagemaker/workflow/function_step.py:40
     38 from sagemaker.workflow.retry import RetryPolicy
     39 from sagemaker.workflow.steps import Step, ConfigurableRetryStep, StepTypeEnum
---> 40 from sagemaker.workflow.step_collections import StepCollection
     41 from sagemaker.workflow.step_outputs import StepOutput
     42 from sagemaker.workflow.utilities import trim_request_dict, load_step_compilation_context

File /opt/conda/lib/python3.10/site-packages/sagemaker/workflow/step_collections.py:29
     27 from sagemaker.workflow.step_outputs import StepOutput
     28 from sagemaker.workflow.steps import Step, CreateModelStep, TransformStep
---> 29 from sagemaker.workflow._utils import _RegisterModelStep, _RepackModelStep
     30 from sagemaker.workflow.retry import RetryPolicy
     31 from sagemaker.utils import update_container_with_inference_params, format_tags

File /opt/conda/lib/python3.10/site-packages/sagemaker/workflow/_utils.py:35
     28 from sagemaker.session import get_create_model_package_request, get_model_package_args
     29 from sagemaker.workflow.steps import (
     30     StepTypeEnum,
     31     TrainingStep,
     32     Step,
     33     ConfigurableRetryStep,
     34 )
---> 35 from sagemaker.utils import _save_model, download_file_from_url, format_tags
     36 from sagemaker.workflow.retry import RetryPolicy
     37 from sagemaker.workflow.utilities import trim_request_dict

ImportError: cannot import name 'format_tags' from 'sagemaker.utils' (/opt/conda/lib/python3.10/site-packages/sagemaker/utils.py)

To reproduce Running in Sagemaker Studio Jupyter Lab instance:

• Open notebook
• Upgrade sagemaker: pip install -U sagemaker[local]
• Run code

Expected behavior It works

Screenshots or logs N/A

System information A description of your system. Please provide:

• SageMaker Python SDK version: ~~2.207.1~~
• Framework name (eg. PyTorch) or algorithm (eg. KMeans): N/A
• Framework version: N/A
• Python version: Python 3.10.13
• CPU or GPU: N/A
• Custom Docker image (Y/N): N

Additional context How is there no unit tests to make sure that everything at least imports successfully?

[Rizhiy]

NVM, I forgot to reload notebook, it works in latest version. Leaving this open here until default package version in Studio gets updated.

[akrishna1995]

Thanks @Rizhiy - are you launching studio via Jupyterlab ? I can talk to the team internally to see what version of the sdk is pinned currently

[Rizhiy]

I'm launching lab from within SM Studio, yes.

[nmadan]

Leaving this open here until default package version in Studio gets updated.

Hi @Rizhiy, we regularly update the sagemaker pysdk version in the default Studio image. If you need newer changes that the ones deployed, you can always upgrade manually pip install -U sagemaker. Please feel free to reopen this sim if you need anything else