def _generate_input_data_config()

in src/sagemaker/remote_function/job.py [0:0]

• 73 lines of code
• 5 McCabe index (conditional complexity)

def _generate_input_data_config(job_settings: _JobSettings, s3_base_uri: str):
    """Generates input data config"""
    from sagemaker.workflow.utilities import load_step_compilation_context

    step_compilation_context = load_step_compilation_context()

    bootstrap_scripts_s3uri = _prepare_and_upload_runtime_scripts(
        spark_config=job_settings.spark_config,
        s3_base_uri=s3_base_uri,
        s3_kms_key=job_settings.s3_kms_key,
        sagemaker_session=job_settings.sagemaker_session,
        use_torchrun=job_settings.use_torchrun,
        use_mpirun=job_settings.use_mpirun,
    )

    input_data_config = [
        dict(
            ChannelName=RUNTIME_SCRIPTS_CHANNEL_NAME,
            DataSource={
                "S3DataSource": {
                    "S3Uri": bootstrap_scripts_s3uri,
                    "S3DataType": "S3Prefix",
                }
            },
        )
    ]

    local_dependencies_path = RuntimeEnvironmentManager().snapshot(job_settings.dependencies)

    if step_compilation_context:
        with _tmpdir() as tmp_dir:
            script_and_dependencies_s3uri = _prepare_dependencies_and_pre_execution_scripts(
                local_dependencies_path=local_dependencies_path,
                pre_execution_commands=job_settings.pre_execution_commands,
                pre_execution_script_local_path=job_settings.pre_execution_script,
                s3_base_uri=s3_base_uri,
                s3_kms_key=job_settings.s3_kms_key,
                sagemaker_session=job_settings.sagemaker_session,
                tmp_dir=tmp_dir,
            )

            if script_and_dependencies_s3uri:
                input_data_config.append(
                    dict(
                        ChannelName=SCRIPT_AND_DEPENDENCIES_CHANNEL_NAME,
                        DataSource={
                            "S3DataSource": {
                                "S3Uri": script_and_dependencies_s3uri,
                                "S3DataType": "S3Prefix",
                            }
                        },
                    )
                )

    user_workspace_s3uri = _prepare_and_upload_workspace(
        local_dependencies_path=local_dependencies_path,
        include_local_workdir=job_settings.include_local_workdir,
        pre_execution_commands=job_settings.pre_execution_commands,
        pre_execution_script_local_path=job_settings.pre_execution_script,
        s3_base_uri=s3_base_uri,
        s3_kms_key=job_settings.s3_kms_key,
        sagemaker_session=job_settings.sagemaker_session,
        custom_file_filter=job_settings.custom_file_filter,
    )

    if user_workspace_s3uri:
        input_data_config.append(
            dict(
                ChannelName=(
                    REMOTE_FUNCTION_WORKSPACE
                    if not step_compilation_context
                    else step_compilation_context.pipeline_build_time
                ),
                DataSource={
                    "S3DataSource": {
                        "S3Uri": user_workspace_s3uri,
                        "S3DataType": "S3Prefix",
                    }
                },
            )
        )

    return input_data_config