forked from flyteorg/flytekit
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Make MatchableResources configurable through flyte-cli (flyteorg#118)
- Loading branch information
Showing
6 changed files
with
399 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -2,4 +2,4 @@ | |
|
||
import flytekit.plugins | ||
|
||
__version__ = '0.9.0b0' | ||
__version__ = '0.9.0b1' |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,185 @@ | ||
from flyteidl.admin import matchable_resource_pb2 as _matchable_resource | ||
from flytekit.models import common as _common | ||
|
||
|
||
class ClusterResourceAttributes(_common.FlyteIdlEntity): | ||
|
||
def __init__(self, attributes): | ||
""" | ||
Custom resource attributes which will be applied in cluster resource creation (e.g. quotas). | ||
Dict keys are the *case-sensitive* names of variables in templatized resource files. | ||
Dict values should be the custom values which get substituted during resource creation. | ||
:param dict[Text, Text] attributes: Applied in cluster resource creation (e.g. quotas). | ||
""" | ||
self._attributes = attributes | ||
|
||
@property | ||
def attributes(self): | ||
""" | ||
Custom resource attributes which will be applied in cluster resource management | ||
:rtype: dict[Text, Text] | ||
""" | ||
return self._attributes | ||
|
||
def to_flyte_idl(self): | ||
""" | ||
:rtype: flyteidl.admin.matchable_resource_pb2.ClusterResourceAttributes | ||
""" | ||
return _matchable_resource.ClusterResourceAttributes( | ||
attributes=self.attributes, | ||
) | ||
|
||
@classmethod | ||
def from_flyte_idl(cls, pb2_object): | ||
""" | ||
:param flyteidl.admin.matchable_resource_pb2.ClusterResourceAttributes pb2_object: | ||
:rtype: ClusterResourceAttributes | ||
""" | ||
return cls( | ||
attributes=pb2_object.attributes, | ||
) | ||
|
||
|
||
class ExecutionQueueAttributes(_common.FlyteIdlEntity): | ||
|
||
def __init__(self, tags): | ||
""" | ||
Tags used for assigning execution queues for tasks matching a project, domain and optionally, workflow. | ||
:param list[Text] tags: | ||
""" | ||
self._tags = tags | ||
|
||
@property | ||
def tags(self): | ||
""" | ||
:rtype: list[Text] | ||
""" | ||
return self._tags | ||
|
||
def to_flyte_idl(self): | ||
""" | ||
:rtype: flyteidl.admin.matchable_resource_pb2.ExecutionQueueAttributes | ||
""" | ||
return _matchable_resource.ExecutionQueueAttributes( | ||
tags=self.tags, | ||
) | ||
|
||
@classmethod | ||
def from_flyte_idl(cls, pb2_object): | ||
""" | ||
:param flyteidl.admin.matchable_resource_pb2.ExecutionQueueAttributes pb2_object: | ||
:rtype: ExecutionQueueAttributes | ||
""" | ||
return cls( | ||
tags=pb2_object.tags, | ||
) | ||
|
||
|
||
class ExecutionClusterLabel(_common.FlyteIdlEntity): | ||
|
||
def __init__(self, value): | ||
""" | ||
Label value to determine where the execution will be run | ||
:param Text value: | ||
""" | ||
self._value = value | ||
|
||
@property | ||
def value(self): | ||
""" | ||
:rtype: Text | ||
""" | ||
return self._value | ||
|
||
def to_flyte_idl(self): | ||
""" | ||
:rtype: flyteidl.admin.matchable_resource_pb2.ExecutionClusterLabel | ||
""" | ||
return _matchable_resource.ExecutionClusterLabel( | ||
value=self.value, | ||
) | ||
|
||
@classmethod | ||
def from_flyte_idl(cls, pb2_object): | ||
""" | ||
:param flyteidl.admin.matchable_resource_pb2.ExecutionClusterLabel pb2_object: | ||
:rtype: ExecutionClusterLabel | ||
""" | ||
return cls( | ||
value=pb2_object.value, | ||
) | ||
|
||
|
||
class MatchingAttributes(_common.FlyteIdlEntity): | ||
def __init__(self, cluster_resource_attributes=None, execution_queue_attributes=None, execution_cluster_label=None): | ||
""" | ||
At most one target from cluster_resource_attributes, execution_queue_attributes or execution_cluster_label | ||
can be set. | ||
:param ClusterResourceAttributes cluster_resource_attributes: | ||
:param ExecutionQueueAttributes execution_queue_attributes: | ||
:param ExecutionClusterLabel execution_cluster_label: | ||
""" | ||
if cluster_resource_attributes: | ||
if execution_queue_attributes or execution_cluster_label: | ||
raise ValueError("Only one target can be set") | ||
elif execution_queue_attributes and execution_cluster_label: | ||
raise ValueError("Only one target can be set") | ||
|
||
self._cluster_resource_attributes = cluster_resource_attributes | ||
self._execution_queue_attributes = execution_queue_attributes | ||
self._execution_cluster_label = execution_cluster_label | ||
|
||
@property | ||
def cluster_resource_attributes(self): | ||
""" | ||
Custom resource attributes which will be applied in cluster resource creation (e.g. quotas). | ||
:rtype: ClusterResourceAttributes | ||
""" | ||
return self._cluster_resource_attributes | ||
|
||
@property | ||
def execution_queue_attributes(self): | ||
""" | ||
Tags used for assigning execution queues for tasks. | ||
:rtype: ExecutionQueueAttributes | ||
""" | ||
return self._execution_queue_attributes | ||
|
||
@property | ||
def execution_cluster_label(self): | ||
""" | ||
Label value to determine where the execution will be run. | ||
:rtype: ExecutionClusterLabel | ||
""" | ||
return self._execution_cluster_label | ||
|
||
def to_flyte_idl(self): | ||
""" | ||
:rtype: flyteidl.admin.matchable_resource_pb2.MatchingAttributes | ||
""" | ||
return _matchable_resource.MatchingAttributes( | ||
cluster_resource_attributes=self.cluster_resource_attributes.to_flyte_idl() if | ||
self.cluster_resource_attributes else None, | ||
execution_queue_attributes=self.execution_queue_attributes.to_flyte_idl() if self.execution_queue_attributes | ||
else None, | ||
execution_cluster_label=self.execution_cluster_label.to_flyte_idl() if self.execution_cluster_label | ||
else None, | ||
) | ||
|
||
@classmethod | ||
def from_flyte_idl(cls, pb2_object): | ||
""" | ||
:param flyteidl.admin.matchable_resource_pb2.MatchingAttributes pb2_object: | ||
:rtype: MatchingAttributes | ||
""" | ||
return cls( | ||
cluster_resource_attributes=ClusterResourceAttributes.from_flyte_idl( | ||
pb2_object.cluster_resource_attributes) if pb2_object.HasField("cluster_resource_attributes") else None, | ||
execution_queue_attributes=ExecutionQueueAttributes.from_flyte_idl(pb2_object.execution_queue_attributes) if | ||
pb2_object.HasField("execution_queue_attributes") else None, | ||
execution_cluster_label=ExecutionClusterLabel.from_flyte_idl(pb2_object.execution_cluster_label) if | ||
pb2_object.HasField("execution_cluster_label") else None, | ||
) |
Oops, something went wrong.