diff --git a/litdata/streaming/resolver.py b/litdata/streaming/resolver.py index 9bf5560b..f9dbfe40 100644 --- a/litdata/streaming/resolver.py +++ b/litdata/streaming/resolver.py @@ -23,6 +23,8 @@ try: from lightning_sdk import Machine, Studio + from lightning_sdk.organization import Organization + from lightning_sdk.user import User _LIGHTNING_SDK_AVAILABLE = True except (ImportError, ModuleNotFoundError): @@ -295,6 +297,15 @@ def _resolve_time_template(path: str) -> str: return path.replace("{" + pattern + "}", datetime.datetime.now().strftime(pattern)) +def _resolve_owner_name(studio: "Studio") -> str: + if issubclass(studio.owner, Organization): + return studio.owner.name + if issubclass(studio.owner, User): + return studio.owner + + raise RuntimeError("Neither user nor org provided, but one of them needs to be provided") + + def _execute( name: str, num_nodes: int, @@ -333,8 +344,9 @@ def _execute( ) if not has_printed: cloud_url = os.getenv("LIGHTNING_CLOUD_URL", "https://lightning.ai").replace(":443", "") - job_url = f"{cloud_url}/{studio.owner.name}/{studio._teamspace.name}" - job_url += f"/studios/{studio.name}/app?app_id=litdata&app_tab=Runs&job_name={curr_job.name}" + owner_name = _resolve_owner_name(studio) + job_url = f"{cloud_url}/{owner_name}/{studio._teamspace.name}" + job_url += f"/studios/{studio.name}/app?app_id=data-prep&job_name={curr_job.name}" print(f"Find your job at {job_url}") has_printed = True