Skip to content

Commit

Permalink
fix(benchmark): Allow for specifying path based buckets
Browse files Browse the repository at this point in the history
  • Loading branch information
ADustyOldMuffin committed Aug 26, 2024
1 parent 6148adf commit 86fc1bf
Showing 1 changed file with 10 additions and 1 deletion.
11 changes: 10 additions & 1 deletion examples/benchmark_buffer_size/benchmark.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,11 +128,20 @@
help="The bucket where the models are located (default: tensorized)",
default="tensorized"
)
parser.add_argument(
"--path-based",
action="store_true",
default=False,
help="Use path based buckets"
)
args = parser.parse_args()

model_name: str = args.model

http_uri = f"http://{args.bucket}.{args.s3_endpoint}/{model_name}/model.tensors"
if args.path_based:
http_uri = f"http://{args.s3_endpoint}/{args.bucket}/{model_name}/model.tensors"
else:
http_uri = f"http://{args.bucket}.{args.s3_endpoint}/{model_name}/model.tensors"

https_uri = http_uri.replace("http://", "https://")
s3_uri = f"s3://{args.bucket}/{model_name}/model.tensors"
Expand Down

0 comments on commit 86fc1bf

Please sign in to comment.