jd-worker 1.0.0__py3-none-any.whl

jd/__init__.py

@@ -0,0 +1,14 @@
+"""jd — Job Distributor client package."""
+__version__ = "1.0.0"
+
+from jd.files import jd_get_last_checkpoint, jd_update_checkpoint, jd_upload
+from jd.paths import jd_exp_dir, jd_job_dir, jd_worker_workspace
+
+__all__ = [
+    "jd_upload",
+    "jd_update_checkpoint",
+    "jd_get_last_checkpoint",
+    "jd_job_dir",
+    "jd_worker_workspace",
+    "jd_exp_dir",
+]

jd/files.py

@@ -0,0 +1,236 @@
+"""
+jd.files — worker-side file/checkpoint helpers
+================================================
+These functions are designed to be called **from inside your entry script**
+while a job is running.  jd_worker automatically injects the required
+context (server URL, job ID, experiment ID) as environment variables, so
+you normally do not need to pass those arguments explicitly.
+
+Quick-start
+-----------
+    from jd import jd_upload, jd_job_dir, jd_update_checkpoint, jd_get_last_checkpoint
+
+    out = jd_job_dir() / "metrics.csv"
+    # … write to out …
+    jd_upload(str(out))
+
+    jd_update_checkpoint({"epoch": 5, "state_dict": model.state_dict()})
+    ckpt = jd_get_last_checkpoint()
+    if ckpt:
+        model.load_state_dict(ckpt["state_dict"])
+
+Environment variables (set automatically by jd_worker)
+-------------------------------------------------------
+    JD_SERVER                   — job server base URL, e.g. http://10.0.0.1:8000
+    JD_JOB_ID                   — integer job ID assigned by the server
+    JD_EXP_ID                   — experiment identifier
+    JD_WORKER_JOB_DIR           — absolute …/<parent>/jd_data/<expId>/<job_id>/
+                                  (same as ``--base_path``); prefer ``jd_job_dir()``
+    JD_WORKER_WORKSPACE_ROOT    — absolute ``<parent>/jd_data`` (same as ``jd_worker_workspace()``)
+
+You can override server/job_id via the upload/checkpoint function keyword arguments.
+"""
+
+import io
+import logging
+import os
+import pickle
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+_MAX_BYTES = 100 * 1024 * 1024   # 100 MB
+_TIMEOUT   = 180                  # seconds for upload/download HTTP calls
+
+
+# ── Internal helpers ──────────────────────────────────────────────────────────
+
+def _ctx(job_id, server):
+    """Resolve job_id and server from kwargs → env vars → error."""
+    if job_id is None:
+        raw = os.environ.get("JD_JOB_ID", "")
+        if not raw:
+            raise RuntimeError(
+                "job_id not provided and JD_JOB_ID environment variable is not set. "
+                "Make sure you are calling this function from inside an entry script "
+                "launched by jd_worker."
+            )
+        job_id = int(raw)
+    if server is None:
+        server = os.environ.get("JD_SERVER", "")
+        if not server:
+            raise RuntimeError(
+                "server not provided and JD_SERVER environment variable is not set."
+            )
+    return job_id, server.rstrip("/")
+
+
+def _auth_headers() -> dict:
+    """Return Authorization header if JD_WORKER_TOKEN is set (Hub mode)."""
+    token = os.environ.get("JD_WORKER_TOKEN", "").strip()
+    if token:
+        return {"Authorization": f"Bearer {token}"}
+    return {}
+
+
+def _check_size(data: bytes, label: str) -> None:
+    if len(data) > _MAX_BYTES:
+        mb = len(data) / (1024 ** 2)
+        raise ValueError(
+            f"{label} is {mb:.1f} MB which exceeds the 100 MB limit."
+        )
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def jd_upload(file_path: str, job_id: int = None, server: str = None) -> dict:
+    """
+    Upload a result file (≤ 100 MB) to the server.
+
+    The file is stored in the experiment's job directory as::
+
+        result_v{N}_{timestamp}.<original_ext>
+
+    where N auto-increments across calls so every upload is preserved.
+
+    Parameters
+    ----------
+    file_path : str
+        Local path to the file you want to upload.
+    job_id : int, optional
+        Defaults to the JD_JOB_ID environment variable.
+    server : str, optional
+        Job server base URL. Defaults to JD_SERVER environment variable.
+
+    Returns
+    -------
+    dict
+        ``{"success": True, "filename": "result_v0_…", "version": 0, "size_bytes": …}``
+    """
+    job_id, server = _ctx(job_id, server)
+
+    if not os.path.isfile(file_path):
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    with open(file_path, "rb") as fh:
+        data = fh.read()
+
+    _check_size(data, f"File '{os.path.basename(file_path)}'")
+
+    original_name = os.path.basename(file_path)
+    logger.info(f"[jd_upload] Uploading '{original_name}' ({len(data)/(1024**2):.2f} MB) …")
+
+    resp = requests.post(
+        f"{server}/upload",
+        data={"job_id": str(job_id)},
+        files={"file": (original_name, io.BytesIO(data), "application/octet-stream")},
+        headers=_auth_headers(),
+        timeout=_TIMEOUT,
+    )
+    resp.raise_for_status()
+    result = resp.json()
+    logger.info(f"[jd_upload] Saved as '{result.get('filename')}' (version {result.get('version')})")
+    return result
+
+
+def jd_update_checkpoint(obj, job_id: int = None, server: str = None) -> dict:
+    """
+    Serialise *obj* with pickle and upload it as a versioned checkpoint.
+
+    The checkpoint is stored in the experiment's job directory as::
+
+        checkpoint_v{N}_{timestamp}.pt
+
+    Each call creates a new version, so previous checkpoints are never
+    overwritten.  The file uses standard Python pickle serialisation and is
+    compatible with PyTorch state-dicts as well as arbitrary Python objects.
+
+    Parameters
+    ----------
+    obj : any
+        The Python object to checkpoint (e.g. a dict containing
+        ``model.state_dict()`` and optimizer state).
+    job_id : int, optional
+    server : str, optional
+
+    Returns
+    -------
+    dict
+        ``{"success": True, "filename": "checkpoint_v0_…", "version": 0, …}``
+    """
+    job_id, server = _ctx(job_id, server)
+
+    logger.info("[jd_update_checkpoint] Serialising checkpoint …")
+    data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
+    _check_size(data, "Checkpoint")
+
+    logger.info(f"[jd_update_checkpoint] Uploading {len(data)/(1024**2):.2f} MB …")
+    resp = requests.post(
+        f"{server}/checkpoint",
+        data={"job_id": str(job_id)},
+        files={"checkpoint": ("checkpoint.pkl", io.BytesIO(data), "application/octet-stream")},
+        headers=_auth_headers(),
+        timeout=_TIMEOUT,
+    )
+    resp.raise_for_status()
+    result = resp.json()
+    logger.info(f"[jd_update_checkpoint] Saved as '{result.get('filename')}' (version {result.get('version')})")
+    return result
+
+
+def jd_get_last_checkpoint(job_id: int = None, server: str = None):
+    """
+    Download the highest-versioned checkpoint for this job and return it as
+    a Python object — **nothing is written to disk**.
+
+    The server sends the raw pickle bytes; this function deserialises them
+    directly from an in-memory buffer so your script can resume immediately.
+
+    Parameters
+    ----------
+    job_id : int, optional
+    server : str, optional
+
+    Returns
+    -------
+    object
+        The Python object that was passed to ``jd_update_checkpoint``, or
+        ``None`` if no checkpoint exists yet for this job.
+
+    Examples
+    --------
+    ::
+
+        ckpt = jd_get_last_checkpoint()
+        if ckpt is not None:
+            model.load_state_dict(ckpt["model"])
+            start_epoch = ckpt["epoch"] + 1
+    """
+    job_id, server = _ctx(job_id, server)
+
+    logger.info(f"[jd_get_last_checkpoint] Fetching latest checkpoint for job {job_id} …")
+    resp = requests.get(
+        f"{server}/checkpoint/latest",
+        params={"job_id": job_id},
+        headers=_auth_headers(),
+        timeout=_TIMEOUT,
+        stream=True,
+    )
+
+    if resp.status_code == 404:
+        logger.info("[jd_get_last_checkpoint] No checkpoint found.")
+        return None
+
+    resp.raise_for_status()
+
+    # Read directly into a BytesIO buffer — no temp file, stays in memory
+    buf = io.BytesIO()
+    for chunk in resp.iter_content(chunk_size=1024 * 256):
+        buf.write(chunk)
+
+    logger.info(f"[jd_get_last_checkpoint] Received {buf.tell()/(1024**2):.2f} MB. Deserialising …")
+    buf.seek(0)
+    obj = pickle.load(buf)
+    logger.info("[jd_get_last_checkpoint] Checkpoint ready.")
+    return obj

jd/paths.py

@@ -0,0 +1,58 @@
+"""
+Paths injected by jd_worker for entry scripts
+==============================================
+Layout on the worker::
+
+    <parent>/jd_data/<expId>/<job_id>/   ← save job outputs here (jd_job_dir)
+
+``parent`` is ``JD_WORKSPACE_PATH`` if set, otherwise your home directory.
+``jd_worker_workspace()`` returns the resolved ``…/jd_data`` directory.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+
+def jd_job_dir() -> Path:
+    """
+    Absolute directory for **this job's** local files — same as ``--base_path``.
+
+    Equivalent to ``JD_WORKER_JOB_DIR``. Prefer this for CSVs, checkpoints on
+    disk, temp files, etc.
+    """
+    raw = os.environ.get("JD_WORKER_JOB_DIR", "").strip()
+    if not raw:
+        raise RuntimeError(
+            "JD_WORKER_JOB_DIR is not set. Run your script via jd_worker "
+            "(or set JD_WORKER_JOB_DIR to your job sandbox directory)."
+        )
+    return Path(raw).expanduser().resolve()
+
+
+def jd_worker_workspace() -> Path:
+    """
+    Absolute **jd_data root**: ``<parent>/jd_data`` where ``parent`` is
+    ``JD_WORKSPACE_PATH`` or ``~``.
+
+    Set by ``JD_WORKER_WORKSPACE_ROOT`` for the entry-script process.
+    If only ``JD_WORKER_JOB_DIR`` is present, derives
+    ``…/jd_data/<expId>/<job_id>`` → parent.parent as that root.
+    """
+    raw = os.environ.get("JD_WORKER_WORKSPACE_ROOT", "").strip()
+    if raw:
+        return Path(raw).expanduser().resolve()
+    job_raw = os.environ.get("JD_WORKER_JOB_DIR", "").strip()
+    if job_raw:
+        p = Path(job_raw).expanduser().resolve()
+        return p.parent.parent
+    raise RuntimeError(
+        "Cannot resolve worker workspace: set JD_WORKER_WORKSPACE_ROOT or "
+        "run inside jd_worker so JD_WORKER_JOB_DIR is set."
+    )
+
+
+def jd_exp_dir() -> Path:
+    """Absolute ``<parent>/jd_data/<expId>/`` for this run (parent of ``jd_job_dir()``)."""
+    return jd_job_dir().parent

jd/worker.py

@@ -0,0 +1,621 @@
+"""
+jd_worker — Job Distributor Worker CLI
+=======================================
+Requests jobs from a jd server, runs the entry script with the job's
+parameters as CLI flags, sends a heartbeat ping every 57 seconds, and
+reports DONE or ABORTED when the script finishes.
+
+Usage
+-----
+    jd_worker expId=<id> entry_script=<script.py> [options]
+
+Required
+--------
+    expId=<id>              Experiment identifier (must match the server).
+    entry_script=<path>     Python script to run for each job.
+
+Optional
+--------
+    (no workspace CLI)      Local data lives under ``<parent>/jd_data/<expId>/<job_id>/``.
+                            ``parent`` is ``JD_WORKSPACE_PATH`` if set, otherwise
+                            ``~`` (your home directory). So the default root is
+                            ``~/jd_data/``. Passed to the entry script as
+                            ``--base_path``, ``JD_WORKER_JOB_DIR``, and
+                            ``JD_WORKER_WORKSPACE_ROOT`` (= ``.../jd_data``).
+
+                            Note: ``JD_WORKSPACE_PATH`` here is **worker-only**
+                            (parent of ``jd_data``). The job server uses the same
+                            variable name for its own layout — avoid exporting one
+                            value for both in the same shell.
+    server=<url>            Job server base URL or host (default: http://localhost,
+                            env: JD_SERVER). If you omit ``http://`` or ``https://``,
+                            ``http://`` is assumed.
+    port=<N>                Port if not included in server URL
+                            (default: 5000, env: JD_PORT)
+    log_dir=<path>          If set, logs go under <log_dir>/<expId>/; otherwise
+                            under <jd_data>/<expId>/jd_worker_logs/
+                            (env: JD_LOG_DIR)
+    machine_type=<type>     Label for this machine in the dashboard
+                            (default: worker, env: JD_MACHINE_TYPE)
+    process_id=<N>          Numeric ID when running multiple workers on the
+                            same machine  (default: 0)
+    once=true               Exit after completing a single job instead of
+                            looping until no jobs remain.
+
+Examples
+--------
+    jd_worker expId=mnist_tune entry_script=train.py
+
+    # Optional: put jd_data under /scratch/jd_data (parent=/scratch)
+    JD_WORKSPACE_PATH=/scratch jd_worker expId=mnist_tune entry_script=train.py \\
+              server=http://10.0.0.5 port=8000 \\
+              machine_type=gpu_node
+
+    # Run exactly one job:
+    jd_worker expId=mnist_tune entry_script=train.py once=true
+
+Install
+-------
+    pip install -e ./client     # from the repo root
+    # then `jd_worker` is available in whatever env is active
+"""
+
+import logging
+import os
+import platform
+import random
+import signal
+import socket
+import subprocess
+import sys
+import threading
+import time
+from urllib.parse import urlparse, urlunparse
+
+import psutil
+import requests
+
+from jd import __version__
+
+IS_WINDOWS = platform.system() == "Windows"
+PING_INTERVAL = 57  # seconds — intentionally not 60 to avoid racing the idle timeout
+
+# Fixed subdirectory under JD_WORKSPACE_PATH (or home): …/jd_data/<expId>/<job_id>/
+_WORKER_JD_DATA_DIRNAME = "jd_data"
+
+
+# ── Argument parsing ─────────────────────────────────────────────────────────
+
+def _parse_kv(argv: list) -> dict:
+    """Parse key=value positional arguments into a plain dict."""
+    cfg = {}
+    for arg in argv:
+        if '=' in arg:
+            k, v = arg.split('=', 1)
+            cfg[k.strip()] = v.strip()
+        else:
+            cfg[arg.strip()] = 'true'
+    return cfg
+
+
+def _normalize_server_base_url(server_raw: str, port_raw: str) -> str:
+    """
+    Build a base URL that requests/lib can open.
+
+    Host-only values like ``localhost`` must become ``http://localhost:<port>``.
+    Without a scheme, Python requests raises "No connection adapters were found".
+    """
+    s = (server_raw or "").strip().rstrip("/")
+    if not s:
+        s = "http://localhost"
+    if not urlparse(s).scheme:
+        s = "http://" + s.lstrip("/")
+    parsed = urlparse(s)
+    if parsed.port is not None:
+        return s
+    host = parsed.hostname
+    if not host:
+        host = "localhost"
+    netloc = f"{host}:{port_raw}"
+    return urlunparse(
+        (parsed.scheme, netloc, parsed.path or "", "", "", "")
+    ).rstrip("/")
+
+
+def _resolve(cfg: dict) -> dict:
+    """Merge CLI key=value pairs with environment variables and defaults."""
+
+    def get(cli_key, env_key=None, default=None):
+        if cli_key in cfg:
+            return cfg[cli_key]
+        if env_key:
+            val = os.environ.get(env_key, '')
+            if val:
+                return val
+        return default
+
+    server_raw = get('server', 'JD_SERVER', 'http://localhost')
+    port_raw   = str(get('port', 'JD_PORT', '5000')).strip()
+
+    base_url = _normalize_server_base_url(server_raw, port_raw)
+
+    # …/<parent>/jd_data/<expId>/<job_id>/  — parent from env or ~
+    parent = os.environ.get("JD_WORKSPACE_PATH", "").strip()
+    if not parent:
+        parent = os.path.expanduser("~")
+    parent = os.path.abspath(os.path.expanduser(parent))
+    workspace_path = os.path.join(parent, _WORKER_JD_DATA_DIRNAME)
+    os.makedirs(workspace_path, exist_ok=True)
+
+    log_override = None
+    if 'log_dir' in cfg:
+        log_override = os.path.expanduser(cfg['log_dir'].strip())
+    elif os.environ.get('JD_LOG_DIR', '').strip():
+        log_override = os.path.expanduser(os.environ['JD_LOG_DIR'].strip())
+
+    return {
+        'exp_id':           get('expId',        'JD_EXP_ID',       None),
+        'entry_script':     get('entry_script', 'JD_ENTRY_SCRIPT', None),
+        'base_url':         base_url,
+        'workspace_path':   workspace_path,
+        'log_dir_override': log_override,
+        'machine_type':     get('machine_type', 'JD_MACHINE_TYPE', 'worker'),
+        'process_id':       get('process_id',   None,              '0'),
+        'once':             get('once',         'JD_ONCE',         'false').lower() == 'true',
+        # Hub authentication (optional)
+        'hub_url':          get('hub_url',      'JD_HUB_URL',      '').strip().rstrip('/'),
+        'api_key':          get('api_key',      'JD_API_KEY',      '').strip(),
+    }
+
+
+# ── Hub authentication ────────────────────────────────────────────────────────
+
+def _hub_get_worker_token(hub_url: str, api_key: str,
+                           exp_id: str, logger: logging.Logger) -> tuple[str, str] | None:
+    """
+    Obtain a worker JWT and server URL from the Hub.
+
+    Returns (worker_token, server_url) on success, None on failure.
+    The returned server_url overrides any CLI-supplied server address.
+    """
+    endpoint = f"{hub_url}/api/worker/token"
+    try:
+        r = requests.post(
+            endpoint,
+            json={"experiment_name": exp_id},
+            headers={"Authorization": f"Bearer {api_key}"},
+            timeout=30,
+        )
+        if r.status_code == 200:
+            data = r.json()
+            token      = data.get("worker_token", "")
+            server_url = data.get("server_url", "")
+            expires_at = data.get("expires_at", "unknown")
+            logger.info(f"Worker token obtained from Hub (expires: {expires_at})")
+            return token, server_url
+        logger.error(f"Hub token request failed: HTTP {r.status_code} — {r.text[:300]}")
+        return None
+    except Exception as exc:
+        logger.error(f"Hub connection error: {exc}")
+        return None
+
+
+# ── Logging ───────────────────────────────────────────────────────────────────
+
+def _setup_logger(log_dir: str, runner_id: str) -> logging.Logger:
+    os.makedirs(log_dir, exist_ok=True)
+    log_path = os.path.join(log_dir, f"jd_worker_{runner_id}.log")
+
+    logger = logging.getLogger(f'jd_worker.{runner_id}')
+    logger.setLevel(logging.INFO)
+
+    fmt = logging.Formatter("%(asctime)s  %(levelname)-8s  %(message)s")
+    fh  = logging.FileHandler(log_path, encoding='utf-8')
+    fh.setFormatter(fmt)
+    sh  = logging.StreamHandler(sys.stdout)
+    sh.setFormatter(fmt)
+
+    logger.addHandler(fh)
+    logger.addHandler(sh)
+    logger.propagate = False
+    return logger
+
+
+# ── System metrics ────────────────────────────────────────────────────────────
+
+def _collect_metrics(machine_type: str, logger: logging.Logger) -> dict:
+    """Collect a single snapshot of system metrics."""
+    try:
+        cpu_pct          = psutil.cpu_percent(interval=0.1)
+        cpu_phys         = psutil.cpu_count(logical=False) or 1
+        cpu_logi         = psutil.cpu_count(logical=True)  or 1
+
+        try:
+            freq         = psutil.cpu_freq()
+            freq_mhz     = freq.current if freq else 0
+        except (AttributeError, RuntimeError):
+            freq_mhz     = 0
+
+        mem              = psutil.virtual_memory()
+        ram_total_gb     = mem.total     / (1024 ** 3)
+        ram_avail_gb     = mem.available / (1024 ** 3)
+
+        if not IS_WINDOWS:
+            try:
+                la           = os.getloadavg()
+                load1, load5, load15 = la[0], la[1], la[2]
+            except (AttributeError, OSError):
+                load1 = load5 = load15 = 0.0
+            load_per_cpu = load1 / cpu_logi if cpu_logi > 0 else 0.0
+            idle_slots   = max(0, cpu_logi - load1)
+        else:
+            load1 = load5 = load15 = cpu_pct / 100.0 * cpu_logi
+            load_per_cpu = cpu_pct / 100.0
+            idle_slots   = max(0, cpu_logi * (1 - cpu_pct / 100.0))
+
+        try:
+            d0           = psutil.disk_io_counters()
+            time.sleep(0.1)
+            d1           = psutil.disk_io_counters()
+            if d0 and d1:
+                dt           = 0.1
+                ops          = (d1.read_count  - d0.read_count  +
+                                d1.write_count - d0.write_count) / dt
+                bps          = (d1.read_bytes  - d0.read_bytes  +
+                                d1.write_bytes - d0.write_bytes) / dt
+                disk_io_util = max(min(100.0, (ops / 100_000) * 100),
+                                   min(100.0, (bps / 1e9)     * 100))
+            else:
+                disk_io_util = 0.0
+        except (AttributeError, RuntimeError, TypeError):
+            disk_io_util = 0.0
+
+        return {
+            "cpu_util":     round(cpu_pct,        1),
+            "ram_util":     round(mem.percent,     1),
+            "ram_available": round(ram_avail_gb,  15),
+            "ram_total":    round(ram_total_gb,    1),
+            "worker_type":  machine_type,
+            "idle_slots":   int(round(idle_slots)),
+            "load_1min":    round(load1,          10),
+            "load_5min":    round(load5,          10),
+            "load_15min":   round(load15,         10),
+            "load_per_cpu": round(load_per_cpu,   13),
+            "disk_io_util": round(disk_io_util,    2),
+            "cpu_cores":    cpu_phys,
+            "cpu_threads":  cpu_logi,
+            "cpu_freq_mhz": int(round(freq_mhz)) if freq_mhz > 0 else 0,
+        }
+    except Exception as exc:
+        logger.error(f"Metrics collection error: {exc}")
+        return {k: 0 for k in (
+            "cpu_util", "ram_util", "ram_available", "ram_total",
+            "idle_slots", "load_1min", "load_5min", "load_15min",
+            "load_per_cpu", "disk_io_util", "cpu_cores", "cpu_threads",
+            "cpu_freq_mhz",
+        )} | {"worker_type": machine_type}
+
+
+def _averaged_metrics(machine_type: str, logger: logging.Logger,
+                      samples: int = 5, interval: float = 3.0) -> dict:
+    """Collect `samples` snapshots and return their numeric averages."""
+    logger.info(f"Collecting system metrics ({samples} samples × {interval}s)…")
+    snapshots = []
+    for i in range(samples):
+        logger.info(f"  Sample {i+1}/{samples}")
+        snapshots.append(_collect_metrics(machine_type, logger))
+        if i < samples - 1:
+            time.sleep(interval)
+
+    numeric = [
+        "cpu_util", "ram_util", "ram_available", "ram_total",
+        "idle_slots", "load_1min", "load_5min", "load_15min",
+        "load_per_cpu", "disk_io_util", "cpu_freq_mhz",
+    ]
+    result = {"worker_type": machine_type,
+              "cpu_cores":   snapshots[0]["cpu_cores"],
+              "cpu_threads": snapshots[0]["cpu_threads"]}
+    for key in numeric:
+        vals = [s[key] for s in snapshots]
+        avg  = sum(vals) / len(vals)
+        if key in ("idle_slots", "cpu_freq_mhz"):
+            result[key] = int(round(avg))
+        elif key in ("cpu_util", "ram_util", "ram_total"):
+            result[key] = round(avg, 1)
+        elif key == "ram_available":
+            result[key] = round(avg, 15)
+        elif key in ("load_1min", "load_5min", "load_15min"):
+            result[key] = round(avg, 10)
+        elif key == "load_per_cpu":
+            result[key] = round(avg, 13)
+        else:
+            result[key] = round(avg, 2)
+    logger.info("System metrics ready.")
+    return result
+
+
+# ── Server communication ──────────────────────────────────────────────────────
+
+def _auth_headers(worker_token: str) -> dict:
+    """Return Authorization header dict if worker_token is set, else empty."""
+    if worker_token:
+        return {"Authorization": f"Bearer {worker_token}"}
+    return {}
+
+
+def _request_job(url: str, runner_id: str, metrics: dict,
+                 logger: logging.Logger, worker_token: str = ""):
+    """Returns (job_dict, None) on success, (None, reason) on failure."""
+    try:
+        r = requests.post(url, json={"requested_by": runner_id,
+                                     "system_metrics": metrics},
+                          headers=_auth_headers(worker_token),
+                          timeout=30)
+        if r.status_code == 404:
+            return None, 'no_jobs'
+        if r.status_code == 200:
+            return r.json(), None
+        return None, f"HTTP {r.status_code}: {r.text[:200]}"
+    except Exception as exc:
+        return None, str(exc)
+
+
+def _update_status(url: str, job_id: int, status: str,
+                   message: str, logger: logging.Logger,
+                   worker_token: str = "") -> None:
+    try:
+        r = requests.post(url, json={"job_id": job_id,
+                                     "status":  status,
+                                     "message": message},
+                          headers=_auth_headers(worker_token),
+                          timeout=30)
+        if r.status_code == 200:
+            logger.info(f"Job {job_id} → {status}")
+        else:
+            logger.warning(f"Status update failed: HTTP {r.status_code}")
+    except Exception as exc:
+        logger.error(f"Status update error: {exc}")
+
+
+def _ping_loop(url: str, job_id: int,
+               stop_event: threading.Event,
+               logger: logging.Logger,
+               worker_token: str = "") -> None:
+    """Background thread: ping the server every PING_INTERVAL seconds."""
+    while not stop_event.wait(PING_INTERVAL):
+        try:
+            r = requests.post(url, json={"id": job_id},
+                              headers=_auth_headers(worker_token),
+                              timeout=10)
+            if r.status_code == 200:
+                logger.info(f"Ping OK (job {job_id})")
+            else:
+                logger.warning(f"Ping HTTP {r.status_code} (job {job_id})")
+        except Exception as exc:
+            logger.warning(f"Ping error (job {job_id}): {exc}")
+
+
+# ── Main entry point ──────────────────────────────────────────────────────────
+
+def main() -> None:
+    argv = sys.argv[1:]
+
+    # Show help
+    if not argv or any(a in argv for a in ('help', '-h', '--help')):
+        print(__doc__)
+        sys.exit(0)
+
+    kv  = _parse_kv(argv)
+    cfg = _resolve(kv)
+
+    # Validate required arguments
+    errors = []
+    if not cfg['exp_id']:
+        errors.append("expId is required")
+    if not cfg['entry_script']:
+        errors.append("entry_script is required")
+    elif not os.path.isfile(cfg['entry_script']):
+        errors.append(f"entry_script '{cfg['entry_script']}' not found")
+    if errors:
+        for e in errors:
+            print(f"Error: {e}")
+        print("Run `jd_worker help` for usage.")
+        sys.exit(1)
+
+    # Build a unique runner ID visible in the dashboard
+    username  = os.getenv('USER') or os.getenv('USERNAME') or 'user'
+    random.seed(int(time.time() * 1000))
+    suffix    = random.randint(10000, 99999)
+    runner_id = (f"{username}@{socket.gethostname()}"
+                 f"({cfg['machine_type']})_{cfg['process_id']}_{suffix}")
+
+    if cfg['log_dir_override'] is not None:
+        log_dir = os.path.join(cfg['log_dir_override'], cfg['exp_id'])
+    else:
+        log_dir = os.path.join(cfg['workspace_path'], cfg['exp_id'], 'jd_worker_logs')
+    logger  = _setup_logger(log_dir, runner_id)
+
+    # ── Hub authentication (optional) ────────────────────────────────────────
+    worker_token = ""
+    if cfg['hub_url'] and cfg['api_key']:
+        logger.info(f"Hub mode: authenticating via {cfg['hub_url']}")
+        result = _hub_get_worker_token(
+            cfg['hub_url'], cfg['api_key'], cfg['exp_id'], logger
+        )
+        if result is None:
+            logger.error("Failed to obtain worker token from Hub. Exiting.")
+            sys.exit(1)
+        worker_token, hub_server_url = result
+        if hub_server_url:
+            cfg['base_url'] = hub_server_url
+            logger.info(f"Server URL from Hub: {hub_server_url}")
+    elif cfg['hub_url'] or cfg['api_key']:
+        logger.warning(
+            "Both hub_url (JD_HUB_URL) and api_key (JD_API_KEY) must be set "
+            "for Hub authentication. Running in standalone mode."
+        )
+
+    urls = {
+        'request': f"{cfg['base_url']}/request_job",
+        'update':  f"{cfg['base_url']}/update_job_status",
+        'ping':    f"{cfg['base_url']}/ping",
+    }
+
+    logger.info(f"jd_worker v{__version__}  |  runner: {runner_id}")
+    logger.info(f"Server:        {cfg['base_url']}")
+    logger.info(f"Entry script:   {cfg['entry_script']}")
+    logger.info(f"Hub mode:       {'enabled' if worker_token else 'disabled'}")
+    logger.info(f"Local jd_data root: {cfg['workspace_path']} "
+                f"(each job: …/jd_data/<expId>/<job_id>/)")
+    logger.info(f"Ping interval: {PING_INTERVAL}s")
+    if cfg['once']:
+        logger.info("Mode: single job (once=true)")
+
+    # Track the current child process so SIGINT/SIGTERM can clean it up
+    _proc: list = [None]
+
+    def _shutdown(signum=None, frame=None):
+        p = _proc[0]
+        if p and p.poll() is None:
+            logger.info(f"Killing subprocess PID {p.pid}…")
+            try:
+                if IS_WINDOWS:
+                    p.send_signal(signal.CTRL_BREAK_EVENT)
+                else:
+                    os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+            except Exception:
+                pass
+        logger.info("jd_worker shut down.")
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT,  _shutdown)
+    signal.signal(signal.SIGTERM, _shutdown)
+
+    # ── Main job loop ────────────────────────────────────────────────────────
+    while True:
+        job_id = None
+        try:
+            metrics      = _averaged_metrics(cfg['machine_type'], logger)
+            job, reason  = _request_job(urls['request'], runner_id, metrics, logger,
+                                        worker_token=worker_token)
+
+            if job is None:
+                if reason == 'no_jobs':
+                    logger.info("No more jobs available. Exiting.")
+                    break
+                logger.error(f"Job request failed: {reason}. Retrying in 10 s…")
+                time.sleep(10)
+                continue
+
+            job_id = job['job_id']
+            params = job['parameters']
+            logger.info(f"Job {job_id} received  |  params: {params}")
+
+            # Local sandbox for this job — keep worker-side I/O under this directory
+            job_root = os.path.abspath(os.path.join(
+                cfg['workspace_path'], cfg['exp_id'], str(job_id)))
+            os.makedirs(job_root, exist_ok=True)
+
+            # Build subprocess command
+            # Uses the *current* Python interpreter so venv/conda is respected
+            cmd = [sys.executable, cfg['entry_script']]
+            for k, v in params.items():
+                cmd.extend([f"--{k}", str(v)])
+            cmd.extend(['--base_path', job_root])
+            logger.info(f"Job workspace: {job_root}")
+            logger.info(f"Command: {' '.join(cmd)}")
+
+            # Start heartbeat ping thread
+            stop_ping = threading.Event()
+            pinger    = threading.Thread(
+                target=_ping_loop,
+                args=(urls['ping'], job_id, stop_ping, logger, worker_token),
+                daemon=True,
+            )
+            pinger.start()
+
+            # Build child environment: inherit everything + inject JD_ context
+            # so jd_upload / jd_update_checkpoint / jd_get_last_checkpoint work
+            # inside the entry script without requiring explicit arguments.
+            child_env = os.environ.copy()
+            child_env["JD_JOB_ID"]                  = str(job_id)
+            child_env["JD_SERVER"]                  = cfg["base_url"]
+            child_env["JD_EXP_ID"]                  = cfg["exp_id"]
+            child_env["JD_WORKER_JOB_DIR"]          = job_root
+            child_env["JD_WORKER_WORKSPACE_ROOT"]   = cfg["workspace_path"]
+            if worker_token:
+                child_env["JD_WORKER_TOKEN"]        = worker_token
+
+            # Launch the entry script
+            popen_kw = dict(stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                            text=True, env=child_env)
+            if IS_WINDOWS:
+                _proc[0] = subprocess.Popen(
+                    cmd,
+                    creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
+                    **popen_kw,
+                )
+            else:
+                _proc[0] = subprocess.Popen(
+                    cmd,
+                    preexec_fn=os.setsid,
+                    **popen_kw,
+                )
+
+            stdout, stderr = _proc[0].communicate()
+            rc             = _proc[0].returncode
+
+            # Stop heartbeat
+            stop_ping.set()
+            pinger.join(timeout=5)
+
+            if rc == 0:
+                logger.info(f"Job {job_id} finished successfully.")
+                _update_status(
+                    urls['update'], job_id, 'DONE',
+                    f"Completed successfully on {runner_id}.",
+                    logger, worker_token=worker_token,
+                )
+            else:
+                logger.error(f"Job {job_id} failed — exit code {rc}")
+                if stderr.strip():
+                    logger.error(f"STDERR (last 1000 chars):\n{stderr.strip()[-1000:]}")
+
+                abort_msg = (
+                    f"Job failed on {runner_id}. "
+                    f"Exit code {rc}."
+                )
+                snippet = (stderr.strip() or stdout.strip())[-500:]
+                if snippet and any(kw in snippet.lower()
+                                   for kw in ('error', 'exception', 'traceback', 'failed')):
+                    abort_msg += f" Last output: {snippet}"
+                else:
+                    abort_msg += " Check worker logs for details."
+
+                if rc == -9:
+                    abort_msg += " (Process killed — possible OOM or time limit.)"
+
+                _update_status(urls['update'], job_id, 'ABORTED', abort_msg, logger,
+                               worker_token=worker_token)
+
+            _proc[0] = None
+
+        except Exception as exc:
+            logger.exception(f"Unexpected error: {exc}")
+            if job_id is not None:
+                _update_status(
+                    urls['update'], job_id, 'ABORTED',
+                    f"Unexpected exception on {runner_id}: {exc}",
+                    logger, worker_token=worker_token,
+                )
+            break
+
+        if cfg['once']:
+            logger.info("once=true — exiting after one job.")
+            break
+
+        time.sleep(3)   # brief pause before requesting the next job
+
+
+if __name__ == '__main__':
+    main()

jd_worker-1.0.0.dist-info/METADATA

@@ -0,0 +1,56 @@
+Metadata-Version: 2.4
+Name: jd-worker
+Version: 1.0.0
+Summary: Worker CLI for the JobDistributor distributed job system
+License: MIT
+Project-URL: Homepage, https://github.com/NWSL-UCF/job-distributor
+Project-URL: Documentation, https://github.com/NWSL-UCF/job-distributor/blob/main/docs/jd-worker.md
+Project-URL: Bug Tracker, https://github.com/NWSL-UCF/job-distributor/issues
+Keywords: distributed,jobs,worker,hpc,ml
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: System :: Distributed Computing
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: requests
+Requires-Dist: psutil
+
+# Worker client (`jd-worker`)
+
+Workers run **`jd_worker`**, which requests jobs from the job server, executes your **entry script** with parameters as CLI flags, sends heartbeats, and reports **DONE** / **ABORTED**.
+
+Full behaviour, environment variables, local paths (`~/jd_data/…`), and library helpers (`jd_upload`, checkpoints, `jd_job_dir`, …) are documented in **[`docs/jd-worker.md`](../docs/jd-worker.md)**.
+
+---
+
+## Install
+
+From the repo root (editable):
+
+```bash
+cd job-distributor/client
+python3 -m venv venv
+source venv/bin/activate   # Windows: venv\Scripts\activate
+pip install -e .
+```
+
+Or install from GitHub (see `docs/jd-worker.md` for branch / subdirectory).
+
+---
+
+## Run
+
+```bash
+jd_worker expId=<experiment_id> entry_script=<your_script.py>
+```
+
+Optional: `server=`, `port=`, `machine_type=`, `once=true`, etc. — see `jd_worker help` or **`docs/jd-worker.md`**.
+
+---
+
+## Example workload
+
+For an end-to-end ML-style example (MNIST tuning), see  
+[MNIST-parameter-tuning](https://github.com/NWSL-UCF/MNIST-parameter-tuning). Point **`entry_script`** at that repo’s training script and align **`expId`** with the server.

jd_worker-1.0.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+jd/__init__.py,sha256=HvsPfgH7X5ao9MpgJ5brs52CeHF7hgDv9sPmmd-kd9E,363
+jd/files.py,sha256=-66sf4fbTZugtsAinBI5YyE_ctK74k7SdAieI8zISsE,8151
+jd/paths.py,sha256=-FVQYEuSYamYtbMdBq9azQS0bqYz49AKBeCg5ZO3tSU,1918
+jd/worker.py,sha256=QPfMEmixVynTfTaU9OMf42EWc1Hl2q4MADxyKCCR5GU,25166
+jd_worker-1.0.0.dist-info/METADATA,sha256=Xam6aG0K-Gt6e3BTpHY_D_vq-tLaq9AHqTaWuH8dDAA,1925
+jd_worker-1.0.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+jd_worker-1.0.0.dist-info/entry_points.txt,sha256=4yBzE58Z9WQN--Ri1vKEKuqvd25QBgVe-XCea-TLm9Q,45
+jd_worker-1.0.0.dist-info/top_level.txt,sha256=2Yt96k5nsj0T_z76Ppc_ss__sm5nD6LRi6bmXiQFrik,3
+jd_worker-1.0.0.dist-info/RECORD,,

jd_worker-1.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

jd_worker-1.0.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+jd_worker = jd.worker:main

jd_worker-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+jd