schrodinger.seam.transforms.resources module¶

class schrodinger.seam.transforms.resources.LocalExecutionHint¶

Bases: apache_beam.transforms.resources.ResourceHint

urn: Optional[str] = 'seam:local_exec_hint:0.1'¶

encode() → bytes¶

static decode(encoded_value: bytes) → schrodinger.seam.transforms.resources.LocalExecutionHint¶

static get_by_name(name)¶

static get_by_urn(urn)¶

classmethod get_merged_value(outer_value: bytes, inner_value: bytes) → bytes¶: Reconciles values of a hint when the hint specified on a transform is also defined in an outer context, for example on a composite transform, or specified in the transform’s execution environment. Override to specify a custom merging logic.

static is_registered(name)¶

classmethod parse(value: str) → Dict[str, bytes]¶: Describes how to parse the hint. Override to specify a custom parsing logic.

static register_resource_hint(hint_name: str, hint_class: type) → None¶

class schrodinger.seam.transforms.resources.LicenseHint¶

Bases: apache_beam.transforms.resources.ResourceHint

urn: Optional[str] = 'seam:license_hint:0.1'¶

classmethod parse(license_requirements: Dict[int, int]) → Dict[str, bytes]¶: Describes how to parse the hint. Override to specify a custom parsing logic.

static decode(encoded_requirements: bytes) → Dict[int, int]¶

static get_by_name(name)¶

static get_by_urn(urn)¶

classmethod get_merged_value(outer_value: bytes, inner_value: bytes) → bytes¶: Reconciles values of a hint when the hint specified on a transform is also defined in an outer context, for example on a composite transform, or specified in the transform’s execution environment. Override to specify a custom merging logic.

static is_registered(name)¶

static register_resource_hint(hint_name: str, hint_class: type) → None¶

class schrodinger.seam.transforms.resources.NumCpusHint¶

Bases: schrodinger.seam.transforms.resources._IntegerHint

urn: Optional[str] = 'seam:num_cpus_hint:0.1'¶

static get_by_name(name)¶

static get_by_urn(urn)¶

classmethod get_merged_value(outer_value: bytes, inner_value: bytes) → bytes¶: Reconciles values of a hint when the hint specified on a transform is also defined in an outer context, for example on a composite transform, or specified in the transform’s execution environment. Override to specify a custom merging logic.

static is_registered(name)¶

classmethod parse(value: int) → Dict[str, bytes]¶: Describes how to parse the hint. Override to specify a custom parsing logic.

static register_resource_hint(hint_name: str, hint_class: type) → None¶

class schrodinger.seam.transforms.resources.NumGpusHint¶

Bases: schrodinger.seam.transforms.resources._IntegerHint

urn: Optional[str] = 'seam:num_gpus_hint:0.1'¶

static get_by_name(name)¶

static get_by_urn(urn)¶

classmethod get_merged_value(outer_value: bytes, inner_value: bytes) → bytes¶: Reconciles values of a hint when the hint specified on a transform is also defined in an outer context, for example on a composite transform, or specified in the transform’s execution environment. Override to specify a custom merging logic.

static is_registered(name)¶

classmethod parse(value: int) → Dict[str, bytes]¶: Describes how to parse the hint. Override to specify a custom parsing logic.

static register_resource_hint(hint_name: str, hint_class: type) → None¶

class schrodinger.seam.transforms.resources.GpusOptionalHint¶

Bases: apache_beam.transforms.resources.ResourceHint

urn: Optional[str] = 'seam:gpus_optional_hint:0.1'¶

classmethod parse(gpus_optional: Union[bool, str]) → Dict[str, bytes]¶: Describes how to parse the hint. Override to specify a custom parsing logic.

static get_by_name(name)¶

static get_by_urn(urn)¶

classmethod get_merged_value(outer_value: bytes, inner_value: bytes) → bytes¶: Reconciles values of a hint when the hint specified on a transform is also defined in an outer context, for example on a composite transform, or specified in the transform’s execution environment. Override to specify a custom merging logic.

static is_registered(name)¶

static register_resource_hint(hint_name: str, hint_class: type) → None¶

schrodinger.seam.transforms.resources.requires_local_execution(cls_) → Callable¶

A class decorator for PTransforms or DoFns that must be executed locally. This means the decorated class will not be parallelized and executed on remote workers. Instead, it will be executed on the same host that the pipeline is being run on.

DoFn example:

>>> @requires_local_execution
... class MyDoFn(DoFn):
...     def process(self, element):
...         yield element

PTransform definition example:

>>> @requires_local_execution
... class MyPTransform(PTransform):
...     def expand(self, pcoll):
...         return (pcoll | beam.Map(lambda x: x + 1))

schrodinger.seam.transforms.resources.with_license_requirements(license_requirements: Optional[Dict[int, int]] = None) → Callable¶

A class decorator for classses that require a set of Schrödinger licenses to be defined either statically or dynamically. If no license requirements are provided as decorator arguments, it is assumed this decorator is being used to define requirements dynamically.

Static definition example:

>>> @with_license_requirements({license.GLIDE_MAIN: 1}
... class DoGlideFn(DoFn):
...     def process(self, element):
...         return glide.dock(element)

Dynamic definition example:

>>> @with_license_requirements()
... class DoGlideFn(DoFn):
...     def __init__(self, glide_config_file: str, *args, **kwargs):
...         super().__init__(*args, **kwargs)
...         self._glide_config_file = glide_config_file
...
...     def process(self, element):
...         return glide.dock(element)
...
...     def getLicenseHints(self) -> Dict[int, int]
...         reqs = _glide.get_license_requirements(
...             self._glide_config_file)
...         return reqs

Parameters: license_requirements – The license requirements needed to execute code in the wrapped class.

class schrodinger.seam.transforms.resources.LicenseHintRegistrar¶

Bases: apache_beam.pipeline.PipelineVisitor

A pipeline visitor that annotates all nodes with their respective license hints, if any.

Note

Nodes wrapping decorated subclasses of PTransform do not need to be visited because they implement get_resource_hints, which allows beam to register them automatically.

visit_transform(node)¶: Callback for visiting a transform leaf node in the pipeline DAG.

enter_composite_transform(transform_node: apache_beam.pipeline.AppliedPTransform) → None¶: Callback for entering traversal of a composite transform node.

leave_composite_transform(transform_node: apache_beam.pipeline.AppliedPTransform) → None¶: Callback for leaving traversal of a composite transform node.

visit_value(value: apache_beam.pvalue.PValue, producer_node: apache_beam.pipeline.AppliedPTransform) → None¶

Callback for visiting a PValue in the pipeline DAG.

Args:: value: PValue visited (typically a PCollection instance). producer_node: AppliedPTransform object whose transform produced the

pvalue.

class schrodinger.seam.transforms.resources.LocalHostRegistrar(*args, **kwargs)¶

Bases: apache_beam.pipeline.PipelineVisitor

A pipeline visitor that annotates any transforms that must be executed locally.

PTransforms that must be executed locally include:

ParDos whose DoFn were decorated with @requires_local_execution

PTransforms that were decorated with @requires_local_execution

PTransforms that read or write to local files

__init__(*args, **kwargs)¶

enter_composite_transform(transform_node: apache_beam.pipeline.AppliedPTransform) → None¶: Callback for entering traversal of a composite transform node.

visit_transform(node)¶: Callback for visiting a transform leaf node in the pipeline DAG.

leave_composite_transform(transform_node: apache_beam.pipeline.AppliedPTransform) → None¶: Callback for leaving traversal of a composite transform node.

visit_value(value: apache_beam.pvalue.PValue, producer_node: apache_beam.pipeline.AppliedPTransform) → None¶

Callback for visiting a PValue in the pipeline DAG.

Args:: value: PValue visited (typically a PCollection instance). producer_node: AppliedPTransform object whose transform produced the

pvalue.

schrodinger.seam.transforms.resources.with_compute_requirements(num_cpus: Optional[int] = None, num_gpus: Optional[int] = None, min_ram: Optional[str] = None, gpus_optional: Optional[bool] = None) → Callable¶

A class decorator for classes that require a set of compute resources to be defined statically or dynamically. If no compute requirements are provided as decorator arguments, it is assumed this decorator is being used to define requirements dynamically.

Static definition example:

>>> @with_compute_requirements(num_cpus=2, num_gpus=2, min_ram='16MB')
... class DoGlideFn(DoFn):
...     def process(self, element):
...         return glide.dock(element)

Dynamic definition example:

>>> @with_runtime_compute_requirements()
... class DoGlideFn(DoFn):
...     def __init__(self, glide_config_file: str, *args, **kwargs):
...         super().__init__(*args, **kwargs)
...         self._glide_config_file = glide_config_file
...
...     def process(self, element):
...         return glide.dock(element)
...
...     def getComputeHints(self) -> Dict[str, bytes]
...         num_cpus = _glide.get_num_cpus(self._glide_config_file)
...         num_gpus = _glide.get_num_gpus(self._glide_config_file)
...         min_ram = _glide.get_min_ram(self._glide_config_file)
...         gpus_optional = _glide.get_gpus_optional(
...             self._glide_config_file)
...         hints = {}
...         hints.update(NumCpusHint.parse(num_cpus))
...         hints.update(NumGpusHint.parse(num_gpus))
...         hints.update(MinRamHint.parse(min_ram))
...         hints.update(GpusOptionalHint.parse(gpus_optional))
...         return hints

Parameters

num_cpus – The desired number of CPUs to use to run the decorated class.
num_gpus – The desired number of GPUs to use to run the decorated class.
min_ram – The minimum amount of RAM to use to run the decorated class. Example:: ‘1MB’.
gpus_optional – Whether it is optional for the compute host to have GPUs available. If True and the host has GPUs, then the specified num_gpus will be used. If True and the host does not have GPUs, then only CPUs will be used.

class schrodinger.seam.transforms.resources.ComputeHintRegistrar¶

Bases: apache_beam.pipeline.PipelineVisitor

A pipeline visitor that annotates all nodes with their respective compute hints, if any.

Note

Nodes wrapping decorated subclasses of PTransform do not need to be visited because they implement get_resource_hints, which allows beam to register them automatically.

visit_transform(node)¶: Callback for visiting a transform leaf node in the pipeline DAG.

enter_composite_transform(transform_node: apache_beam.pipeline.AppliedPTransform) → None¶: Callback for entering traversal of a composite transform node.

leave_composite_transform(transform_node: apache_beam.pipeline.AppliedPTransform) → None¶: Callback for leaving traversal of a composite transform node.

visit_value(value: apache_beam.pvalue.PValue, producer_node: apache_beam.pipeline.AppliedPTransform) → None¶

Callback for visiting a PValue in the pipeline DAG.

Args:: value: PValue visited (typically a PCollection instance). producer_node: AppliedPTransform object whose transform produced the

pvalue.