VictorPlusC commented on a change in pull request #16741:
URL: https://github.com/apache/beam/pull/16741#discussion_r807379457
##########
File path: sdks/python/apache_beam/runners/interactive/interactive_beam.py
##########
@@ -329,6 +359,101 @@ def record(self, pipeline):
return recording_manager.record_pipeline()
+class Clusters():
+ """An interface for users to modify the pipelines that are being run by the
+ Interactive Environment.
+
+ Methods of the Interactive Beam Clusters class can be accessed via:
+ interactive_beam.clusters
+
+ Example of calling the Interactive Beam clusters describe method::
+ interactive_beam.clusters.describe()
+ """
+ def __init__(self) -> None:
+ """Instantiates default values for Dataproc cluster interactions.
+ """
+ self._default_cluster_name = 'interactive-beam-cluster'
+ self._master_urls = bidict()
+ self._dataproc_cluster_managers = {}
+ self._master_urls_to_pipelines = {}
+
+ def describe(self, pipeline: Optional[beam.Pipeline]=None) -> dict:
+ """Returns a description of the cluster associated to the given pipeline.
+
+ If no pipeline is given then this returns a dictionary of descriptions for
+ all pipelines.
+ """
+
+ description = ie.current_env().describe_all_clusters()
+ if pipeline:
+ return description.get(pipeline, None)
+ return description
+
+ @property
+ def default_cluster_name(self) -> str:
+ """The default name to be used when creating Dataproc clusters.
+
+ Defaults to 'interactive-beam-cluster'.
+ """
+ return self._default_cluster_name
+
+ @default_cluster_name.setter
+ def default_cluster_name(self, value: bool) -> None:
+ """Sets the default name to be used when creating Dataproc clusters.
+
+ Defaults to 'interactive-beam-cluster'.
+
+ Example of assigning a default_cluster_name::
+ interactive_beam.clusters.default_cluster_name = 'my-beam-cluster'
+ """
+ self._default_cluster_name = value
+
+ def cleanup(self, pipeline: beam.Pipeline, forcefully=False) -> None:
+ """Attempt to cleanup the Dataproc Cluster corresponding to the given
pipeline.
+
+ If the cluster is not managed by interactive_beam, a corresponding cluster
+ manager is not detected, and deletion is skipped.
+
+ For clusters managed by Interactive Beam, by default, deletion is skipped
+ if any other pipelines are using the cluster.
+
+ Optionally, the cleanup for a cluster managed by Interactive Beam can be
+ forced, by setting the 'forcefully' parameter to True.
+
+ Example usage of default cleanup::
+ interactive_beam.clusters.cleanup(p)
+
+ Example usage of forceful cleanup::
+ interactive_beam.clusters.cleanup(p, forcefully=True)
+ """
+ cluster_manager = ie.current_env().get_dataproc_cluster_manager(pipeline)
+ if cluster_manager:
+ master_url = cluster_manager.master_url
+ if len(self.get_pipelines_using_master_url(master_url)) > 1:
+ if forcefully:
+ _LOGGER.warning(
+ 'Cluster is currently being used by another cluster, but '
+ 'will be forcefully cleaned up.')
+ cluster_manager.cleanup()
+ self._master_urls_to_pipelines[master_url].remove(str(id(pipeline)))
Review comment:
Thanks for the catch!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
