A small compute campaign

A workspace holds the durable state of a campaign, a job is one input-specific calculation, and a manager drives each job through its runner steps. collect turns finished jobs into structured outputs that can be inspected or stored. This notebook is the in-notebook miniature of the workflow CLI quickstart, using only a local temporary directory and a mock VASP.

import shlex
import sys
import tempfile
from pathlib import Path

from httk.core import DataRecord
from httk.store import Backend, SqlStore
from httk.workflow import TaskManager, Workspace
from httk.workflow.collecting import collect
from httk.workflow.scaffold import new_job

tmpdir = Path(tempfile.mkdtemp(prefix="httk-campaign-"))
package = tmpdir / "campaign"
package.mkdir()
(package / "httk_workflow.toml").write_text(
    """[workflow]
id = 'notebook.campaign'

[workflow.runner]
entry = 'run'
steps = ['prepare', 'run', 'publish']
initial_step = 'prepare'
data_mode = 'transactional'

[workflow.collect]
file = 'collect.py'

[workflow.outputs.energy]
entry_type = '_httk_records'
role = 'energy'
""",
    encoding="utf-8",
)

# A deterministic stand-in for VASP: it reads POSCAR and writes familiar result files.
(package / "mock_vasp.py").write_text(
    """from pathlib import Path

poscar = Path('POSCAR')
label = poscar.read_text(encoding='utf-8').splitlines()[0].strip()
energy = {'Si-A': -10.0, 'Si-B': -10.25, 'Si-C': -10.5}[label]
Path('OUTCAR').write_text(f'mock VASP energy {energy} eV', encoding='utf-8')
Path('OSZICAR').write_text(str(energy), encoding='utf-8')
Path('CONTCAR').write_text(poscar.read_text(encoding='utf-8'), encoding='utf-8')
Path('vasprun.xml').write_text('<mock-vasp/>', encoding='utf-8')
Path('energy.txt').write_text(str(energy), encoding='utf-8')
""",
    encoding="utf-8",
)

# The runner follows prepare -> run -> publish, as a packaged VASP runner does.
(package / "run").write_text(
    """#!/usr/bin/env python3
import shlex
import shutil
from httk.workflow import Runner

workflow = Runner('notebook.campaign')

@workflow.step
def prepare(a):
    shutil.copyfile(a.payload / 'files/POSCAR', a.workdir / 'POSCAR')
    a.advance('run')

@workflow.step
def run(a):
    result = a.run(shlex.split(str(a.setting('vasp.command'))))
    if result.returncode:
        a.fail('mock.failed', 'the mock VASP failed')
        return
    a.advance('publish')

@workflow.step
def publish(a):
    a.put(a.workdir / 'energy.txt', 'campaign/energy.txt')
    a.succeed()

raise SystemExit(workflow.main())
""",
    encoding="utf-8",
)
(package / "run").chmod(0o755)

(package / "collect.py").write_text(
    """from httk.core import DataRecord

def collect(record):
    energy = float((record.workdir / 'energy.txt').read_text(encoding='utf-8'))
    return {'energy': DataRecord.from_value('urn:notebook:energy', 'energy', energy)}
""",
    encoding="utf-8",
)

221

workspace = Workspace.initialize(tmpdir / 'workspace')
mock_command = f'{sys.executable} {package / "mock_vasp.py"}'
workspace.set_setting('vasp.command', mock_command)

poscars = []
for label in ('Si-A', 'Si-B', 'Si-C'):
    poscar = tmpdir / f'POSCAR.{label}'
    poscar.write_text(
        f'{label}\n1.0\n2 0 0\n0 2 0\n0 0 2\nSi\n1\nDirect\n0 0 0\n',
        encoding='utf-8',
    )
    poscars.append(poscar)

jobs = [
    new_job(workspace, package, files={'POSCAR': poscar}, tag=f'case-{index}')
    for index, poscar in enumerate(poscars)
]
assert len(jobs) == 3

with TaskManager(workspace, heartbeat_interval=0.01, maximum_workers=3) as manager:
    manager.run_until_idle(timeout=30)

states = [workspace.find_marker_by_id(job.job_id).kind for job in jobs]
assert states == ['succeeded'] * 3, states
print('job states:', states)

could not seal succeeded job case-0--75f6228b-cb3b-4471-8294-7fc1f6cbd12e: no seal signing key could be resolved from ['project', 'identity']

could not seal succeeded job case-1--f1653408-7622-47d4-97d7-ca8314948692: no seal signing key could be resolved from ['project', 'identity']

could not seal succeeded job case-2--f00b436e-08e1-4c96-b50b-b15fbf7c2a77: no seal signing key could be resolved from ['project', 'identity']

job states: ['succeeded', 'succeeded', 'succeeded']

collected = list(collect(workspace, allow_job_collector=True))
assert len(collected) == len(jobs)
energies = []
for item in collected:
    energy = item.outputs['energy']
    assert isinstance(energy, DataRecord)
    energies.append(energy.value)
actual_tags = sorted(job.tag for job in jobs)
assert actual_tags == ['case-0', 'case-1', 'case-2']
print('campaign tags:', actual_tags)
print('campaign energies:', sorted(energies))
assert sorted(energies) == [-10.5, -10.25, -10.0]

campaign tags: ['case-0', 'case-1', 'case-2']
campaign energies: [-10.5, -10.25, -10.0]

# Store selected collected outputs; the complete CLI `collect --into` path also stores provenance.
database = Backend.sqlite(tmpdir / 'results.sqlite')
store = SqlStore(database, entry_records={})
for item in collected:
    store.save(item.outputs['energy'])

search = store.searcher()
record = search.variable(DataRecord)
search.add(record.name == 'energy')
stored_values = sorted(row.record.value for row in search.results(record=record))
assert stored_values == [-10.5, -10.25, -10.0]
print('stored energies:', stored_values)
database.dispose()

stored energies: [-10.5, -10.25, -10.0]

At scale, the same workspace/job boundary supports remote transfer and managers, partitioning a large run into campaign workspaces, and precheck before submitting expensive work. The local mock keeps this example fast; a deployment replaces only the workspace setting and runner details.