From 77f7c954c5d74940bfeebebeb729a271ccb74674 Mon Sep 17 00:00:00 2001
From: Yurij Mikhalevich <yurij@lightning.ai>
Date: Fri, 19 May 2023 20:44:48 +0400
Subject: [PATCH] feat: cloud dispatch improvement

---
 src/lightning/app/plugin/plugin.py | 5 +++++
 src/lightning/app/runners/cloud.py | 7 ++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/lightning/app/plugin/plugin.py b/src/lightning/app/plugin/plugin.py
index b59553256c009..ecb670cc3447c 100644
--- a/src/lightning/app/plugin/plugin.py
+++ b/src/lightning/app/plugin/plugin.py
@@ -81,6 +81,7 @@ def run_job(self, name: str, app_entrypoint: str, env_vars: Dict[str, str] = {})
         os.environ["LIGHTNING_DISPATCHED"] = "1"
 
         url = runtime.cloudspace_dispatch(
+            organization_id=self.organization_id,
             project_id=self.project_id,
             cloudspace_id=self.cloudspace_id,
             name=name,
@@ -91,10 +92,12 @@ def run_job(self, name: str, app_entrypoint: str, env_vars: Dict[str, str] = {})
 
     def _setup(
         self,
+        organization_id: str,
         project_id: str,
         cloudspace_id: str,
         cluster_id: str,
     ) -> None:
+        self.organization_id = organization_id
         self.project_id = project_id
         self.cloudspace_id = cloudspace_id
         self.cluster_id = cluster_id
@@ -103,6 +106,7 @@ def _setup(
 class _Run(BaseModel):
     plugin_entrypoint: str
     source_code_url: str
+    organization_id: str
     project_id: str
     cloudspace_id: str
     cluster_id: str
@@ -159,6 +163,7 @@ def _run_plugin(run: _Run) -> Dict[str, Any]:
         # Setup and run the plugin
         try:
             plugin._setup(
+                organization_id=run.organization_id,
                 project_id=run.project_id,
                 cloudspace_id=run.cloudspace_id,
                 cluster_id=run.cluster_id,
diff --git a/src/lightning/app/runners/cloud.py b/src/lightning/app/runners/cloud.py
index 6ace6929f13ce..e385a377f507f 100644
--- a/src/lightning/app/runners/cloud.py
+++ b/src/lightning/app/runners/cloud.py
@@ -196,6 +196,7 @@ def cloudspace_dispatch(
         cloudspace_id: str,
         name: str,
         cluster_id: str,
+        organization_id: Optional[str] = None,
     ) -> str:
         """Slim dispatch for creating runs from a cloudspace. This dispatch avoids resolution of some properties
         such as the project and cluster IDs that are instead passed directly.
@@ -218,7 +219,7 @@ def cloudspace_dispatch(
         # Resolution
         root = self._resolve_root()
         repo = self._resolve_repo(root)
-        project = self._resolve_project(project_id=project_id)
+        project = self._resolve_project(project_id=project_id, organization_id=organization_id)
         existing_instances = self._resolve_run_instances_by_name(project_id, name)
         name = self._resolve_run_name(name, existing_instances)
         cloudspace = self._resolve_cloudspace(project_id, cloudspace_id)
@@ -470,9 +471,9 @@ def _resolve_repo(
 
         return LocalSourceCodeDir(path=root, ignore_functions=ignore_functions)
 
-    def _resolve_project(self, project_id: Optional[str] = None) -> V1Membership:
+    def _resolve_project(self, project_id: Optional[str] = None, organization_id: Optional[str] = None) -> V1Membership:
         """Determine the project to run on, choosing a default if multiple projects are found."""
-        return _get_project(self.backend.client, project_id=project_id)
+        return _get_project(self.backend.client, project_id=project_id, organization_id=organization_id)
 
     def _resolve_existing_cloudspaces(self, project_id: str, cloudspace_name: str) -> List[V1CloudSpace]:
         """Lists all the cloudspaces with a name matching the provided cloudspace name."""