Register stages¶
fasthep-flow
provides a set of useful stages for common tasks and, if
fasthep-stages
is installed, a set of stages for common HEP tasks. However,
you may want to write your own stages. This is easy to do, and we encourage you
to share your stages with the community via fasthep-stages-contrib
, which are
also automatically loaded into the scope. Sometimes you may want to use a stage
that cannot be publicly shared, or you may want to use a stage that is not yet
available in fasthep-stages-contrib
. In this case, you can register your own
stage with fasthep-flow
.
Registering a stage¶
To register a stage, you need to point fasthep-flow
to your callable. This can
be a function or a class implementing the __call__
method.
register:
my_namespace::my_stage: my_module.my_stage
Here, my_namespace
is the namespace for the stage, my_stage
is the name of
the stage, and my_module.my_stage
is the path to the callable. Ideally, the
namespace will match your module name, but this is not required. Now, you can
use my_namespace::my_stage
as a type for a stage, e.g.
stages:
- name: my_stage
type: my_namespace::my_stage
Note
There are protected namespaces, which you cannot use like this. These are anything starting with fasthep
.
Registering a stage that can run on GPU¶
If your stage can run on GPU, you can register it with fasthep-flow
so that it
can be used with the process_on: gpu
option. To do this, you need to register
the GPU version of the stage with fasthep-flow
. This is done by adding a gpu
key to the stage definition:
register:
my_namespace::my_stage:
cpu: my_module.my_stage
gpu: my_module.my_stage_gpu
If you do not have a cpu
version of the stage, you can omit the cpu
key. The
stage will fail if process_on: cpu
is used and there is no cpu
version of
the stage or during fall-back from gpu_if_available
.