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 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}"