databricks-sdk-air 0.1.0__py3-none-any.whl

databricks/air/__init__.py

@@ -0,0 +1,50 @@
+"""Databricks AI Runtime Python package."""
+
+import importlib
+import logging
+import os
+from typing import Any, Callable
+
+__version__ = "0.1.0"
+
+logging.getLogger("databricks.sdk").setLevel(logging.WARNING)
+
+__all__ = [
+    "distributed",
+    "ray_launch",
+    "ray_init",
+    "__version__",
+]
+
+_DATABRICKS_INTERNAL_BETA_WARNING_ENV_VAR = "DATABRICKS_INTERNAL_BETA_WARNING_PRINTED"
+_DISTRIBUTED_EXTRA_MESSAGE = (
+    "databricks.air distributed APIs require additional dependencies. "
+    "Install them with `pip install databricks-sdk-air[distributed]`."
+)
+
+
+def _load_distributed_api(module_name: str, attr_name: str) -> Callable[..., Any]:
+    try:
+        module = importlib.import_module(f"databricks.air.{module_name}")
+    except ModuleNotFoundError as e:
+        raise ImportError(_DISTRIBUTED_EXTRA_MESSAGE) from e
+    return getattr(module, attr_name)
+
+
+def distributed(*args: Any, **kwargs: Any) -> Any:
+    return _load_distributed_api("launcher", "distributed")(*args, **kwargs)
+
+
+def ray_init(*args: Any, **kwargs: Any) -> Any:
+    return _load_distributed_api("ray", "ray_init")(*args, **kwargs)
+
+
+def ray_launch(*args: Any, **kwargs: Any) -> Any:
+    return _load_distributed_api("ray", "ray_launch")(*args, **kwargs)
+
+
+_logger = logging.getLogger(__name__)
+
+if not os.environ.get(_DATABRICKS_INTERNAL_BETA_WARNING_ENV_VAR):
+    os.environ[_DATABRICKS_INTERNAL_BETA_WARNING_ENV_VAR] = "1"
+    _logger.warning("Warning: databricks.air is in Beta. The API is subject to change.")