Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 46 additions & 22 deletions cwms/timeseries/timeseries.py
Original file line number Diff line number Diff line change
Expand Up @@ -528,14 +528,20 @@ def store_multi_timeseries_df(
DELETE_INSERT.
override_protection: bool, optional, default is False
A flag to ignore the protected data quality flag when storing data.
multithread: bool, default is false
multithread: bool, default is true
Specifies whether to store chunked time series values using multiple threads.
max_workers: Int, Optional, default is None
It is a number of Threads aka size of pool in concurrent.futures.ThreadPoolExecutor.

Returns
-------
None

Raises
------
RuntimeError
If any series fails to store. The message identifies failed series;
other series may already have been stored successfully.
"""

def store_ts_ids(
Expand All @@ -544,24 +550,21 @@ def store_ts_ids(
office_id: str,
version_date: Optional[datetime] = None,
) -> None:
try:
units = data["units"].iloc[0]
data_json = timeseries_df_to_json(
data=data,
ts_id=ts_id,
units=units,
office_id=office_id,
version_date=version_date,
)
store_timeseries(
data=data_json,
create_as_ltrs=create_as_ltrs,
store_rule=store_rule,
override_protection=override_protection,
multithread=multithread,
)
except Exception as e:
print(f"Error processing {ts_id}: {e}")
units = data["units"].iloc[0]
data_json = timeseries_df_to_json(
data=data,
ts_id=ts_id,
units=units,
office_id=office_id,
version_date=version_date,
)
store_timeseries(
data=data_json,
create_as_ltrs=create_as_ltrs,
store_rule=store_rule,
override_protection=override_protection,
multithread=multithread,
)
return None

required_columns = ["date-time", "value", "ts_id", "units"]
Expand All @@ -577,7 +580,9 @@ def store_ts_ids(
ts_data_all["ts_id"].astype(str) + ":" + ts_data_all["version_date"].astype(str)
).unique()

errors: List[str] = []
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = {}
for unique_tsid in unique_tsids:
ts_id, version_date = unique_tsid.split(":", 1)
if version_date != "NaT":
Expand All @@ -592,9 +597,21 @@ def store_ts_ids(
(ts_data_all["ts_id"] == ts_id) & ts_data_all["version_date"].isna()
]
if not data.empty:
executor.submit(
future = executor.submit(
store_ts_ids, ts_data, ts_id, office_id, version_date_dt
)
futures[future] = unique_tsid

for future in concurrent.futures.as_completed(futures):
try:
future.result()
except Exception as e:
errors.append(f"{futures[future]}: {e}")

if errors:
raise RuntimeError(
f"{len(errors)} time series failed to store:\n" + "\n".join(errors)
)


def chunk_timeseries_data(
Expand Down Expand Up @@ -686,7 +703,14 @@ def store_timeseries(
if len(chunks) == 1 or not multithread:
return api.post(endpoint, data, params)

actual_workers = min(max_workers, len(chunks))
if max_workers <= 0:
raise ValueError("max_workers must be greater than 0")

# A new series must exist before multiple transactions can write its data.
# Complete one normal write first, then retain parallelism for the rest.
_call_with_retry(api.post, endpoint, chunks[0], params)
remaining_chunks = chunks[1:]
actual_workers = min(max_workers, len(remaining_chunks))
logging.debug(
f"Storing {len(chunks)} chunks of timeseries data with {actual_workers} threads"
)
Expand All @@ -698,7 +722,7 @@ def store_timeseries(
with concurrent.futures.ThreadPoolExecutor(max_workers=actual_workers) as executor:
future_to_chunk = {
executor.submit(_call_with_retry, api.post, endpoint, chunk, params): chunk
for chunk in chunks
for chunk in remaining_chunks
}

for future in concurrent.futures.as_completed(future_to_chunk):
Expand Down
90 changes: 90 additions & 0 deletions tests/mock/timeseries/test_concurrent_store.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
import threading

import pandas as pd
import pytest

import cwms.timeseries.timeseries as ts


def test_chunked_store_creates_series_before_parallel_writes(monkeypatch):
created = threading.Event()
lock = threading.Lock()
received = []
initializing = False
data = {
"name": "Test.Stage.Inst.15Minutes.0.Raw",
"office-id": "MVP",
"units": "ft",
"values": [[i, i, 0] for i in range(8)],
}

def post(endpoint, chunk, params):
nonlocal initializing
with lock:
first = not initializing
initializing = True
if first:
# Model the database transaction that creates a new series. Other
# writes cannot use its identifier until that transaction commits.
assert not created.wait(0.1)
created.set()
else:
assert created.is_set(), "concurrent write raced series creation"
assert endpoint == "timeseries"
assert params["store-rule"] == "REPLACE_ALL"
assert chunk["office-id"] == "MVP"
with lock:
received.extend(chunk["values"])

monkeypatch.setattr(ts.api, "post", post)
ts.store_timeseries(data, chunk_size=2, store_rule="REPLACE_ALL")
assert sorted(received) == data["values"]


def test_initial_chunk_failure_stops_remaining_writes(monkeypatch):
calls = []

def post(endpoint, chunk, params):
calls.append(chunk["values"])
raise ValueError("cannot create series")

monkeypatch.setattr(ts.api, "post", post)
with pytest.raises(ValueError, match="cannot create series"):
ts.store_timeseries(
{
"name": "Test.Stage.Inst.15Minutes.0.Raw",
"office-id": "MVP",
"units": "ft",
"values": [[i, i, 0] for i in range(6)],
},
chunk_size=2,
)
assert calls
assert all(chunk == [[0, 0, 0], [1, 1, 0]] for chunk in calls)


def test_multi_store_reports_all_failed_series(monkeypatch):
attempted = []
lock = threading.Lock()
data = pd.DataFrame(
{
"date-time": pd.to_datetime(["2025-01-01"] * 3, utc=True),
"value": [1, 2, 3],
"ts_id": ["good", "bad-one", "bad-two"],
"units": ["ft"] * 3,
}
)

def store(data, **kwargs):
with lock:
attempted.append(data["name"])
if data["name"].startswith("bad"):
raise ValueError("write rejected")

monkeypatch.setattr(ts, "store_timeseries", store)
with pytest.raises(RuntimeError) as error:
ts.store_multi_timeseries_df(data, "MVP")
assert "bad-one" in str(error.value)
assert "bad-two" in str(error.value)
assert "write rejected" in str(error.value)
assert sorted(attempted) == ["bad-one", "bad-two", "good"]
Loading