diff --git a/maestro_worker_python/upload_files.py b/maestro_worker_python/upload_files.py index e329a95..b3cd623 100644 --- a/maestro_worker_python/upload_files.py +++ b/maestro_worker_python/upload_files.py @@ -1,6 +1,5 @@ from __future__ import annotations -import json import logging import tempfile import threading @@ -9,6 +8,7 @@ from os.path import join from typing import List, Any +import orjson import requests @@ -24,11 +24,10 @@ def upload_files(upload_files: List[UploadFile]): threads_upload = [] did_raise_exception = threading.Event() for file_ in upload_files: - t = threading.Thread(target=_upload, args=( - file_, did_raise_exception,)) + t = threading.Thread(target=_upload, args=(file_, did_raise_exception)) threads_upload.append(t) t.start() - + for t in threads_upload: t.join() @@ -40,7 +39,9 @@ def _upload(upload_file: UploadFile, did_raise_exception): logging.info(f"Uploading:{upload_file.file_path}") try: with open(upload_file.file_path, "rb") as data: - response = requests.put(upload_file.signed_url, data=data, headers={"Content-Type": upload_file.file_type}, timeout=300) + response = requests.put( + upload_file.signed_url, data=data, headers={"Content-Type": upload_file.file_type}, timeout=300 + ) response.raise_for_status() logging.info(f"Uploaded {upload_file.file_path}") except Exception as e: @@ -60,8 +61,8 @@ def upload_json_data(upload_data: list[UploadJsonData]): with tempfile.TemporaryDirectory() as tmp_dir: for i, upload in enumerate(upload_data): - with open(join(tmp_dir, f"{i}.json"), "w") as tmp_file: - json.dump(upload.data, tmp_file) + with open(join(tmp_dir, f"{i}.json"), "wb") as tmp_file: + tmp_file.write(orjson.dumps(upload.data, option=orjson.OPT_SERIALIZE_NUMPY)) tmp_file.flush() files_to_upload.append( UploadFile(file_path=tmp_file.name, signed_url=upload.signed_url, file_type="application/json") diff --git a/pyproject.toml b/pyproject.toml index dbcd5f0..4acd8dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "maestro-worker-python" -version = "3.5.1" +version = "3.5.2" description = "Utility to run workers on Moises/Maestro" authors = ["Moises.ai"] license = "MIT" @@ -15,6 +15,7 @@ json-logging = "^1.3.0" uvicorn = "^0.20" sentry-sdk = {extras = ["fastapi"], version = "^1.16.0"} requests = "2.31.0" +orjson = "3.10.5" [tool.poetry.group.dev.dependencies]