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.

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

Display traces via post-run hook #11937

Merged
merged 6 commits into from May 21, 2024

Conversation

daniellok-db
Copy link
Collaborator

@daniellok-db daniellok-db commented May 8, 2024

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

This PR resolves a common complaint that the trace UI does not show up at the end of a cell as expected. To fix this, we register a post-execution hook instead of updating a display handle every time.

I'm keeping the user-facing API as display_traces(), because I feel that it would be confusing to change the name to something like update_traces_to_display(). My main concern is that people will wonder if there's an additional display() function that they have to call, but actually it's handled automatically by the post-execution hook.

How is this PR tested?

  • Existing unit/integration tests
  • New unit/integration tests
  • Manual tests

I'll add some automated tests in a bit

Does this PR require documentation update?

  • No. You can skip the rest of this section.
  • Yes. I've updated:
    • Examples
    • API references
    • Instructions

Release Notes

Is this a user-facing change?

  • No. You can skip the rest of this section.
  • Yes. Give a description of this change to be included in the release notes for MLflow users.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

  • area/artifacts: Artifact stores and artifact logging
  • area/build: Build and test infrastructure for MLflow
  • area/deployments: MLflow Deployments client APIs, server, and third-party Deployments integrations
  • area/docs: MLflow documentation pages
  • area/examples: Example code
  • area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
  • area/models: MLmodel format, model serialization/deserialization, flavors
  • area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
  • area/projects: MLproject format, project running backends
  • area/scoring: MLflow Model server, model deployment tools, Spark UDFs
  • area/server-infra: MLflow Tracking server backend
  • area/tracking: Tracking Service, tracking client APIs, autologging

Interface

  • area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
  • area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
  • area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
  • area/windows: Windows support

Language

  • language/r: R APIs and clients
  • language/java: Java APIs and clients
  • language/new: Proposals for new client languages

Integrations

  • integrations/azure: Azure and Azure ML integrations
  • integrations/sagemaker: SageMaker integrations
  • integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

  • rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
  • rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
  • rn/feature - A new user-facing feature worth mentioning in the release notes
  • rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
  • rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?
  • Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
    Bug fixes, doc updates and new features usually go into minor releases.
  • Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
    Bug fixes and doc updates usually go into patch releases.
  • Yes (this PR will be cherry-picked and included in the next patch release)
  • No (this PR will be included in the next minor release)

Copy link

github-actions bot commented May 8, 2024

Documentation preview for 8a3845a will be available when this CircleCI job
completes successfully.

More info

@daniellok-db daniellok-db marked this pull request as ready for review May 20, 2024 05:47
@github-actions github-actions bot added the rn/none List under Small Changes in Changelogs. label May 20, 2024
@@ -46,40 +88,11 @@ def display_traces(self, traces: List[Trace]):
# this should do nothing if not in an IPython environment
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As mentioned in the PR description, even though this function no longer actually displays anything, I think it makes sense from a developer UX point of view to keep the name as display_traces().

Comment on lines -56 to -85

# if the current ipython exec count has changed, then
# we're in a different cell (or rerendering the current
# cell), so we should create a new display handle.
current_exec_count = get_ipython().execution_count
if self._prev_execution_count != current_exec_count:
self._prev_execution_count = current_exec_count
self.traces_to_display = traces_dict
self.display_handle = None
else:
self.traces_to_display.update(traces_dict)

deduped_trace_list = list(self.traces_to_display.values())
if self.display_handle:
self.display_handle.update(
self.get_mimebundle(deduped_trace_list),
raw=True,
)
else:
self.display_handle = display(
self.get_mimebundle(deduped_trace_list),
display_id=True,
raw=True,
)
self.traces_to_display.update(traces_dict)
except Exception:
# swallow exceptions. this function is called as
# a side-effect in a few other functions (e.g. log_trace,
# get_traces, search_traces), and we don't want to block
# the core functionality if the display fails.
_logger.debug("Failed to display traces", exc_info=True)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

basically moved all this logic into the post-run hook, and removed all logic about tracking execution count, since now we're just guaranteed to execute once per cell

@daniellok-db daniellok-db changed the title [wip] display traces at the end of the cell Display traces via post-run hook May 20, 2024
Copy link
Member

@harupy harupy left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Awesome demo :)

# the core functionality if the display fails.
_logger.debug("Failed to register post-run cell display hook", exc_info=True)

def _display_traces_post_run(self):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def _display_traces_post_run(self):
def _display_traces_post_run(self, result):

can we add a result argument?

https://ipython.readthedocs.io/en/stable/config/callbacks.html#ipython-events

Copy link
Member

@harupy harupy May 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pre_run_cell might be useful to reset self.traces_to_display or initialize the state.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added result! as for resetting state, i'm currently just doing it immediately after displaying, so pre_run_cell should be unnecessary, but we can keep it in mind for future use-cases.

Signed-off-by: Daniel Lok <daniel.lok@databricks.com>
Signed-off-by: Daniel Lok <daniel.lok@databricks.com>
Signed-off-by: Daniel Lok <daniel.lok@databricks.com>
Signed-off-by: Daniel Lok <daniel.lok@databricks.com>
Signed-off-by: Daniel Lok <daniel.lok@databricks.com>
Signed-off-by: Daniel Lok <daniel.lok@databricks.com>
@daniellok-db daniellok-db merged commit 2f22db0 into mlflow:master May 21, 2024
46 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
rn/none List under Small Changes in Changelogs.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants