Short answer: there isn’t a first‑class Prefect integration for Ray Tune or Ray Train today. In Prefect 3.x, the official Ray support is the prefect-ray integration, which provides a RayTaskRunner to parallelize Prefect tasks on a Ray cluster. You can still orchestrate Ray Tune/Train from Prefect, but it’s done by calling them within your tasks/flows rather than via a dedicated integration. Links and current status - Prefect + Ray docs: prefect-ray integration docs - Open/closed items related to Ray in Prefect: - Issue: resource options with RayTaskRunner (older 2.x context) - Issue: RayTaskRunner on remote cluster not picking up PREFECT_API_URL - PR: improve logging when tasks run on remote Ray/Dask workers - I couldn’t find a feature request for a dedicated “Prefect + Ray Tune/Train” integration: - Issues search: issues: ray tune/train - Discussions search: discussions: ray tune How to use Ray Tune/Train with Prefect today - Orchestrate Ray directly from a Prefect task or flow: - Keep Prefect for orchestration, retries, caching, and scheduling. - Let Ray handle HPO/distributed training within that step. Minimal examples - Run Ray Tune within a Prefect flow (local or remote Ray, depending on your Ray config):

from prefect import flow, task
from ray import tune

@task
def run_tune():
    def trainable(config):
        # your training logic here
        acc = 1.0 - config["lr"]  # placeholder metric
        return {"accuracy": acc}

    analysis = tune.run(
        trainable,
        config={"lr": tune.grid_search([0.001, 0.01, 0.1])},
        num_samples=1,
        stop={"training_iteration": 1},
    )
    return analysis.best_result

@flow
def hpo_flow():
    return run_tune()

if __name__ == "__main__":
    hpo_flow()

- If you want Prefect tasks to execute on a Ray cluster while you run Tune inside a task, use RayTaskRunner and ensure worker env has Prefect configured:

from prefect import flow, task
from prefect_ray.task_runners import RayTaskRunner
from ray import tune

@task
def run_tune():
    # same as above
    ...

@flow(
    task_runner=RayTaskRunner(
        address="ray://<your-ray-head>:10001",
        init_kwargs={
            "runtime_env": {
                # ensure Prefect + any deps exist on Ray workers
                "pip": ["prefect>=3", "prefect-ray>=0.4", "ray[tune]"],
                # make sure workers can reach your API if using Prefect Cloud/Server
                "env_vars": {"PREFECT_API_URL": "https://<your-api>/api"},
            }
        },
    )
)
def hpo_flow():
    return run_tune()

Notes and caveats - RayTaskRunner focuses on parallelizing Prefect tasks; it doesn’t manage Tune trials or Train backends natively. - When running on a remote Ray cluster, make sure Prefect and your Prefect API settings are available to Ray workers (see the env_vars/runtime_env pattern above). One of the linked issues tracks an env propagation pitfall. - For Ray Train, use the same pattern: build your Trainer inside a Prefect task and return metrics/checkpoints. Would you like me to open a feature request for a first‑class Ray Tune/Train integration on your behalf, or share a more complete template tailored to your environment (local, K8s, Ray address, storage)?

Got it — you want each Ray Tune config/trial to show up as its own Prefect task. There isn’t a native “Ray Tune trial = Prefect task” integration. In practice you have two viable patterns: Option A — Prefect owns trials (each config = a Prefect task) - You generate the parameter configs in your flow. - Map over them with a Prefect task. - Use RayTaskRunner so each task executes on your Ray cluster (and you can set per‑task CPU/GPU with remote_options). - Pros: every trial is a first‑class Prefect task with individual state, retries, caching, artifacts. - Cons: you won’t get Tune’s advanced schedulers (ASHA, PBT, early stopping across trials) unless you re‑implement coordination logic yourself. Example (Prefect 3.x + prefect-ray)

from prefect import flow, task
from prefect_ray.task_runners import RayTaskRunner
from prefect_ray.context import remote_options

# one training run = one Prefect task
@task
def train_one_config(cfg: dict) -> dict:
    # do your single-trial training here, using GPUs if allocated
    # return metrics + anything needed to choose the best
    # e.g., train(model, cfg) ...
    accuracy = 0.9  # placeholder
    return {"metrics": {"accuracy": accuracy}, "config": cfg}

@flow(
    task_runner=RayTaskRunner(
        address="ray://<your-ray-head>:10001",
        init_kwargs={
            "runtime_env": {
                # ensure workers have what they need
                "pip": ["prefect>=3", "prefect-ray>=0.4", "torch", "numpy"],  # add deps as needed
                # if using Prefect Cloud/Server, propagate settings so Ray workers can call back
                "env_vars": {"PREFECT_API_URL": "https://<your-prefect-api>/api"},
            }
        },
    )
)
def hpo_as_prefect_tasks(n_trials: int = 12):
    # build your search space (grid or random)
    import random
    configs = [
        {"lr": 10 ** random.uniform(-4, -1), "batch_size": random.choice([32, 64, 128])}
        for _ in range(n_trials)
    ]

    # per-trial resource allocation on Ray workers
    with remote_options(num_cpus=4, num_gpus=1):
        futs = train_one_config.map(configs)

    # collect results and pick best
    results = [f.result() for f in futs]
    best = max(results, key=lambda r: r["metrics"]["accuracy"])
    return best

if __name__ == "__main__":
    print(hpo_as_prefect_tasks())

Option B — Ray Tune owns trials (Prefect wraps the whole HPO as one task/flow) - You call Tune (or Train) inside a single Prefect task/flow, and let it launch/manage trials. - Pros: you get Tune’s schedulers, search algorithms, early stopping, trial pruning, etc. - Cons: trials will not appear as individual Prefect tasks (they’re Ray trials), though you can still persist best metrics/artifacts back to Prefect at the end. Example (Ray Tune inside a Prefect task)

from prefect import flow, task
from ray import tune

@task
def run_tune():
    def trainable(cfg):
        # training loop
        # tune.report(metric=...) inside loop if you want schedulers like ASHA
        return {"accuracy": 0.9}

    analysis = tune.run(
        trainable,
        config={"lr": tune.grid_search([1e-4, 1e-3, 1e-2]),
                "batch_size": tune.choice([32, 64, 128])},
        num_samples=1,
        stop={"training_iteration": 5},
    )
    return analysis.best_result

@flow
def hpo_with_tune():
    return run_tune()

Key trade-off - If “each config must be a Prefect task” is a hard requirement, choose Option A. You’ll parallelize via RayTaskRunner and keep full Prefect visibility per trial, but you give up Tune’s adaptive/early‑stopping orchestration. - If you need advanced HPO features (ASHA, PBT, search algorithms), choose Option B and let Tune manage trials end‑to‑end. Caveats and tips - On remote Ray clusters, make sure Prefect is reachable from Ray workers (set PREFECT_API_URL in RayTaskRunner.init_kwargs.runtime_env.env_vars). - If you want per‑trial GPUs, use the remote_options context (as shown) around the submit/map calls.

- Running Ray Train multi‑worker training from inside a RayTaskRunner task can be tricky (nested Ray orchestration). For distributed training per trial, it’s typically cleaner to let Train launch the workers from the driver (Option B) rather than from inside a Ray worker. Docs - Prefect + Ray: prefect-ray docs - Ray Tune: Ray Tune docs - Ray Train: Ray Train docs If you share: - whether you need ASHA/PBT/early stopping - single‑GPU vs multi‑GPU/multi‑node per trial - your environment (local vs K8s, Ray address, Prefect Cloud/Server) …I can tailor a ready‑to‑drop‑in template for your setup.

Great — if you’re using Ray Tune with ASHA and early stopping (Option B), the cleanest way to surface progress in Prefect is to: - Run Tune inside a single Prefect task - Attach a Ray Tune Callback that reports trial progress back to Prefect via logs and artifacts - Use a Prefect progress artifact to show overall HPO progress; optionally publish a summary artifact at the end Below is a ready-to-drop-in pattern. What you get - Live Prefect logs for each trial result (metric, iteration, config) - A Prefect “progress” artifact that increments as trials finish (pruned or completed) - A final Markdown artifact summarizing the best trial Code ``` from prefect import flow, task from prefect.logging import get_run_logger from prefect.artifacts import ( create_progress_artifact, update_progress_artifact, create_markdown_artifact, ) from ray import tune from ray.tune.schedulers import ASHAScheduler # Optional: If you prefer the modern Tuner API: # from ray import tune # from ray.tune import Tuner, TuneConfig # from ray.tune.schedulers import ASHAScheduler # A Ray Tune Callback that logs to Prefect and updates a Prefect progress artifact class PrefectTuneCallback(tune.Callback): def __init__(self, total_trials: int, metric: str = "accuracy", mode: str = "max", description: str | None = None): self.total_trials = max(1, int(total_trials)) self.metric = metric self.mode = mode self.description = description or f"Ray Tune progress ({metric} {mode})" self.logger = None self.progress_id = None self.best_metric = None self.best_config = None def on_experiment_start(self, **info): self.logger = get_run_logger() self.progress_id = create_progress_artifact(0.0, description=self.description) self.logger.info("Started Ray Tune experiment; tracking progress in Prefect.") def _update_best(self, result: dict, config: dict): val = result.get(self.metric) if val is None: return if self.best_metric is None: self.best_metric, self.best_config = val, config return better = (val > self.best_metric) if self.mode == "max" else (val < self.best_metric) if better: self.best_metric, self.best_config = val, config def _update_progress(self, trials): # Count trials that have finished one way or another done = sum(t.status in ("TERMINATED", "ERROR") for t in trials) progress = min(1.0, done / self.total_trials) if self.progress_id: update_progress_artifact(self.progress_id, progress, description=self.description) def on_trial_result(self, iteration, trials, trial, result: dict, **info): # Called whenever a trial reports via tune.report(...) # Log the key metrics and update best-so-far cfg = trial.config or {} metric_val = result.get(self.metric) it = result.get("training_iteration") or result.get("iteration") or iteration self.logger.info( f"Trial {trial.trial_id} iter {it}: {self.metric}={metric_val}, config={cfg}" ) self._update_best(result, cfg) self._update_progress(trials) def on_trial_complete(self, iteration, trials, trial, **info): # Update progress upon completion (including early-stopped trials) self.logger.info(f"Trial {trial.trial_id} completed with status={trial.status}") self._update_progress(trials) def on_trial_error(self, iteration, trials, trial, **info): self.logger.warning(f"Trial {trial.trial_id} errored; status={trial.status}") self._update_progress(trials) def on_experiment_end(self, trials, **info): # Mark progress as complete and publish a summary if self.progress_id: update_progress_artifact(self.progress_id, 1.0, description=self.description) summary_lines = ["# Ray Tune Summary"]

if self.best_metric is not None: summary_lines.append(f"- Best {self.metric} ({self.mode}): {self.best_metric}") summary_lines.append(f"- Best config: `{self.best_config}`") else: summary_lines.append("- No best metric recorded.") # Top few finished trials (optional) finished = [t for t in trials if t.status == "TERMINATED" and getattr(t, "last_result", None)] finished = sorted( finished, key=lambda t: t.last_result.get(self.metric, float("-inf") if self.mode == "max" else float("+inf")), reverse=(self.mode == "max"), )[:5] if finished: summary_lines.append("\n## Top finished trials") for t in finished: mv = t.last_result.get(self.metric) summary_lines.append(f"- {t.trial_id}: {self.metric}={mv}, config=`{t.config}`") create_markdown_artifact("\n".join(summary_lines), description="Ray Tune results") @task def run_tune_with_asha(num_samples: int = 20): """ Runs Ray Tune with ASHA and logs live progress + results into Prefect. """ logger = get_run_logger() metric = "accuracy" mode = "max" # Define your trainable def trainable(config): import math, random, time # Example: pretend we're training and reporting results to Tune for step in range(1, 21): # Simulate some training time.sleep(0.1) # Report a metric that improves (or not) over time acc = 0.5 + 0.5 * (1 - math.exp(-step / 10)) + random.uniform(-0.02, 0.02) tune.report(**{metric: acc, "training_iteration": step}) # ASHA early-stopping scheduler scheduler = ASHAScheduler( max_t=20, # max training iterations per trial grace_period=5, # min iterations before pruning reduction_factor=3 # how aggressively to downsample ) # Search space + Tune configuration param_space = { "lr": tune.loguniform(1e-4, 1e-1), "batch_size": tune.choice([32, 64, 128]), } callback = PrefectTuneCallback(total_trials=num_samples, metric=metric, mode=mode) # You can use either the classic tune.run(...) API or the modern Tuner API. # Classic API: analysis = tune.run( trainable, config=param_space, num_samples=num_samples, scheduler=scheduler, metric=metric, mode=mode, callbacks=[callback], stop={"training_iteration": 20}, ) # Best result summary (also captured by the callback) logger.info(f"Best result: {analysis.best_result}") return analysis.best_result @flow def hpo_flow(num_samples: int = 20): return run_tune_with_asha(num_samples) ``` Notes - The callback runs in the Tune driver process (the same process as your Prefect task), so Prefect logging and artifacts work reliably. - The progress artifact advances as trials finish (including early-stopped ones). We treat “progress” as completed_trials / num_samples. - At the end, you’ll see a Markdown artifact with best metrics and a short leaderboard. If you run this task on a remote Ray cluster via RayTaskRunner - Ensure your Ray workers have Prefect and your Prefect API settings: - Provide runtime_env pip deps and env_vars (e.g., PREFECT_API_URL) in RayTaskRunner - Known pitfall: missing PREFECT_API_URL on workers prevents logs from reaching Prefect (see linked issue above) References - Ray Tune schedulers (ASHA): Ray Tune ASHA docs - Prefect + Ray integration: prefect-ray docs Want me to adapt this to your exact setup (Kubernetes/Cloud, Ray address, GPUs per trial) and include artifact tables for top-K trials?

Sorry, I encountered an error while processing your request. Please try again.

Yes — you can run Ray Tune directly in the flow body (no @task needed). That’s actually a good fit for Option B: the Prefect run acts as the “driver,” Ray handles trials/ASHA/early stopping, and you still log progress and publish artifacts from the flow. Here’s a compact version using the same PrefectTuneCallback pattern, but embedded entirely in the flow: ``` from prefect import flow from prefect.logging import get_run_logger from prefect.artifacts import ( create_progress_artifact, update_progress_artifact, create_markdown_artifact, ) from ray import tune from ray.tune.schedulers import ASHAScheduler class PrefectTuneCallback(tune.Callback): def __init__(self, total_trials, metric="accuracy", mode="max", description=None): self.total_trials = max(1, int(total_trials)) self.metric, self.mode = metric, mode self.description = description or f"Ray Tune progress ({metric} {mode})" self.logger = None self.progress_id = None self.best_metric = None self.best_config = None def on_experiment_start(self, **info): self.logger = get_run_logger() self.progress_id = create_progress_artifact(0.0, description=self.description) self.logger.info("Started Ray Tune experiment; reporting to Prefect.") def _update_best(self, result, config): val = result.get(self.metric) if val is None: return if self.best_metric is None: self.best_metric, self.best_config = val, config else: better = (val > self.best_metric) if self.mode == "max" else (val < self.best_metric) if better: self.best_metric, self.best_config = val, config def _update_progress(self, trials): done = sum(t.status in ("TERMINATED", "ERROR") for t in trials) progress = min(1.0, done / self.total_trials) if self.progress_id: update_progress_artifact(self.progress_id, progress, description=self.description) def on_trial_result(self, iteration, trials, trial, result, **info): cfg = trial.config or {} it = result.get("training_iteration") or result.get("iteration") or iteration self.logger.info(f"Trial {trial.trial_id} iter {it}: {self.metric}={result.get(self.metric)}, config={cfg}") self._update_best(result, cfg) self._update_progress(trials) def on_trial_complete(self, iteration, trials, trial, **info): self.logger.info(f"Trial {trial.trial_id} completed with status={trial.status}") self._update_progress(trials) def on_trial_error(self, iteration, trials, trial, **info): self.logger.warning(f"Trial {trial.trial_id} errored; status={trial.status}") self._update_progress(trials) def on_experiment_end(self, trials, **info): if self.progress_id: update_progress_artifact(self.progress_id, 1.0, description=self.description) lines = ["# Ray Tune Summary"] if self.best_metric is not None: lines += [f"- Best {self.metric} ({self.mode}): {self.best_metric}", f"- Best config: `{self.best_config}`"] else: lines.append("- No best metric recorded.") create_markdown_artifact("\n".join(lines), description="Ray Tune results") @flow def hpo_flow(num_samples: int = 20, ray_address: str | None = None): logger = get_run_logger() # Optional: connect to an existing Ray cluster if ray_address: import ray ray.init(address=ray_address) metric, mode = "accuracy", "max" def trainable(config): # Replace with your real training loop; use tune.report(...) periodically import math, random, time for step in range(1, 21): time.sleep(0.1) acc = 0.5 + 0.5 * (1 - math.exp(-step / 10)) + random.uniform(-0.02, 0.02) tune.report(**{metric: acc, "training_iteration": step})

scheduler = ASHAScheduler(max_t=20, grace_period=5, reduction_factor=3) param_space = {"lr": tune.loguniform(1e-4, 1e-1), "batch_size": tune.choice([32, 64, 128])} callback = PrefectTuneCallback(total_trials=num_samples, metric=metric, mode=mode) analysis = tune.run( trainable, config=param_space, num_samples=num_samples, scheduler=scheduler, metric=metric, mode=mode, callbacks=[callback], stop={"training_iteration": 20}, ) logger.info(f"Best result: {analysis.best_result}") return analysis.best_result ``` Notes - Running this in the flow body ensures get_run_logger and artifact APIs attach to the current Prefect run. No @task is required. - If you’re using a remote Ray cluster, the callback runs in the flow’s driver process (your Prefect flow container), so Prefect logs/artifacts do not depend on Ray worker env. You do not need PREFECT_API_URL on Ray workers in this pattern. - If you do want Ray worker logs in Prefect, that’s separate; they’ll stay in Ray unless you forward them yourself. Docs for reference: - Ray ASHA scheduler: Ray Tune ASHA docs - Prefect + Ray overview: prefect-ray docs Want me to adapt this for your exact Ray address/Kubernetes setup or add a top‑K trial leaderboard artifact?

Yes — that’s a solid pattern. Since the Ray Tune callbacks run in the driver (your flow process), you can: - Run Tune inside the flow body (Option B) - Set the flow’s task runner to ThreadPoolTaskRunner - From the callbacks, submit lightweight Prefect tasks to “track” trials (start/result/complete) so each trial shows up as task runs in the UI - Keep artifacts/logging in the driver so they’re reliable and don’t depend on Ray worker env Here’s a template that: - Uses ThreadPoolTaskRunner so tracking tasks run concurrently and don’t block Tune - Submits a task on trial start/complete and every Nth iteration to avoid flooding - Still uses a Prefect progress artifact for overall HPO progress ``` from prefect import flow, task from prefect.task_runners import ThreadPoolTaskRunner from prefect.logging import get_run_logger from prefect.artifacts import create_progress_artifact, update_progress_artifact from ray import tune from ray.tune.schedulers import ASHAScheduler # --- Tracking tasks (show up as Prefect task runs) --- @task(name="trial-start", tags=["ray-tune"]) def record_trial_start(trial_id: str, config: dict): logger = get_run_logger() logger.info(f"Ray Tune trial started: {trial_id}, config={config}") @task(name="trial-update", tags=["ray-tune"]) def record_trial_update(trial_id: str, iteration: int, metric_name: str, metric_value: float): logger = get_run_logger() logger.info(f"Trial {trial_id} iter {iteration}: {metric_name}={metric_value}") @task(name="trial-complete", tags=["ray-tune"]) def record_trial_complete(trial_id: str, status: str, final_metrics: dict | None = None): logger = get_run_logger() logger.info(f"Ray Tune trial complete: {trial_id}, status={status}, final={final_metrics}") # --- Ray Tune callback that submits tracking tasks --- class PrefectTrackingCallback(tune.Callback): def __init__(self, total_trials: int, metric="accuracy", mode="max", log_every_iters: int = 5): self.total_trials = max(1, int(total_trials)) self.metric, self.mode = metric, mode self.log_every_iters = max(1, int(log_every_iters)) self.progress_id = None self.best_metric = None self.best_config = None self._futures = [] # hold Prefect task futures so we can wait at the end def on_experiment_start(self, **info): # overall HPO progress bar in Prefect self.progress_id = create_progress_artifact(0.0, description=f"Ray Tune ({self.metric} {self.mode})") def _update_best(self, result: dict, config: dict): val = result.get(self.metric) if val is None: return if self.best_metric is None: self.best_metric, self.best_config = val, config else: better = (val > self.best_metric) if self.mode == "max" else (val < self.best_metric) if better: self.best_metric, self.best_config = val, config def _update_overall_progress(self, trials): done = sum(t.status in ("TERMINATED", "ERROR") for t in trials) p = min(1.0, done / self.total_trials) if self.progress_id: update_progress_artifact(self.progress_id, p, description=f"Ray Tune ({self.metric} {self.mode})") def on_trial_start(self, iteration, trials, trial, **info): # one Prefect task per trial start self._futures.append(record_trial_start.submit(trial.trial_id, trial.config)) def on_trial_result(self, iteration, trials, trial, result: dict, **info): # throttle updates to avoid overwhelming the UI it = result.get("training_iteration") or result.get("iteration") or iteration if it % self.log_every_iters == 0: mv = result.get(self.metric) self._futures.append(record_trial_update.submit(trial.trial_id, it, self.metric, mv if mv is not None else float("nan"))) self._update_best(result, trial.config) self._update_overall_progress(trials)

def on_trial_complete(self, iteration, trials, trial, **info): final = getattr(trial, "last_result", None) or {} self._futures.append(record_trial_complete.submit(trial.trial_id, trial.status, final)) self._update_overall_progress(trials) def on_trial_error(self, iteration, trials, trial, **info): self._futures.append(record_trial_complete.submit(trial.trial_id, trial.status, None)) self._update_overall_progress(trials) def on_experiment_end(self, trials, **info): # finish progress and wait for all tracking tasks if self.progress_id: update_progress_artifact(self.progress_id, 1.0, description=f"Ray Tune ({self.metric} {self.mode})") for f in self._futures: try: f.result() # ensure tasks have finished before the flow returns except Exception: pass # --- Flow runs Ray Tune (driver), tasks run concurrently in threads --- @flow(task_runner=ThreadPoolTaskRunner(max_workers=16)) def hpo_flow(num_samples: int = 20, ray_address: str | None = None): logger = get_run_logger() if ray_address: import ray ray.init(address=ray_address) metric, mode = "accuracy", "max" def trainable(config): # Replace with your real training loop. Call tune.report periodically. import math, random, time for step in range(1, 41): time.sleep(0.1) acc = 0.5 + 0.5 * (1 - math.exp(-step / 10)) + random.uniform(-0.02, 0.02) tune.report(**{metric: acc, "training_iteration": step}) scheduler = ASHAScheduler(max_t=40, grace_period=5, reduction_factor=3) param_space = {"lr": tune.loguniform(1e-4, 1e-1), "batch_size": tune.choice([32, 64, 128])} callback = PrefectTrackingCallback(total_trials=num_samples, metric=metric, mode=mode, log_every_iters=5) analysis = tune.run( trainable, config=param_space, num_samples=num_samples, scheduler=scheduler, metric=metric, mode=mode, callbacks=[callback], stop={"training_iteration": 40}, ) logger.info(f"Best result: {analysis.best_result}") return analysis.best_result ``` Why this works - Callbacks are invoked in the driver process; task.submit is non-blocking and schedules small “tracking tasks” to run on the ThreadPoolTaskRunner. - Tune continues coordinating ASHA/early stopping without being blocked by I/O to Prefect. - You still get a single flow run that encapsulates the entire HPO, plus many short-lived task runs for visibility per trial, and a single overall progress artifact. Tips - Throttle task submissions (e.g., every 5 iterations) to avoid thousands of tiny tasks if trials report frequently. - If you only need logs/artifacts (not tasks), keep everything in the callback and skip task.submit entirely. - If you want per‑trial progress bars, keep a map {trial_id -> artifact_id} in the callback and update a Prefect progress artifact per trial; just be mindful of API call volume. If you share your expected trial count and report frequency, I can suggest sensible throttling and max_workers settings.

Short answer: With the callback + ThreadPoolTaskRunner pattern I shared earlier, you’ll get multiple small tasks per trial (start/update/complete). If you want exactly one Prefect task per Ray trial that shows “Running” for the duration of the trial and switches to “Completed” when that trial finishes, use this variant: - Start one long‑lived “tracker” task at trial start. - Keep a per‑trial Event in the driver. - When Ray Tune signals trial completion (or early stop), set the Event and the tracker task returns. Template ``` from prefect import flow, task from prefect.task_runners import ThreadPoolTaskRunner from prefect.logging import get_run_logger from prefect.artifacts import create_progress_artifact, update_progress_artifact from ray import tune from ray.tune.schedulers import ASHAScheduler import threading # --- In‑process registry for trial lifecycle signaling --- _TRACK = { "events": {}, # trial_id -> threading.Event "finals": {}, # trial_id -> final metrics dict "futures": {}, # trial_id -> Prefect future for the tracker task "lock": threading.Lock(), } def _ensure_event(trial_id: str) -> threading.Event: with _TRACK["lock"]: ev = _TRACK["events"].get(trial_id) if ev is None: ev = threading.Event() _TRACK["events"][trial_id] = ev return ev # --- Single, long‑lived task per trial --- @task(name="trial-tracker", tags=["ray-tune"]) def track_trial_lifecycle(trial_id: str, config: dict): logger = get_run_logger() logger.info(f"Ray Tune trial started: {trial_id}, config={config}") ev = _ensure_event(trial_id) # Wait until the callback signals completion/early-stop ev.wait() with _TRACK["lock"]: final = _TRACK["finals"].pop(trial_id, None) _TRACK["events"].pop(trial_id, None) logger.info(f"Ray Tune trial finished: {trial_id}, final={final}") return {"trial_id": trial_id, "final": final, "config": config} # --- Ray Tune callback that coordinates the tracking tasks --- class OneTaskPerTrialCallback(tune.Callback): def __init__(self, total_trials: int, metric="accuracy", mode="max"): self.total_trials = max(1, int(total_trials)) self.metric, self.mode = metric, mode self.progress_id = None def on_experiment_start(self, **info): self.progress_id = create_progress_artifact( 0.0, description=f"Ray Tune ({self.metric} {self.mode})" ) def _update_overall_progress(self, trials): done = sum(t.status in ("TERMINATED", "ERROR") for t in trials) p = min(1.0, done / self.total_trials) if self.progress_id: update_progress_artifact(self.progress_id, p, description=f"Ray Tune ({self.metric} {self.mode})") def on_trial_start(self, iteration, trials, trial, **info): # Pre-create event and launch the tracker task; it will block until we signal completion _ensure_event(trial.trial_id) fut = track_trial_lifecycle.submit(trial.trial_id, trial.config) with _TRACK["lock"]: _TRACK["futures"][trial.trial_id] = fut def on_trial_complete(self, iteration, trials, trial, **info): # Record final metrics and release the tracker task with _TRACK["lock"]: _TRACK["finals"][trial.trial_id] = getattr(trial, "last_result", None) or {} ev = _TRACK["events"].get(trial.trial_id) if ev: ev.set() self._update_overall_progress(trials) def on_trial_error(self, iteration, trials, trial, **info): # Signal completion even on error so the tracker task can end with _TRACK["lock"]: _TRACK["finals"][trial.trial_id] = {} ev = _TRACK["events"].get(trial.trial_id) if ev: ev.set() self._update_overall_progress(trials) def on_experiment_end(self, trials, **info): # Ensure all tracker tasks have finished with _TRACK["lock"]:

futures = list(_TRACK["futures"].values()) _TRACK["futures"].clear() for f in futures: try: f.result() except Exception: pass if self.progress_id: update_progress_artifact(self.progress_id, 1.0, description=f"Ray Tune ({self.metric} {self.mode})") # --- Flow runs Tune; tracker tasks run in threads and mirror each trial lifecycle --- @flow(task_runner=ThreadPoolTaskRunner(max_workers=64)) def hpo_flow(num_samples: int = 20, ray_address: str | None = None): logger = get_run_logger() if ray_address: import ray ray.init(address=ray_address) metric, mode = "accuracy", "max" def trainable(config): # Your real training loop; report periodically so ASHA can make decisions import math, random, time for step in range(1, 41): time.sleep(0.1) acc = 0.5 + 0.5 * (1 - math.exp(-step / 10)) + random.uniform(-0.02, 0.02) tune.report(**{metric: acc, "training_iteration": step}) scheduler = ASHAScheduler(max_t=40, grace_period=5, reduction_factor=3) param_space = {"lr": tune.loguniform(1e-4, 1e-1), "batch_size": tune.choice([32, 64, 128])} cb = OneTaskPerTrialCallback(total_trials=num_samples, metric=metric, mode=mode) analysis = tune.run( trainable, config=param_space, num_samples=num_samples, scheduler=scheduler, metric=metric, mode=mode, callbacks=[cb], stop={"training_iteration": 40}, ) logger.info(f"Best result: {analysis.best_result}") return analysis.best_result ``` What you’ll see - Exactly one Prefect task run per Ray trial. - Each task will appear “Running” for the life of the trial and flip to “Completed” as soon as that trial finishes or is pruned/errored. - An overall progress artifact for the HPO. Important considerations - Set max_workers high enough to cover your max concurrent active trials; each tracker task holds a thread while the trial runs. - If you’d rather avoid blocking threads, a simpler alternative is to submit a single completion task in on_trial_complete (one task per trial, but short-lived at the end). - Because all tracking runs in the flow’s driver process, you don’t need PREFECT_API_URL on Ray workers for these tasks/artifacts/logs to appear in Prefect. If you share your expected concurrency (max active trials) I can suggest a safe max_workers and any throttling tweaks.