nemo-evaluator-launcher 0.1.11__py3-none-any.whl → 0.1.13__py3-none-any.whl

nemo_evaluator_launcher/executors/slurm/executor.py

@@ -174,10 +174,11 @@ class SlurmExecutor(BaseExecutor):
             for idx, (slurm_job_id, remote_runsub_path) in enumerate(
                 zip(slurm_job_ids, remote_runsub_paths)
             ):
+                job_id = generate_job_id(invocation_id, idx)
                 db.write_job(
                     job=JobData(
                         invocation_id=invocation_id,
-                        job_id=generate_job_id(invocation_id, idx),
+                        job_id=job_id,
                         timestamp=time.time(),
                         executor="slurm",
                         data={
@@ -204,7 +205,7 @@ class SlurmExecutor(BaseExecutor):
         """
         db = ExecutionDB()
 
-        # If id looks like an invocation_id (no dot), get all jobs for it
+        # If id looks like an invocation_id
         if "." not in id:
             jobs = db.get_jobs(id)
             if not jobs:
@@ -605,20 +606,27 @@ def _create_slurm_sbatch_script(
         s += "kill $SERVER_PID  # terminate the server to finish gracefully\n\n"
 
     # auto-export
-    if cfg.execution.get("auto_export", {}).get("destinations", []):
-        s += _generate_auto_export_section(cfg, job_id)
+    ae_cfg = cfg.execution.get("auto_export")
+    destinations: list = []
+    if isinstance(ae_cfg, list):
+        destinations = list(ae_cfg)
+    elif isinstance(ae_cfg, dict) or isinstance(ae_cfg, DictConfig):
+        destinations = list(ae_cfg.get("destinations", []) or [])
+
+    if destinations:
+        export_env = dict(cfg.execution.get("env_vars", {}).get("export", {}) or {})
+        s += _generate_auto_export_section(cfg, job_id, destinations, export_env)
 
     return s
 
 
 def _generate_auto_export_section(
     cfg: DictConfig,
-    job_id: str,  # Complete job_id string
+    job_id: str,
+    destinations: list,
+    export_env: dict,
 ) -> str:
     """Generate simple auto-export section for sbatch script."""
-    auto_export_config = cfg.execution.get("auto_export", {})
-    destinations = auto_export_config.get("destinations", [])
-
     if not destinations:
         return ""
 
@@ -626,18 +634,65 @@ def _generate_auto_export_section(
     s += "EVAL_EXIT_CODE=$?\n"
     s += "if [ $EVAL_EXIT_CODE -eq 0 ]; then\n"
     s += "    echo 'Evaluation completed successfully. Starting auto-export...'\n"
-    s += "    set +e\n"  # per exporter failure allowed
+    s += "    set +e\n"
     s += "    set +x\n"
+    s += "    set +u\n"
     s += '    cd "$TASK_DIR/artifacts"\n'
-    auto_export_cfg = OmegaConf.to_container(
-        cfg.execution.get("auto_export", {}), resolve=True
+
+    # Work with DictConfig; convert only for YAML at the end
+    exec_type = (
+        cfg.execution.type
+        if hasattr(cfg.execution, "type")
+        else cfg.execution.get("type", "slurm")
     )
-    yaml_str = yaml.safe_dump(
-        {"execution": {"auto_export": auto_export_cfg}}, sort_keys=False
+    eval_tasks = (
+        list(cfg.evaluation.tasks)
+        if hasattr(cfg, "evaluation") and hasattr(cfg.evaluation, "tasks")
+        else list((cfg.get("evaluation", {}) or {}).get("tasks", []) or [])
     )
+    export_block = cfg.get("export", {}) or {}
+
+    payload = {
+        "execution": {
+            "auto_export": {
+                "destinations": list(destinations),
+                **({"env_vars": dict(export_env)} if export_env else {}),
+            },
+            "type": exec_type,
+        },
+        "evaluation": {"tasks": eval_tasks},
+    }
+    if export_block:
+        # Convert just this block to plain for YAML
+        payload["export"] = (
+            OmegaConf.to_object(export_block)
+            if OmegaConf.is_config(export_block)
+            else dict(export_block)
+        )
+
+    # Final YAML (single conversion at the end)
+    payload_clean = OmegaConf.to_container(OmegaConf.create(payload), resolve=True)
+    yaml_str = yaml.safe_dump(payload_clean, sort_keys=False)
     s += "    cat > export_config.yml << 'EOF'\n"
     s += yaml_str
     s += "EOF\n"
+
+    # write launcher config as config.yml for exporters (no core command)
+    submitted_yaml = yaml.safe_dump(
+        OmegaConf.to_container(cfg, resolve=True), sort_keys=False
+    )
+    s += "    cat > config.yml << 'EOF'\n"
+    s += submitted_yaml
+    s += "EOF\n"
+
+    # Export host only env before running auto export
+    for k, v in (export_env or {}).items():
+        if isinstance(v, str) and re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", v):
+            s += f'    export {k}="${{{v}}}"\n'
+        else:
+            esc = str(v).replace('"', '\\"')
+            s += f'    export {k}="{esc}"\n'
+
     for dest in destinations:
         s += f"    echo 'Exporting to {dest}...'\n"
         s += f"    nemo-evaluator-launcher export {job_id} --dest {dest} || echo 'Export to {dest} failed'\n"
@@ -656,7 +711,9 @@ def _open_master_connection(
     socket: str,
 ) -> str | None:
     ssh_command = f"ssh -MNf -S {socket} {username}@{hostname}"
-    completed_process = subprocess.run(args=shlex.split(ssh_command))
+    completed_process = subprocess.run(
+        args=shlex.split(ssh_command), capture_output=True
+    )
     if completed_process.returncode == 0:
         return socket
     return None
@@ -694,12 +751,17 @@ def _make_remote_execution_output_dir(
     ssh_command.append(f"{username}@{hostname}")
     ssh_command.append(mkdir_command)
     ssh_command = " ".join(ssh_command)
-    completed_process = subprocess.run(args=shlex.split(ssh_command))
+    completed_process = subprocess.run(
+        args=shlex.split(ssh_command), capture_output=True
+    )
     if completed_process.returncode != 0:
+        error_msg = (
+            completed_process.stderr.decode("utf-8")
+            if completed_process.stderr
+            else "Unknown error"
+        )
         raise RuntimeError(
-            "failed to make a remote execution output dir\n{}".format(
-                completed_process.stderr.decode("utf-8")
-            )
+            "failed to make a remote execution output dir\n{}".format(error_msg)
         )
 
 
@@ -725,13 +787,16 @@ def _rsync_upload_rundirs(
     remote_destination_str = f"{username}@{hostname}:{remote_target}"
     local_sources_str = " ".join(map(str, local_sources))
     rsync_upload_command = f"rsync -qcaz {local_sources_str} {remote_destination_str}"
-    completed_process = subprocess.run(args=shlex.split(rsync_upload_command))
+    completed_process = subprocess.run(
+        args=shlex.split(rsync_upload_command), capture_output=True
+    )
     if completed_process.returncode != 0:
-        raise RuntimeError(
-            "failed to upload local sources\n{}".format(
-                completed_process.stderr.decode("utf-8")
-            )
+        error_msg = (
+            completed_process.stderr.decode("utf-8")
+            if completed_process.stderr
+            else "Unknown error"
         )
+        raise RuntimeError("failed to upload local sources\n{}".format(error_msg))
 
 
 def _sbatch_remote_runsubs(
@@ -757,10 +822,9 @@ def _sbatch_remote_runsubs(
         args=shlex.split(ssh_command), capture_output=True
     )
     if completed_process.returncode != 0:
+        error_msg = completed_process.stderr.decode("utf-8")
         raise RuntimeError(
-            "failed to submit sbatch scripts for execution\n{}".format(
-                completed_process.stderr.decode("utf-8")
-            )
+            "failed to submit sbatch scripts for execution\n{}".format(error_msg)
         )
 
     sbatch_output = completed_process.stdout.decode("utf-8")

nemo_evaluator_launcher/exporters/base.py

@@ -70,6 +70,15 @@ class BaseExporter(ABC):
 
     def get_job_paths(self, job_data: JobData) -> Dict[str, Any]:
         """Get result paths based on executor type from job metadata."""
+        # Special case: remote executor artifacts accessed locally (remote auto-export)
+        if job_data.data.get("storage_type") == "remote_local":
+            output_dir = Path(job_data.data["output_dir"])
+            return {
+                "artifacts_dir": output_dir / "artifacts",
+                "logs_dir": output_dir / "logs",
+                "storage_type": "remote_local",
+            }
+
         if job_data.executor == "local":
             output_dir = Path(job_data.data["output_dir"])
             return {

nemo_evaluator_launcher/exporters/gsheets.py

@@ -15,6 +15,7 @@
 #
 """Google Sheets evaluation results exporter."""
 
+import os
 import shutil
 import tempfile
 from pathlib import Path
@@ -89,28 +90,38 @@ class GSheetsExporter(BaseExporter):
             }
 
         try:
+            # Load exporter config from the first job (supports job-embedded config and CLI overrides)
+            first_job = next(iter(jobs.values()))
+            gsheets_config = extract_exporter_config(first_job, "gsheets", self.config)
+
             # Connect to Google Sheets
-            service_account_file = self.config.get("service_account_file")
-            spreadsheet_name = self.config.get(
+            service_account_file = gsheets_config.get("service_account_file")
+            spreadsheet_name = gsheets_config.get(
                 "spreadsheet_name", "NeMo Evaluator Launcher Results"
             )
 
             if service_account_file:
-                gc = gspread.service_account(filename=service_account_file)
+                gc = gspread.service_account(
+                    filename=os.path.expanduser(service_account_file)
+                )
             else:
                 gc = gspread.service_account()
 
             # Get or create spreadsheet
+            spreadsheet_id = gsheets_config.get("spreadsheet_id")
             try:
-                sh = gc.open(spreadsheet_name)
+                if spreadsheet_id:
+                    sh = gc.open_by_key(spreadsheet_id)
+                else:
+                    sh = gc.open(spreadsheet_name)
                 logger.info(f"Opened existing spreadsheet: {spreadsheet_name}")
             except gspread.SpreadsheetNotFound:
+                if spreadsheet_id:
+                    raise  # Can't create with explicit ID
                 sh = gc.create(spreadsheet_name)
                 logger.info(f"Created new spreadsheet: {spreadsheet_name}")
-                sh.share("", perm_type="anyone", role="reader")
 
             worksheet = sh.sheet1
-
             # Extract metrics from ALL jobs first to determine headers
             all_job_metrics = {}
             results = {}
@@ -226,16 +237,23 @@ class GSheetsExporter(BaseExporter):
                 )
 
                 if service_account_file:
-                    gc = gspread.service_account(filename=service_account_file)
+                    gc = gspread.service_account(
+                        filename=os.path.expanduser(service_account_file)
+                    )
                 else:
                     gc = gspread.service_account()
 
                 # Get or create spreadsheet
+                spreadsheet_id = gsheets_config.get("spreadsheet_id")
                 try:
-                    sh = gc.open(spreadsheet_name)
+                    if spreadsheet_id:
+                        sh = gc.open_by_key(spreadsheet_id)
+                    else:
+                        sh = gc.open(spreadsheet_name)
                 except gspread.SpreadsheetNotFound:
+                    if spreadsheet_id:
+                        raise  # Can't create with explicit ID
                     sh = gc.create(spreadsheet_name)
-                    sh.share("", perm_type="anyone", role="reader")
 
                 worksheet = sh.sheet1
 

nemo_evaluator_launcher/exporters/local.py

@@ -74,6 +74,9 @@ class LocalExporter(BaseExporter):
             # Stage artifacts per storage type
             if paths["storage_type"] == "local_filesystem":
                 exported_files = self._copy_local_artifacts(paths, job_export_dir, cfg)
+            elif paths["storage_type"] == "remote_local":
+                # Same as local_filesystem (we're on the remote machine, accessing locally)
+                exported_files = self._copy_local_artifacts(paths, job_export_dir, cfg)
             elif paths["storage_type"] == "remote_ssh":
                 exported_files = ssh_download_artifacts(
                     paths, job_export_dir, cfg, None
@@ -125,6 +128,8 @@ class LocalExporter(BaseExporter):
                     logger.warning(f"Failed to create {fmt} summary: {e}")
                     msg += " (summary failed)"
 
+            meta["output_dir"] = str(job_export_dir.resolve())
+
             return ExportResult(
                 success=True, dest=str(job_export_dir), message=msg, metadata=meta
             )

nemo_evaluator_launcher/exporters/mlflow.py

@@ -15,6 +15,7 @@
 #
 """Evaluation results exporter for MLflow tracking."""
 
+import os
 import shutil
 import tempfile
 from pathlib import Path
@@ -37,6 +38,7 @@ from nemo_evaluator_launcher.exporters.registry import register_exporter
 from nemo_evaluator_launcher.exporters.utils import (
     extract_accuracy_metrics,
     extract_exporter_config,
+    get_artifact_root,
     get_available_artifacts,
     get_benchmark_info,
     get_task_name,
@@ -100,6 +102,21 @@ class MLflowExporter(BaseExporter):
             # Extract config using common utility
             mlflow_config = extract_exporter_config(job_data, "mlflow", self.config)
 
+            # resolve tracking_uri with fallbacks
+            tracking_uri = mlflow_config.get("tracking_uri")
+            if not tracking_uri:
+                tracking_uri = os.getenv("MLFLOW_TRACKING_URI")
+            # allow env var name
+            if tracking_uri and "://" not in tracking_uri:
+                tracking_uri = os.getenv(tracking_uri, tracking_uri)
+
+            if not tracking_uri:
+                return ExportResult(
+                    success=False,
+                    dest="mlflow",
+                    message="tracking_uri is required (set export.mlflow.tracking_uri or MLFLOW_TRACKING_URI)",
+                )
+
             # Extract metrics
             log_metrics = mlflow_config.get("log_metrics", [])
             accuracy_metrics = extract_accuracy_metrics(
@@ -112,12 +129,6 @@ class MLflowExporter(BaseExporter):
                 )
 
             # Set up MLflow
-            tracking_uri = mlflow_config.get("tracking_uri")
-            if not tracking_uri:
-                return ExportResult(
-                    success=False, dest="mlflow", message="tracking_uri is required"
-                )
-
             tracking_uri = tracking_uri.rstrip("/")
             mlflow.set_tracking_uri(tracking_uri)
 
@@ -253,37 +264,91 @@ class MLflowExporter(BaseExporter):
         try:
             # Use LocalExporter to get files locally first
             temp_dir = tempfile.mkdtemp(prefix="mlflow_artifacts_")
-            local_exporter = LocalExporter({"output_dir": temp_dir})
+            local_exporter = LocalExporter(
+                {
+                    "output_dir": temp_dir,
+                    "copy_logs": mlflow_config.get(
+                        "log_logs", mlflow_config.get("copy_logs", False)
+                    ),
+                    "only_required": mlflow_config.get("only_required", True),
+                    "format": mlflow_config.get("format", None),
+                    "log_metrics": mlflow_config.get("log_metrics", []),
+                    "output_filename": mlflow_config.get("output_filename", None),
+                }
+            )
             local_result = local_exporter.export_job(job_data)
 
             if not local_result.success:
                 logger.error(f"Failed to download artifacts: {local_result.message}")
                 return []
 
-            artifacts_dir = Path(local_result.dest) / "artifacts"
-            logged_names = []
+            base_dir = Path(local_result.dest)
+            artifacts_dir = base_dir / "artifacts"
+            logs_dir = base_dir / "logs"
+            logged_names: list[str] = []
 
-            task_name = get_task_name(job_data)
-            artifact_path = task_name
+            artifact_path = get_artifact_root(job_data)  # "<harness>.<benchmark>"
 
             # Log config at root level
-            with tempfile.TemporaryDirectory() as tmpdir:
-                cfg_file = Path(tmpdir) / "config.yaml"
-                with cfg_file.open("w") as f:
-                    yaml.dump(
-                        job_data.config or {},
-                        f,
-                        default_flow_style=False,
-                        sort_keys=False,
-                    )
-                mlflow.log_artifact(str(cfg_file))
-
-            # Then log results files
-            for fname in get_available_artifacts(artifacts_dir):
-                file_path = artifacts_dir / fname
-                if file_path.exists():
-                    mlflow.log_artifact(str(file_path), artifact_path=artifact_path)
-                    logged_names.append(fname)
+            cfg_logged = False
+            for fname in ("config.yml", "run_config.yml"):
+                p = artifacts_dir / fname
+                if p.exists():
+                    mlflow.log_artifact(str(p))
+                    cfg_logged = True
+                    break
+            if not cfg_logged:
+                with tempfile.TemporaryDirectory() as tmpdir:
+                    cfg_file = Path(tmpdir) / "config.yaml"
+                    with cfg_file.open("w") as f:
+                        yaml.dump(
+                            job_data.config or {},
+                            f,
+                            default_flow_style=False,
+                            sort_keys=False,
+                        )
+                    mlflow.log_artifact(str(cfg_file))
+
+            files_to_upload: list[Path] = []
+            if mlflow_config.get("only_required", True):
+                for fname in get_available_artifacts(artifacts_dir):
+                    p = artifacts_dir / fname
+                    if p.exists():
+                        files_to_upload.append(p)
+            else:
+                for p in artifacts_dir.iterdir():
+                    if p.is_file():
+                        files_to_upload.append(p)
+
+            for fpath in files_to_upload:
+                rel = fpath.relative_to(artifacts_dir).as_posix()
+                parent = os.path.dirname(rel)
+                mlflow.log_artifact(
+                    str(fpath),
+                    artifact_path=f"{artifact_path}/artifacts/{parent}".rstrip("/"),
+                )
+                logged_names.append(rel)
+
+            # Optionally upload logs under "<harness.task>/logs"
+            if mlflow_config.get("log_logs", False) and logs_dir.exists():
+                for p in logs_dir.rglob("*"):
+                    if p.is_file():
+                        mlflow.log_artifact(
+                            str(p),
+                            artifact_path=f"{artifact_path}/logs",
+                        )
+                        logged_names.append(f"logs/{p.name}")
+
+            # Debug summary of what we uploaded
+            logger.info(
+                f"MLflow upload summary: files={len(logged_names)}, only_required={mlflow_config.get('only_required', True)}, log_logs={mlflow_config.get('log_logs', False)}"
+            )
+            if logger.isEnabledFor(10):  # DEBUG
+                try:
+                    preview = "\n  - " + "\n  - ".join(sorted(logged_names)[:50])
+                    logger.debug(f"Uploaded files preview (first 50):{preview}")
+                except Exception:
+                    pass
 
             # cleanup temp
             shutil.rmtree(temp_dir)
@@ -312,6 +377,18 @@ class MLflowExporter(BaseExporter):
             # Extract config using common utility
             mlflow_config = extract_exporter_config(first_job, "mlflow", self.config)
 
+            # resolve tracking_uri with fallbacks
+            tracking_uri = mlflow_config.get("tracking_uri") or os.getenv(
+                "MLFLOW_TRACKING_URI"
+            )
+            if tracking_uri and "://" not in tracking_uri:
+                tracking_uri = os.getenv(tracking_uri, tracking_uri)
+            if not tracking_uri:
+                return {
+                    "success": False,
+                    "error": "tracking_uri is required (set export.mlflow.tracking_uri or MLFLOW_TRACKING_URI)",
+                }
+
             # Collect metrics from ALL jobs
             all_metrics = {}
             for job_id, job_data in jobs.items():
@@ -328,10 +405,6 @@ class MLflowExporter(BaseExporter):
                 }
 
             # Set up MLflow
-            tracking_uri = mlflow_config.get("tracking_uri")
-            if not tracking_uri:
-                return {"success": False, "error": "tracking_uri is required"}
-
             tracking_uri = tracking_uri.rstrip("/")
             mlflow.set_tracking_uri(tracking_uri)
 

nemo_evaluator_launcher/exporters/utils.py

@@ -148,15 +148,12 @@ def extract_exporter_config(
     """Extract and merge exporter configuration from multiple sources."""
     config = {}
 
-    # Get config from dedicated field
+    # root-level `export.<exporter-name>`
     if job_data.config:
-        execution_config = job_data.config.get("execution", {})
-        auto_export_config = execution_config.get("auto_export", {})
-        exporter_configs = auto_export_config.get("configs", {})
-        yaml_config = exporter_configs.get(exporter_name, {})
-
-        # No conversion needed
-        config.update(yaml_config)
+        export_block = (job_data.config or {}).get("export", {})
+        yaml_config = (export_block or {}).get(exporter_name, {})
+        if yaml_config:
+            config.update(yaml_config)
 
     # From webhook metadata (if triggered by webhook)
     if "webhook_metadata" in job_data.data:
@@ -167,8 +164,6 @@ def extract_exporter_config(
             "source_artifact": f"{webhook_data.get('artifact_name', 'unknown')}:{webhook_data.get('artifact_version', 'unknown')}",
             "config_source": webhook_data.get("config_file", "unknown"),
         }
-
-        # For W&B specifically, extract run info if available
         if exporter_name == "wandb" and webhook_data.get("webhook_source") == "wandb":
             wandb_specific = {
                 "entity": webhook_data.get("entity"),
@@ -176,10 +171,9 @@ def extract_exporter_config(
                 "run_id": webhook_data.get("run_id"),
             }
             webhook_config.update({k: v for k, v in wandb_specific.items() if v})
-
         config.update(webhook_config)
 
-    # Constructor config: allows CLI overrides
+    # allows CLI overrides
     if constructor_config:
         config.update(constructor_config)
 
@@ -269,6 +263,14 @@ def get_container_from_mapping(job_data: JobData) -> str:
         return None
 
 
+def get_artifact_root(job_data: JobData) -> str:
+    """Get artifact root from job data."""
+    bench = get_benchmark_info(job_data)
+    h = bench.get("harness", "unknown")
+    b = bench.get("benchmark", get_task_name(job_data))
+    return f"{h}.{b}"
+
+
 # =============================================================================
 # GITLAB DOWNLOAD
 # =============================================================================
@@ -288,91 +290,6 @@ def download_gitlab_artifacts(
         Dictionary mapping artifact names to local file paths
     """
     raise NotImplementedError("Downloading from gitlab is not implemented")
-    # TODO: rework this logic
-    # pipeline_id = paths["pipeline_id"]
-    # project_id = paths["project_id"]
-    # gitlab_token = os.getenv("GITLAB_TOKEN")
-    #
-    # if not gitlab_token:
-    #     raise RuntimeError(
-    #         "GITLAB_TOKEN environment variable required for GitLab remote downloads"
-    #     )
-    #
-    # # GitLab API endpoint for artifacts
-    # base_url = "TODO: replace"
-    # artifacts_url = "TODO: replace"
-    #
-    # headers = {"Private-Token": gitlab_token}
-    # downloaded_artifacts = {}
-    #
-    # try:
-    #     # Get pipeline jobs
-    #     response = requests.get(artifacts_url, headers=headers, timeout=30)
-    #     response.raise_for_status()
-    #     jobs = response.json()
-    #
-    #     for job in jobs:
-    #         if job.get("artifacts_file"):
-    #             job_id = job["id"]
-    #             job_name = job.get("name", f"job_{job_id}")
-    #             artifacts_download_url = (
-    #                 f"{base_url}/api/v4/projects/{project_id}/jobs/{job_id}/artifacts"
-    #             )
-    #
-    #             logger.info(f"Downloading artifacts from job: {job_name}")
-    #
-    #             # Download job artifacts
-    #             response = requests.get(
-    #                 artifacts_download_url, headers=headers, timeout=300
-    #             )
-    #             response.raise_for_status()
-    #
-    #             if extract_specific:
-    #                 # Extract specific files from ZIP
-    #                 with tempfile.NamedTemporaryFile(
-    #                     suffix=".zip", delete=False
-    #                 ) as temp_zip:
-    #                     temp_zip.write(response.content)
-    #                     temp_zip_path = temp_zip.name
-    #
-    #                 try:
-    #                     with zipfile.ZipFile(temp_zip_path, "r") as zip_ref:
-    #                         # Create artifacts directory
-    #                         artifacts_dir = export_dir / "artifacts"
-    #                         artifacts_dir.mkdir(parents=True, exist_ok=True)
-    #
-    #                         # Extract to be logged artifacts
-    #                         for member in zip_ref.namelist():
-    #                             filename = Path(member).name
-    #                             if filename in get_relevant_artifacts():
-    #                                 # Extract the file
-    #                                 source = zip_ref.open(member)
-    #                                 target_path = artifacts_dir / filename
-    #                                 with open(target_path, "wb") as f:
-    #                                     f.write(source.read())
-    #                                 source.close()
-    #
-    #                                 downloaded_artifacts[filename] = target_path
-    #                                 logger.info(f"Extracted: {filename}")
-    #                 finally:
-    #                     os.unlink(temp_zip_path)
-    #             else:
-    #                 # Save as ZIP files (original behavior)
-    #                 artifacts_zip = export_dir / f"job_{job_id}_artifacts.zip"
-    #                 with open(artifacts_zip, "wb") as f:
-    #                     f.write(response.content)
-    #
-    #                 downloaded_artifacts[f"job_{job_id}_artifacts.zip"] = artifacts_zip
-    #                 logger.info(f"Downloaded: {artifacts_zip.name}")
-    #
-    # except requests.RequestException as e:
-    #     logger.error(f"GitLab API request failed: {e}")
-    #     raise RuntimeError(f"GitLab API request failed: {e}")
-    # except Exception as e:
-    #     logger.error(f"GitLab remote download failed: {e}")
-    #     raise RuntimeError(f"GitLab remote download failed: {e}")
-    #
-    # return downloaded_artifacts
 
 
 # =============================================================================
@@ -522,16 +439,16 @@ def ssh_download_artifacts(
 
 def _get_artifacts_dir(paths: Dict[str, Any]) -> Path:
     """Get artifacts directory from paths."""
-    if paths["storage_type"] == "local_filesystem":
-        return paths["artifacts_dir"]
-    elif paths["storage_type"] == "gitlab_ci_local":
-        return paths["artifacts_dir"]
-    elif paths["storage_type"] == "remote_ssh":
-        return None
-    else:
-        logger.error(f"Unsupported storage type: {paths['storage_type']}")
+    storage_type = paths.get("storage_type")
+
+    # For SSH-based remote access, artifacts aren't available locally yet
+    if storage_type == "remote_ssh":
         return None
 
+    # For all local access (local_filesystem, remote_local, gitlab_ci_local)
+    # return the artifacts_dir from paths
+    return paths.get("artifacts_dir")
+
 
 def _extract_metrics_from_results(results: dict) -> Dict[str, float]:
     """Extract metrics from a 'results' dict (with optional 'groups'/'tasks')."""