MITLibraries / timdex-pipeline-lambdas / 21599578295

92%

main: 92%

Pull #334

github

Add embeddings-create and embeddings-load steps to pipeline

Why these changes are being introduced:

Now that we have AWS Batch pipelines that can generate embeddings, and TIM is prepared
to load them into Opensearch, we need the pipeline lambda to prepare commands for both
running the AWS Batch job to create embeddings, and then the TIM command to load the
embeddings.

This dovetails with work in https://mitlibraries.atlassian.net/browse/USE-215 which
has proposed updates to the StepFunction.   There are two new pipeline lambda
invocations in the StepFunction that will utilize the two new allowed 'next-step'
values introduced in this commit.

How this addresses that need:
* Add "embeddings-create" and "embeddings-load" as valid steps in config
* Add SKIP_EMBEDDINGS_SOURCES config for sources that don't need embeddings (alma, gisogm)
* Add generate_embeddings_create_command() which determines compute env (cpu vs gpu-spot)
  based on record count threshold
* Add generate_embeddings_load_command() for TIM bulk-update-embeddings command
* Add handlers for both new steps in format_input.py
* Update handle_load() to flow into embeddings-create instead of end
* Add run_id and embeddings fields to ResultPayload
* Add unit tests for new functionality

Side effects of this change:
* Pipeline will now continue to embeddings steps after load completes

Relevant ticket(s):
* https://mitlibraries.atlassian.net/browse/USE-140

31 of 50 new or added lines in 3 files covered. (62.0%)

378 of 411 relevant lines covered (91.97%)

0.92 hits per line