diff --git a/.nextchanges/cli/ssh-jupyter-pip-magics.md b/.nextchanges/cli/ssh-jupyter-pip-magics.md new file mode 100644 index 00000000000..92fdf57595c --- /dev/null +++ b/.nextchanges/cli/ssh-jupyter-pip-magics.md @@ -0,0 +1 @@ +* Fix `%pip` in Jupyter notebooks opened over `databricks ssh connect` on serverless environment versions 5 and 6 and on dedicated clusters running recent Databricks Runtime versions, which fell back to the stock IPython magic instead of the Databricks one. ([#7263](https://github.com/databricks/cli/pull/7263)) diff --git a/experimental/ssh/internal/server/jupyter-init.py b/experimental/ssh/internal/server/jupyter-init.py index 055a3966584..ea535e10d1d 100644 --- a/experimental/ssh/internal/server/jupyter-init.py +++ b/experimental/ssh/internal/server/jupyter-init.py @@ -151,25 +151,53 @@ def _register_common_magics(): @_log_exceptions def _register_pip_magics(): """Register the pip magic command parser with IPython.""" + import inspect + from dbruntime import UserNamespaceInitializer from dbruntime.DatasetInfo import UserNamespaceDict from dbruntime.PipMagicOverrides import PipMagicOverrides user_namespace_initializer = UserNamespaceInitializer.getOrCreate() entry_point = user_namespace_initializer.get_spark_entry_point() - user_ns = UserNamespaceDict( - user_namespace_initializer.get_namespace_globals(), - entry_point.getDriverConf(), - entry_point, - ) ip = get_ipython() - try: - # Older DBRs - pip_magic = PipMagicOverrides(entry_point, ip.user_ns["sc"]._conf, user_ns) - except Exception: - # Newer DBRs - pip_magic = PipMagicOverrides(entry_point, user_ns, ip) + params = inspect.signature(PipMagicOverrides).parameters + if "logger" in params: + # Runtimes that take a logger (serverless environment 5+ and recent Databricks Runtime + # versions) are migrating off Py4J, so build the magic from the driver connection the way + # dbruntime's own kernel does: on serverless environment 6 the entry point is None. The + # constructor gained keyword arguments over several runtime releases, so pass only the ones + # this runtime accepts. Values are looked up lazily because the attributes behind + # newer arguments may not exist on runtimes that don't accept them. + driver_connection = user_namespace_initializer.get_driver_connection() + driver_conf = driver_connection.runtime_integration_client.getDriverConf() + user_ns = UserNamespaceDict(user_namespace_initializer.get_namespace_globals(), driver_conf, entry_point) + optional_kwargs = { + "orig_pip_magic": lambda: ip.magics_manager.magics["line"]["pip"], + "orig_conda_magic": lambda: ip.magics_manager.magics["line"]["conda"], + "driver_conf": lambda: driver_conf, + "environment_resource_client": lambda: driver_connection.environment_resource_client, + "sc_conf": lambda: user_namespace_initializer.db_connection.sc_conf, + } + pip_magic = PipMagicOverrides( + entry_point, + user_ns, + ip, + logger=driver_connection.observability_client.getLogger(), + **{name: get_value() for name, get_value in optional_kwargs.items() if name in params}, + ) + else: + user_ns = UserNamespaceDict( + user_namespace_initializer.get_namespace_globals(), + entry_point.getDriverConf(), + entry_point, + ) + try: + # Older DBRs + pip_magic = PipMagicOverrides(entry_point, ip.user_ns["sc"]._conf, user_ns) + except Exception: + # Newer DBRs + pip_magic = PipMagicOverrides(entry_point, user_ns, ip) ip.register_magics(pip_magic)