Mlflow implementation (#2)

Reviewed-on: #2

src/commands/infra.py

@@ -77,7 +77,8 @@ def setup(
     if outputs.get("SageMakerRoleArn"):
         CONSOLE.print(f"[green]✓[/green] IAM role:  {outputs['SageMakerRoleArn']}")
     if cfg.mlflow.mode is MlflowMode.create and outputs.get("MlflowTrackingServerArn"):
-        CONSOLE.print(f"[green]✓[/green] MLflow:    {outputs['MlflowTrackingServerArn']}")
+        mlflow_name = outputs.get("MlflowTrackingServerName", cfg.managed_mlflow_tracking_server_name)
+        CONSOLE.print(f"[green]✓[/green] MLflow:    {mlflow_name}")
     elif cfg.mlflow.mode is MlflowMode.existing:
         CONSOLE.print(f"[green]✓[/green] MLflow:    {cfg.mlflow.tracking_server_name}")
     CONSOLE.print("\n[bold green]Infrastructure ready.[/bold green]")
@@ -102,7 +103,7 @@ def status(config: str = CONFIG_OPT) -> None:
         if cfg.mlflow.mode is not MlflowMode.disabled:
             table.add_row(
                 "MLflow",
-                cfg.mlflow.tracking_server_name or "-",
+                cfg.effective_mlflow_tracking_server_name or "-",
                 "[red]unknown[/red]",
                 "-",
             )
@@ -126,7 +127,7 @@ def status(config: str = CONFIG_OPT) -> None:
     if cfg.mlflow.mode is MlflowMode.create:
         table.add_row(
             "MLflow",
-            cfg.mlflow.tracking_server_name or "-",
+            outputs.get("MlflowTrackingServerName", cfg.managed_mlflow_tracking_server_name),
             "[green]managed[/green]",
             outputs.get("MlflowTrackingServerArn", outputs.get("MlflowArtifactUri", "-")),
         )
@@ -149,6 +150,35 @@ def status(config: str = CONFIG_OPT) -> None:
     CONSOLE.print(table)
 
 
+@app.command(name="mlflow-url")
+def mlflow_url(config: str = CONFIG_OPT) -> None:
+    """Print a presigned URL for the configured MLflow tracking server."""
+    cfg = load_cfg(config)
+    tracking_server_name = cfg.effective_mlflow_tracking_server_name
+    if not tracking_server_name:
+        CONSOLE.print("[red]MLflow is disabled in config.yaml.[/red]")
+        raise typer.Exit(1)
+
+    try:
+        url = mlflow.create_presigned_tracking_server_url(
+            cfg.aws.region,
+            cfg.aws.profile,
+            tracking_server_name,
+        )
+    except Exception as e:
+        CONSOLE.print("[yellow]Could not create a SageMaker MLflow UI URL.[/yellow]")
+        CONSOLE.print(f"Tracking server: [cyan]{tracking_server_name}[/cyan]")
+        CONSOLE.print(f"Reason: {e}")
+        CONSOLE.print(
+            "This command can create presigned URLs only for MLflow tracking servers managed by "
+            "Amazon SageMaker. If this is an external MLflow server, open it with that server's own URL."
+        )
+        raise typer.Exit(1)
+
+    CONSOLE.print(f"MLflow tracking server: [cyan]{tracking_server_name}[/cyan]")
+    CONSOLE.print(f"MLflow UI: {url}")
+
+
 @app.command()
 def destroy(
     config: str = CONFIG_OPT,
@@ -209,6 +239,7 @@ def _role_name(configured_name: str, role_arn: str) -> str:
         return role_arn.rsplit("/", 1)[-1]
     return "-"
 
+
 def _destroy_account_id(config_path: str, cfg: Config) -> str:
     config_dir = str(Path(config_path).parent)
     state = read_infra_state(config_dir)