Concurrent runs creating a new MLflow experiment race into the Default experiment (id 0) with silent artifact-upload failures

[physicistphil]

Summary

When multiple ergoExo runs start concurrently and target an MLflow experiment that does not yet exist (e.g. the first batch of a SLURM job array, or any parallel scan), there is a race in experiment creation. The race loser silently logs its run into the Default experiment (id 0) instead of the intended experiment, and — because (on our tracking server) the Default experiment's artifact root is a server-local path — all artifact uploads for that run fail. Params and metrics are still logged (to the wrong experiment), so the failure is easy to miss: the run simply appears "missing" from the target experiment.

Impact

• A run vanishes from its intended experiment and ends up in Default (id 0).
• That run has no artifacts (plots, config, NetCDF, …) — uploads fail (for us: PermissionError: [Errno 13] Permission denied: '/mlfs', because Default's artifact root is server-local while named experiments use S3).
• The failure is silent: robust_log_artifacts swallows the exception after retries and the run still ends FINISHED. Nothing surfaces unless someone counts runs in the target experiment.
• Only the first concurrent batch into a brand-new experiment is affected; once the experiment exists, later runs resolve it fine.

Root cause

Two interacting pieces:

1. adept — adept/_base_.py:256 discards the return value of set_experiment and relies on its process-global side effect:

mlflow.set_experiment(cfg["mlflow"]["experiment"])   # return value discarded
with mlflow.start_run(run_name=cfg["mlflow"]["run"], ...) as mlflow_run:
    ...

2. MLflow set_experiment (observed in 3.12.0, mlflow/tracking/fluent.py) is not race-safe in a way that interacts badly with the above:

with _experiment_lock:                  # threading.Lock -> per-process only
    if experiment_id is None:
        experiment = client.get_experiment_by_name(experiment_name)
        if not experiment:
            ...
            try:
                experiment_id = client.create_experiment(experiment_name)
            except MlflowException as e:
                if e.error_code == "RESOURCE_ALREADY_EXISTS":
                    return client.get_experiment_by_name(experiment_name)   # <-- early return
                raise
            ...
# reached only on the non-race path:
global _active_experiment_id
_active_experiment_id = experiment.experiment_id        # <-- skipped on race
MLFLOW_EXPERIMENT_ID.set(_active_experiment_id)

On the concurrent-create race the loser hits RESOURCE_ALREADY_EXISTS and returns early — before the lines that set the global active experiment. (The _experiment_lock is a threading.Lock, so it gives no protection across separate processes/nodes.)

Combined effect: the race-loser's process never sets an active experiment, so the subsequent start_run() falls back to the default experiment id "0".

Reproduction

1. Ensure experiment E does not exist.
2. Launch ≥2 ergoExo runs concurrently in separate processes/nodes, all with cfg["mlflow"]["experiment"] = "E", started close enough that both observe E missing.
3. One run creates E and logs there; the other logs into experiment 0, and its artifact uploads fail if experiment 0's artifact store isn't writable from the client.

Real-world evidence

• A 6-task SLURM array (one ergoExo run each) into a new experiment. SLURM logs show two tasks logging Experiment with name '...' does not exist. Creating a new experiment. in the same second. One landed in the named experiment; the other landed in experiment 0 with repeated Permission denied: '/mlfs' → Failed to log artifacts after 5 attempts., yet finished FINISHED with full params/metrics and zero artifacts.
• A later array into the (now-existing) same experiment did not hit this at all.

Secondary issue: silent, lossy artifact-upload failure

adept/utils.py:268 robust_log_artifacts:

for attempt in range(retries):
    try:
        mlflow.log_artifacts(directory); break
    except Exception as e:
        print(f"Attempt {attempt + 1} failed: {e}"); time.sleep(delay)
else:
    print(f"Failed to log artifacts after {retries} attempts.")   # swallowed; run continues

After exhausting retries it only prints — the run continues and finishes FINISHED. The artifacts were staged in a tempfile.TemporaryDirectory (_base_.py:254) that is deleted on context exit regardless of upload success, so on failure the staged artifacts are lost (recoverable only if the solver also persists raw outputs elsewhere).

Suggested fixes

1. Don't rely on set_experiment's global; use an explicit id. In _base_.py, resolve the experiment to an id once via a get-or-create helper that swallows RESOURCE_ALREADY_EXISTS, and pass experiment_id=... directly to start_run. This sidesteps the MLflow early-return bug regardless of concurrency.
2. Make experiment creation race-tolerant (the get-or-create helper above), and/or pre-create the experiment a single time before fanning out parallel runs.
3. Make artifact-upload failure loud. After exhausting retries, raise (so the orchestrator can mark the run FAILED) or set the run status to FAILED and tag it — so a failed upload can't masquerade as a clean run.
4. (Optional) On upload failure, copy the staged temp dir to a persistent location instead of letting TemporaryDirectory delete it.

Environment

• adept commit ba10baa (as run)
• mlflow 3.12.0
• Tracking server behind a reverse proxy (/ajax-api/2.0); Default experiment artifact root /mlfs/mlruns/0 (server-local), named experiments on S3.

[physicistphil]

Claude generated above. Let me know what the best fix would be; from here on out I'll just create the experiments ahead of time