Lazy load unnecessary modules to improve start speed (#525)

docs/source/custom_event.rst

@@ -65,7 +65,8 @@ Custom Event
 
     .. code-block:: python
 
-        from viztracer import VizTracer, VizCounter
+        from viztracer import VizTracer
+        from viztracer.vizcounter import VizCounter
         tracer = VizTracer()
         tracer.start()
         counter = VizCounter(tracer, "counter name")

docs/source/extra_log.rst

@@ -217,7 +217,8 @@ data as Instant Events.
 
 .. code-block:: python
 
-    from viztracer import VizTracer, VizLoggingHandler
+    from viztracer import VizTracer
+    from viztracer.vizlogging import VizLoggingHandler
 
     tracer = VizTracer()
     handler = VizLoggingHandler()

docs/source/global_tracer.rst

@@ -26,14 +26,16 @@ You can do things like:
 
 .. code-block:: python
 
-    from viztracer import VizLoggingHandler, get_tracer
+    from viztracer import get_tracer
+    from viztracer.vizlogging import VizLoggingHandler
 
     handler = VizLoggingHandler()
 
     handler.setTracer(get_tracer())
 
 .. code-block:: python
 
-    from viztracer import get_tracer, VizObject
+    from viztracer import get_tracer
+    from viztracer.vizobject import VizObject
 
     obj = VizObject(get_tracer(), "my variable")

example/src/gradient_descent.py

@@ -5,7 +5,8 @@
 
 import numpy
 
-from viztracer import VizCounter, VizTracer
+from viztracer import VizTracer
+from viztracer.vizcounter import VizCounter
 
 # List of input, output pairs
 train_data = (

example/src/logging_integration.py

@@ -1,6 +1,7 @@
 import logging
 
-from viztracer import VizLoggingHandler, get_tracer
+from viztracer import get_tracer
+from viztracer.vizlogging import VizLoggingHandler
 
 
 def fib(n):

pyproject.toml

@@ -36,8 +36,8 @@ Documentation = "https://viztracer.readthedocs.io"
 full = ["orjson"]
 
 [project.scripts]
-viztracer = "viztracer:main"
-vizviewer = "viztracer:viewer_main"
+viztracer = "viztracer.main:main"
+vizviewer = "viztracer.viewer:viewer_main"
 
 [tool.setuptools.dynamic]
 version = {attr = "viztracer.__version__"}

src/viztracer/__init__.py

@@ -6,24 +6,14 @@
 
 from .cellmagic import load_ipython_extension
 from .decorator import ignore_function, log_sparse, trace_and_save
-from .main import main
-from .viewer import viewer_main
-from .vizcounter import VizCounter
-from .vizlogging import VizLoggingHandler
-from .vizobject import VizObject
 from .viztracer import VizTracer, get_tracer
 
 __all__ = [
     "__version__",
-    "main",
-    "viewer_main",
     "VizTracer",
     "ignore_function",
     "trace_and_save",
     "log_sparse",
     "get_tracer",
-    "VizCounter",
-    "VizObject",
-    "VizLoggingHandler",
     "load_ipython_extension",
 ]

src/viztracer/main.py

@@ -5,7 +5,6 @@
 import atexit
 import base64
 import builtins
-import configparser
 import io
 import json
 import multiprocessing.util  # type: ignore
@@ -23,7 +22,6 @@
 from typing import Any, Optional, Union
 
 from . import __version__
-from .attach_process.add_code_to_python_process import run_python_code  # type: ignore
 from .code_monkey import CodeMonkey
 from .patch import install_all_hooks
 from .report_builder import ReportBuilder
@@ -166,6 +164,7 @@ def create_parser(self) -> argparse.ArgumentParser:
     def load_config_file(self, filename: str = ".viztracerrc") -> argparse.Namespace:
         ret = argparse.Namespace()
         if os.path.exists(filename):
+            import configparser
             cfg_parser = configparser.ConfigParser()
             cfg_parser.read(filename)
             if "default" not in cfg_parser:
@@ -531,6 +530,8 @@ def _check_attach_availability(self) -> tuple[bool, Optional[str]]:
         return True, None
 
     def attach(self) -> VizProcedureResult:
+        from .attach_process.add_code_to_python_process import run_python_code  # type: ignore
+
         pid = self.options.attach
         interval = self.options.t
 
@@ -570,6 +571,8 @@ def attach(self) -> VizProcedureResult:
         return True, None
 
     def uninstall(self) -> VizProcedureResult:
+        from .attach_process.add_code_to_python_process import run_python_code  # type: ignore
+
         pid = self.options.uninstall
 
         success, err_msg = self._check_attach_availability()

src/viztracer/modules/snaptrace.c

@@ -34,7 +34,7 @@ PyObject* trio_lowlevel_module = NULL;
 PyObject* sys_module = NULL;
 PyObject* sys_monitoring_missing = NULL;
 
-static PyObject* curr_task_getters[2] = {0};
+PyObject* curr_task_getters[2] = {0};
 
 // =============================================================================
 // Utility function
@@ -2142,11 +2142,6 @@ PyInit_snaptrace(void)
 
     threading_module = PyImport_ImportModule("threading");
     multiprocessing_module = PyImport_ImportModule("multiprocessing");
-    asyncio_module = PyImport_ImportModule("asyncio");
-    asyncio_tasks_module = PyImport_AddModule("asyncio.tasks");
-    if (PyObject_HasAttrString(asyncio_tasks_module, "current_task")) {
-        curr_task_getters[0] = PyObject_GetAttrString(asyncio_tasks_module, "current_task");
-    }
     if ((trio_module = PyImport_ImportModule("trio"))) {
         trio_lowlevel_module = PyImport_AddModule("trio.lowlevel");
         curr_task_getters[1] = PyObject_GetAttrString(trio_lowlevel_module, "current_task");

src/viztracer/modules/snaptrace_member.c

@@ -1,6 +1,13 @@
+// Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+// For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
+
 #include "pythoncapi_compat.h"
 #include "snaptrace.h"
 
+extern PyObject* asyncio_module;
+extern PyObject* asyncio_tasks_module;
+extern PyObject* curr_task_getters[2];
+
 // ================================================================
 // Tracer members
 // ================================================================
@@ -368,6 +375,14 @@ Tracer_log_async_setter(TracerObject* self, PyObject* value, void* closure)
     }
 
     if (value == Py_True) {
+        // Lazy import asyncio because it's slow
+        if (asyncio_module == NULL) {
+            asyncio_module = PyImport_ImportModule("asyncio");
+            asyncio_tasks_module = PyImport_AddModule("asyncio.tasks");
+            if (PyObject_HasAttrString(asyncio_tasks_module, "current_task")) {
+                curr_task_getters[0] = PyObject_GetAttrString(asyncio_tasks_module, "current_task");
+            }
+        }
         SET_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC);
     } else {
         UNSET_FLAG(self->check_flags, SNAPTRACE_LOG_ASYNC);