Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix test_vdb_upload_pipe #1381

Merged
merged 2 commits into from
Nov 22, 2023
Merged

Fix test_vdb_upload_pipe #1381

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
423e52b
Revert accidental change during merge conflicts
dagardner-nv Nov 22, 2023
10e9b9f
When caching is disabled web-scraper uses requests not requests_cache…
dagardner-nv Nov 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 10 additions & 9 deletions tests/llm/test_vdb_upload_pipe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,8 +57,9 @@ def _run_pipeline(config: Config,

pipe = LinearPipeline(config)

pipe.set_source(RSSSourceStage(config, feed_input=rss_files, batch_size=128, run_indefinitely=False))
pipe.add_stage(web_scraper_stage_mod.WebScraperStage(config, chunk_size=MODEL_FEA_LENGTH))
pipe.set_source(
RSSSourceStage(config, feed_input=rss_files, batch_size=128, run_indefinitely=False, enable_cache=False))
pipe.add_stage(web_scraper_stage_mod.WebScraperStage(config, chunk_size=MODEL_FEA_LENGTH, enable_cache=False))
pipe.add_stage(DeserializeStage(config))

pipe.add_stage(
Expand All @@ -70,7 +71,7 @@ def _run_pipeline(config: Config,
column='page_content'))

pipe.add_stage(
TritonInferenceStage(config, model_name='all-MiniLM-L6-v2', server_url='test:0000', force_convert_inputs=True))
TritonInferenceStage(config, model_name='test-model', server_url='test:0000', force_convert_inputs=True))

pipe.add_stage(
WriteToVectorDBStage(config,
Expand All @@ -89,10 +90,10 @@ def _run_pipeline(config: Config,
os.path.join(TEST_DIRS.examples_dir, 'llm/common/utils.py'),
os.path.join(TEST_DIRS.examples_dir, 'llm/common/web_scraper_stage.py')
])
@mock.patch('requests_cache.CachedSession')
@mock.patch('requests.Session')
@mock.patch('tritonclient.grpc.InferenceServerClient')
def test_vdb_upload_pipe(mock_triton_client: mock.MagicMock,
mock_cache_session: mock.MagicMock,
mock_requests_session: mock.MagicMock,
config: Config,
dataset: DatasetManager,
milvus_server_uri: str,
Expand Down Expand Up @@ -134,17 +135,17 @@ def test_vdb_upload_pipe(mock_triton_client: mock.MagicMock,
async_infer = mk_async_infer(inf_results)
mock_triton_client.async_infer.side_effect = async_infer

# Mock requests_cache, since we are feeding the RSSSourceStage with a local file it won't be using the
# requests_cache lib, only web_scraper_stage.py will use it.
# Mock requests, since we are feeding the RSSSourceStage with a local file it won't be using the
# requests lib, only web_scraper_stage.py will use it.
def mock_get_fn(url: str):
mock_response = mock.MagicMock()
mock_response.ok = True
mock_response.status_code = 200
mock_response.text = web_responses[url]
return mock_response

mock_cache_session.return_value = mock_cache_session
mock_cache_session.get.side_effect = mock_get_fn
mock_requests_session.return_value = mock_requests_session
mock_requests_session.get.side_effect = mock_get_fn

(utils_mod, web_scraper_stage_mod) = import_mod
collection_name = "test_vdb_upload_pipe"
Expand Down
4 changes: 2 additions & 2 deletions tests/tests_data/service/milvus_rss_data.json
View file
Open in desktop
Git LFS file not shown
Loading