Refactor ArtifactManager and restart TaskManager #529

gabegma · 2023-04-03T18:42:28Z

Fixes #521

Description:

ArtifactManager is now a real Singleton (1 per process). 3 are created -> 1 per worker and 1 in the main thread.
- By extension, fewer DMs are created. Different tasks will access the same DM.
The client is now restarted after the startup tasks. This empties the memory, creating new ArtifactManagers. This avoids that the pipeline and the sentence encoder are kept in memory after the startup.

Checklist:

You should check all boxes before the PR is ready. If a box does not apply, check it to acknowledge it.

ISSUE NUMBER. You linked the issue number (Ex: Resolve #XXX).
PRE-COMMIT. You ran pre-commit on all commits, or else, you
ran pre-commit run --all-files at the end.
USER CHANGES. The changes are added to CHANGELOG.md and the documentation, if they impact
our users.
DEV CHANGES.
- Update the documentation if this PR changes how to develop/launch on the app.
- Update the README files and our wiki for any big design decisions, if relevant.
- Add unit tests, docstrings, typing and comments for complex sections.

gabegma · 2023-04-03T20:46:55Z

azimuth/app.py

@@ -297,6 +296,32 @@ def create_app() -> FastAPI:
    return app


+def load_dataset_split_managers_from_config(


I needed to move this function because of circular import issues.
One big change that I made here is to use an ArtifactManager in the main thread. I was afraid that multiple dms were created otherwise.

azimuth/modules/base_classes/artifact_manager.py

gabegma · 2023-04-03T20:49:26Z

azimuth/routers/export.py

+    dm: DatasetSplitManager,
+    results: List[List[PerturbedUtteranceResult]],
+    pipeline_index: int,
+    config: AzimuthConfig,


This was causing an error in a test. We shouldn't rely on the config of the dataset split manager, because it only depends on the project_hash.

azimuth/startup.py

gabegma · 2023-04-03T20:52:07Z

tests/test_task_manager.py

@@ -63,29 +63,6 @@ def get_module_data(simple_text_config):
    )


-def test_clearing_cache(tiny_text_config):


I'm not sure if we need a new test to replace this one - but it didn't look relevant anymore. Replacing task_manager.clear_worker_cache() by task_manager.restart() would work, but then the test takes a lot of time (~50 seconds), and I'm not sure it is testing something relevant.

I don't really understand the assertions any([m and d for m, d in cached.values()]). I leave it up to you! Saving 50 seconds does sound tempting if it's useless.

azimuth/task_manager.py

azimuth/modules/base_classes/artifact_manager.py

azimuth/config.py

azimuth/modules/base_classes/artifact_manager.py

JosephMarinier

LGTM! I don't understand all the details, but I guess we'll confirm that everything works by using it!

* Add logging to help with config update debugging * Restart client to free memory * Refactor artifact manager to real singleton and remove clear_cache * Adapt based on comments

gabegma changed the title ~~Refactor artifact manager and restart TaskManager~~ Refactor ArtifactManager and restart TaskManager Apr 3, 2023

gabegma force-pushed the ggm/refactor-artifact-manager branch from 6fd26a1 to 0fc4e4a Compare April 3, 2023 19:58

gabegma changed the base branch from main to ggm/refactor-dm April 3, 2023 20:23

gabegma mentioned this pull request Apr 3, 2023

Restart task manager #528

Closed

4 tasks

gabegma commented Apr 3, 2023

View reviewed changes

azimuth/modules/base_classes/artifact_manager.py Show resolved Hide resolved

gabegma commented Apr 3, 2023

View reviewed changes

azimuth/startup.py Show resolved Hide resolved

gabegma commented Apr 3, 2023

View reviewed changes

gabegma marked this pull request as ready for review April 3, 2023 21:01

gabegma self-assigned this Apr 4, 2023

gabegma force-pushed the ggm/refactor-dm branch 2 times, most recently from 4402ef6 to 5f81471 Compare April 5, 2023 19:12

gabegma force-pushed the ggm/refactor-artifact-manager branch from 0fc4e4a to a3f8412 Compare April 5, 2023 19:13

gabegma force-pushed the ggm/refactor-dm branch from 1b501bb to 750cca7 Compare April 7, 2023 21:57

gabegma force-pushed the ggm/refactor-artifact-manager branch from a3f8412 to 228f0ac Compare April 7, 2023 21:58

gabegma force-pushed the ggm/refactor-dm branch 2 times, most recently from a93803f to 4d33219 Compare April 9, 2023 01:41

gabegma force-pushed the ggm/refactor-artifact-manager branch from 228f0ac to 4095a73 Compare April 9, 2023 01:44

gabegma force-pushed the ggm/refactor-dm branch from 4d33219 to 69b7949 Compare April 13, 2023 20:41

gabegma force-pushed the ggm/refactor-artifact-manager branch from 4095a73 to 3ecfa72 Compare April 13, 2023 20:58

gabegma force-pushed the ggm/refactor-dm branch from 69b7949 to 99178a8 Compare April 17, 2023 21:34

JosephMarinier reviewed Apr 17, 2023

View reviewed changes

azimuth/task_manager.py Outdated Show resolved Hide resolved

JosephMarinier reviewed Apr 18, 2023

View reviewed changes

azimuth/modules/base_classes/artifact_manager.py Show resolved Hide resolved

JosephMarinier reviewed Apr 18, 2023

View reviewed changes

azimuth/config.py Show resolved Hide resolved

Base automatically changed from ggm/refactor-dm to main April 18, 2023 03:05

gabegma added 3 commits April 17, 2023 23:06

Add logging to help with config update debugging

b60a743

Restart client to free memory

47b044a

Refactor artifact manager to real singleton and remove clear_cache

c5fbc3c

gabegma force-pushed the ggm/refactor-artifact-manager branch from 3ecfa72 to c5fbc3c Compare April 18, 2023 03:06

JosephMarinier reviewed Apr 18, 2023

View reviewed changes

azimuth/modules/base_classes/artifact_manager.py Show resolved Hide resolved

Adapt based on comments

0bedba9

JosephMarinier approved these changes Apr 18, 2023

View reviewed changes

gabegma merged commit 9ba2f9b into main Apr 18, 2023

gabegma deleted the ggm/refactor-artifact-manager branch April 18, 2023 21:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor ArtifactManager and restart TaskManager #529

Refactor ArtifactManager and restart TaskManager #529

gabegma commented Apr 3, 2023 •

edited

Loading

gabegma Apr 3, 2023

gabegma Apr 3, 2023

gabegma Apr 3, 2023 •

edited

Loading

JosephMarinier Apr 18, 2023

JosephMarinier left a comment

		@@ -297,6 +296,32 @@ def create_app() -> FastAPI:
		return app


		def load_dataset_split_managers_from_config(

		@@ -63,29 +63,6 @@ def get_module_data(simple_text_config):
		)


		def test_clearing_cache(tiny_text_config):

Refactor ArtifactManager and restart TaskManager #529

Refactor ArtifactManager and restart TaskManager #529

Conversation

gabegma commented Apr 3, 2023 • edited Loading

Description:

Checklist:

gabegma Apr 3, 2023

Choose a reason for hiding this comment

gabegma Apr 3, 2023

Choose a reason for hiding this comment

gabegma Apr 3, 2023 • edited Loading

Choose a reason for hiding this comment

JosephMarinier Apr 18, 2023

Choose a reason for hiding this comment

JosephMarinier left a comment

Choose a reason for hiding this comment

gabegma commented Apr 3, 2023 •

edited

Loading

gabegma Apr 3, 2023 •

edited

Loading