fix: artifacts with dynamic names (#706)

Cahllagerfeld · web-flow · commit c868b7ff0db7 · 2024-12-04T09:36:11.000+01:00
diff --git a/src/app/runs/[id]/_Tabs/Configuration/CodeCollapsible.tsx b/src/app/runs/[id]/_Tabs/Configuration/CodeCollapsible.tsx
@@ -1,23 +1,22 @@
 import ChevronDown from "@/assets/icons/chevron-down.svg?react";
+import { Codesnippet } from "@/components/CodeSnippet";
+import { getRunSnippet } from "@/lib/code-snippets";
+import { PipelineRun } from "@/types/pipeline-runs";
 import {
 	CollapsibleContent,
 	CollapsibleHeader,
 	CollapsiblePanel,
 	CollapsibleTrigger
 } from "@zenml-io/react-component-library";
-import { PipelineRun } from "@/types/pipeline-runs";
 import { useState } from "react";
-import { Codesnippet } from "@/components/CodeSnippet";
 
 type Props = {
 	runId: PipelineRun["id"];
 };
 
 export function CodeCollapsible({ runId }: Props) {
 	const [open, setOpen] = useState(true);
-	const runCode = `from zenml.client import Client
-run = Client().get_pipeline_run('${runId}')
-config = run.config`;
+	const runCode = getRunSnippet(runId);
 
 	return (
 		<CollapsiblePanel open={open} onOpenChange={setOpen}>
diff --git a/src/app/runs/[id]/_Tabs/Configuration/index.tsx b/src/app/runs/[id]/_Tabs/Configuration/index.tsx
@@ -44,6 +44,11 @@ export function ConfigurationTab() {
 				// eslint-disable-next-line @typescript-eslint/no-explicit-any
 				data={(data.metadata?.config.settings as { [key: string]: any })?.resources || {}}
 			/>
+			<NestedCollapsible
+				isInitialOpen
+				title="Substitutions"
+				data={data.metadata?.config.substitutions}
+			/>
 			<NestedCollapsible isInitialOpen title="Extra" data={data.metadata?.config.extra} />
 		</div>
 	);
diff --git a/src/app/runs/[id]/useDag.tsx b/src/app/runs/[id]/useDag.tsx
@@ -37,7 +37,11 @@ export function useDag() {
 	}, [pipelineDeployment.data?.metadata?.step_configurations]);
 
 	const realNodes = useMemo(() => {
-		return extractExistingNodes((pipelineRun.data?.metadata?.steps as StepDict) ?? {});
+		return extractExistingNodes(
+			(pipelineRun.data?.metadata?.steps as StepDict) ?? {},
+			(pipelineRun.data?.metadata?.step_substitutions as Record<string, Record<string, string>>) ||
+				{}
+		);
 	}, [pipelineRun.data?.metadata?.steps]);
 
 	const onDagreLayout = useCallback(() => {
diff --git a/src/app/settings/secrets/columns.tsx b/src/app/settings/secrets/columns.tsx
@@ -2,6 +2,7 @@ import LockIcon from "@/assets/icons/Lock.svg?react";
 import { CopyButton } from "@/components/CopyButton";
 import { DisplayDate } from "@/components/DisplayDate";
 import { InlineAvatar } from "@/components/InlineAvatar";
+import { getSecretSnippet } from "@/lib/code-snippets";
 import { getUsername } from "@/lib/user";
 import { routes } from "@/router/routes";
 import { SecretNamespace } from "@/types/secret";
@@ -17,9 +18,7 @@ export const secretsColumns: ColumnDef<SecretNamespace>[] = [
 		header: "Secret",
 		accessorFn: (row) => row.name,
 		cell: ({ getValue, row }) => {
-			const code = `from zenml.client import Client
-secret = Client().get_secret("${row.original.name}")
-`;
+			const code = getSecretSnippet(row.original.name);
 
 			return (
 				<div className="flex items-center space-x-2">
diff --git a/src/app/stacks/StackList.tsx b/src/app/stacks/StackList.tsx
@@ -7,16 +7,17 @@ import { routes } from "@/router/routes";
 import { useQuery } from "@tanstack/react-query";
 import { Button, DataTable, Skeleton } from "@zenml-io/react-component-library";
 import { Link } from "react-router-dom";
-import { getStackColumns } from "./columns";
 import { useStacklistQueryParams } from "./service";
 import { StackListQueryParams } from "../../types/stack";
+import { useStackColumns } from "./columns";
 
 type Props = {
 	fixedQueryParams?: StackListQueryParams;
 };
 
 export function StackList({ fixedQueryParams = {} }: Props) {
 	const queryParams = useStacklistQueryParams();
+	const columns = useStackColumns();
 	const { refetch, data } = useQuery({
 		...stackQueries.stackList({ ...queryParams, sort_by: "desc:updated", ...fixedQueryParams }),
 		throwOnError: true
@@ -43,7 +44,7 @@ export function StackList({ fixedQueryParams = {} }: Props) {
 				<div className="flex flex-col items-center gap-5">
 					<div className="w-full">
 						{data ? (
-							<DataTable columns={getStackColumns()} data={data.items} />
+							<DataTable columns={columns} data={data.items} />
 						) : (
 							<Skeleton className="h-[500px] w-full" />
 						)}
diff --git a/src/app/stacks/columns.tsx b/src/app/stacks/columns.tsx
@@ -7,10 +7,11 @@ import { Stack } from "@/types/stack";
 import { User } from "@/types/user";
 import { ColumnDef } from "@tanstack/react-table";
 import { Avatar, AvatarFallback } from "@zenml-io/react-component-library";
+import { useMemo } from "react";
 import { StackActionsMenu } from "./ActionsDropdown";
 
-export function getStackColumns(): ColumnDef<Stack>[] {
-	return [
+export function useStackColumns(): ColumnDef<Stack>[] {
+	const cols: ColumnDef<Stack>[] = [
 		{
 			id: "name",
 			header: "Stack",
@@ -67,4 +68,6 @@ export function getStackColumns(): ColumnDef<Stack>[] {
 			}
 		}
 	];
+
+	return useMemo(() => cols, []);
 }
diff --git a/src/components/artifacts/JsonVisualization.tsx b/src/components/artifacts/JsonVisualization.tsx
@@ -0,0 +1,23 @@
+import { isString } from "@/lib/type-guards";
+import { Codesnippet } from "../CodeSnippet";
+import { Props } from "./Visualization";
+
+export function JSONVisualization({ content }: Props) {
+	const json = parseJSON(content);
+	return (
+		<Codesnippet
+			fullWidth
+			highlightCode
+			code={isString(json) ? json : JSON.stringify(json, null, 2)}
+		/>
+	);
+}
+
+function parseJSON(content: string): unknown {
+	try {
+		const parsedJSON = JSON.parse(content);
+		return parsedJSON;
+	} catch (e) {
+		return content;
+	}
+}
diff --git a/src/components/artifacts/Visualization.tsx b/src/components/artifacts/Visualization.tsx
@@ -12,6 +12,7 @@ import { HTMLVisualization } from "./HtmlVisualization";
 import { ImageVisualization } from "./ImageVisualization";
 import { InfoBox } from "../Infobox";
 import { useArtifactLoadConfirmationContext } from "@/context/VisualizationConfirmationContext";
+import { JSONVisualization } from "./JsonVisualization";
 
 const CSVVisualization = lazy(() => import("./CsvVizualization"));
 const MarkdownVisualization = lazy(() => import("./MarkdownVisualization"));
@@ -126,6 +127,7 @@ export function Visualization({ artifactVersionId, artifactName }: Visualization
 				{data.type === "html" && <HTMLVisualization content={data.value} />}
 				{data.type === "markdown" && <MarkdownVisualization content={data.value} />}
 				{data.type === "csv" && <CSVVisualization content={data.value} />}
+				{data.type === "json" && <JSONVisualization content={data.value} />}
 			</div>
 		</div>
 	);
@@ -165,14 +167,16 @@ function DownloadButton({
 		image: "png",
 		html: "html",
 		markdown: "md",
-		csv: "csv"
+		csv: "csv",
+		json: "json"
 	} as const;
 
 	const typeMap = {
 		image: "image/png",
 		html: "text/html",
 		markdown: "text/markdown",
-		csv: "text/csv"
+		csv: "text/csv",
+		json: "application/json"
 	} as const;
 
 	function prepareImagedownload() {
diff --git a/src/components/artifacts/artifact-node-sheet/DetailCards.tsx b/src/components/artifacts/artifact-node-sheet/DetailCards.tsx
@@ -7,9 +7,12 @@ import { ExecutionStatusIcon, getExecutionStatusTagColor } from "@/components/Ex
 import { InlineAvatar } from "@/components/InlineAvatar";
 import { Key, KeyValue, Value } from "@/components/KeyValue";
 import { useArtifactVersion } from "@/data/artifact-versions/artifact-version-detail-query";
+import { componentQueries } from "@/data/components";
 import { usePipelineRun } from "@/data/pipeline-runs/pipeline-run-detail-query";
 import { useStepDetail } from "@/data/steps/step-detail-query";
+import { getArtifactVersionSnippet } from "@/lib/code-snippets";
 import { routes } from "@/router/routes";
+import { useQuery } from "@tanstack/react-query";
 import {
 	Skeleton,
 	Tag,
@@ -21,8 +24,6 @@ import {
 import { Link } from "react-router-dom";
 import { Codesnippet } from "../../CodeSnippet";
 import { CollapsibleCard } from "../../CollapsibleCard";
-import { useQuery } from "@tanstack/react-query";
-import { componentQueries } from "../../../data/components";
 
 type Props = {
 	artifactVersionId: string;
@@ -251,16 +252,14 @@ export function DataCard({ artifactVersionId }: Props) {
 }
 
 export function CodeCard({ artifactVersionId }: Props) {
-	function returnConfigSchema(id: string) {
-		return `from zenml.client import Client
-
-step = Client().get_run_step(${id})
-config = step.config`;
-	}
-
 	return (
 		<CollapsibleCard initialOpen title="Code">
-			<Codesnippet fullWidth highlightCode wrap code={returnConfigSchema(artifactVersionId)} />
+			<Codesnippet
+				fullWidth
+				highlightCode
+				wrap
+				code={getArtifactVersionSnippet(artifactVersionId)}
+			/>
 		</CollapsibleCard>
 	);
 }
diff --git a/src/components/dag-visualizer/ArtifactNode.tsx b/src/components/dag-visualizer/ArtifactNode.tsx
@@ -1,4 +1,5 @@
 import Copy from "@/assets/icons/copy.svg?react";
+import { getArtifactVersionSnippet } from "@/lib/code-snippets";
 import { ArtifactVersion } from "@/types/artifact-versions";
 import { NodeProps, useStore } from "reactflow";
 import { ArtifactIcon } from "../ArtifactIcon";
@@ -51,10 +52,7 @@ export function ArtifactNode({ data, selected }: NodeProps<ArtifactVersion & { n
 					</div>
 					<CopyNodeButton
 						className="h-4 w-4 shrink-0 rounded-sm hover:bg-primary-100 active:bg-primary-200"
-						code={`from zenml.client import Client
-							
-artifact = Client().get_artifact_version("${data.id}")
-loaded_artifact = artifact.load()`}
+						code={getArtifactVersionSnippet(data.id)}
 						type="artifact"
 					>
 						<Copy className="h-3 w-3 fill-primary-400" />
diff --git a/src/components/dag-visualizer/StepNode.tsx b/src/components/dag-visualizer/StepNode.tsx
@@ -1,4 +1,5 @@
 import Copy from "@/assets/icons/copy.svg?react";
+import { getStepSnippet } from "@/lib/code-snippets";
 import { calculateTimeDifference } from "@/lib/dates";
 import { Step } from "@/types/steps";
 import { clsx } from "clsx";
@@ -49,9 +50,7 @@ export function StepNode({ data, selected }: NodeProps<Step>) {
 						<p className="truncate font-semibold">{data.name}</p>
 						<CopyNodeButton
 							className="h-4 w-4 shrink-0 rounded-sm hover:bg-theme-surface-secondary active:bg-neutral-300"
-							code={`from zenml.client import Client
-step = Client().get_run_step("${data.id}")
-config = step.config`}
+							code={getStepSnippet(data.id)}
 							type="step"
 						>
 							<Copy className="h-3 w-3 fill-theme-text-tertiary" />
diff --git a/src/components/dag-visualizer/layout/helper.ts b/src/components/dag-visualizer/layout/helper.ts
@@ -48,9 +48,13 @@ export function mergeRealAndPlacehodlerData({
 	});
 
 	placeholderArtifacts.forEach((placeholderArtifact) => {
-		const existingRealArtifacts = realArtifacts.filter((realArtifact) =>
-			realArtifact.placeholderId.startsWith(placeholderArtifact.id)
-		);
+		const existingRealArtifacts = realArtifacts.filter((realArtifact) => {
+			const substitutedId = substitutePlaceholder(
+				placeholderArtifact.id,
+				realArtifact.substitutions || {}
+			);
+			return realArtifact.placeholderId.startsWith(substitutedId);
+		});
 		if (existingRealArtifacts.length === 0) {
 			finalNodes.push({
 				...placeholderArtifact,
@@ -125,3 +129,12 @@ export function mergeRealAndPlacehodlerData({
 
 	return { nodes: finalNodes, edges: styledEdges };
 }
+
+function substitutePlaceholder(initString: string, substitutuions: Record<string, string>) {
+	const regex = /\{([^}]+)\}/g;
+
+	return initString.replace(regex, (match, group) => {
+		const key = group.trim();
+		return substitutuions[key] || match;
+	});
+}
diff --git a/src/components/dag-visualizer/layout/real-data.ts b/src/components/dag-visualizer/layout/real-data.ts
@@ -3,13 +3,16 @@ import { RealNode, ZenEdge } from "@/types/pipeline-runs";
 import { StepDict } from "@/types/steps";
 import { addEdge } from "./helper";
 
-export function extractExistingNodes(stepConfig: StepDict) {
-	const nodes = extractNodes(stepConfig);
+export function extractExistingNodes(
+	stepConfig: StepDict,
+	substitutions: Record<string, Record<string, string>>
+) {
+	const nodes = extractNodes(stepConfig, substitutions);
 	const edges = extractEdges(stepConfig);
 	return { nodes, edges };
 }
 
-function extractNodes(stepConfig: StepDict) {
+function extractNodes(stepConfig: StepDict, substitutions: Record<string, Record<string, string>>) {
 	const nodes: RealNode[] = [];
 
 	Object.keys(stepConfig).forEach((stepName) => {
@@ -32,6 +35,7 @@ function extractNodes(stepConfig: StepDict) {
 					id: version.id,
 					placeholderId,
 					type: "artifact",
+					substitutions: substitutions[stepName] || {},
 					data: { ...version, name: outputName }
 				});
 			});
diff --git a/src/components/stacks/Sheet/index.tsx b/src/components/stacks/Sheet/index.tsx
@@ -147,9 +147,11 @@ function ComponentListItem({ component }: ComponentListItemProps) {
 					</div>
 				</div>
 			</div>
-			<ComponentBadge type={component.body?.type || "orchestrator"}>
-				{snakeCaseToTitleCase(component.body?.type || "")}
-			</ComponentBadge>
+			<Link to={routes.components.overview}>
+				<ComponentBadge type={component.body?.type || "orchestrator"}>
+					{snakeCaseToTitleCase(component.body?.type || "")}
+				</ComponentBadge>
+			</Link>
 		</Box>
 	);
 }
diff --git a/src/components/steps/step-sheet/ConfigurationTab.tsx b/src/components/steps/step-sheet/ConfigurationTab.tsx
@@ -4,6 +4,7 @@ import { CollapsibleCard } from "@/components/CollapsibleCard";
 import { usePipelineBuild } from "@/data/pipeline-builds/all-pipeline-builds-query";
 import { usePipelineRun } from "@/data/pipeline-runs/pipeline-run-detail-query";
 import { useStepDetail } from "@/data/steps/step-detail-query";
+import { getStepSnippet } from "@/lib/code-snippets";
 import { AnyDict } from "@/types/common";
 import { BuildItemMap } from "@/types/pipeline-builds";
 import { Skeleton } from "@zenml-io/react-component-library";
@@ -82,22 +83,21 @@ export function StepConfigTab({ stepId }: Props) {
 				title="Resources"
 				data={(data.metadata?.config.settings as { [key: string]: any })?.resources || {}}
 			/>
+			<NestedCollapsible
+				isInitialOpen
+				title="Substitutions"
+				data={data.metadata?.config.substitutions}
+			/>
 			<NestedCollapsible isInitialOpen title="Extra" data={data.metadata?.config.extra} />
 		</div>
 	);
 }
 
 function CodeSnippetCard({ id }: { id: string }) {
-	function returnConfigSchema(id: string) {
-		return `from zenml.client import Client
-artifact = Client().get_pipeline_run('${id}')
-config = run.config()`;
-	}
-
 	return (
 		<CollapsibleCard initialOpen title="Code">
 			<h2 className="mb-2  text-text-md text-theme-text-secondary">Get config programmatically</h2>
-			<Codesnippet fullWidth highlightCode wrap code={returnConfigSchema(id)} />
+			<Codesnippet fullWidth highlightCode wrap code={getStepSnippet(id)} />
 		</CollapsibleCard>
 	);
 }
diff --git a/src/lib/code-snippets.ts b/src/lib/code-snippets.ts
@@ -0,0 +1,27 @@
+export function getArtifactVersionSnippet(id: string) {
+	return `from zenml.client import Client
+
+artifact = Client().get_artifact_version("${id}")
+data = artifact.load()`;
+}
+
+export function getStepSnippet(id: string) {
+	return `from zenml.client import Client
+
+step = Client().get_run_step("${id}")
+config = step.config`;
+}
+
+export function getRunSnippet(id: string) {
+	return `from zenml.client import Client
+
+run = Client().get_pipeline_run("${id}")
+config = run.config`;
+}
+
+export function getSecretSnippet(id: string) {
+	return `from zenml.client import Client
+
+secret = Client().get_secret("${id}")
+`;
+}
diff --git a/src/lib/constants.ts b/src/lib/constants.ts
diff --git a/src/types/core.ts b/src/types/core.ts
diff --git a/src/types/pipeline-runs.ts b/src/types/pipeline-runs.ts

Original file line number	Diff line number	Diff line change
`@@ -7,10 +7,11 @@ import { Stack } from "@/types/stack";`
`7`	`7`	`import { User } from "@/types/user";`
`8`	`8`	`import { ColumnDef } from "@tanstack/react-table";`
`9`	`9`	`import { Avatar, AvatarFallback } from "@zenml-io/react-component-library";`
	`10`	`+import { useMemo } from "react";`
`10`	`11`	`import { StackActionsMenu } from "./ActionsDropdown";`
`11`	`12`
`12`		`-export function getStackColumns(): ColumnDef<Stack>[] {`
`13`		`- return [`
	`13`	`+export function useStackColumns(): ColumnDef<Stack>[] {`
	`14`	`+ const cols: ColumnDef<Stack>[] = [`
`14`	`15`	`{`
`15`	`16`	`id: "name",`
`16`	`17`	`header: "Stack",`
`@@ -67,4 +68,6 @@ export function getStackColumns(): ColumnDef<Stack>[] {`
`67`	`68`	`}`
`68`	`69`	`}`
`69`	`70`	`];`
	`71`	`+`
	`72`	`+ return useMemo(() => cols, []);`
`70`	`73`	`}`