Merge branch 'master' into feature/modules

wandb · Dec 7, 2024 · 5f517d3 · 5f517d3
2 parents 7bce780 + 059f388
commit 5f517d3
Show file tree

Hide file tree

Showing 5 changed files with 209 additions and 123 deletions.
diff --git a/tests/trace/test_trace_server_common.py b/tests/trace/test_trace_server_common.py
@@ -1,4 +1,5 @@
 from weave.trace_server.trace_server_common import (
+    DynamicBatchProcessor,
     LRUCache,
     get_nested_key,
     set_nested_key,
@@ -54,3 +55,26 @@ def test_lru_cache():
     cache["c"] = 10
     assert cache["c"] == 10
     assert cache["d"] == 4
+
+
+def test_dynamic_batch_processor():
+    # Initialize processor with:
+    # - initial batch size of 2
+    # - max size of 8
+    # - growth factor of 2
+    processor = DynamicBatchProcessor(initial_size=2, max_size=8, growth_factor=2)
+
+    test_data = range(15)
+
+    batches = list(processor.make_batches(iter(test_data)))
+
+    # Expected batch sizes: 2, 4, 8, 1
+    assert batches[0] == [0, 1]
+    assert batches[1] == [2, 3, 4, 5]
+    assert batches[2] == [6, 7, 8, 9, 10, 11, 12, 13]
+    assert batches[3] == [14]
+    assert len(batches) == 4
+
+    # Verify all items were processed
+    flattened = [item for batch in batches for item in batch]
+    assert flattened == list(range(15))
diff --git a/weave-js/src/components/FancyPage/FancyPageMenu.tsx b/weave-js/src/components/FancyPage/FancyPageMenu.tsx
@@ -60,7 +60,6 @@ export const FancyPageMenu = ({
               return null;
             }
             const linkProps = {
-              key: menuItem.slug,
               to: menuItem.isDisabled
                 ? {}
                 : {
@@ -76,7 +75,7 @@ export const FancyPageMenu = ({
               },
             };
             return (
-              <Link {...linkProps}>
+              <Link key={menuItem.slug} {...linkProps}>
                 <DropdownMenu.Item {...menuItemProps}>
                   <Icon name={menuItem.iconName} />
                   {menuItem.nameTooltip || menuItem.name}

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/CallsTable.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/CallsTable.tsx
@@ -85,7 +85,10 @@ import {TraceCallSchema} from '../wfReactInterface/traceServerClientTypes';
 import {traceCallToUICallSchema} from '../wfReactInterface/tsDataModelHooks';
 import {EXPANDED_REF_REF_KEY} from '../wfReactInterface/tsDataModelHooksCallRefExpansion';
 import {objectVersionNiceString} from '../wfReactInterface/utilities';
-import {CallSchema} from '../wfReactInterface/wfDataModelHooksInterface';
+import {
+  CallSchema,
+  OpVersionSchema,
+} from '../wfReactInterface/wfDataModelHooksInterface';
 import {CallsCharts} from './CallsCharts';
 import {CallsCustomColumnMenu} from './CallsCustomColumnMenu';
 import {
@@ -743,71 +746,13 @@ export const CallsTable: FC<{
         <TailwindContents>
           <RefreshButton onClick={() => calls.refetch()} />
           {!hideOpSelector && (
-            <div className="flex-none">
-              <ListItem
-                sx={{minWidth: 190, width: 320, height: 32, padding: 0}}>
-                <FormControl fullWidth sx={{borderColor: MOON_200}}>
-                  <Autocomplete
-                    PaperComponent={paperProps => (
-                      <StyledPaper {...paperProps} />
-                    )}
-                    sx={{
-                      '& .MuiOutlinedInput-root': {
-                        height: '32px',
-                        '& fieldset': {
-                          borderColor: MOON_200,
-                        },
-                        '&:hover fieldset': {
-                          borderColor: `rgba(${TEAL_300}, 0.48)`,
-                        },
-                      },
-                      '& .MuiOutlinedInput-input': {
-                        height: '32px',
-                        padding: '0 14px',
-                        boxSizing: 'border-box',
-                      },
-                    }}
-                    size="small"
-                    // Temp disable multiple for simplicity - may want to re-enable
-                    // multiple
-                    limitTags={1}
-                    disabled={Object.keys(frozenFilter ?? {}).includes(
-                      'opVersions'
-                    )}
-                    value={selectedOpVersionOption}
-                    onChange={(event, newValue) => {
-                      if (newValue === ALL_TRACES_OR_CALLS_REF_KEY) {
-                        setFilter({
-                          ...filter,
-                          opVersionRefs: [],
-                        });
-                      } else {
-                        setFilter({
-                          ...filter,
-                          opVersionRefs: newValue ? [newValue] : [],
-                        });
-                      }
-                    }}
-                    renderInput={renderParams => (
-                      <StyledTextField
-                        {...renderParams}
-                        sx={{maxWidth: '350px'}}
-                      />
-                    )}
-                    getOptionLabel={option => {
-                      return opVersionOptions[option]?.title ?? 'loading...';
-                    }}
-                    disableClearable={
-                      selectedOpVersionOption === ALL_TRACES_OR_CALLS_REF_KEY
-                    }
-                    groupBy={option => opVersionOptions[option]?.group}
-                    options={Object.keys(opVersionOptions)}
-                    popupIcon={<Icon name="chevron-down" />}
-                    clearIcon={<Icon name="close" />}
-                  />
-                </FormControl>
-              </ListItem>
-            </div>
+            <OpSelector
+              frozenFilter={frozenFilter}
+              filter={filter}
+              setFilter={setFilter}
+              selectedOpVersionOption={selectedOpVersionOption}
+              opVersionOptions={opVersionOptions}
+            />
           )}
           {filterModel && setFilterModel && (
             <FilterPanel
@@ -1088,6 +1033,90 @@ export const CallsTable: FC<{
   );
 };
 
+const OpSelector = ({
+  frozenFilter,
+  filter,
+  setFilter,
+  selectedOpVersionOption,
+  opVersionOptions,
+}: {
+  frozenFilter: WFHighLevelCallFilter | undefined;
+  filter: WFHighLevelCallFilter;
+  setFilter: (state: WFHighLevelCallFilter) => void;
+  selectedOpVersionOption: string;
+  opVersionOptions: Record<
+    string,
+    {
+      title: string;
+      ref: string;
+      group: string;
+      objectVersion?: OpVersionSchema;
+    }
+  >;
+}) => {
+  const frozenOpFilter = Object.keys(frozenFilter ?? {}).includes('opVersions');
+  const handleChange = useCallback(
+    (event: any, newValue: string | null) => {
+      if (newValue === ALL_TRACES_OR_CALLS_REF_KEY) {
+        setFilter({
+          ...filter,
+          opVersionRefs: [],
+        });
+      } else {
+        setFilter({
+          ...filter,
+          opVersionRefs: newValue ? [newValue] : [],
+        });
+      }
+    },
+    [filter, setFilter]
+  );
+
+  return (
+    <div className="flex-none">
+      <ListItem sx={{minWidth: 190, width: 320, height: 32, padding: 0}}>
+        <FormControl fullWidth sx={{borderColor: MOON_200}}>
+          <Autocomplete
+            PaperComponent={paperProps => <StyledPaper {...paperProps} />}
+            sx={{
+              '& .MuiOutlinedInput-root': {
+                height: '32px',
+                '& fieldset': {
+                  borderColor: MOON_200,
+                },
+                '&:hover fieldset': {
+                  borderColor: `rgba(${TEAL_300}, 0.48)`,
+                },
+              },
+              '& .MuiOutlinedInput-input': {
+                height: '32px',
+                padding: '0 14px',
+                boxSizing: 'border-box',
+              },
+            }}
+            size="small"
+            limitTags={1}
+            disabled={frozenOpFilter}
+            value={selectedOpVersionOption}
+            onChange={handleChange}
+            renderInput={renderParams => (
+              <StyledTextField {...renderParams} sx={{maxWidth: '350px'}} />
+            )}
+            getOptionLabel={option => opVersionOptions[option]?.title ?? ''}
+            disableClearable={
+              selectedOpVersionOption === ALL_TRACES_OR_CALLS_REF_KEY
+            }
+            groupBy={option => opVersionOptions[option]?.group}
+            options={Object.keys(opVersionOptions)}
+            popupIcon={<Icon name="chevron-down" />}
+            clearIcon={<Icon name="close" />}
+          />
+        </FormControl>
+      </ListItem>
+    </div>
+  );
+};
+
 const ButtonDivider = () => (
   <div className="h-24 flex-none border-l-[1px] border-moon-250"></div>
 );