schrodinger.seam.transforms.resources module

This module provides decorators for defining PTransforms and DoFns that require local execution and/or Schrödinger license requirements.

class schrodinger.seam.transforms.resources.requires_local_execution(func=None)

Bases: object

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

DoFn example:

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

PTransform definition example:

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

__init__(func=None)

schrodinger.seam.transforms.resources.with_compute_requirements(num_cpus: int = None, num_gpus: int = None, min_ram: str = 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:

>>> from apache_beam import DoFn
>>> @with_compute_requirements(num_gpus=1)
... class DoGlideFn(DoFn):
...     def process(self, element):
...         return element

Dynamic definition example:

>>> from schrodinger.seam.transforms._resources import ComputeRequirements
>>> @with_compute_requirements()
... class DoGlideFn(DoFn):
...
...     def process(self, element):
...         return element
...
...     def getComputeHints(self) -> ComputeRequirements:
...         return ComputeRequirements(
...             num_cpus=None,
...             num_gpus=1,
...             min_ram=None,
...         )

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. Only 1 GPU per DoFn is supported right now.

• min_ram – The minimum amount of RAM to use to run the decorated class. Example:: ‘1MB’.

schrodinger.seam.transforms.resources.requires_min_ram(min_ram: str) → Callable

A decorator for PTransforms or DoFns that require a minimum amount of RAM to run.

Example usage:

>>> import apache_beam as beam
>>> from apache_beam import DoFn
>>> @requires_min_ram("1MB")
... class MyDoFn(DoFn):
...     def process(self, element):
...         yield element
>>>
>>> # For composite transforms where all subtransforms require a minimum amount of RAM
>>> @requires_min_ram("1MB")
... class MyPTransform(beam.PTransform):
...     def expand(self, pcoll):
...         return (pcoll | beam.Map(lambda x: x + 1))

class schrodinger.seam.transforms.resources.requires_gpu(func=None)

Bases: object

A convenience class decorator for DoFns requiring a single GPU.

DoFn example:

>>> from apache_beam import DoFn
>>> @requires_gpu
... class MyDoFn(DoFn):
...     def process(self, element):
...         yield element

__init__(func=None)

schrodinger.seam.transforms.resources.with_license_requirements(license_requirements: 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:

>>> from apache_beam import DoFn
>>> from schrodinger.utils import license
>>> @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.