Replace S3 subprocess/multiprocessing with direct boto3 + ThreadPoolExecutor

.github/workflows/metaflow.s3_tests.minio.yml

@@ -18,10 +18,23 @@ jobs:
     name: metaflow.s3.minio / Python ${{ matrix.ver }} on ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu-22.04]
-        ver: ['3.8', '3.9', '3.10', '3.11', '3.12']
-
+        ver: ['3.11']
+
+    timeout-minutes: 90
+
+    env:
+      AWS_ACCESS_KEY_ID: rootuser
+      AWS_SECRET_ACCESS_KEY: rootpass123
+      AWS_DEFAULT_REGION: us-east-1
+      METAFLOW_S3_TEST_ROOT: s3://metaflow-test/metaflow/
+      METAFLOW_DATASTORE_SYSROOT_S3: s3://metaflow-test/metaflow/
+      AWS_ENDPOINT_URL_S3: http://localhost:9000
+      MINIO_TEST: "1"
+      METAFLOW_S3_TRANSIENT_RETRY_COUNT: "7"
+
     steps:
     - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8  # v6.0.1
       with:
@@ -31,32 +44,40 @@ jobs:
       uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6.2.0
       with:
         python-version: ${{ matrix.ver }}
-    - name: Install Python ${{ matrix.ver }} dependencies
+    - name: Install dependencies
       run: |
         python3 -m pip install --upgrade pip
-        python3 -m pip install . kubernetes tox numpy pytest click boto3 requests pylint pytest-benchmark
-    - name: Start MinIO development environment
+        python3 -m pip install . pytest click boto3 requests numpy pytest-benchmark
+    - name: Start MinIO
       run: |
-          echo "Starting environment in the background..."
-          MINIKUBE_CPUS=2 metaflow-dev all-up &
-          # Give time to spin up. Adjust as needed:
-          sleep 150
+        docker run -d --name minio \
+          -p 9000:9000 \
+          -e MINIO_ROOT_USER=rootuser \
+          -e MINIO_ROOT_PASSWORD=rootpass123 \
+          minio/minio server /data
+        for i in $(seq 1 30); do
+          if curl -sf http://localhost:9000/minio/health/live; then
+            echo "MinIO is ready"
+            break
+          fi
+          echo "Waiting for MinIO... ($i/30)"
+          sleep 2
+        done
+    - name: Create test bucket
+      run: |
+        python3 -c "
+        import boto3
+        s3 = boto3.client('s3',
+            endpoint_url='http://localhost:9000',
+            aws_access_key_id='rootuser',
+            aws_secret_access_key='rootpass123')
+        s3.create_bucket(Bucket='metaflow-test')
+        print('Bucket metaflow-test created')
+        "
     - name: Execute tests
       run: |
-        cat <<EOF | metaflow-dev shell
-        # Set MinIO environment variables
-        export AWS_ACCESS_KEY_ID=rootuser
-        export AWS_SECRET_ACCESS_KEY=rootpass123
-        export AWS_DEFAULT_REGION=us-east-1
-        export METAFLOW_S3_TEST_ROOT=s3://metaflow-test/metaflow/
-        export METAFLOW_DATASTORE_SYSROOT_S3=s3://metaflow-test/metaflow/
-        export AWS_ENDPOINT_URL_S3=http://localhost:9000
-        export MINIO_TEST=1
-
-        # Run the same test command as the original workflow
         cd test/data
-        PYTHONPATH=\$(pwd)/../../ python3 -m pytest --benchmark-skip -s -v
-        EOF
-    - name: Tear down environment
-      run: |
-        metaflow-dev down
+        PYTHONPATH=$(pwd)/../../ python3 -m pytest --benchmark-skip -s -v --tb=long
+    - name: Stop MinIO
+      if: always()
+      run: docker rm -f minio || true

metaflow/metaflow_config.py

@@ -158,6 +158,8 @@
     "S3_CLIENT_RETRY_CONFIG", {"max_attempts": 10, "mode": "adaptive"}
 )
 
+S3_DIRECT_BOTO3 = from_conf("S3_DIRECT_BOTO3", True)
+
 # Threshold to start printing warnings for an AWS retry
 RETRY_WARNING_THRESHOLD = 3
 

metaflow/plugins/datatools/s3/s3.py

@@ -16,6 +16,7 @@
 from metaflow.metaflow_current import current
 from metaflow.metaflow_config import (
     DATATOOLS_S3ROOT,
+    S3_DIRECT_BOTO3,
     S3_RETRY_COUNT,
     S3_TRANSIENT_RETRY_COUNT,
     S3_LOG_TRANSIENT_RETRIES,
@@ -1449,6 +1450,14 @@ def _jitter_sleep(
     # and url_unquote.
     def _read_many_files(self, op, prefixes_and_ranges, **options):
         prefixes_and_ranges = list(prefixes_and_ranges)
+        if S3_DIRECT_BOTO3:
+            from .s3op_boto import S3DirectClient
+
+            direct = S3DirectClient(
+                self._s3_client, self._tmpdir, self._s3_inject_failures
+            )
+            yield from direct.read_many(op, prefixes_and_ranges, **options)
+            return
         with NamedTemporaryFile(
             dir=self._tmpdir,
             mode="wb",
@@ -1489,6 +1498,13 @@ def _read_many_files(self, op, prefixes_and_ranges, **options):
 
     def _put_many_files(self, url_info, overwrite):
         url_info = list(url_info)
+        if S3_DIRECT_BOTO3:
+            from .s3op_boto import S3DirectClient
+
+            direct = S3DirectClient(
+                self._s3_client, self._tmpdir, self._s3_inject_failures
+            )
+            return direct.put_objects(url_info, overwrite)
         url_dicts = [
             dict(
                 chain([("local", os.path.realpath(local)), ("url", url)], info.items())