[Python] Snowflake SPCS Multi Model Training - 403 Errors After some time

I am attempting to run a long running training job leveraging snowflakes distributed multi model training functionality. Using a sample of my training data (~100k rows) this process completes with no issues. However when expanding to the entire training set (~150M rows) I start to run into 403 errors across the later processed partitions in my dataset. These errors occur around the 6hr mark of training, partitions that were trained earlier have no issues.

The following exceptions are thrown which are the same for every failed partition

An exception was raised from a task of operator "ReadResultSetDataSource->SplitBlocks(2)". Dataset execution will now abort. To ignore this exception and continue, set DataContext.max_errored_blocks.
Traceback (most recent call last):
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/data/_internal/execution/streaming_executor_state.py", line 472, in process_completed_tasks
bytes_read = task.on_data_ready(
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/data/_internal/execution/interfaces/physical_operator.py", line 138, in on_data_ready
raise ex from None
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/data/_internal/execution/interfaces/physical_operator.py", line 134, in on_data_ready
ray.get(block_ref)
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
return fn(*args, **kwargs)
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
return func(*args, **kwargs)
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/_private/worker.py", line 2849, in get
values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/_private/worker.py", line 937, in get_objects
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(RuntimeError): [36mray::ReadResultSetDataSource->SplitBlocks(2)()[39m (pid=58742, ip=**.***.**.**)
snowflake.connector.network.RetryRequest: 290403: HTTP 403: Forbidden


[36mray::ReadResultSetDataSource->SplitBlocks(2)()[39m (pid=58742, ip=**.***.**.**)
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/data/_internal/execution/operators/map_operator.py", line 552, in _map_task
for b_out in map_transformer.apply_transform(iter(blocks), ctx):
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/data/_internal/execution/operators/map_transformer.py", line 601, in __call__
for block in blocks:
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/data/_internal/execution/operators/map_transformer.py", line 532, in __call__
for data in iter:
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/data/_internal/execution/operators/map_transformer.py", line 377, in __call__
yield from self._block_fn(input, ctx)
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/data/_internal/planner/plan_read_op.py", line 106, in do_read
yield from read_task()
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/ray/data/datasource/datasource.py", line 168, in __call__
yield from result
File "/opt/python/cpython-3.10-linux-x86_64-gnu/lib/python3.10/site-packages/implementations/ray_data_ingester.py", line 255, in _read_fn
raise RuntimeError(
RuntimeError: Unable to refresh batches: invalid query_id


I have tried setting the session keep alive to true

SESSION.sql("ALTER SESSION SET CLIENT_SESSION_KEEP_ALIVE = TRUE").collect()


as well as passing a reference to a table and dataframe in memory when training with no luck.

trainer = ManyModelTraining(
train_func=train_model,
stage_name=MMT_STAGE,
serde=ModelOnlySerde(),
)

run = trainer.run(
partition_by="PARTITION_ID",
#snowpark_dataframe=SESSION.table(<table_ref using (<Select statement).collect())
snowpark_dataframe=SESSION.table(<table with training data>),
run_id=RUN_ID,
on_existing_artifacts="overwrite",
)

run.wait()
print("MMT training complete. Run ID:", RUN_ID)
scale_cluster(1)
SESSION.close()


Any ideas?

UPDATE:

Have been working w/ Snowflake support to no success. The initial issue was due to cached queries expiring after 6hrs which is why the later partitions failed. I've tried the suggestion to persist training data as files in a stage which bypasses the cached query limit. Again, I'm seeing no issues when training on a sample of data but the process fails with no error logs and no other indication of failure after some time anywhere between 4 and 8 hours. I've even taken the full data of certain partitions that were failing as a sample and these have completed with no issues when training in isolation.

Continue reading...