feat(plugin): store source app (#17892)

requirements/app/base.txt

@@ -1,4 +1,4 @@
-lightning-cloud >=0.5.34
+lightning-cloud >=0.5.37
 packaging
 typing-extensions >=4.0.0, <=4.4.0
 deepdiff >=5.7.0, <6.3.1

src/lightning/app/plugin/plugin.py

@@ -42,6 +42,7 @@ def __init__(self) -> None:
  self.project_id = ""
  self.cloudspace_id = ""
  self.cluster_id = ""
+ self.source_app = ""
 
  def run(self, *args: str, **kwargs: str) -> Optional[List[_Action]]:
  """Override with the logic to execute on the cloudspace."""
@@ -85,6 +86,7 @@ def run_job(self, name: str, app_entrypoint: str, env_vars: Dict[str, str] = {})
  cloudspace_id=self.cloudspace_id,
  name=name,
  cluster_id=self.cluster_id,
+ source_app=self.source_app,
  )
  # Return a relative URL so it can be used with the NavigateTo action.
  return url.replace(constants.get_lightning_cloud_url(), "")
@@ -94,7 +96,9 @@ def _setup(
  project_id: str,
  cloudspace_id: str,
  cluster_id: str,
+ source_app: str,
  ) -> None:
+ self.source_app = source_app
  self.project_id = project_id
  self.cloudspace_id = cloudspace_id
  self.cluster_id = cluster_id
@@ -107,6 +111,7 @@ class _Run(BaseModel):
  cloudspace_id: str
  cluster_id: str
  plugin_arguments: Dict[str, str]
+ source_app: str
 
 
 def _run_plugin(run: _Run) -> Dict[str, Any]:
@@ -162,6 +167,7 @@ def _run_plugin(run: _Run) -> Dict[str, Any]:
  project_id=run.project_id,
  cloudspace_id=run.cloudspace_id,
  cluster_id=run.cluster_id,
+ source_app=run.source_app,
  )
  actions = plugin.run(**run.plugin_arguments) or []
  return {"actions": [action.to_spec().to_dict() for action in actions]}

src/lightning/app/runners/cloud.py

@@ -196,6 +196,7 @@ def cloudspace_dispatch(
  cloudspace_id: str,
  name: str,
  cluster_id: str,
+ source_app: 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.
@@ -260,6 +261,7 @@ def cloudspace_dispatch(
  V1LightningappInstanceState.RUNNING,
  queue_server_type,
  env_vars,
+ source_app=source_app,
  )
 
  return self._get_app_url(project, run_instance, "logs" if run.is_headless else "web-ui")
@@ -978,6 +980,7 @@ def _api_create_run_instance(
  queue_server_type: Optional[V1QueueServerType] = None,
  env_vars: Optional[List[V1EnvVar]] = None,
  auth: Optional[V1LightningAuth] = None,
+ source_app: Optional[str] = None,
  ) -> Externalv1LightningappInstance:
  """Create a new instance of the given run with the given specification."""
  return self.backend.client.cloud_space_service_create_lightning_run_instance(
@@ -991,6 +994,7 @@ def _api_create_run_instance(
  queue_server_type=queue_server_type,
  env=env_vars,
  auth=auth,
+ source_app=source_app,
  ),
  )
 

tests/tests_app/plugin/test_plugin.py

@@ -84,6 +84,7 @@ def run(self):
  cloudspace_id="any",
  cluster_id="any",
  plugin_arguments={},
+ source_app="any",
  ),
  "Error downloading plugin source:",
  None,
@@ -97,6 +98,7 @@ def run(self):
  cloudspace_id="any",
  cluster_id="any",
  plugin_arguments={},
+ source_app="any",
  ),
  "Error extracting plugin source:",
  None,
@@ -110,6 +112,7 @@ def run(self):
  cloudspace_id="any",
  cluster_id="any",
  plugin_arguments={},
+ source_app="any",
  ),
  "Error loading plugin:",
  "plugin.py",
@@ -123,6 +126,7 @@ def run(self):
  cloudspace_id="any",
  cluster_id="any",
  plugin_arguments={},
+ source_app="any",
  ),
  "Error running plugin:",
  "plugin.py",
@@ -203,6 +207,7 @@ def test_run_job(mock_requests, mock_cloud_runtime, mock_plugin_server, plugin_s
  cloudspace_id="test_cloudspace_id",
  cluster_id="test_cluster_id",
  plugin_arguments={"name": "test_name", "entrypoint": "test_entrypoint"},
+ source_app="test_source_app",
  )
 
  mock_app = mock.MagicMock()
@@ -230,6 +235,7 @@ def test_run_job(mock_requests, mock_cloud_runtime, mock_plugin_server, plugin_s
  cloudspace_id=body.cloudspace_id,
  name="test_name",
  cluster_id=body.cluster_id,
+ source_app=body.source_app,
  )