replace blocking file I/O with async aiofiles for better perf (#4520)
Co-authored-by: Shuchang Zheng <wintonzheng0325@gmail.com>
This commit is contained in:
@@ -4,6 +4,7 @@ from datetime import datetime
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import BinaryIO
|
from typing import BinaryIO
|
||||||
|
|
||||||
|
import aiofiles
|
||||||
import structlog
|
import structlog
|
||||||
|
|
||||||
from skyvern.config import settings
|
from skyvern.config import settings
|
||||||
@@ -55,8 +56,9 @@ class LocalStorage(BaseStorage):
|
|||||||
if not file_path.exists():
|
if not file_path.exists():
|
||||||
return []
|
return []
|
||||||
try:
|
try:
|
||||||
with open(file_path) as f:
|
async with aiofiles.open(file_path, "r") as f:
|
||||||
return [line.strip() for line in f.readlines() if line.strip()]
|
lines = await f.readlines()
|
||||||
|
return [line.strip() for line in lines if line.strip()]
|
||||||
except Exception:
|
except Exception:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
@@ -122,8 +124,8 @@ class LocalStorage(BaseStorage):
|
|||||||
if WINDOWS:
|
if WINDOWS:
|
||||||
file_path = file_path.with_name(_windows_safe_filename(file_path.name))
|
file_path = file_path.with_name(_windows_safe_filename(file_path.name))
|
||||||
self._create_directories_if_not_exists(file_path)
|
self._create_directories_if_not_exists(file_path)
|
||||||
with open(file_path, "wb") as f:
|
async with aiofiles.open(file_path, "wb") as f:
|
||||||
f.write(data)
|
await f.write(data)
|
||||||
except Exception:
|
except Exception:
|
||||||
LOG.exception(
|
LOG.exception(
|
||||||
"Failed to store artifact locally.",
|
"Failed to store artifact locally.",
|
||||||
@@ -150,8 +152,8 @@ class LocalStorage(BaseStorage):
|
|||||||
file_path = None
|
file_path = None
|
||||||
try:
|
try:
|
||||||
file_path = parse_uri_to_path(artifact.uri)
|
file_path = parse_uri_to_path(artifact.uri)
|
||||||
with open(file_path, "rb") as f:
|
async with aiofiles.open(file_path, "rb") as f:
|
||||||
return f.read()
|
return await f.read()
|
||||||
except Exception:
|
except Exception:
|
||||||
LOG.exception(
|
LOG.exception(
|
||||||
"Failed to retrieve local artifact.",
|
"Failed to retrieve local artifact.",
|
||||||
@@ -174,8 +176,8 @@ class LocalStorage(BaseStorage):
|
|||||||
Path(f"{get_skyvern_temp_dir()}/{organization_id}").mkdir(parents=True, exist_ok=True)
|
Path(f"{get_skyvern_temp_dir()}/{organization_id}").mkdir(parents=True, exist_ok=True)
|
||||||
file_path = Path(f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}")
|
file_path = Path(f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}")
|
||||||
try:
|
try:
|
||||||
with open(file_path, "rb") as f:
|
async with aiofiles.open(file_path, "rb") as f:
|
||||||
return f.read()
|
return await f.read()
|
||||||
except Exception:
|
except Exception:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
@@ -481,8 +483,8 @@ class LocalStorage(BaseStorage):
|
|||||||
"""Download a user-uploaded file from local filesystem."""
|
"""Download a user-uploaded file from local filesystem."""
|
||||||
try:
|
try:
|
||||||
file_path = parse_uri_to_path(uri)
|
file_path = parse_uri_to_path(uri)
|
||||||
with open(file_path, "rb") as f:
|
async with aiofiles.open(file_path, "rb") as f:
|
||||||
return f.read()
|
return await f.read()
|
||||||
except Exception:
|
except Exception:
|
||||||
LOG.exception("Failed to read local file", uri=uri)
|
LOG.exception("Failed to read local file", uri=uri)
|
||||||
return None
|
return None
|
||||||
|
|||||||
Reference in New Issue
Block a user