Sagemaker processing job with PySpark and Step Functions

Question

this is my problem: I have to run a Sagemaker processing job using custom code written in PySpark. I've used the Sagemaker SDK by running these commands:

spark_processor = sagemaker.spark.processing.PySparkProcessor(
        base_job_name="spark-preprocessor",
        framework_version="2.4",
        role=role_arn,
        instance_count=2,
        instance_type="ml.m5.xlarge",
        max_runtime_in_seconds=1800,
    )

    spark_processor.run(
        submit_app="processing.py",
        arguments=['s3_input_bucket', bucket_name,
                   's3_input_file_path', file_path
                   ]
    )

Now I have to automate the workflow by using Step Functions. For this purpose, I've written a lambda function to do that but I receive the following error:

{
  "errorMessage": "Unable to import module 'lambda_function': No module named 'sagemaker'",
  "errorType": "Runtime.ImportModuleError"
}

This is my lambda function:

import sagemaker

def lambda_handler(event, context):
    spark_processor = sagemaker.spark.processing.PySparkProcessor(
        base_job_name="spark-preprocessor",
        framework_version="2.4",
        role=role_arn,
        instance_count=2,
        instance_type="ml.m5.xlarge",
        max_runtime_in_seconds=1800,
    )

    spark_processor.run(
        submit_app="processing.py",
        arguments=['s3_input_bucket', event["bucket_name"],
                   's3_input_file_path', event["file_path"]
                   ]
    )

My question is: How can I create a step in my state machine for running a PySpark code using Sagemaker processing?

Thank you

user_1242446 · Answer

sagemaker sdk is not installed by default in the lambda container environment: you should include it in the lambda zip that you upload to s3.

There are various ways to do this, one of the easiest is to deploy your lambda with Serverless Application Model (SAM) cli. In this case it might be enough to place sagemaker in the requirements.txt placed in the folder that contains your lambda code, and SAM will ensure that the dependency is included in the zip.

Alternatively you can create the zip manually with pip install sagemaker -t lambda_folder but you should execute this command in an Amazon Linux OS, for example with an EC2 with the appropriate image or in a Docker container. Search for "python dependencies in aws lambda" for more info.

Sagemaker processing job with PySpark and Step Functions

1 Answers