Always show console link, conditional show studio link

Author: mollyheamazon

sagemaker-train/src/sagemaker/train/common_utils/trainer_wait.py

@@ -307,37 +307,39 @@ def get_cached_mlflow_url():
                     header_table.add_row("TrainingJob Name", f"[bold green]{training_job.training_job_name}[/bold green]")
                     header_table.add_row("TrainingJob ARN", f"[dim]{training_job.training_job_arn}[/dim]")
 
-                    # Build links row
-                    links = []
+                    # Build links rows
+                    links_row1 = []
+                    links_row2 = []
                     try:
                         from sagemaker.train.common_utils.metrics_visualizer import (
                             _is_in_studio, get_console_job_url, get_cloudwatch_logs_url, get_studio_url
                         )
+                        console_url = get_console_job_url(training_job.training_job_arn)
+                        if console_url:
+                            links_row1.append(f"[bright_blue underline][link={console_url}]🔗 Training Job (Console)[/link][/bright_blue underline]")
                         if _is_in_studio():
                             studio_url = get_studio_url(training_job)
                             if studio_url:
-                                links.append(f"[bright_blue underline][link={studio_url}]🔗 Training Job (Studio)[/link][/bright_blue underline]")
-                        else:
-                            console_url = get_console_job_url(training_job.training_job_arn)
-                            if console_url:
-                                links.append(f"[bright_blue underline][link={console_url}]🔗 Training Job[/link][/bright_blue underline]")
+                                links_row1.append(f"[bright_blue underline][link={studio_url}]🔗 Training Job (Studio)[/link][/bright_blue underline]")
                         cw_url = get_cloudwatch_logs_url(training_job.training_job_arn)
                         if cw_url:
-                            links.append(f"[bright_blue underline][link={cw_url}]🔗 CloudWatch Logs[/link][/bright_blue underline]")
+                            links_row2.append(f"[bright_blue underline][link={cw_url}]🔗 CloudWatch Logs[/link][/bright_blue underline]")
                     except Exception:
                         pass
                     if has_mlflow_config:
                         cached_url = get_cached_mlflow_url()
                         if cached_url:
-                            links.append(f"[bright_blue underline][link={cached_url}]🔗 MLflow Experiment[/link][/bright_blue underline]")
+                            links_row2.append(f"[bright_blue underline][link={cached_url}]🔗 MLflow Experiment[/link][/bright_blue underline]")
                         elif mlflow_link_cache['error']:
                             header_table.add_row("MLflow Experiment", f"[red]{mlflow_link_cache['error']}[/red]")
                     if has_mlflow_config:
                         exp_name = training_job.mlflow_config.mlflow_experiment_name if hasattr(training_job, 'mlflow_config') else None
                         if exp_name and not _is_unassigned_attribute(exp_name):
                             header_table.add_row("MLflow Experiment", f"{exp_name}")
-                    if links:
-                        header_table.add_row("Links", " | ".join(links))
+                    if links_row1:
+                        header_table.add_row("Links", " | ".join(links_row1))
+                    if links_row2:
+                        header_table.add_row("" if links_row1 else "Links", " | ".join(links_row2))
 
                     status_table = Table(show_header=False, box=None, padding=(0, 1))
                     status_table.add_column("Property", style="cyan bold", width=20)

sagemaker-train/src/sagemaker/train/evaluate/execution.py

@@ -1051,45 +1051,51 @@ def wait(
                     if job_arn_entries:
                         links_table = Table(show_header=True, header_style="bold magenta", box=None, padding=(0, 1))
                         links_table.add_column("Step", style="cyan", width=20)
-                        links_table.add_column("Job Link", style="dim")
-                        links_table.add_column("Logs", style="dim")
-                        links_table.add_column("Job ARN", style="dim", overflow="fold")
+                        links_table.add_column("Console", style="dim")
                         from sagemaker.train.common_utils.metrics_visualizer import (
                             _is_in_studio, _parse_job_arn, _get_studio_base_url,
                             get_console_job_url, get_cloudwatch_logs_url,
                         )
                         in_studio = _is_in_studio()
                         studio_base = _get_studio_base_url(region) if in_studio else ""
+                        if in_studio:
+                            links_table.add_column("Studio", style="dim")
+                        links_table.add_column("Logs", style="dim")
+                        links_table.add_column("Job ARN", style="dim", overflow="fold")
                         studio_path_map = {
                             "training-job/": "jobs/train/",
                             "processing-job/": "jobs/processing/",
                             "transform-job/": "jobs/transform/",
                         }
                         for entry in job_arn_entries:
-                            job_link = ""
+                            console_link = ""
                             logs_link = ""
+                            studio_link = ""
                             try:
                                 arn = entry['job_arn']
+                                url = get_console_job_url(arn)
+                                if url:
+                                    console_link = f"[bright_blue underline][link={url}]🔗 link[/link][/bright_blue underline]"
+                                cw_url = get_cloudwatch_logs_url(arn)
+                                if cw_url:
+                                    logs_link = f"[bright_blue underline][link={cw_url}]🔗 logs[/link][/bright_blue underline]"
                                 if in_studio and studio_base:
                                     parsed = _parse_job_arn(arn)
                                     if parsed:
                                         _, resource = parsed
                                         for prefix, path in studio_path_map.items():
                                             if resource.startswith(prefix):
                                                 job_name = resource.split("/", 1)[1]
-                                                url = f"{studio_base}/{path}{job_name}"
-                                                job_link = f"[bright_blue underline][link={url}]🔗 link[/link][/bright_blue underline]"
+                                                s_url = f"{studio_base}/{path}{job_name}"
+                                                studio_link = f"[bright_blue underline][link={s_url}]🔗 studio[/link][/bright_blue underline]"
                                                 break
-                                else:
-                                    url = get_console_job_url(arn)
-                                    if url:
-                                        job_link = f"[bright_blue underline][link={url}]🔗 link[/link][/bright_blue underline]"
-                                cw_url = get_cloudwatch_logs_url(arn)
-                                if cw_url:
-                                    logs_link = f"[bright_blue underline][link={cw_url}]🔗 logs[/link][/bright_blue underline]"
                             except Exception:
                                 pass
-                            links_table.add_row(entry['step_name'], job_link, logs_link, entry['job_arn'])
+                            row = [entry['step_name'], console_link]
+                            if in_studio:
+                                row.append(studio_link)
+                            row.extend([logs_link, entry['job_arn']])
+                            links_table.add_row(*row)
                         content_parts.append(Text(""))
                         content_parts.append(Text("Job ARNs", style="bold magenta"))
                         content_parts.append(links_table)