Tiltfile

# Tilt configuration for running Metaflow on a local Kubernetes stack
#
# Usage:
#   Start the development environment:
#     $ tilt up
#   Stop and clean up:
#     $ tilt down

version_settings(constraint='>=0.22.2')
allow_k8s_contexts('minikube')

# ---------------------------------------------------------------------------
# Version configuration
# ---------------------------------------------------------------------------
ARGO_WORKFLOWS_HELM_CHART_VERSION = os.getenv("ARGO_WORKFLOWS_HELM_CHART_VERSION", "0.45.2")
ARGO_WORKFLOWS_IMAGE_TAG = os.getenv("ARGO_WORKFLOWS_IMAGE_TAG", "v3.6.0")
AIRFLOW_HELM_CHART_VERSION = os.getenv("AIRFLOW_HELM_CHART_VERSION", "1.15.0")
AIRFLOW_IMAGE_TAG = os.getenv("AIRFLOW_IMAGE_TAG", "2.10.4")
ARGO_EVENTS_HELM_CHART_VERSION = os.getenv("ARGO_EVENTS_HELM_CHART_VERSION", "2.4.8")
ARGO_EVENTS_IMAGE_TAG = os.getenv("ARGO_EVENTS_IMAGE_TAG", "v1.9.2")
JOBSET_VERSION = os.getenv("JOBSET_VERSION", "v0.8.2")

# ---------------------------------------------------------------------------
# Component dependency graph
# ---------------------------------------------------------------------------
components = {
    "metadata-service": ["postgresql"],
    "ui": ["postgresql", "minio"],
    "minio": [],
    "postgresql": [],
    "argo-workflows": [],
    "localbatch": ["minio"],
    "ddb-local": [],
    "sfn-local": ["ddb-local"],
    "airflow": ["postgresql"],
    "argo-events": ["argo-workflows"],
    "jobset": [],
}

# ---------------------------------------------------------------------------
# Resolve requested components with dependencies
# ---------------------------------------------------------------------------
services_env = os.getenv("SERVICES", "metadata-service,argo-workflows").strip().lower()

if services_env == "all":
    requested_components = list(components.keys())
else:
    requested_components = [s.strip() for s in services_env.split(",") if s.strip()]

load('ext://helm_resource', 'helm_resource', 'helm_repo')
load('ext://helm_remote', 'helm_remote')

def resolve(component, resolved=None):
    if resolved == None:
        resolved = []
    if component in resolved:
        return resolved
    if component in components:
        for dep in components[component]:
            resolve(dep, resolved)
    resolved.append(component)
    return resolved

valid_components = []
for component in components.keys():
    if component not in valid_components:
        valid_components.append(component)
for deps in components.values():
    for dep in deps:
        if dep not in valid_components:
            valid_components.append(dep)

enabled_components = []
for component in requested_components:
    if component not in valid_components:
        fail("Unknown component: " + component)
    for result in resolve(component):
        if result not in enabled_components:
            enabled_components.append(result)

if config.tilt_subcommand == 'up':
    print("\n📦 Components to install:")
    for component in enabled_components:
        print("• " + component)
        if component in components and components[component]:
            print("  ↳ requires: " + ", ".join(components[component]))

# ---------------------------------------------------------------------------
# Load component setup functions
# ---------------------------------------------------------------------------
load('./tilt/_result.tiltfile', 'new_result')
load('./tilt/minio.tiltfile', 'setup_minio')
load('./tilt/postgresql.tiltfile', 'setup_postgresql')
load('./tilt/argo_workflows.tiltfile', 'setup_argo_workflows')
load('./tilt/metadata_service.tiltfile', 'setup_metadata_service')
load('./tilt/ui.tiltfile', 'setup_ui')
load('./tilt/localbatch.tiltfile', 'setup_localbatch')
load('./tilt/ddb_local.tiltfile', 'setup_ddb_local')
load('./tilt/sfn_local.tiltfile', 'setup_sfn_local')
load('./tilt/airflow.tiltfile', 'setup_airflow')
load('./tilt/argo_events.tiltfile', 'setup_argo_events')
load('./tilt/jobset.tiltfile', 'setup_jobset')

_SETUP = {
    "minio": setup_minio,
    "postgresql": setup_postgresql,
    "argo-workflows": setup_argo_workflows,
    "metadata-service": setup_metadata_service,
    "ui": setup_ui,
    "localbatch": setup_localbatch,
    "ddb-local": setup_ddb_local,
    "sfn-local": setup_sfn_local,
    "airflow": setup_airflow,
    "argo-events": setup_argo_events,
    "jobset": setup_jobset,
}

# ---------------------------------------------------------------------------
# Build context and run each enabled component
# ---------------------------------------------------------------------------
ctx = struct(
    enabled_components=enabled_components,
    components=components,
    versions=struct(
        argo_workflows_chart=ARGO_WORKFLOWS_HELM_CHART_VERSION,
        argo_workflows_image=ARGO_WORKFLOWS_IMAGE_TAG,
        airflow_chart=AIRFLOW_HELM_CHART_VERSION,
        airflow_image=AIRFLOW_IMAGE_TAG,
        argo_events_chart=ARGO_EVENTS_HELM_CHART_VERSION,
        argo_events_image=ARGO_EVENTS_IMAGE_TAG,
        jobset=JOBSET_VERSION,
    ),
)

metaflow_config = {
    "METAFLOW_KUBERNETES_NAMESPACE": "default",
    "METAFLOW_KUBERNETES_JOB_TTL_SECONDS_AFTER_FINISHED": 300,
}
aws_config = ""
shell_env = {}
config_resources = []
k8s_secrets = []

for name in enabled_components:
    if name in _SETUP:
        r = _SETUP[name](ctx)
        metaflow_config.update(r["config"])
        config_resources.extend(r["config_resources"])
        if r["aws_config"]:
            aws_config = r["aws_config"]
        shell_env.update(r["shell_env"])
        k8s_secrets.extend(r["k8s_secrets"])

if k8s_secrets:
    metaflow_config["METAFLOW_KUBERNETES_SECRETS"] = ",".join(k8s_secrets)

# ---------------------------------------------------------------------------
# Write config files
# ---------------------------------------------------------------------------
def write_config_files():
    metaflow_json = encode_json(metaflow_config)
    cmd = '''cat > .devtools/config_local.json <<EOF
%s
EOF
''' % (metaflow_json)
    if aws_config and aws_config.strip():
        cmd += '''cat > .devtools/aws_config <<EOF
%s
EOF
''' % (aws_config.strip())
    if shell_env:
        env_lines = "\n".join(['export %s="%s"' % (k, v) for k, v in shell_env.items()])
        cmd += '''cat > .devtools/env_local <<EOF
%s
EOF
''' % (env_lines)
    return cmd

# Detect node architecture for @conda/@pypi emulation
cmd = '''
ARCH=$(kubectl get nodes -o jsonpath='{.items[0].status.nodeInfo.architecture}')
case "$ARCH" in
  arm64)   echo linux-aarch64 ;;
  amd64)   echo linux-64 ;;
  *)       echo linux-64 ;;
esac
'''
metaflow_config["METAFLOW_KUBERNETES_CONDA_ARCH"] = str(local(cmd)).strip()

local_resource(
    name="generate-configs",
    cmd=write_config_files(),
    resource_deps=config_resources,
)

# ---------------------------------------------------------------------------
# Python dev dependencies
# ---------------------------------------------------------------------------
local_resource(
    name="install-dev-deps",
    cmd="python -m pip install -r requirements-devstack.txt",
    deps=["requirements-devstack.txt"],
    labels=["setup"],
)