Enabled debug logging for every job task run through a file, which is accessible from both workspace UI and Databricks CLI (#426)

This PR improves debuggability of UCX.

## More concise error messaging from job run
<img width="1390" alt="image"
src="https://github.com/databrickslabs/ucx/assets/259697/be4dc82c-3724-4684-b8cd-a77a05f24d67">

## Viewing logs via Databricks CLI
<img width="929" alt="image"
src="https://github.com/databrickslabs/ucx/assets/259697/b1901d03-1c9a-4c63-81c7-501f0f3e4e48">

## Added `.ucx/logs` folder in the installation
<img width="313" alt="image"
src="https://github.com/databrickslabs/ucx/assets/259697/a396b637-308a-4791-9cc5-dbd0494431c4">

where every executes workflow has a folder per job run, with ability to
get back to Job Run in Workspace UI
<img width="1092" alt="image"
src="https://github.com/databrickslabs/ucx/assets/259697/1bf60990-eab6-4866-b0b5-a75c8a44674e">

## Viewing full debug logs via workspace UI

<img width="1218" alt="image"
src="https://github.com/databrickslabs/ucx/assets/259697/0d1a84bb-0af9-4cbc-b662-77c1715ff4b1">

Closes #86

---------

Co-authored-by: Serge Smertin <[email protected]>

Author: FastLee

src/databricks/labs/ucx/framework/logger.py

@@ -54,10 +54,11 @@ def format(self, record: logging.LogRecord):  # noqa: A003
         return f"{self.GRAY}{ts}{self.RESET} {level} {color_marker}[{name}] {msg}{self.RESET}"
 
 
-def _install():
+def _install(level="DEBUG"):
     for h in logging.root.handlers:
         logging.root.removeHandler(h)
     console_handler = logging.StreamHandler(sys.stderr)
     console_handler.setFormatter(NiceFormatter())
-    console_handler.setLevel("DEBUG")
+    console_handler.setLevel(level)
     logging.root.addHandler(console_handler)
+    return console_handler

src/databricks/labs/ucx/framework/tasks.py

@@ -29,9 +29,10 @@ def decorator(func):
         def wrapper(*args, **kwargs):
             # Perform any task-specific logic here
             # For example, you can log when the task is started and completed
-            print(f"Task '{workflow}' is starting...")
+            logger = logging.getLogger(func.__name__)
+            logger.info(f"Task '{workflow}' is starting...")
             result = func(*args, **kwargs)
-            print(f"Task '{workflow}' is completed!")
+            logger.info(f"Task '{workflow}' is completed!")
             return result
 
         deps = []
@@ -76,18 +77,60 @@ def trigger(*argv):
     if "config" not in args:
         msg = "no --config specified"
         raise KeyError(msg)
-
     task_name = args.get("task", "not specified")
+    # `{{parent_run_id}}` is the run of entire workflow, whereas `{{run_id}}` is the run of a task
+    workflow_run_id = args.get("parent_run_id", "unknown_run_id")
+    job_id = args.get("job_id")
     if task_name not in _TASKS:
         msg = f'task "{task_name}" not found. Valid tasks are: {", ".join(_TASKS.keys())}'
         raise KeyError(msg)
 
     current_task = _TASKS[task_name]
     print(current_task.doc)
 
-    _install()
-
-    cfg = WorkspaceConfig.from_file(Path(args["config"]))
-    logging.getLogger("databricks").setLevel(cfg.log_level)
-
-    current_task.fn(cfg)
+    config_path = Path(args["config"])
+    cfg = WorkspaceConfig.from_file(config_path)
+
+    # see https://docs.python.org/3/howto/logging-cookbook.html
+    databricks_logger = logging.getLogger("databricks")
+    databricks_logger.setLevel(logging.DEBUG)
+
+    ucx_logger = logging.getLogger("databricks.labs.ucx")
+    ucx_logger.setLevel(logging.DEBUG)
+
+    log_path = config_path.parent / "logs" / current_task.workflow / f"run-{workflow_run_id}"
+    log_path.mkdir(parents=True, exist_ok=True)
+
+    log_file = log_path / f"{task_name}.log"
+    file_handler = logging.FileHandler(log_file.as_posix())
+    log_format = "%(asctime)s %(levelname)s [%(name)s] {%(threadName)s} %(message)s"
+    log_formatter = logging.Formatter(fmt=log_format, datefmt="%H:%M:%S")
+    file_handler.setFormatter(log_formatter)
+    file_handler.setLevel(logging.DEBUG)
+
+    console_handler = _install(cfg.log_level)
+    databricks_logger.removeHandler(console_handler)
+    databricks_logger.addHandler(file_handler)
+
+    ucx_logger.info(f"See debug logs at {log_file}")
+
+    log_readme = log_path.joinpath("README.md")
+    if not log_readme.exists():
+        # this may race when run from multiple tasks, but let's accept the risk for now.
+        with log_readme.open(mode="w") as f:
+            f.write(f"# Logs for the UCX {current_task.workflow} workflow\n")
+            f.write("This folder contains UCX log files.\n\n")
+            f.write(f"See the [{current_task.workflow} job](/#job/{job_id}) and ")
+            f.write(f"[run #{workflow_run_id}](/#job/{job_id}/run/{workflow_run_id})\n")
+
+    try:
+        current_task.fn(cfg)
+    except BaseException as error:
+        log_file_for_cli = str(log_file).lstrip("/Workspace")
+        cli_command = f"databricks workspace export /{log_file_for_cli}"
+        ucx_logger.error(f"Task crashed. Execute `{cli_command}` locally to troubleshoot with more details. {error}")
+        databricks_logger.debug("Task crash details", exc_info=error)
+        file_handler.flush()
+        raise
+    finally:
+        file_handler.close()

src/databricks/labs/ucx/install.py

@@ -29,7 +29,11 @@
 TAG_STEP = "step"
 TAG_APP = "App"
 NUM_USER_ATTEMPTS = 10  # number of attempts user gets at answering a question
-
+EXTRA_TASK_PARAMS = {
+    "job_id": "{{job_id}}",
+    "run_id": "{{run_id}}",
+    "parent_run_id": "{{parent_run_id}}",
+}
 DEBUG_NOTEBOOK = """
 # Databricks notebook source
 # MAGIC %md
@@ -69,26 +73,13 @@
 
 # COMMAND ----------
 
-import logging
-import databricks.labs.ucx.runtime
-
-from pathlib import Path
-from databricks.labs.ucx.__about__ import __version__
-from databricks.labs.ucx.config import WorkspaceConfig
-from databricks.labs.ucx.framework.tasks import _TASKS
-from databricks.labs.ucx.framework.logger import _install
-from databricks.sdk import WorkspaceClient
-
-task_name = dbutils.widgets.get('task')
-current_task = _TASKS[task_name]
-
-_install()
-print('UCX version: ' + __version__)
-logging.getLogger("databricks").setLevel('DEBUG')
-
-cfg = WorkspaceConfig.from_file(Path("/Workspace{config_file}"))
+from databricks.labs.ucx.runtime import main
 
-current_task.fn(cfg)
+main(f'--config=/Workspace{config_file}',
+     f'--task=' + dbutils.widgets.get('task'),
+     f'--job_id=' + dbutils.widgets.get('job_id'),
+     f'--run_id=' + dbutils.widgets.get('run_id'),
+     f'--parent_run_id=' + dbutils.widgets.get('parent_run_id'))
 """
 
 logger = logging.getLogger(__name__)
@@ -325,6 +316,7 @@ def _create_jobs(self):
         self._deployed_steps = self._deployed_steps()
         desired_steps = {t.workflow for t in _TASKS.values()}
         wheel_runner = None
+
         if self._override_clusters:
             wheel_runner = self._upload_wheel_runner(remote_wheel)
         for step_name in desired_steps:
@@ -518,9 +510,8 @@ def _apply_cluster_overrides(settings: dict[str, any], overrides: dict[str, str]
                 job_task.job_cluster_key = None
             if job_task.python_wheel_task is not None:
                 job_task.python_wheel_task = None
-                job_task.notebook_task = jobs.NotebookTask(
-                    notebook_path=wheel_runner, base_parameters={"task": job_task.task_key}
-                )
+                params = {"task": job_task.task_key} | EXTRA_TASK_PARAMS
+                job_task.notebook_task = jobs.NotebookTask(notebook_path=wheel_runner, base_parameters=params)
         return settings
 
     def _job_task(self, task: Task, dbfs_path: str) -> jobs.Task:
@@ -555,7 +546,12 @@ def _job_notebook_task(self, jobs_task: jobs.Task, task: Task) -> jobs.Task:
             notebook_task=jobs.NotebookTask(
                 notebook_path=remote_notebook,
                 # ES-872211: currently, we cannot read WSFS files from Scala context
-                base_parameters={"inventory_database": self._current_config.inventory_database},
+                base_parameters={
+                    "inventory_database": self._current_config.inventory_database,
+                    "task": task.name,
+                    "config": f"/Workspace{self._config_file}",
+                }
+                | EXTRA_TASK_PARAMS,
             ),
         )
 
@@ -566,7 +562,7 @@ def _job_wheel_task(self, jobs_task: jobs.Task, task: Task, dbfs_path: str) -> j
             python_wheel_task=jobs.PythonWheelTask(
                 package_name="databricks_labs_ucx",
                 entry_point="runtime",  # [project.entry-points.databricks] in pyproject.toml
-                named_parameters={"task": task.name, "config": f"/Workspace{self._config_file}"},
+                named_parameters={"task": task.name, "config": f"/Workspace{self._config_file}"} | EXTRA_TASK_PARAMS,
             ),
         )
 

src/databricks/labs/ucx/runtime.py

@@ -249,12 +249,14 @@ def destroy_schema(cfg: WorkspaceConfig):
     RuntimeBackend().execute(f"DROP DATABASE {cfg.inventory_database} CASCADE")
 
 
-def main():
-    trigger(*sys.argv)
+def main(*argv):
+    if len(argv) == 0:
+        argv = sys.argv
+    trigger(*argv)
 
 
 if __name__ == "__main__":
     if "DATABRICKS_RUNTIME_VERSION" not in os.environ:
         msg = "Only intended to run in Databricks Runtime"
         raise SystemExit(msg)
-    main()
+    main(*sys.argv)

tests/integration/test_installation.py

@@ -44,6 +44,32 @@ def test_destroying_non_existing_schema_fails_with_correct_message(ws, sql_backe
     assert "cannot be found" in str(failure.value)
 
 
+def test_logs_are_available(ws, sql_backend, env_or_skip, make_random):
+    default_cluster_id = env_or_skip("TEST_DEFAULT_CLUSTER_ID")
+    ws.clusters.ensure_cluster_is_running(default_cluster_id)
+
+    install = WorkspaceInstaller.run_for_config(
+        ws,
+        WorkspaceConfig(
+            inventory_database=f"ucx_{make_random(4)}",
+            instance_pool_id=env_or_skip("TEST_INSTANCE_POOL_ID"),
+            groups=GroupsConfig(auto=True),
+            log_level="INFO",
+        ),
+        prefix=make_random(4),
+        override_clusters={
+            "main": default_cluster_id,
+        },
+    )
+
+    with pytest.raises(OperationFailed):
+        install.run_workflow("destroy-schema")
+        assert True
+
+    workflow_run_logs = list(ws.workspace.list(f"{install._install_folder}/logs"))
+    assert len(workflow_run_logs) == 1
+
+
 def test_jobs_with_no_inventory_database(
     ws,
     sql_backend,