Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

STYLE: Replace thread_local variables with m_TensorStoreData data member #71

Merged
merged 1 commit into from
Sep 3, 2024
Merged

STYLE: Replace thread_local variables with m_TensorStoreData data member #71

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
6 changes: 5 additions & 1 deletion include/itkOMEZarrNGFFImageIO.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@


#include <fstream>
#include <memory> // For unique_ptr.
#include <string>
#include <vector>
#include "itkImageIOBase.h"
Expand Down Expand Up @@ -182,7 +183,7 @@ class IOOMEZarrNGFF_EXPORT OMEZarrNGFFImageIO : public ImageIOBase

protected:
OMEZarrNGFFImageIO();
~OMEZarrNGFFImageIO() override = default;
~OMEZarrNGFFImageIO() override;

void
PrintSelf(std::ostream & os, Indent indent) const override;
Expand Down Expand Up @@ -246,6 +247,9 @@ class IOOMEZarrNGFF_EXPORT OMEZarrNGFFImageIO : public ImageIOBase
char m_EmptyZip[m_EmptyZipSize] = "PK\x05\x06"; // the rest is filled with zeroes
const BufferInfo m_EmptyZipBufferInfo{ m_EmptyZip, m_EmptyZipSize };
const std::string m_EmptyZipFileName = MakeMemoryFileName(m_EmptyZipBufferInfo);

struct TensorStoreData;
const std::unique_ptr<TensorStoreData> m_TensorStoreData;
};
} // end namespace itk

Expand Down
57 changes: 32 additions & 25 deletions src/itkOMEZarrNGFFImageIO.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -374,9 +374,14 @@ MakeKVStoreHTTPDriverSpec(nlohmann::json & spec, const std::string & fullPath)

} // namespace

thread_local tensorstore::Context tsContext = tensorstore::Context::Default();
struct OMEZarrNGFFImageIO::TensorStoreData
{
tensorstore::Context tsContext{ tensorstore::Context::Default() };
tensorstore::TensorStore<> store{};
};

OMEZarrNGFFImageIO::OMEZarrNGFFImageIO()
: m_TensorStoreData(std::make_unique<TensorStoreData>())
{
this->AddSupportedWriteExtension(".zarr");
this->AddSupportedWriteExtension(".zr2");
Expand All @@ -395,6 +400,8 @@ OMEZarrNGFFImageIO::OMEZarrNGFFImageIO()
this->Self::SetCompressionLevel(2);
}

OMEZarrNGFFImageIO::~OMEZarrNGFFImageIO() = default;


void
OMEZarrNGFFImageIO::PrintSelf(std::ostream & os, Indent indent) const
Expand All @@ -407,7 +414,7 @@ OMEZarrNGFFImageIO::PrintSelf(std::ostream & os, Indent indent) const

// JSON file path, e.g. "C:/Dev/ITKIOOMEZarrNGFF/v0.4/cyx.ome.zarr/.zgroup"
void
writeJson(nlohmann::json json, std::string path, std::string driver)
writeJson(nlohmann::json json, std::string path, std::string driver, tensorstore::Context& tsContext)
{
auto attrs_store = tensorstore::Open<nlohmann::json, 0>(
{ { "driver", "json" }, { "kvstore", { { "driver", driver }, { "path", path } } } },
Expand All @@ -428,7 +435,7 @@ writeJson(nlohmann::json json, std::string path, std::string driver)

// JSON file path, e.g. "C:/Dev/ITKIOOMEZarrNGFF/v0.4/cyx.ome.zarr/.zattrs"
bool
jsonRead(const std::string path, nlohmann::json & result, std::string driver)
jsonRead(const std::string path, nlohmann::json & result, std::string driver, tensorstore::Context& tsContext)
{
// Reading JSON via TensorStore allows it to be in the cloud
nlohmann::json readSpec = { { "driver", "json" }, { "kvstore", { { "driver", driver }, { "path", path } } } };
Expand Down Expand Up @@ -466,15 +473,15 @@ OMEZarrNGFFImageIO::CanReadFile(const char * filename)
{
std::string driver = getKVstoreDriver(filename);
nlohmann::json json;
if (!jsonRead(std::string(filename) + "/.zgroup", json, driver))
if (!jsonRead(std::string(filename) + "/.zgroup", json, driver, m_TensorStoreData->tsContext))
{
return false;
}
if (json.at("zarr_format").get<int>() != 2)
{
return false; // unsupported zarr format
}
if (!jsonRead(std::string(filename) + "/.zattrs", json, driver))
if (!jsonRead(std::string(filename) + "/.zattrs", json, driver, m_TensorStoreData->tsContext))
{
return false;
}
Expand All @@ -491,8 +498,6 @@ OMEZarrNGFFImageIO::CanReadFile(const char * filename)
// return this->HasSupportedWriteExtension(filename, true);
}

thread_local tensorstore::TensorStore<> store; // initialized by ReadImageInformation/ReadArrayMetadata

void
OMEZarrNGFFImageIO::ReadArrayMetadata(std::string path, std::string driver)
{
Expand All @@ -503,15 +508,15 @@ OMEZarrNGFFImageIO::ReadArrayMetadata(std::string path, std::string driver)
}

auto openFuture = tensorstore::Open(readSpec,
tsContext,
m_TensorStoreData->tsContext,
tensorstore::OpenMode::open,
tensorstore::RecheckCached{ false },
tensorstore::ReadWriteMode::read);
TS_EVAL_CHECK(openFuture);
store = openFuture.value();
auto shape_span = store.domain().shape();
m_TensorStoreData->store = openFuture.value();
auto shape_span = m_TensorStoreData->store.domain().shape();

tensorstore::DataType dtype = store.dtype();
tensorstore::DataType dtype = m_TensorStoreData->store.dtype();
this->SetComponentType(tensorstoreToITKComponentType(dtype));

std::vector<int64_t> dims(shape_span.rbegin(), shape_span.rend()); // convert KJI into IJK
Expand All @@ -534,13 +539,15 @@ OMEZarrNGFFImageIO::ReadArrayMetadata(std::string path, std::string driver)
ImageIORegion
OMEZarrNGFFImageIO::ConfigureTensorstoreIORegion(const ImageIORegion & ioRegion) const
{
const auto storeRank = m_TensorStoreData->store.rank();

// Set up IO region to match known store dimensions
itkAssertOrThrowMacro(m_StoreAxes.size() == store.rank(), "Detected mismatch in axis count and store rank");
ImageIORegion storeRegion(store.rank());
itkAssertOrThrowMacro(storeRegion.GetImageDimension(), store.rank());
itkAssertOrThrowMacro(m_StoreAxes.size() == storeRank, "Detected mismatch in axis count and store rank");
ImageIORegion storeRegion(storeRank);
itkAssertOrThrowMacro(storeRegion.GetImageDimension(), storeRank);
auto storeAxes = this->GetAxesInStoreOrder();

for (size_t storeIndex = 0; storeIndex < store.rank(); ++storeIndex)
for (size_t storeIndex = 0; storeIndex < storeRank; ++storeIndex)
{
auto axisName = storeAxes.at(storeIndex).name;

Expand Down Expand Up @@ -651,12 +658,12 @@ OMEZarrNGFFImageIO::ReadImageInformation()
std::string driver = getKVstoreDriver(this->GetFileName());

const std::string zgroupFilePath(std::string(this->GetFileName()) + "/.zgroup");
bool status = jsonRead(zgroupFilePath, json, driver);
bool status = jsonRead(zgroupFilePath, json, driver, m_TensorStoreData->tsContext);
itkAssertOrThrowMacro(status, ("Failed to read from " + zgroupFilePath));
itkAssertOrThrowMacro(json.at("zarr_format").get<int>() == 2, "Only v2 zarr format is supported"); // only v2 for now

const std::string zattrsFilePath(std::string(this->GetFileName()) + "/.zattrs");
status = jsonRead(zattrsFilePath, json, driver);
status = jsonRead(zattrsFilePath, json, driver, m_TensorStoreData->tsContext);
itkAssertOrThrowMacro(status, ("Failed to read from " + zattrsFilePath));
json = json.at("multiscales")[0]; // multiscales must be present in OME-NGFF
auto version = json.at("version").get<std::string>();
Expand Down Expand Up @@ -732,7 +739,7 @@ OMEZarrNGFFImageIO::Read(void * buffer)
// This comparison needs to be done carefully, we can compare 3D and 6D regions
if (this->GetLargestRegion().GetNumberOfPixels() == m_IORegion.GetNumberOfPixels())
{
itkAssertOrThrowMacro(store.domain().num_elements() == m_IORegion.GetNumberOfPixels(),
itkAssertOrThrowMacro(m_TensorStoreData->store.domain().num_elements() == m_IORegion.GetNumberOfPixels(),
"Detected mismatch between store size and size of largest possible region");
}
else
Expand All @@ -750,7 +757,7 @@ OMEZarrNGFFImageIO::Read(void * buffer)
}

if (const IOComponentEnum componentType{ this->GetComponentType() };
!TryToReadFromStore(supportedPixelTypes, componentType, store, storeIORegion, buffer))
!TryToReadFromStore(supportedPixelTypes, componentType, m_TensorStoreData->store, storeIORegion, buffer))
{
itkExceptionMacro("Unsupported component type: " << GetComponentTypeAsString(componentType));
}
Expand All @@ -771,7 +778,7 @@ OMEZarrNGFFImageIO::WriteImageInformation()

nlohmann::json group;
group["zarr_format"] = 2;
writeJson(group, std::string(this->GetFileName()) + "/.zgroup", driver);
writeJson(group, std::string(this->GetFileName()) + "/.zgroup", driver, m_TensorStoreData->tsContext);

unsigned dim = this->GetNumberOfDimensions();

Expand Down Expand Up @@ -803,7 +810,7 @@ OMEZarrNGFFImageIO::WriteImageInformation()

nlohmann::json zattrs;
zattrs["multiscales"] = multiscales;
writeJson(zattrs, std::string(this->GetFileName()) + "/.zattrs", driver);
writeJson(zattrs, std::string(this->GetFileName()) + "/.zattrs", driver, m_TensorStoreData->tsContext);
}


Expand All @@ -812,7 +819,7 @@ OMEZarrNGFFImageIO::Write(const void * buffer)
{
if (m_FileName.substr(m_FileName.size() - 4) == ".zip" || m_FileName.substr(m_FileName.size() - 7) == ".memory")
{
tsContext = tensorstore::Context::Default(); // start with clean zip handles
m_TensorStoreData->tsContext = tensorstore::Context::Default(); // start with clean zip handles
}
this->WriteImageInformation();

Expand All @@ -837,8 +844,8 @@ OMEZarrNGFFImageIO::Write(const void * buffer)

if (!TryToWriteToStore(supportedPixelTypes,
componentType,
store,
tsContext,
m_TensorStoreData->store,
m_TensorStoreData->tsContext,
m_FileName,
MakePath(this->GetDatasetIndex()),
shape,
Expand All @@ -851,7 +858,7 @@ OMEZarrNGFFImageIO::Write(const void * buffer)
{
// Attempt to read a non-existent file from the in-memory zip to close the current one
nlohmann::json temp;
bool wasRead = jsonRead(m_EmptyZipFileName + "/non-existent.json", temp, "zip_memory");
bool wasRead = jsonRead(m_EmptyZipFileName + "/non-existent.json", temp, "zip_memory", m_TensorStoreData->tsContext);
assert(wasRead == false);
}
}
Expand Down
Loading