-
Notifications
You must be signed in to change notification settings - Fork 49
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add helper for thread local variables that can be used to add metadat…
…a to the output stream (#1052) * Add helper for thread local variables that can be used to add metadata to the output stream * Add devtools to help with debugging * Extend llm generate to calculate metadata * fix mocked test
- Loading branch information
1 parent
7e6b626
commit 1f05347
Showing
11 changed files
with
214 additions
and
13 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,11 @@ | ||
from sycamore.data import MetadataDocument | ||
from sycamore.utils.thread_local import ThreadLocalAccess, ADD_METADATA_TO_OUTPUT | ||
|
||
|
||
def add_metadata(**metadata): | ||
ThreadLocalAccess(ADD_METADATA_TO_OUTPUT).get().append(MetadataDocument(**metadata)) | ||
|
||
|
||
# At some point we should define particular forms of metadata like metrics | ||
# Maybe following https://github.com/prometheus/OpenMetrics/blob/main/specification/OpenMetrics.md | ||
# as a structure for the metrics -- too complex for now. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,43 @@ | ||
import logging | ||
import threading | ||
|
||
logger = logging.getLogger(__name__) | ||
# Used to inject metadata into the docset output stream to make it easier to use with | ||
# transforms like map. | ||
ADD_METADATA_TO_OUTPUT = "add_metadata_to_output" | ||
|
||
# Create a thread-local data object | ||
thread_local_data = threading.local() | ||
|
||
|
||
class ThreadLocalAccess: | ||
def __init__(self, var_name): | ||
self.var_name = var_name | ||
|
||
def present(self): | ||
return hasattr(thread_local_data, self.var_name) | ||
|
||
def get(self): | ||
assert hasattr(thread_local_data, self.var_name), f"{self.var_name} not present in TLS" | ||
return getattr(thread_local_data, self.var_name) | ||
|
||
def set(self, value): | ||
assert hasattr(thread_local_data, self.var_name), f"{self.var_name} not present in TLS" | ||
setattr(thread_local_data, self.var_name, value) | ||
|
||
|
||
class ThreadLocal(ThreadLocalAccess): | ||
def __init__(self, var_name, var_value): | ||
self.var_name = var_name | ||
self.var_value = var_value | ||
|
||
def __enter__(self): | ||
assert not hasattr(thread_local_data, self.var_name), f"{self.var_name} already set in TLS" | ||
setattr(thread_local_data, self.var_name, self.var_value) | ||
logger.debug(f"Thread-local variable '{self.var_name}' removed.") | ||
return self | ||
|
||
def __exit__(self, exc_type, exc_val, exc_tb): | ||
assert hasattr(thread_local_data, self.var_name), f"{self.var_name} vanished from TLS" | ||
delattr(thread_local_data, self.var_name) | ||
logger.debug(f"Thread-local variable '{self.var_name}' removed.") |
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters