From 51217f3555c2a7eb5313322dff6d601bb3b5da6a Mon Sep 17 00:00:00 2001 From: Jean-Pierre Welch Date: Thu, 8 Jan 2026 13:13:30 -0500 Subject: [PATCH 1/2] feat: add BUCKET_PREFIX support for S3 uploads - Add BUCKET_PREFIX environment variable to organize uploaded files under a custom prefix/folder path in S3 - When BUCKET_PREFIX is set, uses upload_file_to_bucket with prefix (preserves original filename) - When BUCKET_PREFIX is not set, uses upload_image for backward compatibility (random UUID filename) - Leading/trailing slashes are automatically stripped from BUCKET_PREFIX - Update configuration documentation with the new environment variable --- docs/configuration.md | 5 +++-- handler.py | 12 +++++++++++- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 25bed9e70..44c972f0b 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -39,8 +39,9 @@ Configure these variables **only** if you want the worker to upload generated im | `BUCKET_ENDPOINT_URL` | The full endpoint URL of your S3 bucket. **Must be set to enable S3 upload.** | `https://.s3..amazonaws.com` | | `BUCKET_ACCESS_KEY_ID` | Your AWS access key ID associated with the IAM user that has write permissions to the bucket. Required if `BUCKET_ENDPOINT_URL` is set. | `AKIAIOSFODNN7EXAMPLE` | | `BUCKET_SECRET_ACCESS_KEY` | Your AWS secret access key associated with the IAM user. Required if `BUCKET_ENDPOINT_URL` is set. | `wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY` | +| `BUCKET_PREFIX` | Optional prefix (folder path) to prepend to uploaded files in the bucket. Leading/trailing slashes are automatically stripped. | `outputs/comfyui` | -**Note:** Upload uses the `runpod` Python library helper `rp_upload.upload_image`, which handles creating a unique path within the bucket based on the `job_id`. +**Note:** When `BUCKET_PREFIX` is not set, upload uses `rp_upload.upload_image` for backward compatibility (files are stored as `{job_id}/{random_uuid}.{ext}`). When `BUCKET_PREFIX` is set, upload uses `rp_upload.upload_file_to_bucket` which preserves the original filename (files are stored as `{prefix}/{job_id}/{original_filename}`). ### Example S3 Response @@ -66,4 +67,4 @@ If the S3 environment variables (`BUCKET_ENDPOINT_URL`, `BUCKET_ACCESS_KEY_ID`, } ``` -The `data` field contains the presigned URL to the uploaded image file in your S3 bucket. The path usually includes the job ID. +The `data` field contains the presigned URL to the uploaded image file in your S3 bucket. The path includes the job ID, and optionally the `BUCKET_PREFIX` if configured. For example, with `BUCKET_PREFIX=outputs/comfyui`, the path would be `outputs/comfyui/sync-uuid-string/ComfyUI_00001_.png`. diff --git a/handler.py b/handler.py index 65c8390b0..5276f94fa 100644 --- a/handler.py +++ b/handler.py @@ -724,7 +724,17 @@ def handler(job): ) print(f"worker-comfyui - Uploading {filename} to S3...") - s3_url = rp_upload.upload_image(job_id, temp_file_path) + # Use BUCKET_PREFIX if set, otherwise use original upload_image for backward compatibility + bucket_prefix = os.environ.get("BUCKET_PREFIX", "").strip("/") + if bucket_prefix: + s3_prefix = f"{bucket_prefix}/{job_id}" + s3_url = rp_upload.upload_file_to_bucket( + file_name=filename, + file_location=temp_file_path, + prefix=s3_prefix, + ) + else: + s3_url = rp_upload.upload_image(job_id, temp_file_path) os.remove(temp_file_path) # Clean up temp file print( f"worker-comfyui - Uploaded {filename} to S3: {s3_url}" From b5e26d8a4f881300089fbd0042d9c2b0cc35623e Mon Sep 17 00:00:00 2001 From: Jean-Pierre Welch Date: Mon, 15 Jun 2026 19:24:34 -0400 Subject: [PATCH 2/2] chore: add changeset for bucket prefix --- .changeset/feat-add-bucket-prefix.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/feat-add-bucket-prefix.md diff --git a/.changeset/feat-add-bucket-prefix.md b/.changeset/feat-add-bucket-prefix.md new file mode 100644 index 000000000..03b00be1a --- /dev/null +++ b/.changeset/feat-add-bucket-prefix.md @@ -0,0 +1,5 @@ +--- +"worker-comfyui": minor +--- + +feat: add `BUCKET_PREFIX` env var to organize S3 uploaded files under a custom prefix path. When set, uses `rp_upload.upload_file_to_bucket` and the path becomes `{prefix}/{job_id}/{original_filename}`. When unset, behavior is unchanged. \ No newline at end of file