Getting TrainingJobName from Training step of Sagemaker Pipeline

Question

I am running a sagemaker pipeline with a training step. This whole setup runs from a **Lambda function** while I pass a few parameters to the pipeline.py file. To get the TrainingJobName from the training step, 
my code is `step_train.__dict__['step_args']['TrainingJobName']`
This works fine while I running it in sagemaker notebook but when I execute the same code to get the train job name from lambda, I get this error **[ERROR] TypeError: '_StepArguments' object is not subscriptable**
`[ERROR] TypeError: '_StepArguments' object is not subscriptable
Traceback (most recent call last):
  File "/var/task/lambda_function.py", line 46, in lambda_handler
    pipeline = create_pipeline(validated_api_input)
  File "/tmp/pipeline.py", line 105, in create_pipeline
    "job_name" : training_step['step_args']['TrainingJobName']`

How do I resolve this?

Answer

Hello there,\
Thank you for contacting.
The error specifically comes up when you are trying to use a subscript for a data-type which isn't subscript-able, in your case this is coming when you are treating the object ```_StepArguments``` as a dictionary. A workaround for the same will be to implement a ```__get_item___``` function in your code. A sample for the same is as follows:

```
def __getitem__(self, key):
        return self.__dict__[key]
```
However, as this was working on Sage-maker instance by default, also share the Python version that you are using and if you are using the same version of library in both Lambda function and SageMaker.

Answer

`_StepArguments` is a new concept introduced in newer version of SageMaker Python SDK. If you want to keep you original code, you need to pin the same SDK version in Lambda as in your notebook instance. Normally, I guess you will be installing the latest version but there has been a compatibility-broken change since then.

If you are OK to code changes, you could follow https://docs.aws.amazon.com/sagemaker/latest/dg/build-and-manage-steps.html#build-and-manage-properties to use `properties` of each step, which maps to `Describe*` response. You will still be able to retrieve any job info. This is actually the recommended way.

Getting TrainingJobName from Training step of Sagemaker Pipeline

관련 콘텐츠