Make notebook-native auth work with more configurations of the Databricks Runtime #285

nfx · 2023-08-14T18:29:26Z

This PR adds additional logging and hardening to auth_type='runtime', which performs credentials lookup in the following order:

init_runtime_native_auth for the newest DBR versions.
init_runtime_repl_auth via Databricks REPL context and workspaceUrl.
init_runtime_legacy_auth via IPython context for legacy runtimes and modes.

Every detection step adds more logging on DEBUG level.

This PR adds nightly integration testing for all LTS runtimes, as well as testing the latest runtime with UC data access modes.

Based on publicly-accessible code in https://github.com/mlflow/mlflow/blame/6bd97bde24d78bcfbf6d50c1dd0f4fac2ed6987b/mlflow/utils/databricks_utils.py

…icks Runtime This PR adds additional logging and hardening to `auth_type='runtime'`, which performs credentials lookup in the following order: 1. `init_runtime_native_auth` for the newest DBR versions. 2. `init_runtime_repl_auth` via Databricks REPL context and `workspaceUrl`. 3. `init_runtime_legacy_auth` via IPython context for legacy runtimes and modes. Every detection step adds more logging on `DEBUG` level. Based on publicly-accessible code in https://github.com/mlflow/mlflow/blame/6bd97bde24d78bcfbf6d50c1dd0f4fac2ed6987b/mlflow/utils/databricks_utils.py

nfx · 2023-08-14T18:32:21Z

TODO: add a job to integration test native auth for different runtimes.

dby-tmwctw

LGTM.

mgyucht

Overall LGTM. Can we include somewhere the list of DBR versions/access control modes that this PR enables?

mgyucht · 2023-08-16T17:40:46Z

databricks/sdk/runtime/__init__.py

+
+        def inner() -> Dict[str, str]:
+            ctx = dbutils.notebook().getContext()
+            return {'Authorization': f'Bearer {getattr(ctx, "apiToken")().get()}'}


Is this equivalent?

Suggested change

return {'Authorization': f'Bearer {getattr(ctx, "apiToken")().get()}'}

return {'Authorization': f'Bearer {ctx.apiToken().get()}'}

mgyucht · 2023-08-16T17:41:02Z

databricks/sdk/runtime/__init__.py

+        ctx = dbutils.notebook().getContext()
+        if ctx is None:
+            return None, None
+        host = getattr(ctx, 'apiUrl')().get()


Is this equivalent?

Suggested change

host = getattr(ctx, 'apiUrl')().get()

host = ctx.apiUrl().get()

* Added collection of Databricks Runtime versions used together with Python SDK ([#287](#287)). * Applied attribute transformer when reading in attributes from the environment ([#293](#293)). * Made notebook-native auth work with more configurations of the Databricks Runtime ([#285](#285)). * Added retry in `w.clusters.ensure_cluster_is_running(id)` when cluster is simultaneously started by two different processes. ([#283](#283)). * Set necessary headers when authenticating via Azure CLI ([#290](#290)). * Updated classifier to `Development Status :: 4 - Beta` ([#291](#291)). * Introduced Artifact Allowlist, Securable Tags, and Subentity Tags services. * Introduced DeleteRuns and RestoreRuns methods in the Experiments API. * Introduced the GetSecret method in the Secrets API. * Renamed Auto Maintenance to Predictive Optimization. New Services: * Added [w.artifact_allowlists](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/artifact_allowlists.html) workspace-level service. * Added [w.securable_tags](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/securable_tags.html) workspace-level service. * Added [w.subentity_tags](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/subentity_tags.html) workspace-level service. * Added `databricks.sdk.service.catalog.ArtifactAllowlistInfo` dataclass. * Added `databricks.sdk.service.catalog.ArtifactMatcher` dataclass. * Added `databricks.sdk.service.catalog.ArtifactType` dataclass. * Added `databricks.sdk.service.catalog.GetArtifactAllowlistRequest` dataclass. * Added `databricks.sdk.service.catalog.ListSecurableTagsRequest` dataclass. * Added `databricks.sdk.service.catalog.ListSecurableType` dataclass. * Added `databricks.sdk.service.catalog.ListSubentityTagsRequest` dataclass. * Added `databricks.sdk.service.catalog.MatchType` dataclass. * Added `databricks.sdk.service.catalog.SetArtifactAllowlist` dataclass. * Added `databricks.sdk.service.catalog.TagChanges` dataclass. * Added `databricks.sdk.service.catalog.TagKeyValuePair` dataclass. * Added `databricks.sdk.service.catalog.TagSecurable` dataclass. * Added `databricks.sdk.service.catalog.TagSecurableAssignment` dataclass. * Added `databricks.sdk.service.catalog.TagSecurableAssignmentsList` dataclass. * Added `databricks.sdk.service.catalog.TagSubentity` dataclass. * Added `databricks.sdk.service.catalog.TagSubentityAssignmentsList` dataclass. * Added `databricks.sdk.service.catalog.TagsSubentityAssignment` dataclass. * Added `databricks.sdk.service.catalog.UpdateSecurableType` dataclass. * Added `databricks.sdk.service.catalog.UpdateTags` dataclass. New APIs: * Added `delete_runs()` method for [w.experiments](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/experiments.html) workspace-level service. * Added `restore_runs()` method for [w.experiments](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/experiments.html) workspace-level service. * Added `databricks.sdk.service.ml.DeleteRuns` dataclass. * Added `databricks.sdk.service.ml.DeleteRunsResponse` dataclass. * Added `databricks.sdk.service.ml.RestoreRuns` dataclass. * Added `databricks.sdk.service.ml.RestoreRunsResponse` dataclass. * Added `get_secret()` method for [w.secrets](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/secrets.html) workspace-level service. * Added `databricks.sdk.service.workspace.GetSecretRequest` dataclass. * Added `databricks.sdk.service.workspace.GetSecretResponse` dataclass. Service Renames: * Removed `effective_auto_maintenance_flag` field for `databricks.sdk.service.catalog.CatalogInfo`. * Removed `enable_auto_maintenance` field for `databricks.sdk.service.catalog.CatalogInfo`. * Added `effective_predictive_optimization_flag` field for `databricks.sdk.service.catalog.CatalogInfo`. * Added `enable_predictive_optimization` field for `databricks.sdk.service.catalog.CatalogInfo`. * Removed `databricks.sdk.service.catalog.EffectiveAutoMaintenanceFlag` dataclass. * Removed `databricks.sdk.service.catalog.EffectiveAutoMaintenanceFlagInheritedFromType` dataclass. * Removed `databricks.sdk.service.catalog.EnableAutoMaintenance` dataclass. * Removed `effective_auto_maintenance_flag` field for `databricks.sdk.service.catalog.SchemaInfo`. * Removed `enable_auto_maintenance` field for `databricks.sdk.service.catalog.SchemaInfo`. * Added `effective_predictive_optimization_flag` field for `databricks.sdk.service.catalog.SchemaInfo`. * Added `enable_predictive_optimization` field for `databricks.sdk.service.catalog.SchemaInfo`. * Removed `effective_auto_maintenance_flag` field for `databricks.sdk.service.catalog.TableInfo`. * Removed `enable_auto_maintenance` field for `databricks.sdk.service.catalog.TableInfo`. * Added `effective_predictive_optimization_flag` field for `databricks.sdk.service.catalog.TableInfo`. * Added `enable_predictive_optimization` field for `databricks.sdk.service.catalog.TableInfo`. * Added `databricks.sdk.service.catalog.EffectivePredictiveOptimizationFlag` dataclass. * Added `databricks.sdk.service.catalog.EffectivePredictiveOptimizationFlagInheritedFromType` dataclass. * Added `databricks.sdk.service.catalog.EnablePredictiveOptimization` dataclass. OpenAPI SHA: beff621d7b3e1d59244e2e34fc53a496f310e130, Date: 2023-08-17

nfx added the do-not-merge label Aug 14, 2023

nfx requested review from mgyucht, dby-tmwctw and MrBago August 14, 2023 18:29

nfx added 2 commits August 14, 2023 20:35

fix tests

bc6cb45

quick path to exit

5f1b201

dby-tmwctw approved these changes Aug 14, 2023

View reviewed changes

mgyucht approved these changes Aug 16, 2023

View reviewed changes

nfx added 2 commits August 17, 2023 12:30

wip

f94ca5b

Add integration tests

7a122fc

nfx removed the do-not-merge label Aug 17, 2023

nfx added this pull request to the merge queue Aug 17, 2023

Merged via the queue into main with commit 9aea9a0 Aug 17, 2023

nfx deleted the fix/runtime-auth branch August 17, 2023 13:32

mgyucht mentioned this pull request Aug 17, 2023

Release v0.6.0 #296

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make notebook-native auth work with more configurations of the Databricks Runtime #285

Make notebook-native auth work with more configurations of the Databricks Runtime #285

nfx commented Aug 14, 2023 •

edited

Loading

nfx commented Aug 14, 2023

dby-tmwctw left a comment

mgyucht left a comment

mgyucht Aug 16, 2023

mgyucht Aug 16, 2023

	return {'Authorization': f'Bearer {getattr(ctx, "apiToken")().get()}'}
	return {'Authorization': f'Bearer {ctx.apiToken().get()}'}

	host = getattr(ctx, 'apiUrl')().get()
	host = ctx.apiUrl().get()

Make notebook-native auth work with more configurations of the Databricks Runtime #285

Make notebook-native auth work with more configurations of the Databricks Runtime #285

Conversation

nfx commented Aug 14, 2023 • edited Loading

nfx commented Aug 14, 2023

dby-tmwctw left a comment

Choose a reason for hiding this comment

mgyucht left a comment

Choose a reason for hiding this comment

mgyucht Aug 16, 2023

Choose a reason for hiding this comment

mgyucht Aug 16, 2023

Choose a reason for hiding this comment

nfx commented Aug 14, 2023 •

edited

Loading