0g-orbit 0.1.0 → 0.2.0

package/src/fine-tuning.test.ts

@@ -0,0 +1,299 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest'
+import { NETWORKS } from './networks.js'
+import { FineTuningError } from './errors.js'
+
+// Hoist mocks so they can be referenced in vi.mock factories
+const {
+    mockListModel,
+    mockListService,
+    mockListTask,
+    mockGetTask,
+    mockGetLog,
+    mockCreateTask,
+    mockAcknowledgeModel,
+    mockAcknowledgeProviderSigner,
+    mockStore,
+} = vi.hoisted(() => ({
+    mockListModel: vi.fn(),
+    mockListService: vi.fn(),
+    mockListTask: vi.fn(),
+    mockGetTask: vi.fn(),
+    mockGetLog: vi.fn(),
+    mockCreateTask: vi.fn(),
+    mockAcknowledgeModel: vi.fn(),
+    mockAcknowledgeProviderSigner: vi.fn(),
+    mockStore: vi.fn(),
+}))
+
+vi.mock('@0glabs/0g-serving-broker', () => ({
+    createZGComputeNetworkBroker: vi.fn().mockResolvedValue({
+        fineTuning: {
+            listModel: mockListModel,
+            listService: mockListService,
+            listTask: mockListTask,
+            getTask: mockGetTask,
+            getLog: mockGetLog,
+            createTask: mockCreateTask,
+            acknowledgeModel: mockAcknowledgeModel,
+            acknowledgeProviderSigner: mockAcknowledgeProviderSigner,
+        },
+        ledger: {
+            getLedger: vi.fn(),
+            addLedger: vi.fn(),
+        },
+        inference: {},
+    }),
+}))
+
+vi.mock('./storage.js', () => ({
+    StorageClient: vi.fn().mockImplementation(() => ({
+        store: mockStore,
+    })),
+}))
+
+import { FineTuningClient } from './fine-tuning.js'
+import { StorageClient } from './storage.js'
+
+const mockWallet = {
+    address: '0x1234567890abcdef1234567890abcdef12345678',
+    privateKey: '0xdeadbeef',
+} as any
+const testNetwork = { ...NETWORKS.testnet }
+const mockStorageClient = new StorageClient(testNetwork, mockWallet)
+
+describe('FineTuningClient', () => {
+    let client: FineTuningClient
+
+    beforeEach(() => {
+        vi.clearAllMocks()
+        client = new FineTuningClient(testNetwork, mockWallet, mockStorageClient)
+    })
+
+    describe('lazy broker init', () => {
+        it('initializes broker on first call', async () => {
+            mockListModel.mockResolvedValueOnce([[], []])
+
+            await client.listModels()
+            const { createZGComputeNetworkBroker } = await import('@0glabs/0g-serving-broker')
+            expect(createZGComputeNetworkBroker).toHaveBeenCalledTimes(1)
+        })
+
+        it('reuses broker on subsequent calls', async () => {
+            mockListModel.mockResolvedValue([[], []])
+
+            await client.listModels()
+            await client.listModels()
+            const { createZGComputeNetworkBroker } = await import('@0glabs/0g-serving-broker')
+            expect(createZGComputeNetworkBroker).toHaveBeenCalledTimes(1)
+        })
+
+        it('throws FineTuningError on broker init failure', async () => {
+            const { createZGComputeNetworkBroker } = await import('@0glabs/0g-serving-broker')
+            ;(createZGComputeNetworkBroker as any).mockRejectedValueOnce(
+                new Error('Network unreachable')
+            )
+
+            const freshClient = new FineTuningClient(testNetwork, mockWallet, mockStorageClient)
+            try {
+                await freshClient.listModels()
+                expect.unreachable('Should have thrown')
+            } catch (err) {
+                expect(err).toBeInstanceOf(FineTuningError)
+                expect((err as FineTuningError).message).toContain('Failed to initialize')
+            }
+        })
+    })
+
+    describe('uploadDataset', () => {
+        it('uploads a dataset and returns root + tx hash', async () => {
+            mockStore.mockResolvedValueOnce({
+                root: '0xdataset_root',
+                txHash: '0xtx_hash',
+            })
+
+            const result = await client.uploadDataset('/tmp/dataset.jsonl')
+            expect(result.root).toBe('0xdataset_root')
+            expect(result.txHash).toBe('0xtx_hash')
+            expect(mockStore).toHaveBeenCalledWith('/tmp/dataset.jsonl')
+        })
+
+        it('throws FineTuningError on file error', async () => {
+            mockStore.mockRejectedValueOnce(new Error('ENOENT: no such file'))
+
+            await expect(client.uploadDataset('/tmp/missing.jsonl')).rejects.toThrow(
+                FineTuningError
+            )
+        })
+    })
+
+    describe('createTask', () => {
+        it('creates a fine-tuning task', async () => {
+            mockAcknowledgeProviderSigner.mockResolvedValueOnce(undefined)
+            mockCreateTask.mockResolvedValueOnce('task-001')
+
+            const task = await client.createTask({
+                model: 'Qwen2.5-0.5B-Instruct',
+                dataset: '0xdataset_root',
+                providerAddress: '0xprovider',
+            })
+
+            expect(task.id).toBe('task-001')
+            expect(task.model).toBe('Qwen2.5-0.5B-Instruct')
+            expect(task.dataset).toBe('0xdataset_root')
+            expect(task.provider).toBe('0xprovider')
+            expect(task.status).toBe('init')
+        })
+
+        it('throws FineTuningError on invalid model', async () => {
+            mockAcknowledgeProviderSigner.mockResolvedValueOnce(undefined)
+            mockCreateTask.mockRejectedValueOnce(new Error('Model not found'))
+
+            await expect(
+                client.createTask({
+                    model: 'nonexistent-model',
+                    dataset: '0xdataset_root',
+                    providerAddress: '0xprovider',
+                })
+            ).rejects.toThrow(FineTuningError)
+        })
+
+        it('throws FineTuningError on missing provider', async () => {
+            mockAcknowledgeProviderSigner.mockRejectedValueOnce(
+                new Error('Provider not found')
+            )
+
+            await expect(
+                client.createTask({
+                    model: 'Qwen2.5-0.5B-Instruct',
+                    dataset: '0xdataset_root',
+                    providerAddress: '0xbad_provider',
+                })
+            ).rejects.toThrow(FineTuningError)
+        })
+    })
+
+    describe('getTask', () => {
+        it('returns a task by ID', async () => {
+            mockGetTask.mockResolvedValueOnce({
+                id: 'task-001',
+                preTrainedModelHash: '0xmodel',
+                datasetHash: '0xdata',
+                progress: 'Training',
+                createdAt: '2024-01-01',
+                updatedAt: '2024-01-02',
+            })
+
+            const task = await client.getTask('0xprovider', 'task-001')
+            expect(task.id).toBe('task-001')
+            expect(task.status).toBe('training')
+            expect(task.provider).toBe('0xprovider')
+        })
+
+        it('throws FineTuningError when task not found', async () => {
+            mockGetTask.mockRejectedValueOnce(new Error('No task found'))
+
+            await expect(
+                client.getTask('0xprovider', 'nonexistent')
+            ).rejects.toThrow(FineTuningError)
+        })
+    })
+
+    describe('getTaskLog', () => {
+        it('returns the training log', async () => {
+            mockGetLog.mockResolvedValueOnce('Epoch 1/3: loss=0.5\nEpoch 2/3: loss=0.3')
+
+            const log = await client.getTaskLog('0xprovider', 'task-001')
+            expect(log).toContain('Epoch 1/3')
+        })
+    })
+
+    describe('listTasks', () => {
+        it('returns all tasks for a provider', async () => {
+            mockListTask.mockResolvedValueOnce([
+                { id: 'task-001', preTrainedModelHash: '0xm1', datasetHash: '0xd1', progress: 'Finished' },
+                { id: 'task-002', preTrainedModelHash: '0xm2', datasetHash: '0xd2', progress: 'Training' },
+            ])
+
+            const tasks = await client.listTasks('0xprovider')
+            expect(tasks).toHaveLength(2)
+            expect(tasks[0].status).toBe('finished')
+            expect(tasks[1].status).toBe('training')
+        })
+    })
+
+    describe('downloadModel', () => {
+        it('downloads a model via acknowledgeModel', async () => {
+            mockAcknowledgeModel.mockResolvedValueOnce(undefined)
+
+            await expect(
+                client.downloadModel('0xprovider', 'task-001', '/tmp/model')
+            ).resolves.toBeUndefined()
+
+            expect(mockAcknowledgeModel).toHaveBeenCalledWith(
+                '0xprovider',
+                'task-001',
+                '/tmp/model'
+            )
+        })
+
+        it('throws FineTuningError when task not ready', async () => {
+            mockAcknowledgeModel.mockRejectedValueOnce(
+                new Error('No deliverable found')
+            )
+
+            await expect(
+                client.downloadModel('0xprovider', 'task-001', '/tmp/model')
+            ).rejects.toThrow(FineTuningError)
+        })
+    })
+
+    describe('listModels', () => {
+        it('returns combined standard and customized models', async () => {
+            mockListModel.mockResolvedValueOnce([
+                [['Qwen2.5-0.5B-Instruct', { turbo: '0xhash', description: 'Qwen model' }]],
+                [['my-custom', { description: 'Custom fine-tuned', provider: '0xp1' }]],
+            ])
+
+            const models = await client.listModels()
+            expect(models).toHaveLength(2)
+            expect(models[0].name).toBe('Qwen2.5-0.5B-Instruct')
+            expect(models[1].name).toBe('my-custom')
+            expect(models[1].config.provider).toBe('0xp1')
+        })
+
+        it('returns empty array when no models', async () => {
+            mockListModel.mockResolvedValueOnce([[], []])
+
+            const models = await client.listModels()
+            expect(models).toHaveLength(0)
+        })
+    })
+
+    describe('listProviders', () => {
+        it('returns fine-tuning providers', async () => {
+            mockListService.mockResolvedValueOnce([
+                { provider: '0xprovider1', url: 'https://p1.example.com', models: ['model1'] },
+                { provider: '0xprovider2', url: 'https://p2.example.com', models: [] },
+            ])
+
+            const providers = await client.listProviders()
+            expect(providers).toHaveLength(2)
+            expect(providers[0].address).toBe('0xprovider1')
+            expect(providers[0].url).toBe('https://p1.example.com')
+            expect(providers[0].models).toEqual(['model1'])
+        })
+    })
+
+    describe('error suggestions', () => {
+        it('FineTuningError includes default suggestion', () => {
+            const err = new FineTuningError('Something failed')
+            expect(err.code).toBe('FINE_TUNING_ERROR')
+            expect(err.suggestion).toContain('listModels')
+        })
+
+        it('FineTuningError allows custom suggestion', () => {
+            const err = new FineTuningError('Oops', 'Try this instead')
+            expect(err.suggestion).toBe('Try this instead')
+        })
+    })
+})

package/src/fine-tuning.ts

@@ -0,0 +1,330 @@
+import { Wallet } from 'ethers'
+import {
+    createZGComputeNetworkBroker,
+    type ZGComputeNetworkBroker,
+} from '@0glabs/0g-serving-broker'
+import type { NetworkConfig } from './networks.js'
+import type {
+    DatasetUploadResult,
+    CreateTaskOptions,
+    FineTuneTask,
+    FineTuneModel,
+    FineTuneProvider,
+    FineTuneStatus,
+} from './types.js'
+import { FineTuningError } from './errors.js'
+import { withRetry } from './retry.js'
+import { StorageClient } from './storage.js'
+
+const DEFAULT_TRAINING_PARAMS = {
+    nEpochs: 3,
+    batchSize: 4,
+    learningRate: 5e-5,
+    loraRank: 8,
+    loraAlpha: 16,
+}
+
+export class FineTuningClient {
+    private broker: ZGComputeNetworkBroker | null = null
+    private initialized = false
+    private wallet: Wallet
+    private network: NetworkConfig
+    private storageClient: StorageClient
+
+    constructor(network: NetworkConfig, wallet: Wallet, storageClient: StorageClient) {
+        this.network = network
+        this.wallet = wallet
+        this.storageClient = storageClient
+    }
+
+    private async ensureBroker(): Promise<ZGComputeNetworkBroker> {
+        if (this.broker && this.initialized) return this.broker
+
+        try {
+            this.broker = await createZGComputeNetworkBroker(
+                this.wallet as any,
+                this.network.ledgerContractAddress,
+                this.network.inferenceContractAddress,
+                this.network.fineTuningContractAddress
+            )
+            this.initialized = true
+            return this.broker
+        } catch (err: unknown) {
+            const msg = err instanceof Error ? err.message : String(err)
+            throw new FineTuningError(
+                `Failed to initialize fine-tuning broker: ${msg}`,
+                'Check your network connection and wallet private key. The compute contracts may be temporarily unavailable.'
+            )
+        }
+    }
+
+    // --- Dataset ---
+
+    /**
+     * Upload a dataset to 0G Storage for fine-tuning.
+     * Uses the existing StorageClient for upload, returning the root hash
+     * needed to create a fine-tuning task.
+     */
+    async uploadDataset(filePath: string): Promise<DatasetUploadResult> {
+        return withRetry(
+            async () => {
+                try {
+                    const result = await this.storageClient.store(filePath)
+                    return {
+                        root: result.root,
+                        txHash: result.txHash,
+                    }
+                } catch (err: unknown) {
+                    const msg = err instanceof Error ? err.message : String(err)
+                    throw new FineTuningError(
+                        `Failed to upload dataset: ${msg}`,
+                        'Check that the dataset file exists, is valid JSONL, and you have sufficient OG balance.'
+                    )
+                }
+            },
+            { maxAttempts: 3 }
+        )
+    }
+
+    // --- Tasks ---
+
+    /**
+     * Create a fine-tuning task. Requires a dataset already uploaded to 0G Storage.
+     */
+    async createTask(options: CreateTaskOptions): Promise<FineTuneTask> {
+        const broker = await this.ensureBroker()
+
+        if (!broker.fineTuning) {
+            throw new FineTuningError(
+                'Fine-tuning broker not available.',
+                'Fine-tuning requires a Wallet signer (not JsonRpcSigner).'
+            )
+        }
+
+        const params = { ...DEFAULT_TRAINING_PARAMS, ...options.trainingParams }
+
+        return withRetry(
+            async () => {
+                try {
+                    // Acknowledge the provider signer (required before task creation)
+                    await broker.fineTuning!.acknowledgeProviderSigner(
+                        options.providerAddress
+                    )
+
+                    // Write training params as JSON string (broker expects a file path,
+                    // but the underlying createTask reads the file content — we'll pass
+                    // the JSON string directly via a temp file approach)
+                    const { writeFileSync, unlinkSync } = await import('node:fs')
+                    const { tmpdir } = await import('node:os')
+                    const { join } = await import('node:path')
+                    const { randomBytes } = await import('node:crypto')
+
+                    const tempPath = join(tmpdir(), `orbit-params-${randomBytes(8).toString('hex')}.json`)
+                    writeFileSync(tempPath, JSON.stringify(params))
+
+                    let taskId: string
+                    try {
+                        taskId = await broker.fineTuning!.createTask(
+                            options.providerAddress,
+                            options.model,
+                            options.dataset,
+                            tempPath
+                        )
+                    } finally {
+                        try { unlinkSync(tempPath) } catch { /* ignore cleanup errors */ }
+                    }
+
+                    return {
+                        id: taskId,
+                        model: options.model,
+                        dataset: options.dataset,
+                        provider: options.providerAddress,
+                        status: 'init' as FineTuneStatus,
+                    }
+                } catch (err: unknown) {
+                    const msg = err instanceof Error ? err.message : String(err)
+                    if (msg.includes('User opted not to continue')) {
+                        throw new FineTuningError(msg, 'The provider has pending tasks in queue. Try again later or use a different provider.')
+                    }
+                    throw new FineTuningError(
+                        `Failed to create task: ${msg}`,
+                        'Verify the model name, dataset hash, and provider address. Run orbit.listModels() to see available options.'
+                    )
+                }
+            },
+            { maxAttempts: 2 }
+        )
+    }
+
+    /**
+     * Get the status of a fine-tuning task.
+     */
+    async getTask(providerAddress: string, taskId: string): Promise<FineTuneTask> {
+        const broker = await this.ensureBroker()
+
+        if (!broker.fineTuning) {
+            throw new FineTuningError('Fine-tuning broker not available.')
+        }
+
+        try {
+            const task = await broker.fineTuning.getTask(providerAddress, taskId)
+            return this.mapTask(task, providerAddress)
+        } catch (err: unknown) {
+            const msg = err instanceof Error ? err.message : String(err)
+            if (msg.includes('No task found') || msg.includes('not found')) {
+                throw new FineTuningError(
+                    `Task "${taskId}" not found for provider ${providerAddress}.`,
+                    'Check the task ID and provider address. Run orbit.listTasks() to see your tasks.'
+                )
+            }
+            throw new FineTuningError(`Failed to get task: ${msg}`)
+        }
+    }
+
+    /**
+     * Get the training log for a fine-tuning task.
+     */
+    async getTaskLog(providerAddress: string, taskId: string): Promise<string> {
+        const broker = await this.ensureBroker()
+
+        if (!broker.fineTuning) {
+            throw new FineTuningError('Fine-tuning broker not available.')
+        }
+
+        try {
+            return await broker.fineTuning.getLog(providerAddress, taskId)
+        } catch (err: unknown) {
+            const msg = err instanceof Error ? err.message : String(err)
+            throw new FineTuningError(`Failed to get task log: ${msg}`)
+        }
+    }
+
+    /**
+     * List all fine-tuning tasks for a given provider.
+     */
+    async listTasks(providerAddress: string): Promise<FineTuneTask[]> {
+        const broker = await this.ensureBroker()
+
+        if (!broker.fineTuning) {
+            throw new FineTuningError('Fine-tuning broker not available.')
+        }
+
+        try {
+            const tasks = await broker.fineTuning.listTask(providerAddress)
+            return tasks.map((t) => this.mapTask(t, providerAddress))
+        } catch (err: unknown) {
+            const msg = err instanceof Error ? err.message : String(err)
+            throw new FineTuningError(`Failed to list tasks: ${msg}`)
+        }
+    }
+
+    // --- Models ---
+
+    /**
+     * Download a fine-tuned model. Combines acknowledge + download.
+     * The task must be in 'delivered' status.
+     */
+    async downloadModel(
+        providerAddress: string,
+        taskId: string,
+        outputPath: string
+    ): Promise<void> {
+        const broker = await this.ensureBroker()
+
+        if (!broker.fineTuning) {
+            throw new FineTuningError('Fine-tuning broker not available.')
+        }
+
+        return withRetry(
+            async () => {
+                try {
+                    await broker.fineTuning!.acknowledgeModel(
+                        providerAddress,
+                        taskId,
+                        outputPath
+                    )
+                } catch (err: unknown) {
+                    const msg = err instanceof Error ? err.message : String(err)
+                    if (msg.includes('No deliverable found')) {
+                        throw new FineTuningError(
+                            `Model not ready for task "${taskId}". The task may still be training.`,
+                            'Check task status with orbit.getFineTuneTask(). The task must be in "delivered" status.'
+                        )
+                    }
+                    throw new FineTuningError(
+                        `Failed to download model: ${msg}`,
+                        'Ensure the task is in "delivered" status and the output path is writable.'
+                    )
+                }
+            },
+            { maxAttempts: 2 }
+        )
+    }
+
+    /**
+     * List available base models for fine-tuning.
+     */
+    async listModels(): Promise<FineTuneModel[]> {
+        const broker = await this.ensureBroker()
+
+        if (!broker.fineTuning) {
+            throw new FineTuningError('Fine-tuning broker not available.')
+        }
+
+        try {
+            const [standardModels, customizedModels] = await broker.fineTuning.listModel()
+
+            const models: FineTuneModel[] = []
+
+            for (const [name, config] of standardModels) {
+                models.push({ name, config })
+            }
+
+            for (const [name, config] of customizedModels) {
+                models.push({ name, config })
+            }
+
+            return models
+        } catch (err: unknown) {
+            const msg = err instanceof Error ? err.message : String(err)
+            throw new FineTuningError(`Failed to list models: ${msg}`)
+        }
+    }
+
+    /**
+     * List fine-tuning service providers.
+     */
+    async listProviders(): Promise<FineTuneProvider[]> {
+        const broker = await this.ensureBroker()
+
+        if (!broker.fineTuning) {
+            throw new FineTuningError('Fine-tuning broker not available.')
+        }
+
+        try {
+            const services = await broker.fineTuning.listService()
+            return services.map((s: any) => ({
+                address: s.provider ?? '',
+                url: s.url ?? '',
+                models: s.models ?? [],
+            }))
+        } catch (err: unknown) {
+            const msg = err instanceof Error ? err.message : String(err)
+            throw new FineTuningError(`Failed to list providers: ${msg}`)
+        }
+    }
+
+    // --- Helpers ---
+
+    private mapTask(task: any, providerAddress: string): FineTuneTask {
+        return {
+            id: task.id ?? '',
+            model: task.preTrainedModelHash ?? '',
+            dataset: task.datasetHash ?? '',
+            provider: providerAddress,
+            status: (task.progress?.toLowerCase() ?? 'init') as FineTuneStatus,
+            createdAt: task.createdAt,
+            updatedAt: task.updatedAt,
+        }
+    }
+}

package/src/index.ts

@@ -4,6 +4,7 @@ export { Orbit } from './orbit.js'
 // Clients (for direct access)
 export { StorageClient } from './storage.js'
 export { InferenceClient } from './inference.js'
+export { FineTuningClient } from './fine-tuning.js'
 
 // Types
 export type {
@@ -15,6 +16,12 @@ export type {
     InferOptions,
     ServiceInfo,
     AccountStatus,
+    DatasetUploadResult,
+    CreateTaskOptions,
+    FineTuneTask,
+    FineTuneStatus,
+    FineTuneModel,
+    FineTuneProvider,
 } from './types.js'
 
 // Network config
@@ -34,4 +41,5 @@ export {
     InsufficientBalanceError,
     ProviderNotFoundError,
     TimeoutError,
+    FineTuningError,
 } from './errors.js'