run_job.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
SFMC NBA Orchestrator (zero-arg; auto-train on first run from 3 DEs)
"""

import io
import os
import re
import sys
import stat
import glob
import shutil
import traceback
import subprocess
from datetime import datetime, timezone
import paramiko

# ======================= CONFIG =======================

SFTP_HOST            = os.getenv("SFTP_HOST", "ftp.marketingcloudapps.com")
SFTP_PORT            = int(os.getenv("SFTP_PORT", "22"))
SFTP_USER            = os.getenv("SFTP_USER")
SFTP_PASSWORD        = os.getenv("SFTP_PASSWORD")
SFTP_PKEY            = os.getenv("SFTP_PKEY")

REMOTE_INBOUND_DIR   = os.getenv("SFTP_INBOUND_PATH",  "Export/sfmc_nba")
REMOTE_OUTBOUND_DIR  = os.getenv("SFTP_OUTBOUND_PATH", "Import/sfmc_nba")

REMOTE_ACTIONS_CATALOG_PATH = os.getenv("REMOTE_ACTIONS_CATALOG_PATH", "Config/actions_catalog.csv")
REMOTE_MODEL_DIR            = os.getenv("REMOTE_MODEL_DIR",            "Models/sfmc_nba")

EMAIL_STEM = os.getenv("EMAIL_STEM", "EmailEngagement")
CUST_STEM  = os.getenv("CUST_STEM",  "CustomerActivity")
PROD_STEM  = os.getenv("PROD_STEM",  "ProductCatalog")

BASE_DIR   = os.getenv("BASE_DIR", os.getcwd())
IN_DIR     = os.getenv("IN_DIR",  os.path.join(BASE_DIR, "in"))
OUT_DIR    = os.getenv("OUT_DIR", os.path.join(BASE_DIR, "out"))
ARTIFACTS_DIR = os.getenv("MODEL_ARTIFACTS_LOCAL", os.path.join(BASE_DIR, "artifacts"))

os.makedirs(IN_DIR, exist_ok=True)
os.makedirs(OUT_DIR, exist_ok=True)
os.makedirs(ARTIFACTS_DIR, exist_ok=True)

LOCAL_EMAIL_CSV = os.path.join(IN_DIR, "EmailEngagement_DE.csv")
LOCAL_CUST_CSV  = os.path.join(IN_DIR, "CustomerActivity_DE.csv")
LOCAL_PROD_CSV  = os.path.join(IN_DIR, "ProductCatalog_DE.csv")
LOCAL_CATALOG   = os.path.join(IN_DIR, os.path.basename(REMOTE_ACTIONS_CATALOG_PATH))

TONE = os.getenv("TONE")
NBA_PIPELINE_SCRIPT_PATH = os.getenv("NBA_PIPELINE_SCRIPT_PATH", os.path.join(BASE_DIR, "nba_pipeline.py"))

# ======================= SFTP HELPERS =======================

def _get_sftp_client():
    if not SFTP_USER:
        raise SystemExit("❌ SFTP_USER missing.")
    transport = paramiko.Transport((SFTP_HOST, SFTP_PORT))
    if SFTP_PKEY:
        from paramiko import RSAKey, Ed25519Key, ECDSAKey, DSSKey
        pkey = None
        for cls in (RSAKey, Ed25519Key, ECDSAKey, DSSKey):
            try:
                pkey = cls.from_private_key(io.StringIO(SFTP_PKEY))
                break
            except Exception:
                continue
        if not pkey:
            raise RuntimeError("Invalid SFTP_PKEY.")
        transport.connect(username=SFTP_USER, pkey=pkey)
    else:
        if not SFTP_PASSWORD:
            raise SystemExit("Missing SFTP_PASSWORD or SFTP_PKEY")
        transport.connect(username=SFTP_USER, password=SFTP_PASSWORD)
    return paramiko.SFTPClient.from_transport(transport)

def _ensure_remote_dir(sftp, remote_dir):
    parts = [p for p in remote_dir.replace("\\","/").split("/") if p]
    cur = ""
    for p in parts:
        cur = f"{cur}/{p}" if cur else p
        try: sftp.listdir(cur)
        except IOError:
            try: sftp.mkdir(cur)
            except Exception: pass

def _listdir_attrs(sftp, remote_dir):
    return sftp.listdir_attr(remote_dir)

def _latest_by_stem(sftp, remote_dir, stem):
    regex = re.compile(rf"^\d{{4}}-\d{{2}}-\d{{2}}-{re.escape(stem)}\.csv$", re.IGNORECASE)
    files = _listdir_attrs(sftp, remote_dir)
    matches = [(f.filename, f.st_mtime) for f in files if regex.match((f.filename or ""))]
    if not matches:
        raise FileNotFoundError(f"No files matching '{stem}' in {remote_dir}")
    return max(matches, key=lambda x: x[1])[0]

def _sftp_download(sftp, remote_dir, filename, local_path):
    remote = f"{remote_dir.rstrip('/')}/{filename}"
    print(f"⬇️  {remote} → {local_path}")
    sftp.get(remote, local_path)

def _sftp_download_exact(sftp, remote_path, local_path):
    rp = remote_path.replace("\\","/")
    rdir, rfile = os.path.split(rp)
    try:
        for attr in sftp.listdir_attr(rdir):
            if attr.filename == rfile:
                os.makedirs(os.path.dirname(local_path), exist_ok=True)
                print(f"⬇️  {rp} → {local_path}")
                sftp.get(rp, local_path)
                return True
        print(f"ℹ️  Remote file not found: {remote_path}")
        return False
    except Exception as e:
        print(f"⚠️  Error downloading {remote_path}: {e}")
        return False

def _sftp_download_tree(sftp, remote_dir, local_dir):
    os.makedirs(local_dir, exist_ok=True)
    for entry in sftp.listdir_attr(remote_dir):
        rpath = f"{remote_dir.rstrip('/')}/{entry.filename}"
        lpath = os.path.join(local_dir, entry.filename)
        if stat.S_ISDIR(entry.st_mode):
            _sftp_download_tree(sftp, rpath, lpath)
        else:
            print(f"⬇️  {rpath} → {lpath}")
            sftp.get(rpath, lpath)

def _sftp_upload(sftp, local_path, remote_dir):
    fname = os.path.basename(local_path)
    remote = f"{remote_dir.rstrip('/')}/{fname}"
    print(f"⬆️  {local_path} → {remote}")
    sftp.put(local_path, remote)

def _is_abs_existing(path):
    return bool(path) and os.path.isabs(path) and os.path.exists(path)

# ======================= WORKFLOW =======================

def resolve_inputs_and_resources():
    print(f"Connecting to SFMC SFTP ({SFTP_HOST})...")
    with _get_sftp_client() as sftp:
        print("✅ Connected.")

        email_remote = _latest_by_stem(sftp, REMOTE_INBOUND_DIR, EMAIL_STEM)
        cust_remote  = _latest_by_stem(sftp, REMOTE_INBOUND_DIR, CUST_STEM)
        prod_remote  = _latest_by_stem(sftp, REMOTE_INBOUND_DIR, PROD_STEM)

        _sftp_download(sftp, REMOTE_INBOUND_DIR, email_remote, LOCAL_EMAIL_CSV)
        _sftp_download(sftp, REMOTE_INBOUND_DIR, cust_remote,  LOCAL_CUST_CSV)
        _sftp_download(sftp, REMOTE_INBOUND_DIR, prod_remote,  LOCAL_PROD_CSV)

        if os.path.exists(LOCAL_CATALOG):
            print(f"✅ Using local actions catalog: {LOCAL_CATALOG}")
        elif _is_abs_existing(REMOTE_ACTIONS_CATALOG_PATH) and os.path.isfile(REMOTE_ACTIONS_CATALOG_PATH):
            shutil.copy2(REMOTE_ACTIONS_CATALOG_PATH, LOCAL_CATALOG)
            print(f"📄 Copied local catalog: {REMOTE_ACTIONS_CATALOG_PATH} → {LOCAL_CATALOG}")
        else:
            ok = _sftp_download_exact(sftp, REMOTE_ACTIONS_CATALOG_PATH, LOCAL_CATALOG)
            if not ok:
                raise SystemExit(f"❌ Required actions catalog not found: {REMOTE_ACTIONS_CATALOG_PATH}")

        if not os.listdir(ARTIFACTS_DIR):
            if _is_abs_existing(REMOTE_MODEL_DIR) and os.path.isdir(REMOTE_MODEL_DIR):
                shutil.copytree(REMOTE_MODEL_DIR, ARTIFACTS_DIR, dirs_exist_ok=True)
                print(f"📁 Seeded artifacts from {REMOTE_MODEL_DIR}")
            else:
                # optional: fetch any seed models; harmless if folder absent remotely
                try:
                    _sftp_download_tree(sftp, REMOTE_MODEL_DIR, ARTIFACTS_DIR)
                except Exception:
                    pass

    return LOCAL_CATALOG

def _has_any_model_artifacts(model_dir: str) -> bool:
    if not os.path.isdir(model_dir):
        return False
    for _, _, files in os.walk(model_dir):
        if any(f.startswith("model_") for f in files) or {"encoder.pkl","features.json"}.issubset(set(files)):
            return True
    return False

def auto_train_if_needed(actions_catalog_path: str):
    if _has_any_model_artifacts(ARTIFACTS_DIR):
        print("✅ Model artifacts detected; skipping auto-train.")
        return
    # Train once WITHOUT send log (pipeline will synthesize from 3 DEs)
    cmd = [
        sys.executable, NBA_PIPELINE_SCRIPT_PATH,
        "train",
        "--customer-activity", LOCAL_CUST_CSV,
        "--email-engagement",  LOCAL_EMAIL_CSV,
        "--product-catalog",   LOCAL_PROD_CSV,
        "--actions-catalog",   actions_catalog_path,
        "--model-dir",         ARTIFACTS_DIR,
    ]
    print(f"🧪 No artifacts found → Auto-training from 3 DEs:\n   {' '.join(cmd)}")
    proc = subprocess.run(cmd, capture_output=True, text=True)
    print(proc.stdout)
    if proc.returncode != 0:
        print(proc.stderr, file=sys.stderr)
        raise SystemExit(f"❌ Auto-train failed with exit code {proc.returncode}")

def run_scoring(actions_catalog_path: str):
    cmd = [
        sys.executable, NBA_PIPELINE_SCRIPT_PATH,
        "score",
        "--customer-activity", LOCAL_CUST_CSV,
        "--email-engagement",  LOCAL_EMAIL_CSV,
        "--product-catalog",   LOCAL_PROD_CSV,
        "--actions-catalog",   actions_catalog_path,
        "--model-dir",         ARTIFACTS_DIR,
        "--out-dir",           OUT_DIR,
    ]
    if TONE:
        cmd += ["--tone", TONE]

    print(f"▶️  Running pipeline: {' '.join(cmd)}")
    proc = subprocess.run(cmd, capture_output=True, text=True)
    print(proc.stdout)
    if proc.returncode != 0:
        print(proc.stderr, file=sys.stderr)
        raise SystemExit(f"❌ nba_pipeline.py failed with exit code {proc.returncode}")

def upload_outputs():
    print("\n📤 Uploading results to SFMC Import folder…")
    with _get_sftp_client() as sftp:
        _ensure_remote_dir(sftp, REMOTE_OUTBOUND_DIR)
        files = sorted(glob.glob(os.path.join(OUT_DIR, "*.csv")))
        if not files:
            raise RuntimeError("No output files found in OUT_DIR to upload.")
        for f in files:
            _sftp_upload(sftp, f, REMOTE_OUTBOUND_DIR)
    print("✅ Upload complete.")

def main():
    print(f"🗓  Started {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S %Z')}")
    catalog = resolve_inputs_and_resources()
    print("\n🔍 Checking model artifacts...")
    auto_train_if_needed(catalog)
    print("\n🧠 Scoring…")
    run_scoring(catalog)
    upload_outputs()
    print("\n🎉 Done.")

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n❌ Error: {e}", file=sys.stderr)
        traceback.print_exc()
        sys.exit(1)