(app) Make Logging DEBUG mode lazy (#14464)

.azure/app-cloud-e2e.yml

@@ -108,12 +108,15 @@ jobs:
       displayName: 'Install lightning'
 
     - bash: |
+        rm -rf examples/app_template_jupyterlab || true
         git clone https://github.com/Lightning-AI/LAI-lightning-template-jupyterlab-App examples/app_template_jupyterlab
         cp examples/app_template_jupyterlab/tests/test_template_jupyterlab.py tests/tests_app_examples/test_template_jupyterlab.py
       condition: eq(variables['name'], 'template_jupyterlab')
       displayName: 'Clone Template Jupyter Lab Repo'
 
-    - bash: git clone https://github.com/Lightning-AI/lightning-template-react examples/app_template_react_ui
+    - bash:  |
+        rm -rf examples/app_template_react_ui || true
+        git clone https://github.com/Lightning-AI/lightning-template-react examples/app_template_react_ui
       condition: eq(variables['name'], 'template_react_ui')
       displayName: 'Clone Template React UI Repo'
 
@@ -137,6 +140,7 @@ jobs:
         LIGHTNING_API_KEY: $(LIGHTNING_API_KEY_PROD)
         LIGHTNING_USERNAME: $(LIGHTNING_USERNAME)
         LIGHTNING_CLOUD_URL: $(LIGHTNING_CLOUD_URL_PROD)
+        LIGHTNING_DEBUG: '1'
       displayName: 'Run the tests'
 
     - publish: '$(Build.ArtifactStagingDirectory)/videos'

.gitignore

@@ -174,3 +174,4 @@ our_model.tar
 test.png
 saved_models
 data/
+.shared

examples/app_boring/app_dynamic.py

@@ -64,4 +64,4 @@ class BoringApp(L.LightningFlow):
         return {"name": "Boring Tab", "content": self.dict["dst_w"].url + "/file" if "dst_w" in self.dict else ""}
 
 
-app = L.LightningApp(BoringApp())
+app = L.LightningApp(BoringApp(), debug=True)

examples/app_commands_and_api/app.py

@@ -40,4 +40,4 @@ class FlowCommands(LightningFlow):
         return [Post("/user/command_without_client", self.command_without_client)]
 
 
-app = LightningApp(FlowCommands())
+app = LightningApp(FlowCommands(), debug=True)

examples/app_drive/app.py

@@ -48,4 +48,4 @@ class Flow(L.LightningFlow):
         self._exit("Application End!")
 
 
-app = L.LightningApp(Flow())
+app = L.LightningApp(Flow(), debug=True)

examples/app_template_streamlit_ui/app.py

@@ -45,4 +45,4 @@ class HelloWorld(LightningFlow):
         return [{"name": "StreamLitUI", "content": self.streamlit_ui}]
 
 
-app = LightningApp(HelloWorld())
+app = LightningApp(HelloWorld(), debug=True)

examples/app_v0/app.py

@@ -46,4 +46,4 @@ class V0App(L.LightningFlow):
         return [tab1, tab2, tab3]
 
 
-app = L.LightningApp(V0App())
+app = L.LightningApp(V0App(), debug=True)

src/lightning_app/CHANGELOG.md

@@ -12,6 +12,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Adds `PanelFrontend` to easily create complex UI in Python ([#13531](https://github.com/Lightning-AI/lightning/pull/13531))
 
 
+### Fixed
+
+- Resolved `LightningApp(..., debug=True)` ([#14464](https://github.com/Lightning-AI/lightning/pull/14464))
 
 
 ## [0.6.0] - 2022-08-23

src/lightning_app/api/http_methods.py

@@ -10,6 +10,9 @@ from uuid import uuid4
 from fastapi import FastAPI
 
 from lightning_app.api.request_types import APIRequest, CommandRequest
+from lightning_app.utilities.app_helpers import Logger
+
+logger = Logger(__name__)
 
 
 def _signature_proxy_function():
@@ -51,6 +54,7 @@ class HttpMethod:
         async def _handle_request(*args, **kwargs):
             async def fn(*args, **kwargs):
                 request_id = str(uuid4()).split("-")[0]
+                logger.debug(f"Processing request {request_id} for route: {self.route}")
                 request_queue.put(
                     request_cls(
                         name=self.component_name,
@@ -67,6 +71,8 @@ class HttpMethod:
                     if (time.time() - t0) > self.timeout:
                         raise Exception("The response was never received.")
 
+                logger.debug(f"Processed request {request_id} for route: {self.route}")
+
                 return responses_store.pop(request_id)
 
             return await asyncio.create_task(fn(*args, **kwargs))

src/lightning_app/cli/cmd_init.py

@@ -1,9 +1,10 @@
-import logging
 import os
 import re
 import shutil
 
-logger = logging.getLogger(__name__)
+from lightning_app.utilities.app_helpers import Logger
+
+logger = Logger(__name__)
 
 
 def app(app_name):

src/lightning_app/cli/cmd_install.py

@@ -1,4 +1,3 @@
-import logging
 import os
 import re
 import shutil
@@ -9,8 +8,9 @@ import requests
 from packaging.version import Version
 
 from lightning_app.core.constants import LIGHTNING_APPS_PUBLIC_REGISTRY, LIGHTNING_COMPONENT_PUBLIC_REGISTRY
+from lightning_app.utilities.app_helpers import Logger
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 def gallery_component(name, yes_arg, version_arg, cwd=None):

src/lightning_app/cli/cmd_react_ui_init.py

@@ -1,10 +1,11 @@
-import logging
 import os
 import re
 import shutil
 import subprocess
 
-logger = logging.getLogger(__name__)
+from lightning_app.utilities.app_helpers import Logger
+
+logger = Logger(__name__)
 
 
 def react_ui(dest_dir=None):

src/lightning_app/cli/lightning_cli.py

@@ -1,4 +1,3 @@
-import logging
 import os
 import sys
 from pathlib import Path
@@ -26,6 +25,7 @@ from lightning_app.cli.lightning_cli_list import get_list
 from lightning_app.core.constants import get_lightning_cloud_url
 from lightning_app.runners.runtime import dispatch
 from lightning_app.runners.runtime_type import RuntimeType
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.app_logs import _app_logs_reader
 from lightning_app.utilities.cli_helpers import _arrow_time_callback, _format_input_env_variables
 from lightning_app.utilities.cloud import _get_project
@@ -33,7 +33,7 @@ from lightning_app.utilities.cluster_logs import _cluster_logs_reader
 from lightning_app.utilities.login import Auth
 from lightning_app.utilities.network import LightningClient
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 def get_app_url(runtime_type: RuntimeType, *args) -> str:
@@ -405,7 +405,7 @@ def install():
 def install_app(name, yes, version, overwrite: bool = False):
     if "github.com" in name:
         if version != "latest":
-            logger.warning(
+            logger.warn(
                 f"The provided version {version} isn't the officially supported one. "
                 f"The provided version will be ignored."
             )
@@ -428,7 +428,7 @@ def install_app(name, yes, version, overwrite: bool = False):
 def install_component(name, yes, version):
     if "github.com" in name:
         if version != "latest":
-            logger.warning(
+            logger.warn(
                 f"The provided version {version} isn't the officially supported one. "
                 f"The provided version will be ignored."
             )

src/lightning_app/cli/pl-app-template/core/callbacks.py

@@ -1,11 +1,11 @@
 import inspect
-import logging
 from typing import Any, Dict, TYPE_CHECKING, Union
 
 from core.state import ProgressBarState, TrainerState
 
 import pytorch_lightning as pl
 from lightning.app.storage import Path
+from lightning_app.utilities.app_helpers import Logger
 from pytorch_lightning import Callback
 from pytorch_lightning.callbacks.progress.base import get_standard_metrics
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
@@ -14,7 +14,8 @@ from pytorch_lightning.utilities.parsing import collect_init_args
 if TYPE_CHECKING:
     from core.components.script_runner import ScriptRunner
 
-_log = logging.getLogger(__name__)
+
+_log = Logger(__name__)
 
 
 class PLAppProgressTracker(Callback):

src/lightning_app/components/python/popen.py

@@ -1,4 +1,3 @@
-import logging
 import os
 import signal
 import subprocess
@@ -7,10 +6,10 @@ from pathlib import Path
 from typing import Dict, List, Optional, Union
 
 from lightning_app import LightningWork
-from lightning_app.utilities.app_helpers import _collect_child_process_pids
+from lightning_app.utilities.app_helpers import _collect_child_process_pids, Logger
 from lightning_app.utilities.tracer import Tracer
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 class PopenPythonScript(LightningWork):

src/lightning_app/components/python/tracer.py

@@ -1,4 +1,3 @@
-import logging
 import os
 import signal
 import sys
@@ -8,11 +7,11 @@ from typing import Any, Dict, List, Optional, TypedDict, Union
 from lightning_app import LightningWork
 from lightning_app.storage.drive import Drive
 from lightning_app.storage.payload import Payload
-from lightning_app.utilities.app_helpers import _collect_child_process_pids
+from lightning_app.utilities.app_helpers import _collect_child_process_pids, Logger
 from lightning_app.utilities.packaging.tarfile import clean_tarfile, extract_tarfile
 from lightning_app.utilities.tracer import Tracer
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 class Code(TypedDict):

src/lightning_app/components/serve/serve.py

@@ -1,6 +1,5 @@
 import abc
 import inspect
-import logging
 import os
 import pydoc
 import subprocess
@@ -15,8 +14,9 @@ from starlette.responses import RedirectResponse
 
 from lightning_app import LightningWork
 from lightning_app.components.serve.types import _DESERIALIZER, _SERIALIZER
+from lightning_app.utilities.app_helpers import Logger
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 fastapi_service = FastAPI()

src/lightning_app/components/training.py

@@ -1,4 +1,3 @@
-import logging
 import os
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
@@ -6,8 +5,9 @@ from lightning import CloudCompute
 from lightning_app import LightningFlow, structures
 from lightning_app.components.python import TracerPythonScript
 from lightning_app.storage.path import Path
+from lightning_app.utilities.app_helpers import Logger
 
-_logger = logging.getLogger(__name__)
+_logger = Logger(__name__)
 
 
 class PyTorchLightningScriptRunner(TracerPythonScript):

src/lightning_app/core/api.py

@@ -1,5 +1,4 @@
 import asyncio
-import logging
 import os
 import queue
 import sys
@@ -24,7 +23,7 @@ from lightning_app.api.http_methods import HttpMethod
 from lightning_app.api.request_types import DeltaRequest
 from lightning_app.core.constants import FRONTEND_DIR
 from lightning_app.core.queues import RedisQueue
-from lightning_app.utilities.app_helpers import InMemoryStateStore, StateStore
+from lightning_app.utilities.app_helpers import InMemoryStateStore, Logger, StateStore
 from lightning_app.utilities.enum import OpenAPITags
 from lightning_app.utilities.imports import _is_redis_available, _is_starsessions_available
 
@@ -58,7 +57,7 @@ app_spec: Optional[List] = None
 # In the future, this would be abstracted to support horizontal scaling.
 responses_store = {}
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 # This can be replaced with a consumer that publishes states in a kv-store

src/lightning_app/core/app.py

@@ -11,12 +11,18 @@ from time import time
 from deepdiff import DeepDiff, Delta
 
 import lightning_app
+from lightning_app import _console
 from lightning_app.api.request_types import APIRequest, CommandRequest, DeltaRequest
-from lightning_app.core.constants import FLOW_DURATION_SAMPLES, FLOW_DURATION_THRESHOLD, STATE_ACCUMULATE_WAIT
+from lightning_app.core.constants import (
+    DEBUG_ENABLED,
+    FLOW_DURATION_SAMPLES,
+    FLOW_DURATION_THRESHOLD,
+    STATE_ACCUMULATE_WAIT,
+)
 from lightning_app.core.queues import BaseQueue, SingleProcessQueue
 from lightning_app.frontend import Frontend
 from lightning_app.storage.path import storage_root_dir
-from lightning_app.utilities.app_helpers import _delta_to_app_state_delta, _LightningAppRef
+from lightning_app.utilities.app_helpers import _delta_to_app_state_delta, _LightningAppRef, Logger
 from lightning_app.utilities.commands.base import _process_requests
 from lightning_app.utilities.component import _convert_paths_after_init
 from lightning_app.utilities.enum import AppStage, CacheCallsKeys
@@ -30,7 +36,7 @@ from lightning_app.utilities.warnings import LightningFlowWarning
 if t.TYPE_CHECKING:
     from lightning_app.runners.backends.backend import Backend, WorkManager
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 class LightningApp:
@@ -116,8 +122,13 @@ class LightningApp:
         # is only available after all Flows and Works have been instantiated.
         _convert_paths_after_init(self.root)
 
-        if debug:
-            logging.getLogger().setLevel(logging.DEBUG)
+        # Lazily enable debugging.
+        if debug or DEBUG_ENABLED:
+            if not DEBUG_ENABLED:
+                os.environ["LIGHTNING_DEBUG"] = "2"
+            _console.setLevel(logging.DEBUG)
+
+        logger.debug(f"ENV: {os.environ}")
 
     def get_component_by_name(self, component_name: str):
         """Returns the instance corresponding to the given component name."""
@@ -433,6 +444,8 @@ class LightningApp:
 
             self._has_updated = False
 
+        self._on_run_end()
+
         return True
 
     def _update_layout(self) -> None:
@@ -550,3 +563,8 @@ class LightningApp:
         # disable any flow schedules.
         for flow in self.flows:
             flow._disable_running_schedules()
+
+    def _on_run_end(self):
+        if os.getenv("LIGHTNING_DEBUG") == "2":
+            del os.environ["LIGHTNING_DEBUG"]
+            _console.setLevel(logging.INFO)

src/lightning_app/core/constants.py

@@ -30,7 +30,7 @@ DISABLE_DEPENDENCY_CACHE = bool(int(os.getenv("DISABLE_DEPENDENCY_CACHE", "0")))
 LIGHTNING_CLOUD_PROJECT_ID = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
 LIGHTNING_CREDENTIAL_PATH = os.getenv("LIGHTNING_CREDENTIAL_PATH", str(Path.home() / ".lightning" / "credentials.json"))
 DOT_IGNORE_FILENAME = ".lightningignore"
-
+DEBUG_ENABLED = bool(int(os.getenv("LIGHTNING_DEBUG", "0")))
 LIGHTNING_COMPONENT_PUBLIC_REGISTRY = "https://lightning.ai/v1/components"
 LIGHTNING_APPS_PUBLIC_REGISTRY = "https://lightning.ai/v1/apps"
 

src/lightning_app/core/queues.py

@@ -1,4 +1,3 @@
-import logging
 import multiprocessing
 import pickle
 import queue
@@ -15,12 +14,13 @@ from lightning_app.core.constants import (
     REDIS_WARNING_QUEUE_SIZE,
     STATE_UPDATE_TIMEOUT,
 )
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.imports import _is_redis_available, requires
 
 if _is_redis_available():
     import redis
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 READINESS_QUEUE_CONSTANT = "READINESS_QUEUE"

src/lightning_app/frontend/panel/app_state_comm.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import asyncio
-import logging
 import os
 from threading import Thread
 from typing import Callable
@@ -12,8 +11,9 @@ from typing import Callable
 import websockets
 
 from lightning_app.core.constants import APP_SERVER_PORT
+from lightning_app.utilities.app_helpers import Logger
 
-_logger = logging.getLogger(__name__)
+_logger = Logger(__name__)
 
 _CALLBACKS = []
 _THREAD: Thread = None

src/lightning_app/frontend/panel/app_state_watcher.py

@@ -7,15 +7,15 @@ This is particularly useful for the PanelFrontend but can be used by other Front
 """
 from __future__ import annotations
 
-import logging
 import os
 
 from lightning_app.frontend.panel.app_state_comm import watch_app_state
 from lightning_app.frontend.utils import _get_flow_state
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.imports import _is_param_available, requires
 from lightning_app.utilities.state import AppState
 
-_logger = logging.getLogger(__name__)
+_logger = Logger(__name__)
 
 
 if _is_param_available():

src/lightning_app/frontend/panel/panel_frontend.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import inspect
-import logging
 import os
 import pathlib
 import subprocess
@@ -11,11 +10,12 @@ from typing import Callable, TextIO
 
 from lightning_app.frontend.frontend import Frontend
 from lightning_app.frontend.utils import _get_frontend_environment
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.cloud import is_running_in_cloud
 from lightning_app.utilities.imports import requires
 from lightning_app.utilities.log import get_frontend_logfile
 
-_logger = logging.getLogger("PanelFrontend")
+_logger = Logger(__name__)
 
 
 def has_panel_autoreload() -> bool:

src/lightning_app/runners/cloud.py

@@ -1,5 +1,4 @@
 import fnmatch
-import logging
 import os
 import random
 import string
@@ -44,12 +43,13 @@ from lightning_app.runners.backends.cloud import CloudBackend
 from lightning_app.runners.runtime import Runtime
 from lightning_app.source_code import LocalSourceCodeDir
 from lightning_app.storage import Drive
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.cloud import _get_project
 from lightning_app.utilities.dependency_caching import get_hash
 from lightning_app.utilities.packaging.app_config import AppConfig, find_config_file
 from lightning_app.utilities.packaging.lightning_utils import _prepare_lightning_wheels_and_requirements
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 @dataclass
@@ -331,4 +331,4 @@ class CloudRuntime(Runtime):
                 )
             else:
                 warning_msg += "\nYou can ignore some files or folders by adding them to `.lightningignore`."
-            logger.warning(warning_msg)
+            logger.warn(warning_msg)

src/lightning_app/runners/runtime.py

@@ -1,4 +1,3 @@
-import logging
 import multiprocessing
 import sys
 from dataclasses import dataclass, field
@@ -10,11 +9,12 @@ import lightning_app
 from lightning_app import LightningApp
 from lightning_app.core.constants import APP_SERVER_HOST, APP_SERVER_PORT
 from lightning_app.runners.backends import Backend, BackendType
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.enum import AppStage, CacheCallsKeys, make_status, WorkStageStatus
 from lightning_app.utilities.load_app import load_app_from_file
 from lightning_app.utilities.proxies import WorkRunner
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 def dispatch(

src/lightning_app/source_code/copytree.py

@@ -1,13 +1,13 @@
 import fnmatch
-import logging
 import os
 from pathlib import Path
 from shutil import copy2, copystat, Error
 from typing import Callable, List, Set, Union
 
 from lightning_app.core.constants import DOT_IGNORE_FILENAME
+from lightning_app.utilities.app_helpers import Logger
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 def copytree(
@@ -159,7 +159,7 @@ def _ignore_filename_spell_check(src: Path):
     possible_spelling_mistakes.extend([p.lstrip(".") for p in possible_spelling_mistakes])
     for path in src.iterdir():
         if path.is_file() and path.name in possible_spelling_mistakes:
-            logger.warning(
+            logger.warn(
                 f"Lightning uses `{DOT_IGNORE_FILENAME}` as the ignore file but found {path.name} at "
                 f"{path.parent} instead. If this was a mistake, please rename the file."
             )

src/lightning_app/storage/copier.py

@@ -1,5 +1,4 @@
 import concurrent.futures
-import logging
 import pathlib
 import threading
 from threading import Thread
@@ -12,9 +11,11 @@ import lightning_app
 from lightning_app.core.queues import BaseQueue
 from lightning_app.storage.path import filesystem
 from lightning_app.storage.requests import ExistsRequest, GetRequest
+from lightning_app.utilities.app_helpers import Logger
 
 _PathRequest = Union[GetRequest, ExistsRequest]
-_logger = logging.getLogger(__name__)
+
+_logger = Logger(__name__)
 
 num_workers = 8
 

src/lightning_app/storage/orchestrator.py

@@ -1,4 +1,3 @@
-import logging
 import threading
 import traceback
 from queue import Empty
@@ -8,6 +7,7 @@ from typing import Dict, Optional, Set, TYPE_CHECKING, Union
 from lightning_app.core.queues import BaseQueue
 from lightning_app.storage.path import filesystem, path_to_work_artifact
 from lightning_app.storage.requests import ExistsRequest, ExistsResponse, GetRequest, GetResponse
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.enum import WorkStageStatus
 
 if TYPE_CHECKING:
@@ -16,7 +16,7 @@ if TYPE_CHECKING:
 
 _PathRequest = Union[GetRequest, ExistsRequest]
 _PathResponse = Union[ExistsResponse, GetResponse]
-_logger = logging.getLogger(__name__)
+_logger = Logger(__name__)
 
 
 class StorageOrchestrator(Thread):

src/lightning_app/storage/path.py

@@ -1,5 +1,4 @@
 import hashlib
-import logging
 import os
 import pathlib
 import shutil
@@ -14,6 +13,7 @@ from fsspec.implementations.local import LocalFileSystem
 import lightning_app
 from lightning_app.core.queues import BaseQueue
 from lightning_app.storage.requests import ExistsRequest, ExistsResponse, GetRequest, GetResponse
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.component import _is_flow_context
 from lightning_app.utilities.imports import _is_s3fs_available
 
@@ -27,7 +27,7 @@ if TYPE_CHECKING:
 
 num_workers = 8
 
-_logger = logging.getLogger(__name__)
+_logger = Logger(__name__)
 
 
 class Path(PathlibPath):

src/lightning_app/storage/payload.py

@@ -1,5 +1,4 @@
 import hashlib
-import logging
 import pathlib
 import pickle
 from abc import ABC, abstractmethod
@@ -10,9 +9,10 @@ import lightning_app
 from lightning_app.core.queues import BaseQueue
 from lightning_app.storage.path import filesystem, Path, shared_storage_path
 from lightning_app.storage.requests import ExistsRequest, ExistsResponse, GetRequest, GetResponse
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.component import _is_flow_context
 
-_logger = logging.getLogger(__name__)
+_logger = Logger(__name__)
 
 
 class BasePayload(ABC):

src/lightning_app/testing/testing.py

@@ -1,6 +1,5 @@
 import asyncio
 import json
-import logging
 import os
 import shutil
 import subprocess
@@ -9,6 +8,7 @@ import tempfile
 import time
 import traceback
 from contextlib import contextmanager
+from multiprocessing import Process
 from subprocess import Popen
 from time import sleep
 from typing import Any, Callable, Dict, Generator, List, Optional, Type
@@ -28,6 +28,7 @@ from lightning_app.utilities.app_logs import _app_logs_reader
 from lightning_app.utilities.cloud import _get_project
 from lightning_app.utilities.enum import CacheCallsKeys
 from lightning_app.utilities.imports import _is_playwright_available, requires
+from lightning_app.utilities.logs_socket_api import _LightningLogsSocketAPI
 from lightning_app.utilities.network import _configure_session, LightningClient
 from lightning_app.utilities.proxies import ProxyWorkRun
 
@@ -36,7 +37,49 @@ if _is_playwright_available():
     from playwright.sync_api import HttpCredentials, sync_playwright
 
 
-_logger = logging.getLogger(__name__)
+def _on_error_callback(ws_app, *_):
+    print(traceback.format_exc())
+    ws_app.close()
+
+
+def print_logs(app_id: str):
+    client = LightningClient()
+    project = _get_project(client)
+
+    works = client.lightningwork_service_list_lightningwork(
+        project_id=project.project_id,
+        app_id=app_id,
+    ).lightningworks
+    component_names = ["flow"] + [w.name for w in works]
+
+    rich_colors = list(ANSI_COLOR_NAMES)
+    colors = {c: rich_colors[i + 1] for i, c in enumerate(component_names)}
+
+    max_length = max(len(c.replace("root.", "")) for c in component_names)
+    identifiers = []
+
+    print("################### PRINTING LOGS ###################")
+
+    logs_api_client = _LightningLogsSocketAPI(client.api_client)
+
+    while True:
+        gen = _app_logs_reader(
+            logs_api_client=logs_api_client,
+            project_id=project.project_id,
+            app_id=app_id,
+            component_names=component_names,
+            follow=False,
+            on_error_callback=_on_error_callback,
+        )
+        for log_event in gen:
+            message = log_event.message
+            identifier = f"{log_event.timestamp}{log_event.message}"
+            if identifier not in identifiers:
+                date = log_event.timestamp.strftime("%m/%d/%Y %H:%M:%S")
+                identifiers.append(identifier)
+                color = colors[log_event.component_name]
+                padding = (max_length - len(log_event.component_name)) * " "
+                print(f"[{color}]{log_event.component_name}{padding}[/{color}] {date} {message}")
 
 
 class LightningTestApp(LightningApp):
@@ -178,12 +221,17 @@ def run_app_in_cloud(app_folder: str, app_name: str = "app.py", extra_args: [str
     # 2. Create the right application name.
     basename = app_folder.split("/")[-1]
     PR_NUMBER = os.getenv("PR_NUMBER", None)
+
     TEST_APP_NAME = os.getenv("TEST_APP_NAME", basename)
+    os.environ["TEST_APP_NAME"] = TEST_APP_NAME
+
     if PR_NUMBER:
         name = f"test-{PR_NUMBER}-{TEST_APP_NAME}-" + str(int(time.time()))
     else:
         name = f"test-{TEST_APP_NAME}-" + str(int(time.time()))
 
+    os.environ["LIGHTNING_APP_NAME"] = name
+
     # 3. Disconnect from the App if any.
     Popen("lightning disconnect", shell=True).wait()
 
@@ -191,6 +239,7 @@ def run_app_in_cloud(app_folder: str, app_name: str = "app.py", extra_args: [str
     with tempfile.TemporaryDirectory() as tmpdir:
         env_copy = os.environ.copy()
         env_copy["PACKAGE_LIGHTNING"] = "1"
+        env_copy["LIGHTNING_DEBUG"] = "1"
         shutil.copytree(app_folder, tmpdir, dirs_exist_ok=True)
         # TODO - add -no-cache to the command line.
         process = Popen(
@@ -295,6 +344,26 @@ def run_app_in_cloud(app_folder: str, app_name: str = "app.py", extra_args: [str
             """
         )
 
+        client = LightningClient()
+        project = _get_project(client)
+
+        lightning_apps = [
+            app
+            for app in client.lightningapp_instance_service_list_lightningapp_instances(
+                project.project_id
+            ).lightningapps
+            if app.name == name
+        ]
+
+        if not lightning_apps:
+            return True
+
+        assert len(lightning_apps) == 1
+        app_id = lightning_apps[0].id
+
+        process = Process(target=print_logs, kwargs={"app_id": app_id})
+        process.start()
+
         while True:
             try:
                 with admin_page.context.expect_page() as page_catcher:
@@ -305,15 +374,12 @@ def run_app_in_cloud(app_folder: str, app_name: str = "app.py", extra_args: [str
             except (playwright._impl._api_types.Error, playwright._impl._api_types.TimeoutError):
                 pass
 
-        client = LightningClient()
-        project = _get_project(client)
-        identifiers = []
-        rich_colors = list(ANSI_COLOR_NAMES)
+        print(f"The Lightning Id Name : [bold magenta]{app_id}[/bold magenta]")
+
+        logs_api_client = _LightningLogsSocketAPI(client.api_client)
 
         def fetch_logs(component_names: Optional[List[str]] = None) -> Generator:
             """This methods creates websockets connection in threads and returns the logs to the main thread."""
-            app_id = admin_page.url.split("/")[-1]
-
             if not component_names:
                 works = client.lightningwork_service_list_lightningwork(
                     project_id=project.project_id,
@@ -321,64 +387,54 @@ def run_app_in_cloud(app_folder: str, app_name: str = "app.py", extra_args: [str
                 ).lightningworks
                 component_names = ["flow"] + [w.name for w in works]
 
-            def on_error_callback(ws_app, *_):
-                print(traceback.print_exc())
-                ws_app.close()
-
-            colors = {c: rich_colors[i + 1] for i, c in enumerate(component_names)}
             gen = _app_logs_reader(
-                client=client,
+                logs_api_client=logs_api_client,
                 project_id=project.project_id,
                 app_id=app_id,
                 component_names=component_names,
                 follow=False,
-                on_error_callback=on_error_callback,
+                on_error_callback=_on_error_callback,
             )
-            max_length = max(len(c.replace("root.", "")) for c in component_names)
             for log_event in gen:
-                message = log_event.message
-                identifier = f"{log_event.timestamp}{log_event.message}"
-                if identifier not in identifiers:
-                    date = log_event.timestamp.strftime("%m/%d/%Y %H:%M:%S")
-                    identifiers.append(identifier)
-                    color = colors[log_event.component_name]
-                    padding = (max_length - len(log_event.component_name)) * " "
-                    print(f"[{color}]{log_event.component_name}{padding}[/{color}] {date} {message}")
-                yield message
-
-        # 7. Print your application ID
-        print(
-            f"The Lightning Id Name : [bold magenta]{str(view_page.url).split('.')[0].split('//')[-1]}[/bold magenta]"
-        )
+                yield log_event.message
 
         try:
             yield admin_page, view_page, fetch_logs, name
         except KeyboardInterrupt:
             pass
         finally:
-            print("##################################################")
-            button = admin_page.locator('[data-cy="stop"]')
-            try:
-                button.wait_for(timeout=3 * 1000)
-                button.click()
-            except (playwright._impl._api_types.Error, playwright._impl._api_types.TimeoutError):
-                pass
-            context.close()
-            browser.close()
-
-            list_lightningapps = client.lightningapp_instance_service_list_lightningapp_instances(project.project_id)
-
-            for lightningapp in list_lightningapps.lightningapps:
-                if lightningapp.name != name:
-                    continue
+            has_finished = False
+            while not has_finished:
                 try:
-                    res = client.lightningapp_instance_service_delete_lightningapp_instance(
-                        project_id=project.project_id,
-                        id=lightningapp.id,
+                    button = admin_page.locator('[data-cy="stop"]')
+                    try:
+                        button.wait_for(timeout=3 * 1000)
+                        button.click()
+                    except (playwright._impl._api_types.Error, playwright._impl._api_types.TimeoutError):
+                        pass
+                    context.close()
+                    browser.close()
+
+                    list_lightningapps = client.lightningapp_instance_service_list_lightningapp_instances(
+                        project.project_id
                     )
-                    assert res == {}
-                except ApiException as e:
-                    print(f"Failed to delete {lightningapp.name}. Exception {e}")
+
+                    for lightningapp in list_lightningapps.lightningapps:
+                        if lightningapp.name != name:
+                            continue
+                        try:
+                            res = client.lightningapp_instance_service_delete_lightningapp_instance(
+                                project_id=project.project_id,
+                                id=lightningapp.id,
+                            )
+                            assert res == {}
+                        except ApiException as e:
+                            print(f"Failed to delete {lightningapp.name}. Exception {e}")
+
+                    process.kill()
+                    has_finished = True
+                except Exception:
+                    pass
 
             Popen("lightning disconnect", shell=True).wait()
 

src/lightning_app/utilities/app_helpers.py

@@ -25,7 +25,6 @@ if TYPE_CHECKING:
     from lightning_app.core.flow import LightningFlow
     from lightning_app.utilities.types import Component
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -385,3 +384,34 @@ class LightningJSONEncoder(json.JSONEncoder):
         if callable(getattr(obj, "__json__", None)):
             return obj.__json__()
         return json.JSONEncoder.default(self, obj)
+
+
+class Logger:
+
+    """This class is used to improve the debugging experience."""
+
+    def __init__(self, name: str):
+        self.logger = logging.getLogger(name)
+        self.level = None
+
+    def info(self, msg, *args, **kwargs):
+        self.logger.info(msg, *args, **kwargs)
+
+    def warn(self, msg, *args, **kwargs):
+        self._set_level()
+        self.logger.warn(msg, *args, **kwargs)
+
+    def debug(self, msg, *args, **kwargs):
+        self._set_level()
+        self.logger.debug(msg, *args, **kwargs)
+
+    def error(self, msg, *args, **kwargs):
+        self._set_level()
+        self.logger.error(msg, *args, **kwargs)
+
+    def _set_level(self):
+        """Lazily set the level once set by the users."""
+        # Set on the first from either log, warn, debug or error call.
+        if self.level is None:
+            self.level = logging.DEBUG if bool(int(os.getenv("LIGHTNING_DEBUG", "0"))) else logging.INFO
+            self.logger.setLevel(self.level)

src/lightning_app/utilities/app_logs.py

@@ -10,7 +10,6 @@ from websocket import WebSocketApp
 
 from lightning_app.utilities.log_helpers import _error_callback, _OrderedLogEntry
 from lightning_app.utilities.logs_socket_api import _LightningLogsSocketAPI
-from lightning_app.utilities.network import LightningClient
 
 
 @dataclass
@@ -57,7 +56,7 @@ def _push_log_events_to_read_queue_callback(component_name: str, read_queue: que
 
 
 def _app_logs_reader(
-    client: LightningClient,
+    logs_api_client: _LightningLogsSocketAPI,
     project_id: str,
     app_id: str,
     component_names: List[str],
@@ -66,7 +65,6 @@ def _app_logs_reader(
 ) -> Iterator[_LogEvent]:
 
     read_queue = queue.PriorityQueue()
-    logs_api_client = _LightningLogsSocketAPI(client.api_client)
 
     # We will use a socket per component
     log_sockets = [

src/lightning_app/utilities/load_app.py

@@ -1,5 +1,4 @@
 import inspect
-import logging
 import os
 import sys
 import traceback
@@ -12,7 +11,9 @@ from lightning_app.utilities.exceptions import MisconfigurationException
 if TYPE_CHECKING:
     from lightning_app import LightningApp, LightningFlow, LightningWork
 
-logger = logging.getLogger(__name__)
+from lightning_app.utilities.app_helpers import Logger
+
+logger = Logger(__name__)
 
 
 def load_app_from_file(filepath: str) -> "LightningApp":

src/lightning_app/utilities/log_helpers.py

@@ -1,11 +1,12 @@
-import logging
 from dataclasses import dataclass
 from datetime import datetime
 from json import JSONDecodeError
 
 from websocket import WebSocketApp
 
-logger = logging.getLogger(__name__)
+from lightning_app.utilities.app_helpers import Logger
+
+logger = Logger(__name__)
 
 
 # This is a superclass to inherit log entry classes from it:

src/lightning_app/utilities/login.py

@@ -1,6 +1,5 @@
 import base64
 import json
-import logging
 import os
 import pathlib
 from dataclasses import dataclass
@@ -16,9 +15,10 @@ from starlette.background import BackgroundTask
 from starlette.responses import RedirectResponse
 
 from lightning_app.core.constants import get_lightning_cloud_url, LIGHTNING_CREDENTIAL_PATH
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.network import find_free_network_port
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 class Keys(Enum):
@@ -165,7 +165,7 @@ class AuthServer:
                 auth.save(token=token, username=user_id, user_id=user_id, api_key=key)
                 logger.info("Authentication Successful")
             else:
-                logger.warning(
+                logger.warn(
                     "Authentication Failed. This is most likely because you're using an older version of the CLI. \n"  # noqa E501
                     "Please try to update the CLI or open an issue with this information \n"  # E501
                     f"expected token in {request.query_params.items()}"

src/lightning_app/utilities/network.py

@@ -1,4 +1,3 @@
-import logging
 import socket
 import time
 from functools import wraps
@@ -13,7 +12,9 @@ from requests.adapters import HTTPAdapter
 from requests.exceptions import ConnectionError, ConnectTimeout, ReadTimeout
 from urllib3.util.retry import Retry
 
-logger = logging.getLogger(__name__)
+from lightning_app.utilities.app_helpers import Logger
+
+logger = Logger(__name__)
 
 
 def find_free_network_port() -> int:

src/lightning_app/utilities/packaging/build_config.py

@@ -1,17 +1,17 @@
 import inspect
-import logging
 import os
 import re
 from dataclasses import asdict, dataclass
 from types import FrameType
 from typing import cast, List, Optional, TYPE_CHECKING, Union
 
+from lightning_app.utilities.app_helpers import Logger
+
 if TYPE_CHECKING:
     from lightning_app import LightningWork
     from lightning_app.utilities.packaging.cloud_compute import CloudCompute
 
-
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 def load_requirements(

src/lightning_app/utilities/packaging/lightning_utils.py

@@ -16,9 +16,10 @@ from packaging.version import Version
 from lightning_app import _logger, _PROJECT_ROOT, _root_logger
 from lightning_app.__version__ import version
 from lightning_app.core.constants import PACKAGE_LIGHTNING
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.git import check_github_repository, get_dir_name
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 
 # FIXME(alecmerdler): Use GitHub release artifacts once the `lightning-ui` repo is public

src/lightning_app/utilities/proxies.py

@@ -1,4 +1,3 @@
-import logging
 import os
 import pathlib
 import signal
@@ -37,7 +36,9 @@ if TYPE_CHECKING:
     from lightning_app.core.queues import BaseQueue
 
 
-logger = logging.getLogger(__name__)
+from lightning_app.utilities.app_helpers import Logger
+
+logger = Logger(__name__)
 _state_observer_lock = threading.Lock()
 
 

src/lightning_app/utilities/state.py

@@ -1,6 +1,5 @@
 import enum
 import json
-import logging
 import os
 from copy import deepcopy
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -11,10 +10,10 @@ from requests.exceptions import ConnectionError
 
 from lightning_app.core.constants import APP_SERVER_HOST, APP_SERVER_PORT
 from lightning_app.storage.drive import _maybe_create_drive
-from lightning_app.utilities.app_helpers import AppStatePlugin, BaseStatePlugin
+from lightning_app.utilities.app_helpers import AppStatePlugin, BaseStatePlugin, Logger
 from lightning_app.utilities.network import _configure_session
 
-logger = logging.getLogger(__name__)
+logger = Logger(__name__)
 
 # GLOBAL APP STATE
 _LAST_STATE = None

tests/tests_app/core/test_lightning_app.py

@@ -955,3 +955,23 @@ def test_non_updated_flow(caplog):
         MultiProcessRuntime(app, start_server=False).dispatch()
     assert caplog.messages == ["Hello World"]
     assert app.counter == 3
+
+
+def test_debug_mode_logging():
+    """This test validates the DEBUG messages are collected when activated by the LightningApp(debug=True) and
+    cleanup once finished."""
+
+    from lightning_app.core.app import _console
+
+    app = LightningApp(A4(), debug=True)
+    assert _console.level == logging.DEBUG
+    assert os.getenv("LIGHTNING_DEBUG") == "2"
+
+    MultiProcessRuntime(app, start_server=False).dispatch()
+
+    assert os.getenv("LIGHTNING_DEBUG") is None
+    assert _console.level == logging.INFO
+
+    app = LightningApp(A4())
+    assert _console.level == logging.INFO
+    MultiProcessRuntime(app, start_server=False).dispatch()

tests/tests_app_examples/collect_failures/app.py

@@ -43,4 +43,4 @@ class RootFlow(LightningFlow):
 
 
 if __name__ == "__main__":
-    app = LightningApp(RootFlow())
+    app = LightningApp(RootFlow(), debug=True)

tests/tests_app_examples/conftest.py

@@ -0,0 +1,44 @@
+import os
+
+from lightning_cloud.openapi.rest import ApiException
+
+from lightning_app.utilities.cloud import _get_project
+from lightning_app.utilities.network import LightningClient
+
+
+def pytest_timeout_cancel_timer(item):
+    """This hook deletes the Lightning App when timeout triggers."""
+
+    if item.name.endswith("_example_cloud"):
+        name = os.getenv("LIGHTNING_APP_NAME")
+        print(f"Timeout was triggered. Deleting the App {name}.")
+
+        client = LightningClient()
+        project = _get_project(client)
+
+        lightning_apps = [
+            app
+            for app in client.lightningapp_instance_service_list_lightningapp_instances(
+                project.project_id
+            ).lightningapps
+            if app.name == name
+        ]
+
+        if not lightning_apps:
+            return True
+
+        assert len(lightning_apps) == 1
+
+        lightning_app = lightning_apps[0]
+
+        try:
+            res = client.lightningapp_instance_service_delete_lightningapp_instance(
+                project_id=project.project_id,
+                id=lightning_app.id,
+            )
+            assert res == {}
+
+        except ApiException as e:
+            print(f"Failed to delete {name}. Exception {e}")
+
+    return True

tests/tests_app_examples/custom_work_dependencies/app.py

@@ -49,4 +49,4 @@ class CustomWorkBuildConfigChecker(LightningFlow):
             self._exit()
 
 
-app = LightningApp(CustomWorkBuildConfigChecker())
+app = LightningApp(CustomWorkBuildConfigChecker(), debug=True)

tests/tests_app_examples/idle_timeout/app.py

@@ -68,4 +68,4 @@ class RootFlow(LightningFlow):
             self._exit()
 
 
-app = LightningApp(RootFlow())
+app = LightningApp(RootFlow(), debug=True)

tests/tests_app_examples/test_commands_and_api.py

@@ -9,6 +9,7 @@ from tests_app import _PROJECT_ROOT
 from lightning_app.testing.testing import run_app_in_cloud
 
 
+@pytest.mark.timeout(300)
 @pytest.mark.cloud
 def test_commands_and_api_example_cloud() -> None:
     with run_app_in_cloud(os.path.join(_PROJECT_ROOT, "examples/app_commands_and_api")) as (
@@ -21,16 +22,19 @@ def test_commands_and_api_example_cloud() -> None:
         app_id = admin_page.url.split("/")[-1]
 
         # 2: Connect to the App
-        Popen(f"lightning connect {app_id} -y", shell=True).wait()
+        Popen(f"python -m lightning connect {app_id} -y", shell=True).wait()
 
         # 3: Send the first command with the client
-        cmd = "lightning command with client --name=this"
+        cmd = "python -m lightning command with client --name=this"
         Popen(cmd, shell=True).wait()
 
         # 4: Send the second command without a client
-        cmd = "lightning command without client --name=is"
+        cmd = "python -m lightning command without client --name=is"
         Popen(cmd, shell=True).wait()
 
+        # This prevents some flakyness in the CI. Couldn't reproduce it locally.
+        sleep(5)
+
         # 5: Send a request to the Rest API directly.
         base_url = view_page.url.replace("/view", "").replace("/child_flow", "")
         resp = requests.post(base_url + "/user/command_without_client?name=awesome")

tests/tests_app_examples/test_v0_app.py

@@ -39,6 +39,7 @@ def run_v0_app(fetch_logs, view_page):
         locator = view_page.frame_locator("iframe").locator("div")
         locator.wait_for(timeout=3 * 1000)
         assert text_content in " ".join(locator.all_text_contents())
+        print(f"Validated {button_name}")
         return True
 
     wait_for(view_page, check_content, "TAB_1", "Hello from component A")