Skip to content

Commit

Permalink
Opt-in to uploading profiles from why.log()
Browse files Browse the repository at this point in the history
  • Loading branch information
richard-rogers committed Sep 26, 2024
1 parent 4ad4cc9 commit 0a2e4be
Show file tree
Hide file tree
Showing 7 changed files with 66 additions and 4 deletions.
19 changes: 18 additions & 1 deletion python/whylogs/api/whylabs/session/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
prompt_default_dataset_id,
prompt_org_id,
prompt_session_type,
prompt_upload_on_log,
)
from whylogs.api.whylabs.session.session_types import ApiKeyV1, ApiKeyV2
from whylogs.api.whylabs.session.session_types import InteractiveLogger as il
Expand Down Expand Up @@ -66,10 +67,11 @@ class ConfigVariableName(Enum):
class InitConfig:
whylabs_api_key: Optional[str] = None
allow_anonymous: bool = True
allow_local: bool = False
allow_local: bool = True
default_dataset_id: Optional[str] = None
config_path: Optional[str] = None
force_local: Optional[bool] = None
upload_on_log: bool = False


class SessionConfig:
Expand All @@ -90,8 +92,10 @@ def __init__(self, init_config: Optional[InitConfig] = None) -> None:
if force_interactive:
self.reset_config()
self.session_type = self._determine_session_type_prompt(self._init_config)
self.upload_on_log = self._determine_upload_on_log_prompt()
else:
self.session_type = self._determine_session_type(self._init_config)
self.upload_on_log = self._determine_upload_on_log(self._init_config)

def _init_parser(self) -> None:
try:
Expand Down Expand Up @@ -406,6 +410,19 @@ def _determine_session_type(self, init_config: InitConfig) -> SessionType:
f"interactive environment. See {INIT_DOCS} for instructions on using why.init()."
)

def _determine_upload_on_log_prompt(self) -> bool:
return self.session_type != SessionType.LOCAL and prompt_upload_on_log()

def _determine_upload_on_log(self, init_config: InitConfig) -> bool:
if init_config.force_local or self.session_type == SessionType.LOCAL:
return False

# If we're in an interactive environment then prompt the user to pick upload on log
if is_interractive():
return self._determine_upload_on_log_prompt()

return init_config.upload_on_log or False


_CONFIG_WHYLABS_SECTION = "whylabs"

Expand Down
9 changes: 9 additions & 0 deletions python/whylogs/api/whylabs/session/prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,15 @@ def prompt_session_type(allow_anonymous: bool = True, allow_local: bool = False)
return [SessionType.WHYLABS, SessionType.WHYLABS_ANONYMOUS, SessionType.LOCAL][choice - 1]


def prompt_upload_on_log() -> bool:
options = [
"No. Use an explicit WhyLabsWriter to manage uploads to WhyLabs",
"Yes. Calling why.log() will automatically upload the results to WhyLabs",
]
choice = _get_user_choice("Do you want to automatically upload profiles in why.log()?", options)
return [False, True][choice - 1]


def prompt_default_dataset_id() -> Optional[str]:
try:
sys.stdout.flush()
Expand Down
6 changes: 4 additions & 2 deletions python/whylogs/api/whylabs/session/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -279,7 +279,9 @@ def __create_log_api(self, config: SessionConfig) -> LogApi:
def upload_reference_profiles(self, profile_aliases: Dict[str, ResultSet]) -> Union[UploadResult, NotSupported]:
results: List[str] = []
for alias, profile in profile_aliases.items():
success, ids = profile.writer("whylabs").option(reference_profile_name=alias).write()
success, ids = (
profile.writer("whylabs")._option(call_from_log=True).option(reference_profile_name=alias).write()
)
if success:
ids = ids if isinstance(ids, list) else [(True, ids)]
results.append(*[id for _, id in ids])
Expand All @@ -301,7 +303,7 @@ def upload_reference_profiles(self, profile_aliases: Dict[str, ResultSet]) -> Un
)

def upload_batch_profile(self, profile: ResultSet) -> Union[UploadResult, NotSupported]:
result = profile.writer("whylabs").write()
result = profile.writer("whylabs")._option(call_from_log=True).write()
utc_now = int(datetime.now(timezone.utc).timestamp() * 1000)
timestamps: Set[int] = set() # For generating the viewing url after upload

Expand Down
2 changes: 1 addition & 1 deletion python/whylogs/api/whylabs/session/session_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def is_active() -> bool:
def init(
reinit: bool = False,
allow_anonymous: bool = True,
allow_local: bool = False,
allow_local: bool = True,
whylabs_api_key: Optional[str] = None,
default_dataset_id: Optional[str] = None,
config_path: Optional[str] = None,
Expand Down
4 changes: 4 additions & 0 deletions python/whylogs/api/writer/whylabs_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,10 @@ def option(self, **kwargs) -> Writer: # type: ignore
self._whylabs_client = self._whylabs_client.option(**kwargs)
return self

def _option(self, **kwargs) -> Writer: # type: ignore
self._whylabs_client = self._whylabs_client._option(**kwargs)
return self

def _get_dataset_epoch(
self, view: Union[DatasetProfileView, SegmentedDatasetProfileView], utc_now: Optional[datetime.datetime] = None
) -> int:
Expand Down
22 changes: 22 additions & 0 deletions python/whylogs/api/writer/whylabs_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,7 @@ def __init__(
self._api_config: Optional[Configuration] = None
self._prefer_sync = read_bool_env_var(WHYLOGS_PREFER_SYNC_KEY, False)
self._transaction_id: Optional[str] = None
self._called_from_log = False

_http_proxy = os.environ.get("HTTP_PROXY")
_https_proxy = os.environ.get("HTTPS_PROXY")
Expand Down Expand Up @@ -361,6 +362,19 @@ def option(self, **kwargs) -> "WhyLabsClient": # type: ignore
)
return self

def _option(self, **kwargs) -> "WhyLabsClient": # type: ignore
"""
Parameters
----------
called_from_log: bool Set this to True in the context of a Session logger
"""

called_from_log = kwargs.get("called_from_log")
if called_from_log is not None:
self._called_from_log = called_from_log
return self

def _tag_columns(self, columns: List[str], value: str) -> Tuple[bool, str]:
"""Sets the column as an input or output for the specified dataset.
Expand Down Expand Up @@ -634,6 +648,14 @@ def do_upload(
profile_file: Optional[IO[bytes]] = None,
) -> Tuple[bool, str]:
assert profile_path or profile_file, "Either a file or file path must be specified when uploading profiles"
session = default_init()
config = session.config
if config.upload_on_log and not self._called_from_log:
logger.warning(
"The current session is configured to upload profiles in the why.log() call. "
+ "Uploading profiles explicitely with a WhyLabsWriter may result in redundant uploads."
)

try:
if profile_file:
status, reason = self._put_file(profile_file, upload_url, dataset_timestamp) # type: ignore
Expand Down
8 changes: 8 additions & 0 deletions python/whylogs/api/writer/writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,10 @@ def write(
def option(self, **kwargs: Any) -> "Writer":
return self

@abstractmethod
def _option(self, **kwargs: Any) -> "Writer":
return self


class WriterWrapper:
"""Elide the Writable argument"""
Expand All @@ -127,6 +131,10 @@ def option(self, **kwargs: Any) -> "WriterWrapper":
self._writer = self._writer.option(**kwargs)
return self

def _option(self, **kwargs: Any) -> "WriterWrapper":
self._writer = self._writer._option(**kwargs)
return self


class Writers:
@staticmethod
Expand Down

0 comments on commit 0a2e4be

Please sign in to comment.