It claimed to run the pipeline with no engine involved, and could not: the node bodies it called save and load artifacts, which raise outside a node by design. Each node is now a thin wrapper over a plain function — make_rows, train_curve, score — and __main__ calls those, which is the split the sandbox already demonstrates and the one worth copying. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
"""The part that would be on the GPU."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
|
|
import fluksio
|
|
from fluksio import Port, node
|
|
|
|
|
|
def train_curve(lr: float, epochs: int) -> list[float]:
|
|
"""The loss per epoch, with no Fluksio in it."""
|
|
return [
|
|
math.exp(-lr * epoch * 10) * (1 + 0.05 * (epoch % 3)) / (1 + lr)
|
|
for epoch in range(epochs)
|
|
]
|
|
|
|
|
|
@node(
|
|
requires=["dataset", Port("lr", "float")],
|
|
provides=[
|
|
Port("loss", "float", stream=True),
|
|
Port("weights", "artifact"),
|
|
Port("final_loss", "float"),
|
|
],
|
|
device="gpu",
|
|
device_policy="prefer",
|
|
timeout=600,
|
|
resources={"cpus": 2},
|
|
)
|
|
def fit(dataset, lr, epochs=25):
|
|
"""Train, reporting the loss as it goes.
|
|
|
|
Yielding is the reporting: each one publishes on the `loss` port the
|
|
instant it happens, and the run keeps every value as a series — which is
|
|
why there is no `log_metric()` to call. `device="gpu"` with
|
|
`device_policy="prefer"` sends this to a worker carrying that label when
|
|
one is attached, and runs it here when none is.
|
|
|
|
`resources` is what a real training node needs and the rest of a flow does
|
|
not: two cores held for the length of the call, and a worker started with
|
|
thread limits saying so — otherwise a numerical library sizes itself to
|
|
every core on the machine, and a sweep of these starves the engine. A GPU
|
|
one would say `{"gpus": 1}` and hold the card exclusively.
|
|
"""
|
|
rows = json.loads(open(fluksio.load_artifact(dataset)).read())["rows"]
|
|
loss = 1.0
|
|
for loss in train_curve(lr, epochs):
|
|
yield {"loss": loss}
|
|
weights = json.dumps({"lr": lr, "epochs": epochs, "n": len(rows)}).encode()
|
|
return {
|
|
"weights": fluksio.save_artifact(weights, "weights.json", "application/json"),
|
|
"final_loss": loss,
|
|
}
|