stroblme
1e62e77aed
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.
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,
|
|
}
|