Repository: PrefectHQ/prefect

Issue:
Flow S3ResultHandler Fails for Dask Worker with nthreads > 1
## Description
Specifying S3ResultHandler for a Flow running on Dask worker(s) with nthreads > 1 fails with: `KeyError: 'credential_provider'`, likely due to a race condition in using the global boto3 session (boto3.client) between threads. 

## Expected Behavior
In a multithreaded environment, boto3 recommends creating a session per thread rather than sharing the default boto3 session, i.e. boto3.client. See boto3 documentation at: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/resources.html?highlight=multithreading#multithreading-multiprocessing

This thread in Prefect's Community Slack describes using this session-per-thread approach to successfully fix a similar issue when using boto3 in Prefect tasks: https://prefect-community.slack.com/archives/CM28LL405/p1581434710167100

## Reproduction
A Flow with tasks that can run in parallel (e.g. mapped tasks or different Flow branches) and where the Flow-level result_handler is set to S3ResultHandler should reproduce this behavior.

Full stack trace:
```
February 29th 2020 at 8:09:43am | prefect.CloudTaskRunner
ERROR 
Failed to set task state with error: KeyError('credential_provider')
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/prefect/engine/cloud/task_runner.py", line 117, in call_runner_target_handlers
    cloud_state = prepare_state_for_cloud(new_state)
  File "/opt/conda/lib/python3.7/site-packages/prefect/engine/cloud/utilities.py", line 21, in prepare_state_for_cloud
    res.store_safe_value()
  File "/opt/conda/lib/python3.7/site-packages/prefect/engine/result.py", line 93, in store_safe_value
    value = self.result_handler.write(self.value)
  File "/opt/conda/lib/python3.7/site-packages/prefect/engine/result_handlers/s3_result_handler.py", line 103, in write
    self.client.upload_fileobj(stream, Bucket=self.bucket, Key=uri)
  File "/opt/conda/lib/python3.7/site-packages/prefect/engine/result_handlers/s3_result_handler.py", line 67, in client
    self.initialize_client()
  File "/opt/conda/lib/python3.7/site-packages/prefect/engine/result_handlers/s3_result_handler.py", line 60, in initialize_client
    aws_secret_access_key=aws_secret_access_key,
  File "/opt/conda/lib/python3.7/site-packages/boto3/__init__.py", line 91, in client
    return _get_default_session().client(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/boto3/session.py", line 263, in client
    aws_session_token=aws_session_token, config=config)
  File "/opt/conda/lib/python3.7/site-packages/botocore/session.py", line 823, in create_client
    credentials = self.get_credentials()
  File "/opt/conda/lib/python3.7/site-packages/botocore/session.py", line 428, in get_credentials
    'credential_provider').load_credentials()
  File "/opt/conda/lib/python3.7/site-packages/botocore/session.py", line 923, in get_component
    del self._deferred[name]
KeyError: 'credential_provider'
```

## Environment
We create a long-running Dask cluster where our Dask workers are started with --nprocs 1 --nthreads 3.

(Thanks to @JLouSRM for identifying this issue and capturing log evidence!)



Discussion:
Very interesting - thanks for the issue!  For reference, [this is the boto3 documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/resources.html?highlight=multithreading#multithreading-multiprocessing) referenced in that Slack thread, and Adam said:

> By first creating a session and then creating the client from the session, each thread has a different session.

As an aside, this could also motivate introducing a new [S3FS](https://github.com/dask/s3fs) Result Handler (similar to the discussion in https://github.com/PrefectHQ/prefect/issues/1475).

Resolving patch:
diff --git a/src/prefect/engine/result_handlers/s3_result_handler.py b/src/prefect/engine/result_handlers/s3_result_handler.py
--- a/src/prefect/engine/result_handlers/s3_result_handler.py
+++ b/src/prefect/engine/result_handlers/s3_result_handler.py
@@ -7,6 +7,7 @@
 import cloudpickle
 import pendulum
 
+import prefect
 from prefect.client import Secret
 from prefect.engine.result_handlers import ResultHandler
 
@@ -54,7 +55,10 @@ def initialize_client(self) -> None:
             aws_access_key = aws_credentials["ACCESS_KEY"]
             aws_secret_access_key = aws_credentials["SECRET_ACCESS_KEY"]
 
-        s3_client = boto3.client(
+        # use a new boto session when initializing in case we are in a new thread
+        # see https://boto3.amazonaws.com/v1/documentation/api/latest/guide/resources.html?#multithreading-multiprocessing
+        session = boto3.session.Session()
+        s3_client = session.client(
             "s3",
             aws_access_key_id=aws_access_key,
             aws_secret_access_key=aws_secret_access_key,
@@ -63,8 +67,13 @@ def initialize_client(self) -> None:
 
     @property
     def client(self) -> "boto3.client":
-        if not hasattr(self, "_client"):
+        """
+        Initializes a client if we believe we are in a new thread.
+        We consider ourselves in a new thread if we haven't stored a client yet in the current context.
+        """
+        if not prefect.context.get("boto3client"):
             self.initialize_client()
+            prefect.context["boto3client"] = self._client
         return self._client
 
     @client.setter


Test patch:
diff --git a/tests/engine/result_handlers/test_result_handlers.py b/tests/engine/result_handlers/test_result_handlers.py
--- a/tests/engine/result_handlers/test_result_handlers.py
+++ b/tests/engine/result_handlers/test_result_handlers.py
@@ -166,32 +166,31 @@ def __getstate__(self):
 @pytest.mark.xfail(raises=ImportError, reason="aws extras not installed.")
 class TestS3ResultHandler:
     @pytest.fixture
-    def s3_client(self, monkeypatch):
+    def session(self, monkeypatch):
         import boto3
 
-        client = MagicMock()
-        with patch.dict("sys.modules", {"boto3": MagicMock(client=client)}):
-            yield client
+        session = MagicMock()
+        with patch.dict("sys.modules", {"boto3": MagicMock(session=session)}):
+            yield session
 
-    def test_s3_client_init_uses_secrets(self, s3_client):
+    def test_s3_client_init_uses_secrets(self, session):
         handler = S3ResultHandler(
             bucket="bob", aws_credentials_secret="AWS_CREDENTIALS"
         )
         assert handler.bucket == "bob"
-        assert s3_client.called is False
+        assert session.Session().client.called is False
 
         with prefect.context(
             secrets=dict(AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42))
         ):
             with set_temporary_config({"cloud.use_local_secrets": True}):
                 handler.initialize_client()
-
-        assert s3_client.call_args[1] == {
+        assert session.Session().client.call_args[1] == {
             "aws_access_key_id": 1,
             "aws_secret_access_key": 42,
         }
 
-    def test_s3_client_init_uses_custom_secrets(self, s3_client):
+    def test_s3_client_init_uses_custom_secrets(self, session):
         handler = S3ResultHandler(bucket="bob", aws_credentials_secret="MY_FOO")
 
         with prefect.context(
@@ -201,12 +200,12 @@ def test_s3_client_init_uses_custom_secrets(self, s3_client):
                 handler.initialize_client()
 
         assert handler.bucket == "bob"
-        assert s3_client.call_args[1] == {
+        assert session.Session().client.call_args[1] == {
             "aws_access_key_id": 1,
             "aws_secret_access_key": 999,
         }
 
-    def test_s3_writes_to_blob_prefixed_by_date_suffixed_by_prefect(self, s3_client):
+    def test_s3_writes_to_blob_prefixed_by_date_suffixed_by_prefect(self, session):
         handler = S3ResultHandler(bucket="foo")
 
         with prefect.context(
@@ -215,7 +214,9 @@ def test_s3_writes_to_blob_prefixed_by_date_suffixed_by_prefect(self, s3_client)
             with set_temporary_config({"cloud.use_local_secrets": True}):
                 uri = handler.write("so-much-data")
 
-        used_uri = s3_client.return_value.upload_fileobj.call_args[1]["Key"]
+        used_uri = session.Session().client.return_value.upload_fileobj.call_args[1][
+            "Key"
+        ]
 
         assert used_uri == uri
         assert used_uri.startswith(pendulum.now("utc").format("Y/M/D"))
@@ -229,7 +230,11 @@ def __init__(self, *args, **kwargs):
             def __getstate__(self):
                 raise ValueError("I cannot be pickled.")
 
-        with patch.dict("sys.modules", {"boto3": MagicMock(client=client)}):
+        import boto3
+
+        with patch.dict("sys.modules", {"boto3": MagicMock()}):
+            boto3.session.Session().client = client
+
             with prefect.context(
                 secrets=dict(AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42))
             ):