chore(weave_query): add file ops for artifact memberships (#3735)

* add file ops

* add comment for commithash on portfolio

* remove error

* remove logs

* weave-js lint

* try fixing node test?

* remove path check

* remove files() ops;

* remove extra file ops and delete line deletes

* delete console logs

* temp logging

* remove console logs and fetch artifact via version instead of commit hash for membership files

* lint

weave-js/src/core/ops/domain/artifactMembership.ts

@@ -1,11 +1,18 @@
+import {opFileType} from '@wandb/weave/core';
+import * as TypeHelpers from '@wandb/weave/core/model/helpers';
+import {docType} from '@wandb/weave/core/util/docs';
+import * as _ from 'lodash';
+
 import {artifact} from '../../_external/util/urls';
-import {list} from '../../model';
+import {list, maybe, union} from '../../model';
 import {makeStandardOp} from '../opKinds';
 
 const artifactArgTypes = {
   artifactMembership: 'artifactMembership' as const,
 };
 
+const artifactMembershipArgDescription = `A ${docType('artifactMembership')}`;
+
 export const opArtifactMembershipId = makeStandardOp({
   hidden: true,
   name: 'artifactMembership-id',
@@ -79,3 +86,82 @@ export const opArtifactMembershipLink = makeStandardOp({
     }),
   }),
 });
+
+// Same as opArtifactVersionFile
+export const opArtifactMembershipFile = makeStandardOp({
+  name: 'artifactMembership-file',
+  argTypes: {...artifactArgTypes, path: 'string'},
+  description: `Returns the ${docType('file')} of the ${docType(
+    'artifactMembership'
+  )} for the given path`,
+  argDescriptions: {
+    artifactMembership: artifactMembershipArgDescription,
+    path: `The path of the ${docType('file')}`,
+  },
+  returnValueDescription: `The ${docType('file')} of the ${docType(
+    'artifactMembership'
+  )} for the given path`,
+  returnType: inputTypes => maybe({type: 'file'}),
+  resolver: async (
+    {artifactMembership, path},
+    inputTypes,
+    rawInputs,
+    forwardGraph,
+    forwardOp,
+    context
+  ) => {
+    if (artifactMembership == null) {
+      throw new Error('opArtifactMembershipFile missing artifactMembership');
+    }
+    if (artifactMembership.artifact == null) {
+      throw new Error('opArtifactMembershipFile missing artifact');
+    }
+    try {
+      const result = await context.backend.getArtifactFileMetadata(
+        artifactMembership.artifact.id,
+        path
+      );
+      if (result == null) {
+        return null;
+      }
+      return {artifact: artifactMembership.artifact, path};
+    } catch (e) {
+      console.warn('Error loading artifact from membership', {
+        err: e,
+        artifact: artifactMembership.artifact,
+        path,
+      });
+      return null;
+    }
+  },
+  resolveOutputType: async (inputTypes, node, executableNode, client) => {
+    const fileTypeNode = opFileType({
+      file: executableNode as any,
+    });
+    let fileType = await client.query(fileTypeNode);
+
+    if (fileType == null) {
+      return 'none';
+    }
+
+    // The standard op pattern does not handle returning types as arrays. This
+    // should be merged into standard op, but since opFileType is the only "type"
+    // op, and only used here, just keeping it simple.
+    // This is for Weave0
+    if (_.isArray(fileType)) {
+      fileType = union(fileType.map(t => (t == null ? 'none' : t)));
+    }
+
+    // This is a Weave1 hack. Weave1's type refinement ops (like file-type) return
+    // tagged/mapped results. But the Weave0 opKinds framework is going to do
+    // the same wrapping to the result we return here. So we unwrap the Weave1
+    // result and let it be rewrapped.
+    if (TypeHelpers.isTaggedValue(fileType)) {
+      fileType = TypeHelpers.taggedValueValueType(fileType);
+    }
+    if (TypeHelpers.isList(fileType)) {
+      fileType = fileType.objectType;
+    }
+    return fileType ?? 'none';
+  },
+});

weave-js/src/core/ops/domain/gql.ts

@@ -1440,6 +1440,13 @@ export const toGqlField = (
         ]),
       },
     ];
+  } else if (forwardOp.op.name === 'artifactMembership-file') {
+    return [
+      {
+        name: 'artifact',
+        fields: gqlBasicField('id'),
+      },
+    ];
   } else if (forwardOp.op.name === 'artifact-memberships') {
     return [
       {

weave_query/weave_query/artifact_wandb.py

@@ -994,7 +994,7 @@ def uri_path(self) -> str:
 
     @property
     def resolved_artifact_uri(self) -> "WeaveWBArtifactURI":
-        if self.version and likely_commit_hash(self.version):
+        if self.version and (likely_commit_hash(self.version) or is_valid_version_index(self.version)):
             return self
         if self._resolved_artifact_uri is None:
             path = f"{self.entity_name}/{self.project_name}/{self.name}"

weave_query/weave_query/ops_domain/artifact_membership_ops.py

@@ -1,6 +1,9 @@
+import logging
 import urllib
+import typing
 
 from weave_query import weave_types as types
+from weave_query import artifact_fs, artifact_wandb
 from weave_query.api import op
 from weave_query.gql_op_plugin import wb_gql_op_plugin
 from weave_query.ops_domain import wb_domain_types as wdt
@@ -9,6 +12,27 @@
     gql_prop_op,
 )
 
+static_art_membership_file_gql = """
+            versionIndex
+            artifactCollection {
+                id 
+                name 
+                defaultArtifactType {
+                    id 
+                    name 
+                    project {
+                        id 
+                        name 
+                        entity {
+                            id 
+                            name
+                        }
+                    }
+                }
+            }
+            """
+
+
 # Section 1/6: Tag Getters
 # None
 
@@ -71,25 +95,7 @@
 @op(
     name="artifactMembership-link",
     plugins=wb_gql_op_plugin(
-        lambda inputs, inner: """
-            versionIndex
-            artifactCollection {
-                id 
-                name 
-                defaultArtifactType {
-                    id 
-                    name 
-                    project {
-                        id 
-                        name 
-                        entity {
-                            id 
-                            name
-                        }
-                    }
-                }
-            }
-        """,
+        lambda inputs, inner: static_art_membership_file_gql,
     ),
 )
 def artifact_membership_link(
@@ -103,3 +109,49 @@ def artifact_membership_link(
         f"{urllib.parse.quote(artifactMembership['artifactCollection']['name'])}"
         f"/v{artifactMembership['versionIndex']}",
     )
+
+
+def _artifact_membership_to_wb_artifact(artifactMembership: wdt.ArtifactCollectionMembership):
+    type_name = artifactMembership["artifactCollection"]["defaultArtifactType"]["name"]
+    collection_name = artifactMembership["artifactCollection"]["name"]
+
+    # This is valid because the commitHash for portfolios is always null. So we will leverage
+    # the artifact's membership in its source collection to fetch it via the commitHash in
+    # downstream paths
+    version = f"v{artifactMembership['versionIndex']}"
+    entity_name = artifactMembership["artifactCollection"]['defaultArtifactType']['project']['entity']['name']
+    project_name = artifactMembership["artifactCollection"]['defaultArtifactType']['project']['name']
+    uri = artifact_wandb.WeaveWBArtifactURI(
+        collection_name, version, entity_name, project_name
+    )
+    return artifact_wandb.WandbArtifact(
+        name=collection_name,
+        type=type_name,
+        uri=uri,
+    )
+
+# Same as the artifactVersion-_file_refine_output_type op
+@op(
+    name="artifactMembership-_file_refine_output_type",
+    hidden=True,
+    output_type=types.TypeType(),
+    plugins=wb_gql_op_plugin(lambda inputs, inner: static_art_membership_file_gql),
+)
+def _file_refine_output_type(artifactMembership: wdt.ArtifactCollectionMembership, path: str):
+    art_local = _artifact_membership_to_wb_artifact(artifactMembership)
+    return types.TypeRegistry.type_of(art_local.path_info(path))
+
+
+# Same as the artifactVersion-file op
+@op(
+    name="artifactMembership-file",
+    refine_output_type=_file_refine_output_type,
+    plugins=wb_gql_op_plugin(lambda inputs, inner: static_art_membership_file_gql),
+)
+def file_(
+        artifactMembership: wdt.ArtifactCollectionMembership, path: str
+) -> typing.Union[
+    None, artifact_fs.FilesystemArtifactFile  # , artifact_fs.FilesystemArtifactDir
+]:
+    art_local = _artifact_membership_to_wb_artifact(artifactMembership)
+    return art_local.path_info(path)  # type: ignore