ygg 0.1.33__py3-none-any.whl → 0.1.35__py3-none-any.whl

{ygg-0.1.33.dist-info → ygg-0.1.35.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.33
+Version: 0.1.35
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 License:                                  Apache License

{ygg-0.1.33.dist-info → ygg-0.1.35.dist-info}/RECORD

@@ -1,15 +1,15 @@
-ygg-0.1.33.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+ygg-0.1.35.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
 yggdrasil/__init__.py,sha256=PfH7Xwt6uue6oqe6S5V8NhDJcVQClkKrBE1KXhdelZc,117
-yggdrasil/version.py,sha256=-w_wSS_K7mBDaL2ciGoQ17Ab9V3-ElRGvc4UASK1hFI,22
+yggdrasil/version.py,sha256=dvUFqQgabIeithBNW7swqwxg6R59T0-i289dJYgomYQ,22
 yggdrasil/databricks/__init__.py,sha256=skctY2c8W-hI81upx9F_PWRe5ishL3hrdiTuizgDjdw,152
 yggdrasil/databricks/compute/__init__.py,sha256=NvdzmaJSNYY1uJthv1hHdBuNu3bD_-Z65DWnaJt9yXg,289
-yggdrasil/databricks/compute/cluster.py,sha256=PTA3OOl2s23NGrMHOCa4ues-cx6mb5p7LtsI0bl_oa8,38088
-yggdrasil/databricks/compute/execution_context.py,sha256=_PWn-Jjb-qRtI7HGm7IOUAKoIpmivNiBZTddxpZNBaM,22142
-yggdrasil/databricks/compute/remote.py,sha256=735VSSkdGwIKXZ2F_SgiPiH-dUh7bheHBc4zaB6fDp8,2265
+yggdrasil/databricks/compute/cluster.py,sha256=yx3xEJ5Vgg-IPeyYxxdf74_9DSBPHvg2-FYs8ptJrl0,40404
+yggdrasil/databricks/compute/execution_context.py,sha256=nxrNXoarq_JAB-Cpj0udHhq2jx-DmMbRWJdAezLrPis,22347
+yggdrasil/databricks/compute/remote.py,sha256=nEN_Fr1Ouul_iKOf4B5QjEGscYAcl7nHjGsl2toRzrU,2874
 yggdrasil/databricks/jobs/__init__.py,sha256=snxGSJb0M5I39v0y3IR-uEeSlZR248cQ_4DJ1sYs-h8,154
 yggdrasil/databricks/jobs/config.py,sha256=9LGeHD04hbfy0xt8_6oobC4moKJh4_DTjZiK4Q2Tqjk,11557
 yggdrasil/databricks/sql/__init__.py,sha256=y1n5yg-drZ8QVZbEgznsRG24kdJSnFis9l2YfYCsaCM,234
-yggdrasil/databricks/sql/engine.py,sha256=BO2lweaL63etVJ9jia14JJvGVGZ6f0X6XsZ3QlA6_Po,38146
+yggdrasil/databricks/sql/engine.py,sha256=kUFBddJJQC0AgDqH0l7GFs7d_Ony5rc8fOv4inLU6Vw,41051
 yggdrasil/databricks/sql/exceptions.py,sha256=Jqd_gT_VyPL8klJEHYEzpv5eHtmdY43WiQ7HZBaEqSk,53
 yggdrasil/databricks/sql/statement_result.py,sha256=VlHXhTcvTVya_2aJ-uUfUooZF_MqQuOZ8k7g6PBDhOM,17227
 yggdrasil/databricks/sql/types.py,sha256=5G-BM9_eOsRKEMzeDTWUsWW5g4Idvs-czVCpOCrMhdA,6412
@@ -18,7 +18,7 @@ yggdrasil/databricks/workspaces/filesytem.py,sha256=Z8JXU7_XUEbw9fpTQT1avRQKi-IA
 yggdrasil/databricks/workspaces/io.py,sha256=Tdde4LaGNJNT50R11OkEYZyNacyIW9QrOXMAicAlIr4,32208
 yggdrasil/databricks/workspaces/path.py,sha256=-XnCD9p42who3DAwnITVE1KyrZUSoXDKHA8iZi-7wk4,47743
 yggdrasil/databricks/workspaces/path_kind.py,sha256=Xc319NysH8_6E9C0Q8nCxDHYG07_SnzyUVKHe0dNdDQ,305
-yggdrasil/databricks/workspaces/workspace.py,sha256=RixYZYhASeKPPPdP6JPN1XQC8M_rC0N-oBheasF--1k,23183
+yggdrasil/databricks/workspaces/workspace.py,sha256=MW-BEyldROqbX9SBbDspvlys_zehJjK5YgM3sGLfW-g,23382
 yggdrasil/dataclasses/__init__.py,sha256=6SdfIyTsoM4AuVw5TW4Q-UWXz41EyfsMcpD30cmjbSM,125
 yggdrasil/dataclasses/dataclass.py,sha256=fKokFUnqe4CmXXGMTdF4XDWbCUl_c_-se-UD48L5s1E,6594
 yggdrasil/libs/__init__.py,sha256=ulzk-ZkFUI2Pfo93YKtO8MBsEWtRZzLos7HAxN74R0w,168
@@ -30,8 +30,10 @@ yggdrasil/libs/extensions/__init__.py,sha256=mcXW5Li3Cbprbs4Ci-b5A0Ju0wmLcfvEiFu
 yggdrasil/libs/extensions/polars_extensions.py,sha256=RTkGi8llhPJjX7x9egix7-yXWo2X24zIAPSKXV37SSA,12397
 yggdrasil/libs/extensions/spark_extensions.py,sha256=E64n-3SFTDgMuXwWitX6vOYP9ln2lpGKb0htoBLEZgc,16745
 yggdrasil/pyutils/__init__.py,sha256=tl-LapAc71TV7RMgf2ftKwrzr8iiLOGHeJgA3RvO93w,293
-yggdrasil/pyutils/callable_serde.py,sha256=prxzYRrjR6-mZ9i1rIWakWL0w1JHPnTKBO_RR1NMacg,20992
+yggdrasil/pyutils/callable_serde.py,sha256=euY7Kiy04i1tpWKuB0b2qQ1FokLC3nq0cv7PObWYUBE,21809
+yggdrasil/pyutils/equality.py,sha256=Xyf8D1dLUCm3spDEir8Zyj7O4US_fBJwEylJCfJ9slI,3080
 yggdrasil/pyutils/exceptions.py,sha256=ssKNm-rjhavHUOZmGA7_1Gq9tSHDrb2EFI-cnBuWgng,3388
+yggdrasil/pyutils/expiring_dict.py,sha256=q9gb09-2EUN-jQZumUw5BXOQGYcj1wb85qKtQlciSxg,5825
 yggdrasil/pyutils/modules.py,sha256=B7IP99YqUMW6-DIESFzBx8-09V1d0a8qrIJUDFhhL2g,11424
 yggdrasil/pyutils/parallel.py,sha256=ubuq2m9dJzWYUyKCga4Y_9bpaeMYUrleYxdp49CHr44,6781
 yggdrasil/pyutils/python_env.py,sha256=tuglnjdqHQjNh18qDladVoSEOjCD0RcnMEPYJ0tArOs,50985
@@ -53,8 +55,8 @@ yggdrasil/types/cast/registry.py,sha256=_zdFGmUBB7P-e_LIcJlOxMcxAkXoA-UXB6HqLMgT
 yggdrasil/types/cast/spark_cast.py,sha256=_KAsl1DqmKMSfWxqhVE7gosjYdgiL1C5bDQv6eP3HtA,24926
 yggdrasil/types/cast/spark_pandas_cast.py,sha256=BuTiWrdCANZCdD_p2MAytqm74eq-rdRXd-LGojBRrfU,5023
 yggdrasil/types/cast/spark_polars_cast.py,sha256=btmZNHXn2NSt3fUuB4xg7coaE0RezIBdZD92H8NK0Jw,9073
-ygg-0.1.33.dist-info/METADATA,sha256=de-ntaFUDrFbT6rErjGnnZfPxN4DFBmxgokYLakRRcA,19204
-ygg-0.1.33.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ygg-0.1.33.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
-ygg-0.1.33.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
-ygg-0.1.33.dist-info/RECORD,,
+ygg-0.1.35.dist-info/METADATA,sha256=qSh-Cd0LtjU_CV9Je0E1ns1AnuRbIUPMIpvVJoJtIUA,19204
+ygg-0.1.35.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ygg-0.1.35.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
+ygg-0.1.35.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.35.dist-info/RECORD,,

yggdrasil/databricks/compute/cluster.py

@@ -24,6 +24,8 @@ from .execution_context import ExecutionContext
 from ..workspaces.workspace import WorkspaceService, Workspace
 from ... import retry, CallableSerde
 from ...libs.databrickslib import databricks_sdk
+from ...pyutils.equality import dicts_equal, dict_diff
+from ...pyutils.expiring_dict import ExpiringDict
 from ...pyutils.modules import PipIndexSettings
 from ...pyutils.python_env import PythonEnv
 
@@ -45,6 +47,31 @@ else:  # pragma: no cover - runtime fallback when SDK is missing
 __all__ = ["Cluster"]
 
 
+NAME_ID_CACHE: dict[str, ExpiringDict] = {}
+
+
+def set_cached_cluster_name(
+    host: str,
+    cluster_name: str,
+    cluster_id: str
+) -> None:
+    existing = NAME_ID_CACHE.get(host)
+
+    if not existing:
+        existing = NAME_ID_CACHE[host] = ExpiringDict(default_ttl=60)
+
+    existing[cluster_name] = cluster_id
+
+
+def get_cached_cluster_id(
+    host: str,
+    cluster_name: str,
+) -> str:
+    existing = NAME_ID_CACHE.get(host)
+
+    return existing.get(cluster_name) if existing else None
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -84,6 +111,7 @@ class Cluster(WorkspaceService):
     
     _details: Optional["ClusterDetails"] = dataclasses.field(default=None, repr=False)
     _details_refresh_time: float = dataclasses.field(default=0, repr=False)
+    _system_context: Optional[ExecutionContext] = dataclasses.field(default=None, repr=False)
 
     # host → Cluster instance
     _env_clusters: ClassVar[Dict[str, "Cluster"]] = {}
@@ -98,10 +126,11 @@ class Cluster(WorkspaceService):
         """Return the current cluster name."""
         return self.cluster_name
 
-    def __post_init__(self):
-        """Initialize cached details after dataclass construction."""
-        if self._details is not None:
-            self.details = self._details
+    @property
+    def system_context(self):
+        if self._system_context is None:
+            self._system_context = self.context(language=Language.PYTHON)
+        return self._system_context
 
     def is_in_databricks_environment(self):
         """Return True when running on a Databricks runtime."""
@@ -233,9 +262,8 @@ class Cluster(WorkspaceService):
         Returns:
             The updated PythonEnv instance.
         """
-        with self.context() as c:
-            m = c.remote_metadata
-            version_info = m.version_info
+        m = self.system_context.remote_metadata
+        version_info = m.version_info
 
         python_version = ".".join(str(_) for _ in version_info)
 
@@ -258,7 +286,7 @@ class Cluster(WorkspaceService):
         )
 
         return target
-    
+
     @property
     def details(self):
         """Return cached cluster details, refreshing when needed."""
@@ -300,21 +328,6 @@ class Cluster(WorkspaceService):
             return details.state
         return State.UNKNOWN
 
-    def get_state(self, max_delay: float = None):
-        """Return the cluster state with a custom refresh delay.
-
-        Args:
-            max_delay: Maximum age in seconds before refresh.
-
-        Returns:
-            The current cluster state.
-        """
-        details = self.fresh_details(max_delay=max_delay)
-
-        if details is not None:
-            return details.state
-        return State.UNKNOWN
-
     @property
     def is_running(self):
         """Return True when the cluster is running."""
@@ -323,7 +336,10 @@ class Cluster(WorkspaceService):
     @property
     def is_pending(self):
         """Return True when the cluster is starting, resizing, or terminating."""
-        return self.state in  (State.PENDING, State.RESIZING, State.RESTARTING, State.TERMINATING)
+        return self.state in (
+            State.PENDING, State.RESIZING, State.RESTARTING,
+            State.TERMINATING
+        )
 
     @property
     def is_error(self):
@@ -507,45 +523,51 @@ class Cluster(WorkspaceService):
     ):
         pip_settings = PipIndexSettings.default_settings()
 
-        if kwargs:
-            details = ClusterDetails(**{
-                **details.as_shallow_dict(),
-                **kwargs
-            })
+        new_details = ClusterDetails(**{
+            **details.as_shallow_dict(),
+            **kwargs
+        })
+
+        default_tags = self.workspace.default_tags()
 
-        if details.custom_tags is None:
-            details.custom_tags = self.workspace.default_tags()
+        if new_details.custom_tags is None:
+            new_details.custom_tags = default_tags
+        elif default_tags:
+            new_tags = new_details.custom_tags.copy()
+            new_tags.update(default_tags)
 
-        if details.cluster_name is None:
-            details.cluster_name = self.workspace.current_user.user_name
+            new_details.custom_tags = new_tags
 
-        if details.spark_version is None or python_version:
-            details.spark_version = self.latest_spark_version(
+        if new_details.cluster_name is None:
+            new_details.cluster_name = self.workspace.current_user.user_name
+
+        if new_details.spark_version is None or python_version:
+            new_details.spark_version = self.latest_spark_version(
                 photon=False, python_version=python_version
             ).key
 
-        if details.single_user_name:
-            if not details.data_security_mode:
-                details.data_security_mode = DataSecurityMode.DATA_SECURITY_MODE_DEDICATED
+        if new_details.single_user_name:
+            if not new_details.data_security_mode:
+                new_details.data_security_mode = DataSecurityMode.DATA_SECURITY_MODE_DEDICATED
 
-        if not details.node_type_id:
-            details.node_type_id = "rd-fleet.xlarge"
+        if not new_details.node_type_id:
+            new_details.node_type_id = "rd-fleet.xlarge"
 
-        if getattr(details, "virtual_cluster_size", None) is None and details.num_workers is None and details.autoscale is None:
-            if details.is_single_node is None:
-                details.is_single_node = True
+        if getattr(new_details, "virtual_cluster_size", None) is None and new_details.num_workers is None and new_details.autoscale is None:
+            if new_details.is_single_node is None:
+                new_details.is_single_node = True
 
-        if details.is_single_node is not None and details.kind is None:
-            details.kind = Kind.CLASSIC_PREVIEW
+        if new_details.is_single_node is not None and new_details.kind is None:
+            new_details.kind = Kind.CLASSIC_PREVIEW
 
         if pip_settings.extra_index_urls:
-            if details.spark_env_vars is None:
-                details.spark_env_vars = {}
+            if new_details.spark_env_vars is None:
+                new_details.spark_env_vars = {}
             str_urls = " ".join(pip_settings.extra_index_urls)
-            details.spark_env_vars["UV_EXTRA_INDEX_URL"] = details.spark_env_vars.get("UV_INDEX", str_urls)
-            details.spark_env_vars["PIP_EXTRA_INDEX_URL"] = details.spark_env_vars.get("PIP_EXTRA_INDEX_URL", str_urls)
+            new_details.spark_env_vars["UV_EXTRA_INDEX_URL"] = new_details.spark_env_vars.get("UV_INDEX", str_urls)
+            new_details.spark_env_vars["PIP_EXTRA_INDEX_URL"] = new_details.spark_env_vars.get("PIP_EXTRA_INDEX_URL", str_urls)
 
-        return details
+        return new_details
 
     def create_or_update(
         self,
@@ -637,8 +659,6 @@ class Cluster(WorkspaceService):
         Returns:
             The updated Cluster instance.
         """
-        self.install_libraries(libraries=libraries, wait_timeout=None, raise_error=False)
-
         existing_details = {
             k: v
             for k, v in self.details.as_shallow_dict().items()
@@ -651,10 +671,23 @@ class Cluster(WorkspaceService):
             if k in _EDIT_ARG_NAMES
         }
 
-        if update_details != existing_details:
+        same = dicts_equal(
+            existing_details,
+            update_details,
+            keys=_EDIT_ARG_NAMES,
+            treat_missing_as_none=True,
+            float_tol=0.0,  # set e.g. 1e-6 if you have float-y stuff
+        )
+
+        if not same:
+            diff = {
+                k: v[1]
+                for k, v in dict_diff(existing_details, update_details, keys=_EDIT_ARG_NAMES).items()
+            }
+
             logger.debug(
                 "Updating %s with %s",
-                self, update_details
+                self, diff
             )
 
             self.wait_for_status()
@@ -665,6 +698,8 @@ class Cluster(WorkspaceService):
                 self
             )
 
+        self.install_libraries(libraries=libraries, wait_timeout=None, raise_error=False)
+
         return self
 
     def list_clusters(self) -> Iterator["Cluster"]:
@@ -704,6 +739,12 @@ class Cluster(WorkspaceService):
         if not cluster_name and not cluster_id:
             raise ValueError("Either name or cluster_id must be provided")
 
+        if not cluster_id:
+            cluster_id = get_cached_cluster_id(
+                host=self.workspace.safe_host,
+                cluster_name=cluster_name
+            )
+
         if cluster_id:
             try:
                 details = self.clusters_client().get(cluster_id=cluster_id)
@@ -716,10 +757,13 @@ class Cluster(WorkspaceService):
                 workspace=self.workspace, cluster_id=details.cluster_id, _details=details
             )
 
-        cluster_name_cf = cluster_name.casefold()
-
         for cluster in self.list_clusters():
-            if cluster_name_cf == cluster.details.cluster_name.casefold():
+            if cluster_name == cluster.details.cluster_name:
+                set_cached_cluster_name(
+                    host=self.workspace.safe_host,
+                    cluster_name=cluster.cluster_name,
+                    cluster_id=cluster.cluster_id
+                )
                 return cluster
 
         if raise_error:
@@ -728,16 +772,18 @@ class Cluster(WorkspaceService):
 
     def ensure_running(
         self,
+        wait_timeout: Optional[dt.timedelta] = dt.timedelta(minutes=20)
     ) -> "Cluster":
         """Ensure the cluster is running.
 
         Returns:
             The current Cluster instance.
         """
-        return self.start()
+        return self.start(wait_timeout=wait_timeout)
 
     def start(
         self,
+        wait_timeout: Optional[dt.timedelta] = dt.timedelta(minutes=20)
     ) -> "Cluster":
         """Start the cluster if it is not already running.
 
@@ -748,8 +794,15 @@ class Cluster(WorkspaceService):
 
         if not self.is_running:
             logger.info("Starting %s", self)
-            self.details = self.clusters_client().start_and_wait(cluster_id=self.cluster_id)
-            return self.wait_installed_libraries()
+
+            if wait_timeout:
+                self.details = (
+                    self.clusters_client()
+                    .start_and_wait(cluster_id=self.cluster_id, timeout=wait_timeout)
+                )
+                self.wait_installed_libraries(timeout=wait_timeout)
+            else:
+                self.clusters_client().start(cluster_id=self.cluster_id)
 
         return self
 
@@ -812,6 +865,7 @@ class Cluster(WorkspaceService):
         env_keys: Optional[List[str]] = None,
         timeout: Optional[dt.timedelta] = None,
         result_tag: Optional[str] = None,
+        context: Optional[ExecutionContext] = None,
     ):
         """Execute a command or callable on the cluster.
 
@@ -823,11 +877,14 @@ class Cluster(WorkspaceService):
             env_keys: Optional environment variable names to pass.
             timeout: Optional timeout for execution.
             result_tag: Optional result tag for parsing output.
+            context: ExecutionContext to run or create new one
 
         Returns:
             The decoded result from the execution context.
         """
-        return self.context(language=language).execute(
+        context = self.system_context if context is None else context
+
+        return context.execute(
             obj=obj,
             args=args,
             kwargs=kwargs,
@@ -849,6 +906,7 @@ class Cluster(WorkspaceService):
         timeout: Optional[dt.timedelta] = None,
         result_tag: Optional[str] = None,
         force_local: bool = False,
+        context: Optional[ExecutionContext] = None,
         **options
     ):
         """
@@ -875,16 +933,28 @@ class Cluster(WorkspaceService):
             timeout: Optional timeout for remote execution.
             result_tag: Optional tag for parsing remote output.
             force_local: force local execution
+            context: ExecutionContext to run or create new one
             **options: Additional execution options passed through.
 
         Returns:
             A decorator or wrapped function that executes remotely.
         """
+        if force_local or self.is_in_databricks_environment():
+            # Support both @ws.remote and @ws.remote(...)
+            if _func is not None and callable(_func):
+                return _func
+
+            def identity(x):
+                return x
+
+            return identity
+
+        context = self.system_context if context is None else context
+
         def decorator(func: Callable):
             if force_local or self.is_in_databricks_environment():
                 return func
 
-            context = self.context(language=language or Language.PYTHON)
             serialized = CallableSerde.from_callable(func)
 
             @functools.wraps(func)
@@ -1075,7 +1145,7 @@ class Cluster(WorkspaceService):
                     "Waiting %s to install libraries timed out" % self
                 )
 
-            time.sleep(10)
+            time.sleep(5)
             statuses = list(self.installed_library_statuses())
 
         return self
@@ -1111,7 +1181,7 @@ class Cluster(WorkspaceService):
                 )
 
                 with open(value, mode="rb") as f:
-                    target_path.write_bytes(f.read())
+                    target_path.open().write_all_bytes(f.read())
 
                 value = str(target_path)
             elif "." in value and not "/" in value:

yggdrasil/databricks/compute/execution_context.py

@@ -78,8 +78,8 @@ class ExecutionContext:
     language: Optional["Language"] = None
     context_id: Optional[str] = None
 
-    _was_connected: Optional[bool] = None
-    _remote_metadata: Optional[RemoteMetadata] = None
+    _was_connected: Optional[bool] = dc.field(default=None, repr=False)
+    _remote_metadata: Optional[RemoteMetadata] = dc.field(default=None, repr=False)
 
     _lock: threading.RLock = dc.field(default_factory=threading.RLock, init=False, repr=False)
 
@@ -367,6 +367,8 @@ print(json.dumps(meta))"""
             args=args,
             kwargs=kwargs,
             result_tag=result_tag,
+            env_keys=env_keys,
+            env_variables=env_variables
         ) if not command else command
 
         raw_result = self.execute_command(
@@ -382,8 +384,9 @@ print(json.dumps(meta))"""
             module_name = module_name.group(1) if module_name else None
             module_name = module_name.split(".")[0]
 
-            if module_name:
+            if module_name and "yggdrasil" not in module_name:
                 self.close()
+
                 self.cluster.install_libraries(
                     libraries=[module_name],
                     raise_error=True,
@@ -442,7 +445,7 @@ print(json.dumps(meta))"""
             module_name = module_name.group(1) if module_name else None
             module_name = module_name.split(".")[0]
 
-            if module_name:
+            if module_name and "yggdrasil" not in module_name:
                 self.close()
                 self.cluster.install_libraries(
                     libraries=[module_name],

yggdrasil/databricks/compute/remote.py

@@ -2,11 +2,12 @@
 
 import datetime as dt
 import logging
+import os
 from typing import (
     Callable,
     Optional,
     TypeVar,
-    List, TYPE_CHECKING,
+    List, TYPE_CHECKING, Union,
 )
 
 if TYPE_CHECKING:
@@ -14,15 +15,26 @@ if TYPE_CHECKING:
 
 from ..workspaces.workspace import Workspace
 
+
+__all__ = [
+    "databricks_remote_compute"
+]
+
+
 ReturnType = TypeVar("ReturnType")
 
 logger = logging.getLogger(__name__)
 
 
+def identity(x):
+    return x
+
+
 def databricks_remote_compute(
+    _func: Optional[Callable] = None,
     cluster_id: Optional[str] = None,
     cluster_name: Optional[str] = None,
-    workspace: Optional[Workspace] = None,
+    workspace: Optional[Union[Workspace, str]] = None,
     cluster: Optional["Cluster"] = None,
     timeout: Optional[dt.timedelta] = None,
     env_keys: Optional[List[str]] = None,
@@ -32,6 +44,7 @@ def databricks_remote_compute(
     """Return a decorator that executes functions on a remote cluster.
 
     Args:
+        _func: function to decorate
         cluster_id: Optional cluster id to target.
         cluster_name: Optional cluster name to target.
         workspace: Workspace instance or host string for lookup.
@@ -45,13 +58,19 @@ def databricks_remote_compute(
         A decorator that runs functions on the resolved Databricks cluster.
     """
     if force_local or Workspace.is_in_databricks_environment():
-        def identity(x):
-            return x
+        return identity if _func is None else _func
 
-        return identity
+    if workspace is None:
+        workspace = os.getenv("DATABRICKS_HOST")
 
-    if isinstance(workspace, str):
-        workspace = Workspace(host=workspace)
+    if workspace is None:
+        return identity if _func is None else _func
+
+    if not isinstance(workspace, Workspace):
+        if isinstance(workspace, str):
+            workspace = Workspace(host=workspace).connect(clone=False)
+        else:
+            raise ValueError("Cannot initialize databricks workspace with %s" % type(workspace))
 
     if cluster is None:
         if cluster_id or cluster_name:
@@ -62,16 +81,15 @@ def databricks_remote_compute(
         else:
             cluster = workspace.clusters().replicated_current_environment(
                 workspace=workspace,
-                cluster_name=cluster_name
+                cluster_name=cluster_name,
+                single_user_name=workspace.current_user.user_name
             )
 
+    cluster.ensure_running(wait_timeout=None)
+
     return cluster.execution_decorator(
+        _func=_func,
         env_keys=env_keys,
         timeout=timeout,
         **options
     )
-
-
-__all__ = [
-    "databricks_remote_compute",
-]