Use of `source_dir` argument for `FrameworkProcessor` disables step caching functionality

[francisco-camargo]

Describe the bug
I am attempting to find a Processor that allows me to specify requirements.txt per container/step and allow containers to import custom code from scripts I write besides the one specified by the code argument. To this end, it seems like sagemaker.processing.FrameworkProcessor is a good option.

However, when I use the source_dir argument (which enables me to import custom code), it seems to prevent step caching from working.

To reproduce

Here is the code such that step caching does work:

I define

BASE_DIR = os.path.dirname(os.path.realpath(__file__))

and the processing step is defined as

    sklearn_processor = FrameworkProcessor(
        estimator_cls=SKLearn,
        framework_version="0.23-1",
        instance_type=processing_instance_type,
        instance_count=processing_instance_count,
        base_job_name=base_job_name,
        sagemaker_session=pipeline_session,
        role=role,
        )
    step_args = sklearn_processor.run(
        outputs=[
            ProcessingOutput(output_name="train", source="/opt/ml/processing/train"),
            ProcessingOutput(output_name="validation", source="/opt/ml/processing/validation"),
            ProcessingOutput(output_name="test", source="/opt/ml/processing/test"),
        ],
        code=os.path.join(BASE_DIR, "preprocess.py"),
        arguments=[
            '--input-data', input_data,
            '--random-seed', random_seed,
            '--train-fraction', train_fraction,
            '--validation-fraction', validation_fraction,
            ],
        )
    step_process = ProcessingStep(
        name=step_name,
        step_args=step_args,
        cache_config=CacheConfig(enable_caching=True, expire_after="T12h"),
        )

This works as expected; takes about 5 minutes to run the first time and then only a couple seconds afterward.

If I switch to

        code="preprocess.py",
        source_dir=BASE_DIR,

Caching no longer works and it takes 5 minutes to execute the pipeline every time.

Note: if I instead switch to

        code=os.path.join(BASE_DIR, "preprocess.py"),
        source_dir=BASE_DIR,

The SageMaker Pipeline fails

System information
A description of your system. Please provide:

• SageMaker Python SDK version:
  sagemaker-2.196.0
• Framework name (eg. PyTorch) or algorithm (eg. KMeans):
• Framework version:
• Python version:
  3.11
• CPU or GPU:
• Custom Docker image (Y/N):

Additional context
Add any other context about the problem here.

[qidewenwhen]

Hi @francisco-camargo , thanks for reaching out!

For ProcessingStep, cache hit is determined by the hash calculated as per the rules below:

1. if source_dir is supplied -> calculate the hash via everything in the source_dir (which includes the code file) and dependencies. See here.
2. if only code is supplied -> calculate the hash via the contents in the code file and dependencies.

Given that in your case, the source_dir is BASE_DIR = os.path.dirname(os.path.realpath(__file__)), which is the current work space right? I guess there might be file changes in this BASE_DIR, which led to the cache miss.

Could you retry by creating a separate folder for source_dir, putting "preprocess.py" and the custom code/script needed import there and making sure the contents in source_dir and dependencies (if any) won't change?
And feel free to let us know if the issue persists.
In this case, if possible, please provide us a zip file with the code defining the pipeline and its steps and the files in source_dir to quickly reproduce the issue. Thanks!

Additionally, for this attempt you've mentioned:

Note: if I instead switch to
code=os.path.join(BASE_DIR, "preprocess.py"),
source_dir=BASE_DIR,
The SageMaker Pipeline fails

This is not a correct usage. IIRC, if source_dir is supplied, the code must be a relative path pointing to a file under source_dir, that's why the error was thrown.

[francisco-camargo]

Hi @qidewenwhen, thanks for the help, your suggestion to create a separate folder to use as source_dir worked! The caching now works as expected.

code = 'preprocess.py'
source_dir = './pipelines/my_project/preprocess/'

[francisco-camargo]

Here is what I ended up going with. It allows me to specify both requirements.txt for the preprocessing step and allows me to import code from helper_functions.py, and step caching works.

The folder structure:

pipelines/my_project/
    __init__.py
    pipeline.py
    preprocess/
        __init__.py
        step_processing.py
        preprocess.py
        requirements.txt
        helper_functions.py

Within step_processing.py, which calls FrameworkProcessor(), I use the following configuration

source_dir = './pipelines/tapeout_forecast_pipeline/preprocess/'
code = 'preprocess.py'
dependencies = ['./pipelines/my_project/preprocess/requirements.txt']

[francisco-camargo]

I tried to remove having to explicitly mention /my_project/ in either defining source_dir or dependencies.

If I use 'os.path.join(BASE_DIR, 'preprocess'), caching no longer works (I executed the pipeline three times without changing anything).

Similarly, if I want to remove explicitly mentioning /my_project/ in defining dependencies I run into trouble. If I ignore the caching issue and I try to use either BASE_DIR or source_dir I can't get a successful execution of the pipeline.