mirror of
https://github.com/hatchet-dev/hatchet.git
synced 2025-12-31 13:49:48 -06:00
8e80faf2d6
* api changes * doc changes * move docs * generated * generate * pkg * backmerge main * revert to main * revert main * race? * remove go tests
55 lines
1.2 KiB
Python
55 lines
1.2 KiB
Python
from hatchet_sdk import Context, EmptyModel, Hatchet, WorkerLabelComparator
|
|
from hatchet_sdk.labels import DesiredWorkerLabel
|
|
|
|
hatchet = Hatchet(debug=True)
|
|
|
|
|
|
# > AffinityWorkflow
|
|
|
|
affinity_worker_workflow = hatchet.workflow(name="AffinityWorkflow")
|
|
|
|
|
|
@affinity_worker_workflow.task(
|
|
desired_worker_labels={
|
|
"model": DesiredWorkerLabel(value="fancy-ai-model-v2", weight=10),
|
|
"memory": DesiredWorkerLabel(
|
|
value=256,
|
|
required=True,
|
|
comparator=WorkerLabelComparator.LESS_THAN,
|
|
),
|
|
},
|
|
)
|
|
|
|
|
|
|
|
# > AffinityTask
|
|
async def step(input: EmptyModel, ctx: Context) -> dict[str, str | None]:
|
|
if ctx.worker.labels().get("model") != "fancy-ai-model-v2":
|
|
ctx.worker.upsert_labels({"model": "unset"})
|
|
# DO WORK TO EVICT OLD MODEL / LOAD NEW MODEL
|
|
ctx.worker.upsert_labels({"model": "fancy-ai-model-v2"})
|
|
|
|
return {"worker": ctx.worker.id()}
|
|
|
|
|
|
|
|
|
|
def main() -> None:
|
|
|
|
# > AffinityWorker
|
|
worker = hatchet.worker(
|
|
"affinity-worker",
|
|
slots=10,
|
|
labels={
|
|
"model": "fancy-ai-model-v2",
|
|
"memory": 512,
|
|
},
|
|
workflows=[affinity_worker_workflow],
|
|
)
|
|
worker.start()
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|